From 888b50de12679b3c087fc6a8a175f01cd56d7c09 Mon Sep 17 00:00:00 2001 From: DK09876 Date: Tue, 10 Feb 2026 01:38:15 -0700 Subject: [PATCH] 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 --- hindsight-api/hindsight_api/api/mcp.py | 33 +++++++++++- hindsight-api/hindsight_api/mcp_tools.py | 16 ++++-- hindsight-api/tests/test_mcp_routing.py | 66 ++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 6 deletions(-) diff --git a/hindsight-api/hindsight_api/api/mcp.py b/hindsight-api/hindsight_api/api/mcp.py index 82fc7cac..279fc2bb 100644 --- a/hindsight-api/hindsight_api/api/mcp.py +++ b/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) diff --git a/hindsight-api/hindsight_api/mcp_tools.py b/hindsight-api/hindsight_api/mcp_tools.py index db4dfdd9..3d652e56 100644 --- a/hindsight-api/hindsight_api/mcp_tools.py +++ b/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: diff --git a/hindsight-api/tests/test_mcp_routing.py b/hindsight-api/tests/test_mcp_routing.py index 78198561..02b5008d 100644 --- a/hindsight-api/tests/test_mcp_routing.py +++ b/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