feat(host-agent): wrap tool_handlers with busy check + status mapping

2026-07-21 14:09:42 +08:00
parent cf8affe4d7
commit b73db01626
2 changed files with 462 additions and 0 deletions
@@ -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)
@@ -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