510 lines
21 KiB
Python
510 lines
21 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import socket
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
from . import __version__
|
|
from .agent_core import AgentSession, AgentToolRegistry
|
|
from .config import AgentConfig
|
|
from .controller_discovery import resolve_controller_base_url
|
|
from .http_client import HttpClientError, request_json
|
|
from .identity import deployment_info, host_name, machine_fingerprint, public_id
|
|
from .local_service import LocalService
|
|
from .model_config import (
|
|
RuntimeRestartManager,
|
|
load_applied_model_config,
|
|
public_model_config_summary,
|
|
save_applied_model_config,
|
|
)
|
|
from .state import load_state, save_state
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
|
|
|
|
def local_ip() -> str:
|
|
try:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
|
|
sock.connect(("8.8.8.8", 80))
|
|
return sock.getsockname()[0]
|
|
except OSError:
|
|
return "127.0.0.1"
|
|
|
|
|
|
class ControlledAgent:
|
|
def __init__(self, config: AgentConfig):
|
|
self.config = AgentConfig(
|
|
root_dir=config.root_dir,
|
|
controller_url=config.controller_url,
|
|
controller_discovery_enabled=config.controller_discovery_enabled,
|
|
controller_discovery_ports=config.controller_discovery_ports,
|
|
controller_discovery_timeout_seconds=config.controller_discovery_timeout_seconds,
|
|
state_path=config.state_path,
|
|
model_config_path=config.model_config_path,
|
|
runtime_state_path=config.runtime_state_path,
|
|
memory_path=config.memory_path,
|
|
display_name=config.display_name,
|
|
deployment_type=config.deployment_type,
|
|
heartbeat_interval_seconds=config.heartbeat_interval_seconds,
|
|
task_poll_interval_seconds=config.task_poll_interval_seconds,
|
|
controller_retry_interval_seconds=config.controller_retry_interval_seconds,
|
|
request_timeout_seconds=config.request_timeout_seconds,
|
|
local_service_enabled=config.local_service_enabled,
|
|
local_service_host=config.local_service_host,
|
|
local_service_port=config.local_service_port,
|
|
runtime_command=config.runtime_command,
|
|
)
|
|
self.controller_url: str | None = None
|
|
self.last_connection_error: str | None = None
|
|
self.controller_retry_count = 0
|
|
self.runtime_started = False
|
|
self.registration_confirmed = False
|
|
self.fingerprint = machine_fingerprint()
|
|
self.public_id = public_id()
|
|
self.host = host_name()
|
|
self.ip = local_ip()
|
|
self.started_at = utc_now()
|
|
self.last_heartbeat_at: str | None = None
|
|
self.last_task_claim_at: str | None = None
|
|
self.current_status = "idle"
|
|
self.local_service: LocalService | None = None
|
|
self.deployment = deployment_info(config.deployment_type)
|
|
self.display_name = config.display_name or self.host
|
|
self.state = load_state(config.state_path)
|
|
self.model_config_state = load_applied_model_config(config.model_config_path)
|
|
self.runtime_manager = RuntimeRestartManager(config.runtime_state_path, config.runtime_command)
|
|
self.tool_registry = AgentToolRegistry()
|
|
print(
|
|
"[agent] startup config: "
|
|
f"controller={self.config.controller_url} "
|
|
f"local_service={self.config.local_service_host}:{self.config.local_service_port} "
|
|
f"state={self.config.state_path} "
|
|
f"model_config={self.config.model_config_path}"
|
|
)
|
|
|
|
def _url(self, path: str) -> str:
|
|
if not self.controller_url:
|
|
raise RuntimeError("Controller is not connected yet")
|
|
return f"{self.controller_url}{path}"
|
|
|
|
def resolve_controller(self) -> str:
|
|
resolved = resolve_controller_base_url(
|
|
configured_base_url=self.config.controller_url,
|
|
discovery_enabled=self.config.controller_discovery_enabled,
|
|
discovery_ports=self.config.controller_discovery_ports,
|
|
timeout_seconds=self.config.controller_discovery_timeout_seconds,
|
|
)
|
|
if self.controller_url != resolved:
|
|
print(f"[agent] controller connected: {resolved}")
|
|
self.controller_url = resolved
|
|
self.last_connection_error = None
|
|
self.controller_retry_count = 0
|
|
return resolved
|
|
|
|
def ensure_registered(self) -> None:
|
|
current_public_id = self.state.get("public_id")
|
|
has_credentials = self.state.get("agent_id") and self.state.get("token")
|
|
if has_credentials and current_public_id == self.public_id:
|
|
if not self.registration_confirmed:
|
|
print(
|
|
f"Using saved registration: agent_id={self.state['agent_id']} "
|
|
f"public_id={self.public_id}"
|
|
)
|
|
self.registration_confirmed = True
|
|
return
|
|
|
|
payload = {
|
|
"public_id": self.public_id,
|
|
"display_name": self.display_name,
|
|
"deployment_type": self.deployment["deployment_type"],
|
|
"machine_fingerprint": self.fingerprint,
|
|
"reported_host": self.host,
|
|
"reported_ip": self.ip,
|
|
"reported_log_web_port": self.config.local_service_port
|
|
if self.config.local_service_enabled
|
|
else None,
|
|
"agent_version": __version__,
|
|
"metadata": {
|
|
"detection_reason": self.deployment["detection_reason"],
|
|
"system": self.deployment["system"],
|
|
},
|
|
}
|
|
print(f"Registering controlled agent at {self.controller_url} ...")
|
|
response = request_json(
|
|
"POST",
|
|
self._url("/api/agents/register"),
|
|
payload,
|
|
timeout=self.config.request_timeout_seconds,
|
|
)
|
|
self.state = {
|
|
"agent_id": response["agent_id"],
|
|
"public_id": response["public_id"],
|
|
"token": response["token"],
|
|
"registered_at": response["registered_at"],
|
|
"controller_url": self.controller_url,
|
|
}
|
|
save_state(self.config.state_path, self.state)
|
|
self.registration_confirmed = True
|
|
if response.get("token_rotated"):
|
|
print(
|
|
"Recovered registration with a new token: "
|
|
f"agent_id={response['agent_id']} state={self.config.state_path}"
|
|
)
|
|
else:
|
|
print(f"Registered: agent_id={response['agent_id']} state={self.config.state_path}")
|
|
|
|
def heartbeat_payload(self) -> dict[str, Any]:
|
|
return {
|
|
"status": self.current_status,
|
|
"display_name": self.display_name,
|
|
"deployment_type": self.deployment["deployment_type"],
|
|
"machine_fingerprint": self.fingerprint,
|
|
"reported_host": self.host,
|
|
"reported_ip": self.ip,
|
|
"reported_log_web_port": self.config.local_service_port
|
|
if self.config.local_service_enabled
|
|
else None,
|
|
"agent_version": __version__,
|
|
"config_revision": int(self.model_config_state.get("version") or 0),
|
|
"runtime_generation": self.runtime_manager.generation,
|
|
"started_at": self.started_at,
|
|
"sent_at": utc_now(),
|
|
"metadata": {
|
|
"detection_reason": self.deployment["detection_reason"],
|
|
"system": self.deployment["system"],
|
|
},
|
|
}
|
|
|
|
def send_heartbeat(self) -> None:
|
|
agent_id = int(self.state["agent_id"])
|
|
response = request_json(
|
|
"POST",
|
|
self._url(f"/api/agents/{agent_id}/heartbeat"),
|
|
self.heartbeat_payload(),
|
|
token=str(self.state["token"]),
|
|
timeout=self.config.request_timeout_seconds,
|
|
)
|
|
self.last_heartbeat_at = response.get("received_at")
|
|
self.apply_model_config_update(response.get("model_config"))
|
|
print(f"Heartbeat OK: {response.get('received_at')}")
|
|
|
|
def apply_model_config_update(self, model_config: Any) -> None:
|
|
if not isinstance(model_config, dict):
|
|
return
|
|
version = int(model_config.get("version") or 0)
|
|
current_version = int(self.model_config_state.get("version") or 0)
|
|
if version <= current_version:
|
|
return
|
|
payload = {
|
|
"version": version,
|
|
"updated_at": model_config.get("updated_at"),
|
|
"applied_at": utc_now(),
|
|
"value": model_config.get("value") or {},
|
|
}
|
|
save_applied_model_config(self.config.model_config_path, payload)
|
|
self.model_config_state = payload
|
|
restart_state = self.runtime_manager.restart_for_model_config(version)
|
|
print(
|
|
"[agent] applied model config "
|
|
f"v{version}; runtime restart generation={restart_state.get('generation')}"
|
|
)
|
|
|
|
def ensure_runtime_started(self) -> None:
|
|
state = self.runtime_manager.ensure_started(int(self.model_config_state.get("version") or 0))
|
|
if state:
|
|
print(
|
|
"[agent] runtime started "
|
|
f"generation={state.get('generation')} mode={state.get('mode')}"
|
|
)
|
|
|
|
def submit_child_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
|
response = request_json(
|
|
"POST",
|
|
self._url("/api/tasks"),
|
|
payload,
|
|
timeout=self.config.request_timeout_seconds,
|
|
)
|
|
return response.get("task") or {}
|
|
|
|
def task_model_config_state(self) -> dict[str, Any]:
|
|
state = dict(self.model_config_state)
|
|
value = dict(state.get("value") or {})
|
|
if self.state.get("agent_id") and self.state.get("token"):
|
|
value["controller_model_proxy_url"] = self._url(
|
|
f"/api/agents/{int(self.state['agent_id'])}/model/chat-json"
|
|
)
|
|
value["controller_model_proxy_token"] = str(self.state["token"])
|
|
value["controller_model_proxy_enabled"] = True
|
|
state["value"] = value
|
|
return state
|
|
|
|
def execute_task(self, task: dict[str, Any]) -> dict[str, Any]:
|
|
session = AgentSession(
|
|
task=task,
|
|
model_config_state=self.task_model_config_state(),
|
|
memory_path=self.config.memory_path,
|
|
request_timeout_seconds=self.config.request_timeout_seconds,
|
|
submit_task=self.submit_child_task,
|
|
)
|
|
return session.run()
|
|
|
|
def task_execution_error_result(self, task: dict[str, Any], exc: Exception, started: float) -> dict[str, Any]:
|
|
payload = task.get("payload") if isinstance(task.get("payload"), dict) else {}
|
|
model_state = self.task_model_config_state()
|
|
model_value = model_state.get("value") if isinstance(model_state.get("value"), dict) else {}
|
|
trace_lines = traceback.format_exc().splitlines()
|
|
message = f"Controlled agent task execution failed: {exc}"
|
|
return {
|
|
"success": False,
|
|
"message": message,
|
|
"task_type": payload.get("task_type") or "manual",
|
|
"agent_framework": "ufo-inspired-lite",
|
|
"p0_mode": True,
|
|
"needs_user_input": True,
|
|
"target_url": payload.get("target_url") or payload.get("url") or "",
|
|
"duration_seconds": round(time.monotonic() - started, 3),
|
|
"action_duration_seconds": 0,
|
|
"model_duration_seconds": 0,
|
|
"model_available": bool(
|
|
model_value.get("controller_model_proxy_url")
|
|
or (model_value.get("vllm_base_url") and model_value.get("vllm_model"))
|
|
),
|
|
"model_errors": [message],
|
|
"exception_type": type(exc).__name__,
|
|
"logs": [
|
|
"Task execution raised an exception before completion.",
|
|
*trace_lines[-25:],
|
|
],
|
|
}
|
|
|
|
def report_task_result(self, task_id: int, result: dict[str, Any]) -> dict[str, Any]:
|
|
agent_id = int(self.state["agent_id"])
|
|
return request_json(
|
|
"POST",
|
|
self._url(f"/api/agents/{agent_id}/tasks/{task_id}/result"),
|
|
result,
|
|
token=str(self.state["token"]),
|
|
timeout=self.config.request_timeout_seconds,
|
|
)
|
|
|
|
def claim_task(self) -> dict[str, Any] | None:
|
|
agent_id = int(self.state["agent_id"])
|
|
response = request_json(
|
|
"POST",
|
|
self._url(f"/api/agents/{agent_id}/tasks/claim"),
|
|
{},
|
|
token=str(self.state["token"]),
|
|
timeout=self.config.request_timeout_seconds,
|
|
)
|
|
task = response.get("task")
|
|
if task:
|
|
self.last_task_claim_at = utc_now()
|
|
print(f"Claimed task: id={task.get('id')} external_id={task.get('external_id')}")
|
|
self.current_status = "running"
|
|
task_started = time.monotonic()
|
|
try:
|
|
try:
|
|
self.send_heartbeat()
|
|
except Exception as exc:
|
|
print(f"[agent] running heartbeat failed: {exc}")
|
|
try:
|
|
result = self.execute_task(task)
|
|
except Exception as exc:
|
|
result = self.task_execution_error_result(task, exc, task_started)
|
|
print(result["message"], file=sys.stderr)
|
|
self.report_task_result(int(task["id"]), result)
|
|
print(
|
|
"Reported task result: "
|
|
f"id={task.get('id')} success={result.get('success')} message={result.get('message')}"
|
|
)
|
|
finally:
|
|
self.current_status = "idle"
|
|
else:
|
|
print("No ready task.")
|
|
return task
|
|
|
|
def local_status(self) -> dict[str, Any]:
|
|
return {
|
|
"ok": True,
|
|
"service": "controlled-agent",
|
|
"version": __version__,
|
|
"configured_controller_url": self.config.controller_url,
|
|
"controller_url": self.controller_url,
|
|
"controller_connected": bool(self.controller_url),
|
|
"last_connection_error": self.last_connection_error,
|
|
"public_id": self.public_id,
|
|
"agent_id": self.state.get("agent_id"),
|
|
"deployment_type": self.deployment["deployment_type"],
|
|
"display_name": self.display_name,
|
|
"reported_host": self.host,
|
|
"reported_ip": self.ip,
|
|
"started_at": self.started_at,
|
|
"last_heartbeat_at": self.last_heartbeat_at,
|
|
"last_task_claim_at": self.last_task_claim_at,
|
|
"model_config": public_model_config_summary(self.model_config_state),
|
|
"runtime": self.runtime_manager.state,
|
|
"agent_core": {
|
|
"framework": "ufo-inspired-lite",
|
|
"memory_path": str(self.config.memory_path),
|
|
"memory": self.agent_memory_status(),
|
|
"tools": self.tool_registry.descriptions(),
|
|
},
|
|
"local_service": {
|
|
"enabled": self.config.local_service_enabled,
|
|
"host": self.config.local_service_host,
|
|
"port": self.config.local_service_port,
|
|
},
|
|
}
|
|
|
|
def agent_memory_status(self) -> dict[str, Any]:
|
|
if not self.config.memory_path.exists():
|
|
return {"exists": False, "counts": {}}
|
|
try:
|
|
data = json.loads(self.config.memory_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
return {"exists": True, "error": str(exc), "counts": {}}
|
|
if not isinstance(data, dict):
|
|
return {"exists": True, "error": "Memory file is not an object.", "counts": {}}
|
|
return {
|
|
"exists": True,
|
|
"counts": {
|
|
key: len(value) if isinstance(value, list) else 0
|
|
for key, value in data.items()
|
|
},
|
|
}
|
|
|
|
def start_local_service(self) -> None:
|
|
if not self.config.local_service_enabled:
|
|
return
|
|
self.local_service = LocalService(
|
|
self.config.local_service_host,
|
|
self.config.local_service_port,
|
|
self.local_status,
|
|
)
|
|
self.local_service.start()
|
|
print(
|
|
"Controlled agent local service started at "
|
|
f"http://127.0.0.1:{self.config.local_service_port}/health"
|
|
)
|
|
|
|
def run(self, once: bool = False) -> None:
|
|
self.start_local_service()
|
|
try:
|
|
while True:
|
|
try:
|
|
if not self.controller_url:
|
|
self.resolve_controller()
|
|
self.ensure_registered()
|
|
if not self.runtime_started:
|
|
self.ensure_runtime_started()
|
|
self.runtime_started = True
|
|
self.send_heartbeat()
|
|
self.claim_task()
|
|
if once:
|
|
return
|
|
time.sleep(min(self.config.heartbeat_interval_seconds, self.config.task_poll_interval_seconds))
|
|
except (RuntimeError, HttpClientError, OSError) as exc:
|
|
if once:
|
|
raise
|
|
self.handle_controller_error(exc)
|
|
time.sleep(max(1.0, self.config.controller_retry_interval_seconds))
|
|
finally:
|
|
self.runtime_manager.stop()
|
|
if once and self.local_service:
|
|
self.local_service.stop()
|
|
|
|
def handle_controller_error(self, exc: Exception) -> None:
|
|
message = str(exc)
|
|
self.last_connection_error = message
|
|
self.controller_retry_count += 1
|
|
|
|
if isinstance(exc, HttpClientError) and exc.status in {401, 404}:
|
|
print(
|
|
"[agent] saved controller credentials are no longer valid; "
|
|
"clearing local registration and retrying."
|
|
)
|
|
self.state = {}
|
|
self.registration_confirmed = False
|
|
save_state(self.config.state_path, self.state)
|
|
|
|
if not isinstance(exc, HttpClientError) or exc.status in {0, 401, 404, 409} or exc.status >= 500:
|
|
self.controller_url = None
|
|
|
|
if self.controller_retry_count == 1 or self.controller_retry_count % 10 == 0:
|
|
print(
|
|
"[agent] controller unavailable; agent stays running and will keep scanning. "
|
|
f"retry={self.controller_retry_count} error={message}"
|
|
)
|
|
|
|
|
|
def parse_args(argv: list[str]) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="System test controlled agent")
|
|
parser.add_argument("--once", action="store_true", help="register, heartbeat, claim once, then exit")
|
|
parser.add_argument("--controller-url", help="override CONTROLLER_URL")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
args = parse_args(argv or sys.argv[1:])
|
|
config = AgentConfig.from_env()
|
|
if args.controller_url:
|
|
config = AgentConfig(
|
|
root_dir=config.root_dir,
|
|
controller_url=args.controller_url.rstrip("/"),
|
|
controller_discovery_enabled=config.controller_discovery_enabled,
|
|
controller_discovery_ports=config.controller_discovery_ports,
|
|
controller_discovery_timeout_seconds=config.controller_discovery_timeout_seconds,
|
|
state_path=config.state_path,
|
|
model_config_path=config.model_config_path,
|
|
runtime_state_path=config.runtime_state_path,
|
|
memory_path=config.memory_path,
|
|
display_name=config.display_name,
|
|
deployment_type=config.deployment_type,
|
|
heartbeat_interval_seconds=config.heartbeat_interval_seconds,
|
|
task_poll_interval_seconds=config.task_poll_interval_seconds,
|
|
controller_retry_interval_seconds=config.controller_retry_interval_seconds,
|
|
request_timeout_seconds=config.request_timeout_seconds,
|
|
local_service_enabled=config.local_service_enabled,
|
|
local_service_host=config.local_service_host,
|
|
local_service_port=config.local_service_port,
|
|
runtime_command=config.runtime_command,
|
|
)
|
|
|
|
try:
|
|
agent = ControlledAgent(config)
|
|
agent.run(once=args.once)
|
|
except RuntimeError as exc:
|
|
print(str(exc), file=sys.stderr)
|
|
raise SystemExit(2) from exc
|
|
except OSError as exc:
|
|
print(
|
|
"Local controlled-agent service failed to start. "
|
|
f"Check whether port {config.local_service_port} is already in use. "
|
|
f"Original error: {exc}",
|
|
file=sys.stderr,
|
|
)
|
|
raise SystemExit(2) from exc
|
|
except HttpClientError as exc:
|
|
if exc.status == 409:
|
|
print(
|
|
"Registration conflict: this machine is already registered on the controller, "
|
|
"but the local token was not found. Restore data/agent_state.json or remove the "
|
|
"existing agent record from the controller database.",
|
|
file=sys.stderr,
|
|
)
|
|
print(str(exc), file=sys.stderr)
|
|
raise SystemExit(2) from exc
|
|
except KeyboardInterrupt:
|
|
print("\nControlled agent stopping...")
|
|
except Exception as exc:
|
|
print("Unexpected SystemTest-Agent error:", file=sys.stderr)
|
|
traceback.print_exc()
|
|
raise SystemExit(2) from exc
|