agentic-mobile-control/apps/device-host-agent/host_agent/web/mcp.py

"""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 Context, 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"})
# Name of the wrapper kwarg FastMCP injects the live ``Context`` into.
# We set ``tool.context_kwarg = _CONTEXT_KWARG`` after swapping the tool's
# ``fn`` (see ``build_mcp_server``) so FastMCP passes ``ctx`` into our
# wrapper alongside the validated arguments.
_CONTEXT_KWARG = "ctx"


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


class FastMcpSdkIncompatibilityError(RuntimeError):
    """Raised when the FastMCP SDK layout diverges from what this module
    expects (e.g. ``Tool.fn`` rename or ``Tool.context_kwarg`` removal)."""


# 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(ctx: Context | None = None) -> str:
    """Extract a stable per-MCP-session identifier from the live context.

    The mcp SDK 1.28.1 ``Context`` exposes ``session`` (a long-lived
    ``ServerSession`` instance per Streamable HTTP session). Its Python
    object identity (``id(ctx.session)``) is stable across every tool call
    the same client makes within that session, which is exactly the
    identity the busy tracker needs to renew leases.

    Falls back to ``_TEST_SESSION_ID`` when no Context is supplied (i.e.
    when invoked outside a FastMCP request lifecycle, as ``_call_tool_sync``
    does in tests).
    """
    if ctx is not None:
        session_obj = getattr(ctx, "session", None)
        if session_obj is not None:
            return f"mcp_session:{id(session_obj)}"
    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
        )
        try:
            tool = server._tool_manager._tools[tool_name]  # type: ignore[attr-defined]
            tool.fn = wrapped
            # FastMCP injects the live Context into the kwarg named by
            # ``tool.context_kwarg``. The raw handler doesn't declare one,
            # so the cached value is None; we override it so the wrapper
            # receives the Context via its ``ctx`` kwarg.
            tool.context_kwarg = _CONTEXT_KWARG
        except AttributeError as exc:
            raise FastMcpSdkIncompatibilityError(
                "FastMCP SDK layout changed: cannot swap Tool.fn or set "
                f"context_kwarg (tool={tool_name!r}). Underlying error: {exc}"
            ) from exc

    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:
        ctx = kwargs.pop(_CONTEXT_KWARG, None)
        session_id = _current_session_id(ctx)
        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 (and the live
    FastMCP Context) so tests don't need an MCP client.

    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)