Fix MCP operations not tracked for usage metering (#334)

MCP middleware was discarding tenant_id and api_key_id after authentication.
The authenticate_mcp() call mutated a RequestContext with these fields, but
tools later created a fresh RequestContext without them. This caused
UsageMeteringValidator to see tenant_id="unknown" and skip billing entirely.

Propagate tenant_id and api_key_id via ContextVars (same pattern as bank_id
and api_key) so the RequestContext passed to the memory engine has the full
auth context needed for usage tracking.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

hindsight-api/hindsight_api/api/mcp.py

@@ -43,6 +43,10 @@ _current_bank_id: ContextVar[str | None] = ContextVar("current_bank_id", default
 # Context variable to hold the current API key (for tenant auth propagation)
 _current_api_key: ContextVar[str | None] = ContextVar("current_api_key", default=None)
 
+# Context variables for tenant_id and api_key_id (set by authenticate, used by usage metering)
+_current_tenant_id: ContextVar[str | None] = ContextVar("current_tenant_id", default=None)
+_current_api_key_id: ContextVar[str | None] = ContextVar("current_api_key_id", default=None)
+
 
 def get_current_bank_id() -> str | None:
     """Get the current bank_id from context."""
@@ -54,6 +58,16 @@ def get_current_api_key() -> str | None:
     return _current_api_key.get()
 
 
+def get_current_tenant_id() -> str | None:
+    """Get the current tenant_id from context."""
+    return _current_tenant_id.get()
+
+
+def get_current_api_key_id() -> str | None:
+    """Get the current api_key_id from context."""
+    return _current_api_key_id.get()
+
+
 def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
     """
     Create and configure the Hindsight MCP server.
@@ -73,6 +87,8 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
     config = MCPToolsConfig(
         bank_id_resolver=get_current_bank_id,
         api_key_resolver=get_current_api_key,  # Propagate API key for tenant auth
+        tenant_id_resolver=get_current_tenant_id,  # Propagate tenant_id for usage metering
+        api_key_id_resolver=get_current_api_key_id,  # Propagate api_key_id for usage metering
         include_bank_id_param=multi_bank,
         tools=None if multi_bank else {"retain", "recall", "reflect"},  # Scoped tools for single-bank mode
         retain_fire_and_forget=False,  # HTTP MCP supports sync/async modes
@@ -165,6 +181,8 @@ class MCPMiddleware:
 
         # Authenticate: check legacy MCP_AUTH_TOKEN first, then TenantExtension
         tenant_context = None
+        auth_tenant_id: str | None = None
+        auth_api_key_id: str | None = None
         if MCP_AUTH_TOKEN:
             # Legacy authentication mode - validate against static token
             if not auth_token:
@@ -178,7 +196,11 @@ class MCPMiddleware:
         else:
             # Use TenantExtension.authenticate_mcp() for auth
             try:
-                tenant_context = await self.tenant_extension.authenticate_mcp(RequestContext(api_key=auth_token))
+                auth_context = RequestContext(api_key=auth_token)
+                tenant_context = await self.tenant_extension.authenticate_mcp(auth_context)
+                # Capture tenant_id and api_key_id set by authenticate() for usage metering
+                auth_tenant_id = auth_context.tenant_id
+                auth_api_key_id = auth_context.api_key_id
             except AuthenticationError as e:
                 await self._send_error(send, 401, str(e))
                 return
@@ -233,10 +255,13 @@ class MCPMiddleware:
         # - Header/env bank_id → multi-bank app (bank_id param, all tools)
         target_app = self.single_bank_app if bank_id_from_path else self.multi_bank_app
 
-        # Set bank_id and api_key context
+        # Set bank_id, api_key, tenant_id, and api_key_id context
         bank_id_token = _current_bank_id.set(bank_id)
         # Store the auth token for tenant extension to validate
         api_key_token = _current_api_key.set(auth_token) if auth_token else None
+        # Store tenant_id and api_key_id from authentication for usage metering
+        tenant_id_token = _current_tenant_id.set(auth_tenant_id) if auth_tenant_id else None
+        api_key_id_token = _current_api_key_id.set(auth_api_key_id) if auth_api_key_id else None
         try:
             new_scope = scope.copy()
             new_scope["path"] = new_path
@@ -258,6 +283,10 @@ class MCPMiddleware:
             _current_bank_id.reset(bank_id_token)
             if api_key_token is not None:
                 _current_api_key.reset(api_key_token)
+            if tenant_id_token is not None:
+                _current_tenant_id.reset(tenant_id_token)
+            if api_key_id_token is not None:
+                _current_api_key_id.reset(api_key_id_token)
             if schema_token is not None:
                 _current_schema.reset(schema_token)
 

hindsight-api/hindsight_api/mcp_tools.py

@@ -35,6 +35,12 @@ class MCPToolsConfig:
     # How to resolve API key for tenant auth (optional)
     api_key_resolver: Callable[[], str | None] | None = None
 
+    # How to resolve tenant_id for usage metering (set by MCP middleware after auth)
+    tenant_id_resolver: Callable[[], str | None] | None = None
+
+    # How to resolve api_key_id for usage metering (set by MCP middleware after auth)
+    api_key_id_resolver: Callable[[], str | None] | None = None
+
     # Whether to include bank_id as a parameter on tools (for multi-bank support)
     include_bank_id_param: bool = False
 
@@ -50,13 +56,15 @@ class MCPToolsConfig:
 
 
 def _get_request_context(config: MCPToolsConfig) -> RequestContext:
-    """Create RequestContext with API key from resolver if available.
+    """Create RequestContext with auth details from resolvers.
 
-    This enables tenant auth to work with MCP tools by propagating
-    the Bearer token from the MCP middleware to the memory engine.
+    This enables tenant auth and usage metering to work with MCP tools by propagating
+    the authentication results from the MCP middleware to the memory engine.
     """
     api_key = config.api_key_resolver() if config.api_key_resolver else None
-    return RequestContext(api_key=api_key)
+    tenant_id = config.tenant_id_resolver() if config.tenant_id_resolver else None
+    api_key_id = config.api_key_id_resolver() if config.api_key_id_resolver else None
+    return RequestContext(api_key=api_key, tenant_id=tenant_id, api_key_id=api_key_id)
 
 
 def parse_timestamp(timestamp: str) -> datetime | None:

hindsight-api/tests/test_mcp_routing.py

@@ -143,6 +143,72 @@ async def test_mcp_tools_propagate_api_key(mock_memory):
         _current_api_key.reset(api_key_token)
 
 
+@pytest.mark.asyncio
+async def test_tenant_id_context_variable():
+    """Test that tenant_id and api_key_id context variables work correctly."""
+    from hindsight_api.api.mcp import (
+        get_current_tenant_id, _current_tenant_id,
+        get_current_api_key_id, _current_api_key_id,
+    )
+
+    # Initially None
+    assert get_current_tenant_id() is None
+    assert get_current_api_key_id() is None
+
+    # Set and verify
+    tenant_token = _current_tenant_id.set("org-123")
+    key_id_token = _current_api_key_id.set("key-456")
+    try:
+        assert get_current_tenant_id() == "org-123"
+        assert get_current_api_key_id() == "key-456"
+    finally:
+        _current_tenant_id.reset(tenant_token)
+        _current_api_key_id.reset(key_id_token)
+
+    # Back to None after reset
+    assert get_current_tenant_id() is None
+    assert get_current_api_key_id() is None
+
+
+@pytest.mark.asyncio
+async def test_mcp_tools_propagate_tenant_id_and_api_key_id(mock_memory):
+    """Test that MCP tools propagate tenant_id and api_key_id to RequestContext.
+
+    This is the critical test for usage metering: the UsageMeteringValidator reads
+    request_context.tenant_id to identify the org for billing. Without this,
+    MCP operations get tenant_id="unknown" and billing is skipped entirely.
+    """
+    from hindsight_api.api.mcp import (
+        create_mcp_server,
+        _current_bank_id, _current_api_key,
+        _current_tenant_id, _current_api_key_id,
+    )
+
+    mcp_server = create_mcp_server(mock_memory)
+    tools = mcp_server._tool_manager._tools
+
+    # Set all context vars (simulating what MCPMiddleware does after authenticate_mcp)
+    bank_token = _current_bank_id.set("test-bank")
+    api_key_token = _current_api_key.set("hsk_test_key")
+    tenant_token = _current_tenant_id.set("org-billing-123")
+    key_id_token = _current_api_key_id.set("key-uuid-456")
+    try:
+        retain_tool = tools["retain"]
+        await retain_tool.fn(content="test content", context="test_context", async_processing=False)
+
+        # Verify the RequestContext passed to memory engine has all auth fields
+        mock_memory.retain_batch_async.assert_called_once()
+        request_context = mock_memory.retain_batch_async.call_args.kwargs["request_context"]
+        assert request_context.api_key == "hsk_test_key"
+        assert request_context.tenant_id == "org-billing-123"
+        assert request_context.api_key_id == "key-uuid-456"
+    finally:
+        _current_bank_id.reset(bank_token)
+        _current_api_key.reset(api_key_token)
+        _current_tenant_id.reset(tenant_token)
+        _current_api_key_id.reset(key_id_token)
+
+
 def test_multi_bank_mode_exposes_all_tools(mock_memory):
     """Test that multi-bank mode exposes all tools including bank management."""
     from hindsight_api.api.mcp import create_mcp_server