236 lines
8.5 KiB
Python
236 lines
8.5 KiB
Python
"""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) |