fix(host-agent): use stable ServerSession id for MCP lock identity

The previous _current_session_id() implementation tried to import a
non-existent get_context() helper, so the production code path always
fell through to the empty _TEST_SESSION_ID ContextVar — meaning every
MCP client shared the empty-string identity and there was no per-session
isolation in production.

Use Context.session (the long-lived ServerSession object) as the source
of identity. id(ctx.session) is stable across every tool call the same
client makes within a Streamable HTTP session, which is exactly what the
busy tracker needs to renew leases.

Wire FastMCP to inject the Context into the wrapper by setting
tool.context_kwarg = "ctx" after swapping tool.fn; wrap the swap in a
defensive try/except that surfaces a FastMcpSdkIncompatibilityError on
future SDK layout drift.

Add 4 tests covering the production path: stability across calls in the
same session, isolation between sessions, fallback to _TEST_SESSION_ID
when no Context is supplied, and verification that the registered tool
declares context_kwarg="ctx".

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

apps/device-host-agent/host_agent/web/mcp.py

@@ -22,12 +22,17 @@ 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
+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):
@@ -40,6 +45,11 @@ class McpDeviceBusyError(Exception):
         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
@@ -50,26 +60,23 @@ _TEST_SESSION_ID: contextvars.ContextVar[str] = contextvars.ContextVar(
 )
 
 
-def _current_session_id() -> str:
-    """Extract session_id from the current FastMCP tool-call context.
+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 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.
+    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).
     """
-    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
+    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("")
 
 
@@ -101,7 +108,19 @@ def build_mcp_server(
         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]
+        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
 
@@ -114,7 +133,8 @@ def _wrap_tool(
     status_tracker: AgentStatusTracker,
 ) -> Callable[..., Any]:
     def wrapped(*args: Any, **kwargs: Any) -> Any:
-        session_id = _current_session_id()
+        ctx = kwargs.pop(_CONTEXT_KWARG, None)
+        session_id = _current_session_id(ctx)
         device_id = kwargs.get("device_id")
 
         if tool_name in _STATUS_TOOLS:
@@ -209,7 +229,8 @@ def _call_tool_sync(
     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.
+    ``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