"""Test MCP server routing with dynamic bank_id.""" from unittest.mock import AsyncMock, MagicMock import pytest @pytest.fixture def mock_memory(): """Create a mock MemoryEngine.""" memory = MagicMock() memory.retain_batch_async = AsyncMock() memory.submit_async_retain = AsyncMock(return_value={"operation_id": "test-op-123"}) memory.recall_async = AsyncMock(return_value=MagicMock(results=[])) return memory @pytest.mark.asyncio async def test_mcp_context_variable(): """Test that context variable works correctly.""" from hindsight_api.api.mcp import _current_bank_id, get_current_bank_id # Initially None assert get_current_bank_id() is None # Set and verify token = _current_bank_id.set("test-bank-123") try: assert get_current_bank_id() == "test-bank-123" finally: _current_bank_id.reset(token) # Back to None after reset assert get_current_bank_id() is None @pytest.mark.asyncio async def test_mcp_tools_use_context_bank_id(mock_memory): """Test that MCP tools use bank_id from context.""" from hindsight_api.api.mcp import _current_bank_id, create_mcp_server mcp_server = create_mcp_server(mock_memory) # Get the tools tools = mcp_server._tool_manager._tools assert "retain" in tools assert "recall" in tools token = _current_bank_id.set("context-bank-id") try: retain_tool = tools["retain"] result = await retain_tool.fn(content="test content", context="test_context") assert result["status"] == "accepted" # Verify the memory was called with the context bank_id mock_memory.submit_async_retain.assert_called_once() call_kwargs = mock_memory.submit_async_retain.call_args.kwargs assert call_kwargs["bank_id"] == "context-bank-id" finally: _current_bank_id.reset(token) def test_path_parsing_logic(): """Test the path parsing logic for bank_id extraction.""" def parse_path(path): """Simulate the path parsing logic from MCPMiddleware.""" if not path.startswith("/") or len(path) <= 1: return None, None # Error case parts = path[1:].split("/", 1) if not parts[0]: return None, None # Error case bank_id = parts[0] new_path = "/" + parts[1] if len(parts) > 1 else "/" return bank_id, new_path # Test bank-specific paths bank_id, remaining = parse_path("/my-bank/") assert bank_id == "my-bank" assert remaining == "/" bank_id, remaining = parse_path("/my-bank") assert bank_id == "my-bank" assert remaining == "/" # Test error case - no bank_id bank_id, remaining = parse_path("/") assert bank_id is None # Test with complex bank_id bank_id, remaining = parse_path("/user_12345/") assert bank_id == "user_12345" assert remaining == "/" # Test with additional path after bank_id bank_id, remaining = parse_path("/my-bank/some/path") assert bank_id == "my-bank" assert remaining == "/some/path" @pytest.mark.asyncio async def test_api_key_context_variable(): """Test that API key context variable works correctly.""" from hindsight_api.api.mcp import _current_api_key, get_current_api_key # Initially None assert get_current_api_key() is None # Set and verify token = _current_api_key.set("test-api-key-123") try: assert get_current_api_key() == "test-api-key-123" finally: _current_api_key.reset(token) # Back to None after reset assert get_current_api_key() is None @pytest.mark.asyncio async def test_mcp_tools_propagate_api_key(mock_memory): """Test that MCP tools propagate API key to RequestContext.""" from hindsight_api.api.mcp import _current_api_key, _current_bank_id, create_mcp_server mcp_server = create_mcp_server(mock_memory) tools = mcp_server._tool_manager._tools # Set both bank_id and api_key context bank_token = _current_bank_id.set("test-bank") api_key_token = _current_api_key.set("test-bearer-token") try: retain_tool = tools["retain"] result = await retain_tool.fn(content="test content", context="test_context") assert result["status"] == "accepted" # Verify the memory was called with request_context containing api_key mock_memory.submit_async_retain.assert_called_once() call_kwargs = mock_memory.submit_async_retain.call_args.kwargs assert call_kwargs["request_context"].api_key == "test-bearer-token" finally: _current_bank_id.reset(bank_token) _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 ( _current_api_key_id, _current_tenant_id, get_current_api_key_id, get_current_tenant_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 ( _current_api_key, _current_api_key_id, _current_bank_id, _current_tenant_id, create_mcp_server, ) 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") # Verify the RequestContext passed to memory engine has all auth fields mock_memory.submit_async_retain.assert_called_once() request_context = mock_memory.submit_async_retain.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 and mental models.""" from hindsight_api.api.mcp import create_mcp_server # Create server in multi-bank mode (default) mcp_server = create_mcp_server(mock_memory, multi_bank=True) tools = mcp_server._tool_manager._tools # Core tools assert "retain" in tools assert "recall" in tools assert "reflect" in tools assert "list_banks" in tools assert "create_bank" in tools # Mental model tools assert "list_mental_models" in tools assert "get_mental_model" in tools assert "create_mental_model" in tools assert "update_mental_model" in tools assert "delete_mental_model" in tools assert "refresh_mental_model" in tools def test_single_bank_mode_excludes_bank_management_tools(mock_memory): """Test that single-bank mode only exposes bank-scoped tools.""" from hindsight_api.api.mcp import create_mcp_server # Create server in single-bank mode mcp_server = create_mcp_server(mock_memory, multi_bank=False) tools = mcp_server._tool_manager._tools # Should have bank-scoped tools assert "retain" in tools assert "recall" in tools assert "reflect" in tools # Mental model tools should also be present (they're bank-scoped) assert "list_mental_models" in tools assert "get_mental_model" in tools assert "create_mental_model" in tools assert "update_mental_model" in tools assert "delete_mental_model" in tools assert "refresh_mental_model" in tools # Should NOT have bank management tools assert "list_banks" not in tools assert "create_bank" not in tools def test_multi_bank_mode_tools_have_bank_id_param(mock_memory): """Test that multi-bank mode tools include bank_id parameter.""" import inspect from hindsight_api.api.mcp import create_mcp_server mcp_server = create_mcp_server(mock_memory, multi_bank=True) tools = mcp_server._tool_manager._tools # All bank-scoped tools should have bank_id parameter in multi-bank mode bank_scoped_tools = [ "retain", "recall", "reflect", "list_mental_models", "get_mental_model", "create_mental_model", "update_mental_model", "delete_mental_model", "refresh_mental_model", ] for tool_name in bank_scoped_tools: tool = tools[tool_name] sig = inspect.signature(tool.fn) assert "bank_id" in sig.parameters, f"{tool_name} should have bank_id param in multi-bank mode" def test_single_bank_mode_tools_no_bank_id_param(mock_memory): """Test that single-bank mode tools do NOT include bank_id parameter.""" import inspect from hindsight_api.api.mcp import create_mcp_server mcp_server = create_mcp_server(mock_memory, multi_bank=False) tools = mcp_server._tool_manager._tools # No bank-scoped tool should have bank_id parameter in single-bank mode bank_scoped_tools = [ "retain", "recall", "reflect", "list_mental_models", "get_mental_model", "create_mental_model", "update_mental_model", "delete_mental_model", "refresh_mental_model", ] for tool_name in bank_scoped_tools: tool = tools[tool_name] sig = inspect.signature(tool.fn) assert "bank_id" not in sig.parameters, f"{tool_name} should NOT have bank_id param in single-bank mode" @pytest.mark.asyncio async def test_middleware_handles_both_endpoints(mock_memory): """Test that MCPMiddleware routes to correct server based on URL path.""" from hindsight_api.api.mcp import MCPMiddleware # Create middleware (single instance) middleware = MCPMiddleware(None, mock_memory) # Verify both server instances exist assert middleware.multi_bank_app is not None assert middleware.single_bank_app is not None # Verify they expose different tools multi_bank_tools = middleware.multi_bank_server._tool_manager._tools single_bank_tools = middleware.single_bank_server._tool_manager._tools # Multi-bank should have all tools assert "retain" in multi_bank_tools assert "recall" in multi_bank_tools assert "list_banks" in multi_bank_tools assert "create_bank" in multi_bank_tools assert "list_mental_models" in multi_bank_tools assert "create_mental_model" in multi_bank_tools # Single-bank should only have scoped tools assert "retain" in single_bank_tools assert "recall" in single_bank_tools assert "list_mental_models" in single_bank_tools assert "create_mental_model" in single_bank_tools assert "list_banks" not in single_bank_tools assert "create_bank" not in single_bank_tools def test_global_mcp_enabled_tools_filter_restricts_registered_tools(mock_memory): """Test that global mcp_enabled_tools env setting restricts which tools are registered.""" from unittest.mock import MagicMock, patch from hindsight_api.api.mcp import create_mcp_server mock_cfg = MagicMock() mock_cfg.mcp_enabled_tools = ["retain", "recall"] with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg): mcp_server = create_mcp_server(mock_memory, multi_bank=True) tools = mcp_server._tool_manager._tools assert "retain" in tools assert "recall" in tools assert "reflect" not in tools assert "list_banks" not in tools assert "create_bank" not in tools assert "list_mental_models" not in tools def test_global_mcp_enabled_tools_none_exposes_all_tools(mock_memory): """Test that mcp_enabled_tools=None (default) exposes all tools.""" from unittest.mock import MagicMock, patch from hindsight_api.api.mcp import create_mcp_server mock_cfg = MagicMock() mock_cfg.mcp_enabled_tools = None with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg): mcp_server = create_mcp_server(mock_memory, multi_bank=True) tools = mcp_server._tool_manager._tools assert "retain" in tools assert "recall" in tools assert "reflect" in tools assert "list_banks" in tools assert "create_bank" in tools def test_global_mcp_enabled_tools_intersects_with_single_bank_mode(mock_memory): """Test that global filter intersects with single-bank mode tool set. list_banks is in the global allowlist but NOT in single-bank mode, so it should be absent from the final registered set. """ from unittest.mock import MagicMock, patch from hindsight_api.api.mcp import create_mcp_server mock_cfg = MagicMock() mock_cfg.mcp_enabled_tools = ["retain", "recall", "list_banks"] with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg): mcp_server = create_mcp_server(mock_memory, multi_bank=False) tools = mcp_server._tool_manager._tools assert "retain" in tools assert "recall" in tools assert "list_banks" not in tools # single-bank mode excludes it regardless @pytest.mark.asyncio async def test_routing_logic_from_url_path(): """Test that routing correctly selects server based on URL structure. Simulates the path parsing logic from MCPMiddleware.__call__ after the prefix has been stripped. Any first path segment is treated as a bank_id. """ from hindsight_api.api.mcp import MCPMiddleware # Mock memory mock_memory = MagicMock() # Create middleware middleware = MCPMiddleware(None, mock_memory) # Simulate different URL patterns and verify routing # Path is what remains after stripping the /mcp prefix test_cases = [ # (path_after_prefix_strip, expected_bank_id_from_path, expected_bank_id, description) ("/alice/messages", True, "alice", "Bank ID in path with endpoint"), ("/my-agent-123/", True, "my-agent-123", "Bank ID in path with trailing slash"), ("/sse/", True, "sse", "Bank named 'sse' routes to single-bank"), ("/messages/", True, "messages", "Bank named 'messages' routes to single-bank"), ("/", False, None, "Root path, no bank ID"), ] for path, expected_bank_from_path, expected_bank_id, description in test_cases: bank_id = None bank_id_from_path = False if path.startswith("/") and len(path) > 1: parts = path[1:].split("/", 1) if parts[0]: bank_id = parts[0] bank_id_from_path = True assert bank_id_from_path == expected_bank_from_path, f"Failed for: {description} (path={path})" assert bank_id == expected_bank_id, f"Failed bank_id for: {description} (path={path}, got={bank_id})"