Add SystemTest Agent and Controller
This commit is contained in:
Codex committed
commit
53f523f187
31 files changed
+9150
No files matched your search
@@ -0,0 +1,509 @@
|
||||
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
|
||||
Reference in new issue
Block a user