feat: improve mcp tools based on endpoint (#318)

* feat: improve mcp tools based on endpoint

* feat: improve mcp tools based on endpoint

* test: add integration test for MCP endpoint routing

- Add test_mcp_endpoint_routing.py to verify single-bank vs multi-bank tool exposure
- Verifies /mcp/ exposes all tools with bank_id parameters
- Verifies /mcp/{bank_id}/ only exposes scoped tools without bank_id parameters
- Regression test for issue #317

Related: #317, #318

* test: use StreamableHTTP client for MCP endpoint routing test

Replace httpx AsyncClient SSE parsing with proper MCP StreamableHTTP
client. This correctly tests the MCP server using the actual protocol
that clients will use.

Fixes #317

Author: nicoloboschi

hindsight-api/hindsight_api/api/__init__.py

@@ -72,22 +72,24 @@ def create_app(
 
     # Mount MCP server and chain its lifespan if enabled
     if mcp_app is not None:
-        # Get the MCP app's underlying Starlette app for lifespan access
-        mcp_starlette_app = mcp_app.mcp_app
+        # Get both MCP apps' underlying Starlette apps for lifespan access
+        multi_bank_starlette_app = mcp_app.multi_bank_app
+        single_bank_starlette_app = mcp_app.single_bank_app
 
         # Store the original lifespan
         original_lifespan = app.router.lifespan_context
 
         @asynccontextmanager
         async def chained_lifespan(app_instance: FastAPI):
-            """Chain the MCP lifespan with the main app lifespan."""
-            # Start MCP lifespan first
-            async with mcp_starlette_app.router.lifespan_context(mcp_starlette_app):
-                logger.info("MCP lifespan started")
-                # Then start the original app lifespan
-                async with original_lifespan(app_instance):
-                    yield
-            logger.info("MCP lifespan stopped")
+            """Chain both MCP lifespans with the main app lifespan."""
+            # Start both MCP lifespans (multi-bank and single-bank)
+            async with multi_bank_starlette_app.router.lifespan_context(multi_bank_starlette_app):
+                async with single_bank_starlette_app.router.lifespan_context(single_bank_starlette_app):
+                    logger.info("MCP lifespans started (multi-bank and single-bank)")
+                    # Then start the original app lifespan
+                    async with original_lifespan(app_instance):
+                        yield
+                logger.info("MCP lifespans stopped")
 
         # Replace the app's lifespan with the chained version
         app.router.lifespan_context = chained_lifespan

hindsight-api/hindsight_api/api/mcp.py

@@ -54,12 +54,14 @@ def get_current_api_key() -> str | None:
     return _current_api_key.get()
 
 
-def create_mcp_server(memory: MemoryEngine) -> FastMCP:
+def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
     """
     Create and configure the Hindsight MCP server.
 
     Args:
         memory: MemoryEngine instance (required)
+        multi_bank: If True, expose all tools with bank_id parameters (default).
+                   If False, only expose bank-scoped tools without bank_id parameters.
 
     Returns:
         Configured FastMCP server instance with stateless_http enabled
@@ -71,8 +73,8 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
     config = MCPToolsConfig(
         bank_id_resolver=get_current_bank_id,
         api_key_resolver=get_current_api_key,  # Propagate API key for tenant auth
-        include_bank_id_param=True,  # HTTP MCP supports multi-bank via parameter
-        tools=None,  # All tools
+        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
     )
 
@@ -88,20 +90,33 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
 
 
 class MCPMiddleware:
-    """ASGI middleware that handles authentication and extracts bank_id from header or path.
+    """ASGI middleware that handles authentication and routes to appropriate MCP server.
 
     Authentication:
         1. If HINDSIGHT_API_MCP_AUTH_TOKEN is set (legacy), validates against that token
         2. Otherwise, uses TenantExtension.authenticate_mcp() from the MemoryEngine
            - DefaultTenantExtension: no auth required (local dev)
            - ApiKeyTenantExtension: validates against env var
 
-    Bank ID can be provided via:
-    1. X-Bank-Id header (recommended for Claude Code)
-    2. URL path: /mcp/{bank_id}/
-    3. Environment variable HINDSIGHT_MCP_BANK_ID (fallback default)
+    Two modes based on URL structure:
 
-    For Claude Code, configure with:
+    1. Multi-bank mode (for /mcp/ root endpoint):
+       - Exposes all tools: retain, recall, reflect, list_banks, create_bank
+       - All tools include optional bank_id parameter for cross-bank operations
+       - Bank ID from: X-Bank-Id header or HINDSIGHT_MCP_BANK_ID env var
+
+    2. Single-bank mode (for /mcp/{bank_id}/ endpoints):
+       - Exposes bank-scoped tools only: retain, recall, reflect
+       - No bank_id parameter (comes from URL)
+       - No bank management tools (list_banks, create_bank)
+       - Recommended for agent isolation
+
+    Examples:
+        # Single-bank mode (recommended for agent isolation)
+        claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/ \\
+            --header "Authorization: Bearer <token>"
+
+        # Multi-bank mode (for cross-bank operations)
         claude mcp add --transport http hindsight http://localhost:8888/mcp \\
             --header "X-Bank-Id: my-bank" --header "Authorization: Bearer <token>"
     """
@@ -110,10 +125,23 @@ def __init__(self, app, memory: MemoryEngine):
         self.app = app
         self.memory = memory
         self.tenant_extension = memory._tenant_extension
-        self.mcp_server = create_mcp_server(memory)
-        self.mcp_app = self.mcp_server.http_app(path="/")
-        # Expose the lifespan for the parent app to chain
-        self.lifespan = self.mcp_app.lifespan_handler if hasattr(self.mcp_app, "lifespan_handler") else None
+
+        # Create two server instances:
+        # 1. Multi-bank server (for /mcp/ root endpoint)
+        self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
+        self.multi_bank_app = self.multi_bank_server.http_app(path="/")
+
+        # 2. Single-bank server (for /mcp/{bank_id}/ endpoints)
+        self.single_bank_server = create_mcp_server(memory, multi_bank=False)
+        self.single_bank_app = self.single_bank_server.http_app(path="/")
+
+        # Backward compatibility: expose multi_bank_app as mcp_app
+        self.mcp_app = self.multi_bank_app
+
+        # Expose the lifespan for the parent app to chain (use multi-bank as default)
+        self.lifespan = (
+            self.multi_bank_app.lifespan_handler if hasattr(self.multi_bank_app, "lifespan_handler") else None
+        )
 
     def _get_header(self, scope: dict, name: str) -> str | None:
         """Extract a header value from ASGI scope."""
@@ -125,7 +153,7 @@ def _get_header(self, scope: dict, name: str) -> str | None:
 
     async def __call__(self, scope, receive, send):
         if scope["type"] != "http":
-            await self.mcp_app(scope, receive, send)
+            await self.multi_bank_app(scope, receive, send)
             return
 
         # Extract auth token from header (for tenant auth propagation)
@@ -173,8 +201,13 @@ async def __call__(self, scope, receive, send):
         elif path == "/mcp":
             path = "/"
 
+        # Ensure path has leading slash (needed after stripping mount path)
+        if path and not path.startswith("/"):
+            path = "/" + path
+
         # Try to get bank_id from header first (for Claude Code compatibility)
         bank_id = self._get_header(scope, "X-Bank-Id")
+        bank_id_from_path = False
 
         # MCP endpoint paths that should not be treated as bank_ids
         MCP_ENDPOINTS = {"sse", "messages"}
@@ -187,13 +220,19 @@ async def __call__(self, scope, receive, send):
             if parts[0] and parts[0] not in MCP_ENDPOINTS:
                 # First segment looks like a bank_id
                 bank_id = parts[0]
+                bank_id_from_path = True
                 new_path = "/" + parts[1] if len(parts) > 1 else "/"
 
         # Fall back to default bank_id
         if not bank_id:
             bank_id = DEFAULT_BANK_ID
             logger.debug(f"Using default bank_id: {bank_id}")
 
+        # Select the appropriate MCP app based on how bank_id was provided:
+        # - Path-based bank_id → single-bank app (no bank_id param, scoped tools)
+        # - 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
         bank_id_token = _current_bank_id.set(bank_id)
         # Store the auth token for tenant extension to validate
@@ -206,15 +245,15 @@ async def __call__(self, scope, receive, send):
 
             # Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing
             async def send_wrapper(message):
-                if message["type"] == "http.response.body":
+                if message["type"] == "http.response.body" and bank_id_from_path:
                     body = message.get("body", b"")
                     if body and b"/messages" in body:
                         # Rewrite /messages to /{bank_id}/messages in SSE endpoint event
                         body = body.replace(b"data: /messages", f"data: /{bank_id}/messages".encode())
                         message = {**message, "body": body}
                 await send(message)
 
-            await self.mcp_app(new_scope, receive, send_wrapper)
+            await target_app(new_scope, receive, send_wrapper)
         finally:
             _current_bank_id.reset(bank_id_token)
             if api_key_token is not None:
@@ -242,15 +281,23 @@ async def _send_error(self, send, status: int, message: str):
 
 def create_mcp_app(memory: MemoryEngine):
     """
-    Create an ASGI app that handles MCP requests.
+    Create an ASGI app that handles MCP requests with dynamic tool exposure.
 
     Authentication:
         Uses the TenantExtension from the MemoryEngine (same auth as REST API).
 
-    Bank ID can be provided via:
-    1. X-Bank-Id header: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
-    2. URL path: /mcp/{bank_id}/
-    3. Environment variable HINDSIGHT_MCP_BANK_ID (fallback, default: "default")
+    Two modes based on URL structure:
+
+    1. Single-bank mode (recommended for agent isolation):
+       - URL: /mcp/{bank_id}/
+       - Tools: retain, recall, reflect (no bank_id parameter)
+       - Example: claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/
+
+    2. Multi-bank mode (for cross-bank operations):
+       - URL: /mcp/
+       - Tools: retain, recall, reflect, list_banks, create_bank (all with bank_id parameter)
+       - Bank ID from: X-Bank-Id header or HINDSIGHT_MCP_BANK_ID env var (default: "default")
+       - Example: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
 
     Args:
         memory: MemoryEngine instance

hindsight-api/tests/test_mcp_endpoint_routing.py

@@ -0,0 +1,78 @@
+"""Integration test for MCP endpoint routing.
+
+This test verifies that /mcp/ and /mcp/{bank_id}/ expose different tool sets.
+"""
+
+import httpx
+import pytest
+from mcp.client.session import ClientSession
+from mcp.client.streamable_http import streamable_http_client
+
+
+@pytest.mark.asyncio
+async def test_mcp_endpoint_routing_integration(memory):
+    """Test that multi-bank and single-bank endpoints expose different tools using StreamableHTTP.
+
+    This is a regression test for issue #317 where /mcp/{bank_id}/ was incorrectly
+    exposing all tools (including list_banks) and bank_id parameters.
+    """
+    from hindsight_api.api import create_app
+
+    # Create app with MCP enabled
+    app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
+
+    # Use the app's lifespan context to properly initialize MCP servers
+    async with app.router.lifespan_context(app):
+        # Create an HTTPX client that routes to our ASGI app
+        from httpx import ASGITransport
+
+        async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
+            # Test 1: Multi-bank endpoint /mcp/
+            async with streamable_http_client("http://test/mcp/", http_client=http_client) as (
+                read_stream,
+                write_stream,
+                _,
+            ):
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+                    multi_result = await session.list_tools()
+
+                    multi_tools = {t.name for t in multi_result.tools}
+
+                    # Multi-bank should have all tools including bank management
+                    assert "retain" in multi_tools
+                    assert "recall" in multi_tools
+                    assert "reflect" in multi_tools
+                    assert "list_banks" in multi_tools, "Multi-bank should expose list_banks"
+                    assert "create_bank" in multi_tools, "Multi-bank should expose create_bank"
+
+                    # Multi-bank retain should have bank_id parameter
+                    retain_tool = next((t for t in multi_result.tools if t.name == "retain"), None)
+                    assert retain_tool is not None
+                    multi_params = set(retain_tool.inputSchema.get("properties", {}).keys())
+                    assert "bank_id" in multi_params, "Multi-bank retain should have bank_id parameter"
+
+            # Test 2: Single-bank endpoint /mcp/test-bank/
+            async with streamable_http_client("http://test/mcp/test-bank/", http_client=http_client) as (
+                read_stream,
+                write_stream,
+                _,
+            ):
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+                    single_result = await session.list_tools()
+
+                    single_tools = {t.name for t in single_result.tools}
+
+                    # Single-bank should only have scoped tools (no bank management)
+                    assert "retain" in single_tools
+                    assert "recall" in single_tools
+                    assert "reflect" in single_tools
+                    assert "list_banks" not in single_tools, "Single-bank should NOT expose list_banks"
+                    assert "create_bank" not in single_tools, "Single-bank should NOT expose create_bank"
+
+                    # Single-bank retain should NOT have bank_id parameter
+                    retain_tool = next((t for t in single_result.tools if t.name == "retain"), None)
+                    assert retain_tool is not None
+                    single_params = set(retain_tool.inputSchema.get("properties", {}).keys())
+                    assert "bank_id" not in single_params, "Single-bank retain should NOT have bank_id parameter"