diff --git a/apps/device-host-agent/host_agent/app.py b/apps/device-host-agent/host_agent/app.py index 8dd8053..4800c21 100644 --- a/apps/device-host-agent/host_agent/app.py +++ b/apps/device-host-agent/host_agent/app.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import logging from contextlib import suppress from dataclasses import dataclass @@ -21,6 +22,8 @@ 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 @@ -28,6 +31,7 @@ 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 @@ -201,13 +205,28 @@ def create_application( 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=60.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, @@ -225,6 +244,9 @@ def create_application( 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( @@ -240,6 +262,7 @@ def create_application( 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 ), diff --git a/apps/device-host-agent/tests/test_app.py b/apps/device-host-agent/tests/test_app.py index 9b425eb..685c3de 100644 --- a/apps/device-host-agent/tests/test_app.py +++ b/apps/device-host-agent/tests/test_app.py @@ -7,6 +7,7 @@ from datetime import UTC, datetime, timedelta import httpx import pytest +from starlette.testclient import TestClient from cloud.internal_api.models import ( AssignmentModel, @@ -15,10 +16,23 @@ from cloud.internal_api.models import ( ) from device.manager import DeviceManager from host_agent.app import HostAgentApplication, create_application +from host_agent.assignment import AssignmentExecutor from host_agent.config import HostAgentConfig +from host_agent.execution import create_execution_factories +from host_agent.history import ConsoleHistoryStore from host_agent.identity import HostIdentityStore from host_agent.instance_lock import InstanceAlreadyRunningError +from host_agent.local_account import LocalAccountStore +from host_agent.mcp_lock import McpBusyTracker +from host_agent.mcp_token import McpTokenStore +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 def _free_loopback_port() -> int: @@ -653,6 +667,79 @@ def test_create_application_with_independent_identity_paths_coexist( asyncio.run(app_a.client.aclose()) +def test_create_application_wires_mcp_components(tmp_path, monkeypatch) -> None: + """create_application produces a console app with /mcp mounted (auth-protected) + and persists the host_mcp_token.json file alongside the identity.""" + monkeypatch.chdir(tmp_path) + config = _config() + config_store = DeviceConfigStore(tmp_path / "devices.sqlite3") + identity_store = HostIdentityStore(config.identity_path) + history_store = ConsoleHistoryStore( + tmp_path / "host_console_history.sqlite3", + limit=config.console_history_limit, + ) + metadata_store = TaskMetadataStore(db_path=config.task_progress_db_path) + timeline = Timeline(ArtifactStore(root=config.task_artifact_dir)) + status_tracker = AgentStatusTracker() + + application = create_application( + config=config, + device_config_store=config_store, + identity_store=identity_store, + manager=DeviceManager(), + ) + + # Token file must exist after create_application. + assert (config.identity_path.parent / "host_mcp_token.json").exists() + + # Heartbeat must hold the in-process McpBusyTracker. + assert application.heartbeat.mcp_busy_tracker is not None + + # Build the same console app the production path builds and verify /mcp + # is mounted (responds 401, not 404) without a bearer token. + mcp_token_store = McpTokenStore( + config.identity_path.parent / "host_mcp_token.json" + ) + mcp_token_store.load_or_create() + mcp_busy_tracker = McpBusyTracker(ttl_seconds=60.0) + mcp_server = build_mcp_server( + manager=application.heartbeat.manager, + mcp_busy_tracker=mcp_busy_tracker, + status_tracker=status_tracker, + ) + console_app = create_console_app( + config=config, + manager=application.heartbeat.manager, + config_store=config_store, + local_account_store=LocalAccountStore(config.local_account_path), + identity_store=identity_store, + history_store=history_store, + status_tracker=status_tracker, + session_manager=SessionManager(ttl_seconds=config.console_session_ttl_seconds), + enrollment_client=None, + host_client=application.client, + metadata_store=metadata_store, + timeline=timeline, + executor=AssignmentExecutor( + create_execution_factories( + application.heartbeat.manager, + metadata_store=metadata_store, + timeline=timeline, + host_agent_config=config, + ), + mcp_busy_tracker=mcp_busy_tracker, + ), + mcp_server=mcp_server, + mcp_token_store=mcp_token_store, + mcp_busy_tracker=mcp_busy_tracker, + ) + with TestClient(console_app) as client: + resp = client.post("/mcp/") + assert resp.status_code == 401 # auth required, not 404 + + asyncio.run(application.client.aclose()) + + def test_lock_released_after_run_async_allows_restart(tmp_path, monkeypatch) -> None: monkeypatch.chdir(tmp_path) identity_path = tmp_path / "host_identity.json"