showtan001
dd8df33910
Tests / Test apps.device-host-agent.tests.test_mcp_token.test_load_or_create_concurrent_calls_do_not_corrupt failed
401 lines
16 KiB
Python
401 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from contextlib import suppress
|
|
from dataclasses import dataclass, replace
|
|
|
|
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.conversation import ConversationAgent
|
|
from host_agent.conversation_log import ConversationLogStore
|
|
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.local_client import LocalHostAgentClient
|
|
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 runtime.executor import default_tool_registry
|
|
from runtime.planner_config import load_config as load_planner_config
|
|
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:
|
|
if startup_config.mode == "local":
|
|
resolved_config = replace(
|
|
startup_config, control_plane_url="", host_id="local-host"
|
|
)
|
|
resolved_manager = manager or _configured_device_manager(
|
|
config_store, config=resolved_config, enrollment_client=None
|
|
)
|
|
else:
|
|
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()
|
|
if resolved_config.mode == "local":
|
|
client = LocalHostAgentClient(
|
|
host_id="local-host",
|
|
device_ids=lambda: [device.id for device in resolved_manager.list_devices()],
|
|
)
|
|
else:
|
|
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)
|
|
conversation_log = ConversationLogStore(resolved_config.conversation_log_path) if resolved_config.mode == "local" else None
|
|
executor = AssignmentExecutor(
|
|
create_execution_factories(
|
|
resolved_manager,
|
|
metadata_store=metadata_store,
|
|
timeline=timeline,
|
|
host_agent_config=resolved_config,
|
|
conversation_log=conversation_log,
|
|
),
|
|
mcp_busy_tracker=mcp_busy_tracker,
|
|
)
|
|
mcp_server = build_mcp_server(
|
|
manager=resolved_manager,
|
|
mcp_busy_tracker=mcp_busy_tracker,
|
|
status_tracker=status_tracker,
|
|
)
|
|
planner_config = load_planner_config()
|
|
conversation_agent = (
|
|
ConversationAgent(
|
|
config=planner_config,
|
|
tools=default_tool_registry(manager=resolved_manager),
|
|
event_logger=conversation_log.append if conversation_log else None,
|
|
)
|
|
if resolved_config.ai_planner_transport == "direct"
|
|
else None
|
|
)
|
|
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,
|
|
conversation_agent=conversation_agent,
|
|
conversation_log=conversation_log,
|
|
)
|
|
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 | None,
|
|
) -> 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
|