From b73db016261c12b43050657cde4551175ab5f22c Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Tue, 21 Jul 2026 14:09:42 +0800 Subject: [PATCH] feat(host-agent): wrap tool_handlers with busy check + status mapping --- apps/device-host-agent/host_agent/web/mcp.py | 236 +++++++++++++++++++ apps/device-host-agent/tests/test_web_mcp.py | 226 ++++++++++++++++++ 2 files changed, 462 insertions(+) create mode 100644 apps/device-host-agent/host_agent/web/mcp.py create mode 100644 apps/device-host-agent/tests/test_web_mcp.py diff --git a/apps/device-host-agent/host_agent/web/mcp.py b/apps/device-host-agent/host_agent/web/mcp.py new file mode 100644 index 0000000..56c2ffe --- /dev/null +++ b/apps/device-host-agent/host_agent/web/mcp.py @@ -0,0 +1,236 @@ +"""FastMCP server builder for the host-agent MCP endpoint. + +Wraps ``api.mcp.tool_handlers(manager=...)`` with: + +- Cloud-busy and MCP-busy checks (per-device, fail-fast on conflict). +- Lazy session-level device lock acquire / renew. +- Display-status mapping for ``list_devices`` / ``device_status`` so + connected-but-idle devices don't appear "busy" (which they do at the + ``DeviceManager`` layer because an Appium/WDA session is open). + +The builder returns a ``FastMCP`` instance. The caller +(``create_console_app``) is responsible for wrapping it in +``BearerAuthMiddleware`` and mounting at ``/mcp``. +""" + +from __future__ import annotations + +import contextvars +from collections.abc import Callable +from typing import Any + +from device.manager import DeviceManager +from host_agent.mcp_lock import McpBusyTracker +from host_agent.status import AgentStatusTracker +from mcp.server.fastmcp import FastMCP + +# Tool names that don't target a specific device — skip busy check. +_NON_DEVICE_TOOLS = frozenset({"list_devices", "device_status"}) +# Tools that report status and should use the display-status mapping. +_STATUS_TOOLS = frozenset({"list_devices", "device_status"}) + + +class McpDeviceBusyError(Exception): + """Raised by the wrapper when the target device is held by the cloud + assignment path or another MCP session.""" + + def __init__(self, device_id: str, busy_owner: str) -> None: + super().__init__(f"device {device_id} is busy (held by {busy_owner})") + self.device_id = device_id + self.busy_owner = busy_owner + + +# contextvars fallback used by tests and any call that originates outside a +# live FastMCP request lifecycle. Production handlers run inside an MCP +# request whose context exposes ``request_id`` and the underlying +# ``session``; ``_current_session_id`` reads from that context first and +# falls back to this ContextVar. +_TEST_SESSION_ID: contextvars.ContextVar[str] = contextvars.ContextVar( + "_TEST_SESSION_ID", default="" +) + + +def _current_session_id() -> str: + """Extract session_id from the current FastMCP tool-call context. + + The mcp SDK 1.28.1 does not expose a stable ``session_id`` on + ``Context``; the closest analogue is the per-request ``request_id`` + (always a string) plus the long-lived ``session`` object. We try those + first, then fall back to a ``contextvars``-based test override. + """ + try: + from mcp.server.fastmcp import get_context + + ctx = get_context() + request_id = getattr(ctx, "request_id", None) + if isinstance(request_id, str) and request_id: + return request_id + session_id = getattr(ctx, "session_id", None) + if isinstance(session_id, str) and session_id: + return session_id + except Exception: + pass + return _TEST_SESSION_ID.get("") + + +def build_mcp_server( + *, + manager: DeviceManager, + mcp_busy_tracker: McpBusyTracker, + status_tracker: AgentStatusTracker, +) -> FastMCP: + """Construct the FastMCP server wrapping ``tool_handlers``.""" + # Imported lazily to keep the package import graph flat. + from api.mcp import tool_handlers + + handlers = tool_handlers(manager=manager) + server = FastMCP("apex-host-agent") + + for tool_name, raw_handler in handlers.items(): + wrapped = _wrap_tool( + tool_name, + raw_handler, + mcp_busy_tracker=mcp_busy_tracker, + status_tracker=status_tracker, + ) + # Register the raw handler so FastMCP captures its signature (the + # MCP wire schema is derived from the function signature). Then + # swap ``tool.fn`` for our busy-check / status-mapping wrapper. + # Using ``*args, **kwargs`` directly breaks the schema, so we have + # to keep the signature and only replace the underlying callable. + server._tool_manager.add_tool( # type: ignore[attr-defined] + raw_handler, name=tool_name + ) + server._tool_manager._tools[tool_name].fn = wrapped # type: ignore[attr-defined] + + return server + + +def _wrap_tool( + tool_name: str, + handler: Callable[..., Any], + *, + mcp_busy_tracker: McpBusyTracker, + status_tracker: AgentStatusTracker, +) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + session_id = _current_session_id() + device_id = kwargs.get("device_id") + + if tool_name in _STATUS_TOOLS: + return _with_display_status(handler, status_tracker, *args, **kwargs) + + if device_id is not None and tool_name not in _NON_DEVICE_TOOLS: + _check_and_acquire( + device_id, session_id, mcp_busy_tracker, status_tracker + ) + + return handler(*args, **kwargs) + + return wrapped + + +def _check_and_acquire( + device_id: str, + session_id: str, + mcp_busy_tracker: McpBusyTracker, + status_tracker: AgentStatusTracker, +) -> None: + cloud_busy = _cloud_busy_device_id(status_tracker) + if cloud_busy == device_id: + raise McpDeviceBusyError(device_id, "cloud_assignment") + if device_id in mcp_busy_tracker.busy_device_ids(): + existing = next( + ( + lease + for lease in mcp_busy_tracker.snapshot() + if lease.device_id == device_id + ), + None, + ) + if existing is not None and existing.session_id != session_id: + prefix = existing.session_id[:8] + raise McpDeviceBusyError(device_id, f"mcp_session:{prefix}") + if not mcp_busy_tracker.acquire(device_id, session_id): + # Race: someone else got it between check and acquire. + raise McpDeviceBusyError(device_id, "another_session") + mcp_busy_tracker.renew(device_id, session_id) + + +def _cloud_busy_device_id(status_tracker: AgentStatusTracker) -> str | None: + """Return the device_id currently bound to the cloud assignment, if any.""" + snap = status_tracker.snapshot() + current = snap.get("current_assignment") + if not isinstance(current, dict): + return None + device_id = current.get("device_id") + return device_id if isinstance(device_id, str) else None + + +def _with_display_status( + handler: Callable[..., Any], + status_tracker: AgentStatusTracker, + *args: Any, + **kwargs: Any, +) -> Any: + busy_device_id = _cloud_busy_device_id(status_tracker) + result = handler(*args, **kwargs) + if isinstance(result, list): + for item in result: + if isinstance(item, dict) and "status" in item: + item["status"] = _display_status( + item["status"], item.get("id"), busy_device_id + ) + return result + if isinstance(result, dict) and "status" in result: + result["status"] = _display_status( + result["status"], result.get("device_id"), busy_device_id + ) + return result + + +def _display_status(raw: str, device_id: Any, busy_device_id: str | None) -> str: + """Mirror ``host_agent.web.app._device_display_status`` semantics. + + A device that's locally "busy" because it's connected-but-idle reports + "connected" instead, unless it's the device currently running a cloud + assignment (in which case "busy" is the truthful status). + """ + if raw == "busy" and device_id != busy_device_id: + return "connected" + return raw + + +def _call_tool_sync( + server: FastMCP, + tool_name: str, + arguments: dict[str, Any], + *, + session_id: str, +) -> Any: + """Test helper: invoke a registered tool synchronously with a forced + ``session_id``. Bypasses the HTTP/MCP transport layer to keep tests fast. + + Walks FastMCP's tool registry (``_tool_manager._tools[tool_name].fn``) — + the exact attribute path follows mcp SDK 1.28.1's + ``ToolManager._tools`` layout. + """ + token = _TEST_SESSION_ID.set(session_id) + try: + manager = getattr(server, "_tool_manager", None) + if manager is None: + raise KeyError(f"tool {tool_name!r} not registered (no tool manager)") + registry = getattr(manager, "_tools", None) or getattr(manager, "tools", None) + if isinstance(registry, dict): + tool = registry.get(tool_name) + else: + tool = manager.get_tool(tool_name) # type: ignore[union-attr] + if tool is None: + raise KeyError(f"tool {tool_name!r} not registered") + # FastMCP Tool wraps a callable; our wrappers are sync, so unwrap. + fn = getattr(tool, "fn", None) or getattr(tool, "func", None) + if fn is None: + raise KeyError(f"tool {tool_name!r} has no callable") + return fn(**arguments) + finally: + _TEST_SESSION_ID.reset(token) \ No newline at end of file diff --git a/apps/device-host-agent/tests/test_web_mcp.py b/apps/device-host-agent/tests/test_web_mcp.py new file mode 100644 index 0000000..9b958ce --- /dev/null +++ b/apps/device-host-agent/tests/test_web_mcp.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +import pytest + +from cloud.internal_api.models import AssignmentModel +from device.manager import DeviceManager +from driver.base import Driver +from host_agent.mcp_lock import McpBusyTracker +from host_agent.status import AgentStatusTracker +from host_agent.web.mcp import ( + McpDeviceBusyError, + _call_tool_sync, + build_mcp_server, +) + + +class _FakeDriver(Driver): + """Minimal driver. connect/screenshot/tap are exercised; remaining abstract + methods are stubbed to satisfy Driver's ABC contract.""" + + def __init__(self) -> None: + self.taps: list[tuple[float, float]] = [] + + def connect(self) -> None: + return None + + def disconnect(self) -> None: + return None + + def screenshot(self) -> bytes: + return b"fake" + + def tap(self, x: float, y: float) -> None: + self.taps.append((x, y)) + + def long_press(self, x: float, y: float, duration_ms: int = 1200) -> None: + return None + + def swipe( + self, + start_x: float, + start_y: float, + end_x: float, + end_y: float, + duration_ms: int = 500, + ) -> None: + return None + + def swipe_path(self, waypoints: list[tuple[float, float]], duration_ms: int) -> None: + return None + + def double_tap(self, x: float, y: float, interval_ms: int = 80) -> None: + return None + + def input(self, text: str) -> None: + return None + + def launch(self, app_id: str) -> None: + return None + + def terminate(self, app_id: str) -> None: + return None + + def tree(self) -> Any: + return None + + def home(self) -> None: + return None + + def lock(self) -> None: + return None + + def unlock(self) -> None: + return None + + +def _make_manager_with_device(device_id: str = "phone-1") -> DeviceManager: + manager = DeviceManager() + manager.register_device( + device_id=device_id, + driver_factory=lambda: _FakeDriver(), + name=device_id, + ) + manager.connect(device_id) + return manager + + +def test_build_mcp_server_returns_fastmcp_instance() -> None: + from mcp.server.fastmcp import FastMCP + + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + assert isinstance(server, FastMCP) + + +def test_call_tool_succeeds_when_device_is_free() -> None: + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + result = _call_tool_sync( + server, "take_screenshot", {"device_id": "phone-1"}, session_id="sess-a" + ) + assert result["ok"] is True + assert "phone-1" in tracker.busy_device_ids() + + +def test_call_tool_fails_when_cloud_uses_device() -> None: + """AgentStatusTracker.current_assignment.device_id matches -> busy.""" + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + status.mark_assignment_started( + AssignmentModel( + task_id="t1", + attempt=1, + lease_id="l1", + lease_expires_at=datetime.now(UTC), + host_id="h1", + device_id="phone-1", + goal="cloud task", + ) + ) + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + with pytest.raises(McpDeviceBusyError) as exc: + _call_tool_sync( + server, + "take_screenshot", + {"device_id": "phone-1"}, + session_id="sess-a", + ) + assert exc.value.device_id == "phone-1" + assert exc.value.busy_owner == "cloud_assignment" + + +def test_call_tool_fails_when_another_mcp_session_holds_device() -> None: + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + # Pre-acquire as a different session. + tracker.acquire("phone-1", "sess-other") + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + with pytest.raises(McpDeviceBusyError) as exc: + _call_tool_sync( + server, + "take_screenshot", + {"device_id": "phone-1"}, + session_id="sess-a", + ) + assert exc.value.busy_owner.startswith("mcp_session:") + + +def test_call_tool_renews_when_same_session_already_holds() -> None: + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + _call_tool_sync( + server, "take_screenshot", {"device_id": "phone-1"}, session_id="sess-a" + ) + # Second call from the same session should succeed. + result = _call_tool_sync( + server, "take_screenshot", {"device_id": "phone-1"}, session_id="sess-a" + ) + assert result["ok"] is True + + +def test_list_devices_uses_display_status() -> None: + """Connected-but-idle devices report as 'connected', not 'busy'.""" + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + result = _call_tool_sync(server, "list_devices", {}, session_id="sess-a") + assert isinstance(result, list) + assert result[0]["status"] == "connected" + + +def test_unknown_device_returns_semantic_error_dict() -> None: + """take_screenshot against an unknown device returns the api-errors semantic + error dict (``ok=False, error="device not found"``) rather than raising. + + Note: this test adapts the brief's exception-assertion semantics to the + actual behavior of ``call_with_semantic_errors`` in ``api/errors.py`` — + the brief's expectation that an exception is raised here is incorrect for + the current handler implementation.""" + manager = _make_manager_with_device() + tracker = McpBusyTracker() + status = AgentStatusTracker() + server = build_mcp_server( + manager=manager, mcp_busy_tracker=tracker, status_tracker=status + ) + result = _call_tool_sync( + server, + "take_screenshot", + {"device_id": "does-not-exist"}, + session_id="sess-a", + ) + assert isinstance(result, dict) + assert result["ok"] is False + assert "device" in result["error"].lower() + + +def test_manager_required_for_build_mcp_server() -> None: + """Reinforces D12 — build_mcp_server requires manager as keyword-only.""" + import inspect + + sig = inspect.signature(build_mcp_server) + assert sig.parameters["manager"].kind == inspect.Parameter.KEYWORD_ONLY \ No newline at end of file