Three final-review deviations closed: I1 (session-end release): mcp SDK 1.28.1 exposes no per-session shutdown callback (only a server-level lifespan). Lower the McpBusyTracker default TTL from 60s to 20s and update spec §6.5, Q5/R3, D9, and docs/MCP_INTEGRATION.md concurrency section to document the TTL-only recovery path. 20s is short enough to recover within one 30s heartbeat interval but long enough that an active session does not lose its lease during normal operator pauses. I2 (JSON-RPC error shape): FastMCP Tool.run wraps every non- UrlElicitationRequiredError exception (including McpError with typed ErrorData) into ToolError, which the lowlevel call_tool handler serializes as CallToolResult(isError=true, content=[TextContent(...)]). There is no public path that surfaces JSON-RPC -32000 with structured data.busy_owner from a tool call site. Update spec §7 error matrix and docs/MCP_INTEGRATION.md error table to document the actual wire shape; busy_owner now lives in the text content. I3 (typing): mcp_server: Any = None -> FastMCP | None = None via TYPE_CHECKING, keeping the mcp import lazy (matches precedent elsewhere in the codebase) while adding static type checking at the create_console_app boundary. Tests added (4): - test_default_ttl_is_20_seconds — locks I1's new default TTL - test_default_ttl_recovers_dead_session_within_one_window — locks I1's recovery semantics (lease sweeped on next read after 20s) - test_busy_error_wire_shape_is_calltoolresult_iserror — pins I2's wire envelope via Tool.run + lowlevel Server._make_error_result - test_busy_error_text_includes_cloud_assignment_owner — same for the cloud_assignment busy_owner branch Full non-integration suite: 697 passed / 54 deselected (was 693 / 54). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
370 lines
14 KiB
Python
370 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass
|
|
|
|
import uvicorn
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
from device.manager import DeviceManager
|
|
from host_agent.assignment import AssignmentExecutor
|
|
from host_agent.client import HostAgentClient, HostAgentEnrollmentClient
|
|
from host_agent.config import HostAgentConfig, load_host_agent_config
|
|
from host_agent.dependency_supervisor import DependencySupervisor
|
|
from host_agent.devices import register_local_device
|
|
from host_agent.enrollment import resolve_host_identity
|
|
from host_agent.execution import create_execution_factories
|
|
from host_agent.heartbeat import HeartbeatSynchronizer
|
|
from host_agent.history import ConsoleHistoryStore
|
|
from host_agent.identity import HostIdentityStore
|
|
from host_agent.instance_lock import InstanceLock
|
|
from host_agent.lease import ActiveAssignmentRunner
|
|
from host_agent.local_account import LocalAccountStore
|
|
from host_agent.mcp_lock import McpBusyTracker
|
|
from host_agent.mcp_token import McpTokenStore
|
|
from host_agent.policy_cache import HostPolicyCacheStore
|
|
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
|
|
from host_agent.retention import prune_task_history
|
|
from host_agent.skill_sync import HostAgentSkillSync
|
|
from host_agent.status import AgentStatusTracker
|
|
from host_agent.web.app import create_console_app
|
|
from host_agent.web.auth import SessionManager
|
|
from host_agent.web.mcp import build_mcp_server
|
|
from storage.artifact_store import ArtifactStore
|
|
from storage.device_config import DeviceConfigStore
|
|
from storage.task_metadata import TaskMetadataStore
|
|
from storage.timeline import Timeline
|
|
|
|
|
|
@dataclass
|
|
class HostAgentApplication:
|
|
client: HostAgentClient
|
|
heartbeat: HeartbeatSynchronizer
|
|
processor: AssignmentProcessor
|
|
console_server: uvicorn.Server | None = None
|
|
console_enrollment_client: HostAgentEnrollmentClient | None = None
|
|
dependency_supervisor: DependencySupervisor | None = None
|
|
instance_lock: InstanceLock | None = None
|
|
skill_sync: HostAgentSkillSync | None = None
|
|
|
|
def run(self) -> None:
|
|
asyncio.run(self.run_async())
|
|
|
|
async def run_async(self, stop: asyncio.Event | None = None) -> None:
|
|
stop_requested = stop or asyncio.Event()
|
|
supervisor_stop = asyncio.Event()
|
|
supervisor_task: asyncio.Task[None] | None = None
|
|
if self.dependency_supervisor is not None:
|
|
# Bring up supervised dependencies (probe + spawn/adopt + readiness
|
|
# wait) before the heartbeat loop's first connect_devices() pass,
|
|
# so a supervised Appium is ready before any Driver.connect().
|
|
await self.dependency_supervisor.start()
|
|
supervisor_task = asyncio.create_task(
|
|
self.dependency_supervisor.run(supervisor_stop)
|
|
)
|
|
heartbeat_stop = asyncio.Event()
|
|
heartbeat_task = asyncio.create_task(self.heartbeat.run(heartbeat_stop))
|
|
if self.skill_sync is not None:
|
|
self.skill_sync.start()
|
|
console_task = (
|
|
asyncio.create_task(self.console_server.serve())
|
|
if self.console_server is not None
|
|
else None
|
|
)
|
|
active_processing: asyncio.Task[AssignmentProcessingResult] | None = None
|
|
try:
|
|
while not stop_requested.is_set():
|
|
assignment = await self._claim_until_stopped(stop_requested)
|
|
if assignment is None:
|
|
continue
|
|
active_processing = asyncio.create_task(
|
|
self.processor.process(assignment)
|
|
)
|
|
stopped = asyncio.create_task(stop_requested.wait())
|
|
done, _ = await asyncio.wait(
|
|
{active_processing, stopped},
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
if stopped in done:
|
|
self.processor.request_stop()
|
|
else:
|
|
stopped.cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await stopped
|
|
await asyncio.shield(active_processing)
|
|
active_processing = None
|
|
finally:
|
|
self.processor.request_stop()
|
|
if active_processing is not None:
|
|
with suppress(Exception):
|
|
await asyncio.shield(active_processing)
|
|
heartbeat_stop.set()
|
|
supervisor_stop.set()
|
|
if self.console_server is not None:
|
|
self.console_server.should_exit = True
|
|
try:
|
|
await asyncio.gather(heartbeat_task, return_exceptions=True)
|
|
with suppress(Exception):
|
|
await self.heartbeat.sync_once()
|
|
if supervisor_task is not None:
|
|
await asyncio.gather(supervisor_task, return_exceptions=True)
|
|
with suppress(Exception):
|
|
await self.dependency_supervisor.stop()
|
|
if console_task is not None:
|
|
with suppress(asyncio.CancelledError):
|
|
await asyncio.gather(console_task, return_exceptions=True)
|
|
finally:
|
|
if self.console_enrollment_client is not None:
|
|
self.console_enrollment_client.close()
|
|
if self.skill_sync is not None:
|
|
self.skill_sync.stop()
|
|
await self.client.aclose()
|
|
if self.instance_lock is not None:
|
|
self.instance_lock.release()
|
|
|
|
async def _claim_until_stopped(
|
|
self,
|
|
stop: asyncio.Event,
|
|
) -> AssignmentModel | None:
|
|
claim = asyncio.create_task(self.client.claim())
|
|
stopped = asyncio.create_task(stop.wait())
|
|
try:
|
|
done, _ = await asyncio.wait(
|
|
{claim, stopped},
|
|
return_when=asyncio.FIRST_COMPLETED,
|
|
)
|
|
except asyncio.CancelledError:
|
|
claim.cancel()
|
|
stopped.cancel()
|
|
await asyncio.gather(claim, stopped, return_exceptions=True)
|
|
raise
|
|
if stopped in done:
|
|
claim.cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await claim
|
|
return None
|
|
stopped.cancel()
|
|
with suppress(asyncio.CancelledError):
|
|
await stopped
|
|
return await claim
|
|
|
|
|
|
class _EmbeddedConsoleServer(uvicorn.Server):
|
|
def install_signal_handlers(self) -> None:
|
|
"""Host Agent owns process-level signal handling."""
|
|
|
|
|
|
def create_application(
|
|
*,
|
|
config: HostAgentConfig | None = None,
|
|
manager: DeviceManager | None = None,
|
|
device_config_store: DeviceConfigStore | None = None,
|
|
identity_store: HostIdentityStore | None = None,
|
|
enrollment_client: HostAgentEnrollmentClient | None = None,
|
|
) -> HostAgentApplication:
|
|
startup_config = config or load_host_agent_config()
|
|
instance_lock = InstanceLock(startup_config.identity_path.parent)
|
|
instance_lock.acquire()
|
|
try:
|
|
config_store = device_config_store or DeviceConfigStore()
|
|
resolved_identity_store = identity_store or HostIdentityStore(
|
|
startup_config.identity_path
|
|
)
|
|
owned_enrollment_client = enrollment_client is None
|
|
bootstrap_client = enrollment_client or HostAgentEnrollmentClient(
|
|
startup_config
|
|
)
|
|
try:
|
|
resolved_config = resolve_host_identity(
|
|
startup_config,
|
|
identity_store=resolved_identity_store,
|
|
client=bootstrap_client,
|
|
)
|
|
bootstrap_client.config = resolved_config
|
|
resolved_manager = manager or _configured_device_manager(
|
|
config_store,
|
|
config=resolved_config,
|
|
enrollment_client=bootstrap_client,
|
|
)
|
|
finally:
|
|
if owned_enrollment_client:
|
|
bootstrap_client.close()
|
|
client = HostAgentClient(resolved_config)
|
|
|
|
history_store = ConsoleHistoryStore(
|
|
resolved_config.identity_path.parent / "host_console_history.sqlite3",
|
|
limit=resolved_config.console_history_limit,
|
|
)
|
|
status_tracker = AgentStatusTracker()
|
|
console_enrollment_client: HostAgentEnrollmentClient | None = None
|
|
if resolved_config.enrollment_managed:
|
|
console_enrollment_client = HostAgentEnrollmentClient(resolved_config)
|
|
metadata_store = TaskMetadataStore(
|
|
db_path=resolved_config.task_progress_db_path
|
|
)
|
|
timeline = Timeline(ArtifactStore(root=resolved_config.task_artifact_dir))
|
|
mcp_token_path = resolved_config.identity_path.parent / "host_mcp_token.json"
|
|
mcp_token_existed = mcp_token_path.exists()
|
|
mcp_token_store = McpTokenStore(mcp_token_path)
|
|
mcp_token_store.load_or_create()
|
|
if not mcp_token_existed:
|
|
logging.getLogger(__name__).info(
|
|
"MCP token generated at %s", mcp_token_path
|
|
)
|
|
mcp_busy_tracker = McpBusyTracker(ttl_seconds=20.0)
|
|
executor = AssignmentExecutor(
|
|
create_execution_factories(
|
|
resolved_manager,
|
|
metadata_store=metadata_store,
|
|
timeline=timeline,
|
|
host_agent_config=resolved_config,
|
|
),
|
|
mcp_busy_tracker=mcp_busy_tracker,
|
|
)
|
|
mcp_server = build_mcp_server(
|
|
manager=resolved_manager,
|
|
mcp_busy_tracker=mcp_busy_tracker,
|
|
status_tracker=status_tracker,
|
|
)
|
|
console_app = create_console_app(
|
|
config=resolved_config,
|
|
manager=resolved_manager,
|
|
config_store=config_store,
|
|
local_account_store=LocalAccountStore(resolved_config.local_account_path),
|
|
identity_store=resolved_identity_store,
|
|
history_store=history_store,
|
|
status_tracker=status_tracker,
|
|
session_manager=SessionManager(
|
|
ttl_seconds=resolved_config.console_session_ttl_seconds
|
|
),
|
|
enrollment_client=console_enrollment_client,
|
|
host_client=client,
|
|
metadata_store=metadata_store,
|
|
timeline=timeline,
|
|
executor=executor,
|
|
mcp_server=mcp_server,
|
|
mcp_token_store=mcp_token_store,
|
|
mcp_busy_tracker=mcp_busy_tracker,
|
|
)
|
|
console_server = _EmbeddedConsoleServer(
|
|
uvicorn.Config(
|
|
console_app,
|
|
host=resolved_config.console_bind_host,
|
|
port=resolved_config.console_port,
|
|
log_level="warning",
|
|
)
|
|
)
|
|
|
|
heartbeat = HeartbeatSynchronizer(
|
|
resolved_manager,
|
|
client,
|
|
resolved_config,
|
|
status_tracker=status_tracker,
|
|
mcp_busy_tracker=mcp_busy_tracker,
|
|
on_sync=lambda device_count: history_store.record_heartbeat(
|
|
device_count=device_count
|
|
),
|
|
policy_cache=HostPolicyCacheStore(
|
|
resolved_config.identity_path.parent / "host_governance_policy.json"
|
|
),
|
|
on_policy_sync=lambda revision: history_store.record_policy_sync(
|
|
revision=revision
|
|
),
|
|
)
|
|
active_runner = ActiveAssignmentRunner(client, executor)
|
|
processor = AssignmentProcessor(
|
|
client,
|
|
active_runner,
|
|
status_tracker=status_tracker,
|
|
on_result=lambda assignment, result: _on_assignment_finished(
|
|
history_store,
|
|
metadata_store,
|
|
timeline,
|
|
resolved_config,
|
|
assignment,
|
|
result,
|
|
),
|
|
)
|
|
dependency_supervisor: DependencySupervisor | None = None
|
|
if resolved_config.dependency_supervisor_enabled:
|
|
dependency_supervisor = DependencySupervisor.from_host_agent_config(
|
|
resolved_config
|
|
)
|
|
skill_sync: HostAgentSkillSync | None = None
|
|
if resolved_config.host_id and resolved_config.token:
|
|
skill_sync = HostAgentSkillSync(resolved_config)
|
|
return HostAgentApplication(
|
|
client=client,
|
|
heartbeat=heartbeat,
|
|
processor=processor,
|
|
console_server=console_server,
|
|
console_enrollment_client=console_enrollment_client,
|
|
dependency_supervisor=dependency_supervisor,
|
|
instance_lock=instance_lock,
|
|
skill_sync=skill_sync,
|
|
)
|
|
except BaseException:
|
|
instance_lock.release()
|
|
raise
|
|
|
|
|
|
def _on_assignment_finished(
|
|
history_store: ConsoleHistoryStore,
|
|
metadata_store: TaskMetadataStore,
|
|
timeline: Timeline,
|
|
config: HostAgentConfig,
|
|
assignment: AssignmentModel,
|
|
result: AssignmentProcessingResult,
|
|
) -> None:
|
|
_record_assignment_history(history_store, assignment, result)
|
|
_prune_task_history(metadata_store, timeline, config)
|
|
|
|
|
|
def _record_assignment_history(
|
|
history_store: ConsoleHistoryStore,
|
|
assignment: AssignmentModel,
|
|
result: AssignmentProcessingResult,
|
|
) -> None:
|
|
status = "done" if result.execution.status == "done" else "failed"
|
|
history_store.record_assignment(
|
|
task_id=assignment.task_id,
|
|
attempt=assignment.attempt,
|
|
status=status,
|
|
failure_reason=result.execution.failure_reason if status == "failed" else None,
|
|
device_id=assignment.device_id,
|
|
)
|
|
|
|
|
|
def _prune_task_history(
|
|
metadata_store: TaskMetadataStore,
|
|
timeline: Timeline,
|
|
config: HostAgentConfig,
|
|
) -> None:
|
|
try:
|
|
prune_task_history(metadata_store, timeline, config=config)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _configured_device_manager(
|
|
config_store: DeviceConfigStore,
|
|
*,
|
|
config: HostAgentConfig,
|
|
enrollment_client: HostAgentEnrollmentClient,
|
|
) -> DeviceManager:
|
|
manager = DeviceManager()
|
|
for device in config_store.list():
|
|
register_local_device(
|
|
config_store,
|
|
manager,
|
|
device_id=device["device_id"],
|
|
driver_type=device["driver_type"],
|
|
connection_info=device["connection_info"],
|
|
name=device["name"],
|
|
config=config,
|
|
enrollment_client=enrollment_client,
|
|
)
|
|
return manager
|