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