Initial commit (forked from jamiepine/voicebox)

backend/mcp_server/context.py

@@ -0,0 +1,152 @@
+"""Per-request client identity for MCP calls.
+
+MCP clients identify themselves via an ``X-Voicebox-Client-Id`` HTTP header
+(direct-HTTP clients set it in their MCP config; the stdio shim forwards it
+from the ``VOICEBOX_CLIENT_ID`` env var). Middleware copies the value into a
+ContextVar so tool implementations can read it without plumbing the request
+object through every service call.
+"""
+
+import asyncio
+import ipaddress
+import logging
+from contextvars import ContextVar
+from datetime import datetime, timezone
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import ASGIApp
+
+
+logger = logging.getLogger(__name__)
+
+# Strong refs to in-flight stamp tasks so asyncio.create_task results
+# don't get garbage-collected mid-flight (cf. asyncio.create_task docs).
+_pending_stamps: set[asyncio.Task] = set()
+
+CLIENT_ID_HEADER = "X-Voicebox-Client-Id"
+
+# Tool handlers read this to apply per-client voice bindings.
+current_client_id: ContextVar[str | None] = ContextVar(
+    "current_client_id", default=None
+)
+
+# Remote address of the in-flight request. Used by tools that gate
+# host-filesystem access to loopback callers (see voicebox.transcribe).
+current_remote_addr: ContextVar[str | None] = ContextVar(
+    "current_remote_addr", default=None
+)
+
+
+def request_is_loopback() -> bool:
+    """True when the in-flight request originated on the loopback interface.
+
+    Returns False if no request is in flight or the remote address can't be
+    parsed — callers gating filesystem reads on this should treat that as
+    "deny".
+    """
+    addr = current_remote_addr.get()
+    if not addr:
+        return False
+    try:
+        return ipaddress.ip_address(addr).is_loopback
+    except ValueError:
+        return False
+
+# Endpoints that consume X-Voicebox-Client-Id for its MCP-semantic
+# meaning (per-client profile resolution + per-client default_personality).
+# These are the paths where a stamp into last_seen_at is accurate.
+# Unrelated REST traffic that happens to set the header is intentionally
+# ignored so the Settings UI's "last heard from" column only reflects
+# calls that actually acted on the client's bindings.
+#
+# - /mcp — FastMCP tool calls (voicebox.speak, voicebox.transcribe, …)
+#   and the /mcp/bindings admin surface. The admin surface is never
+#   called with the header in practice (the frontend manages bindings
+#   over plain REST), so the `startswith("/mcp")` match doesn't cause
+#   false stamps.
+# - /speak — REST mirror of voicebox.speak for non-MCP agents (shell
+#   scripts, ACP, A2A). Uses the same per-client binding lookup, so its
+#   callers belong in the last-seen list too.
+_STAMPED_PATH_PREFIXES: tuple[str, ...] = ("/mcp", "/speak")
+
+
+class ClientIdMiddleware(BaseHTTPMiddleware):
+    """Copy X-Voicebox-Client-Id into a ContextVar and stamp last_seen_at
+    for requests that act on the caller's MCP bindings."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        super().__init__(app)
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        client_id = request.headers.get(CLIENT_ID_HEADER)
+        remote_addr = request.client.host if request.client else None
+        client_token = current_client_id.set(client_id)
+        addr_token = current_remote_addr.set(remote_addr)
+        try:
+            response = await call_next(request)
+        finally:
+            current_client_id.reset(client_token)
+            current_remote_addr.reset(addr_token)
+
+        if client_id and _is_stamped_path(request.url.path):
+            _enqueue_stamp(client_id)
+        return response
+
+
+def _enqueue_stamp(client_id: str) -> None:
+    """Fire-and-forget the SQLite write so it doesn't block the response.
+
+    The stamp does sync SQLAlchemy I/O; running it inline on the event loop
+    serialises every MCP request behind the SQLite write and starves SSE
+    streams. ``asyncio.to_thread`` parks it on the default executor while
+    the response goes back to the caller.
+    """
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        # Middleware shouldn't run outside a loop, but if it ever does
+        # (tests, weird wsgi shim), do the write inline rather than drop it.
+        _stamp_last_seen(client_id)
+        return
+    task = loop.create_task(asyncio.to_thread(_stamp_last_seen, client_id))
+    _pending_stamps.add(task)
+    task.add_done_callback(_pending_stamps.discard)
+
+
+def _is_stamped_path(path: str) -> bool:
+    # Require a path boundary so a future ``/speakers`` or ``/mcpfoo``
+    # route doesn't silently inherit the stamp from ``/speak`` / ``/mcp``.
+    return any(path == p or path.startswith(p + "/") for p in _STAMPED_PATH_PREFIXES)
+
+
+def _stamp_last_seen(client_id: str) -> None:
+    """Update or create the MCPClientBinding row for this client_id."""
+    try:
+        from ..database import get_db
+        from ..database.models import MCPClientBinding
+    except Exception:
+        return
+    try:
+        db = next(get_db())
+    except Exception:
+        return
+    try:
+        row = (
+            db.query(MCPClientBinding)
+            .filter(MCPClientBinding.client_id == client_id)
+            .first()
+        )
+        if row is None:
+            row = MCPClientBinding(client_id=client_id)
+            db.add(row)
+        row.last_seen_at = datetime.now(timezone.utc)
+        db.commit()
+    except Exception:
+        logger.debug(
+            "Could not stamp last_seen_at for %s", client_id, exc_info=True
+        )
+        db.rollback()
+    finally:
+        db.close()