fix: mcp server

hindsight-api/hindsight_api/api/__init__.py

@@ -5,7 +5,7 @@ Provides both HTTP REST API and MCP (Model Context Protocol) server.
 """
 
 import logging
-from typing import Optional
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
@@ -45,33 +45,46 @@ def create_app(
         # Both HTTP and MCP
         app = create_app(memory, mcp_api_enabled=True)
     """
+    mcp_app = None
+    mcp_lifespan = None
 
-    # Import and create HTTP API if enabled
-    if http_api_enabled:
-        from .http import create_app as create_http_app
-
-        app = create_http_app(memory=memory, initialize_memory=initialize_memory)
-        logger.info("HTTP REST API enabled")
-    else:
-        # Create minimal FastAPI app
-        app = FastAPI(title="Hindsight API", version="0.0.7")
-        logger.info("HTTP REST API disabled")
-
-    # Mount MCP server if enabled
+    # Create MCP app if enabled
     if mcp_api_enabled:
         try:
             from .mcp import create_mcp_app
 
             # Create MCP app with dynamic bank_id support
-            # Supports: /mcp/{bank_id}/sse (bank-specific SSE endpoint)
-            mcp_app = create_mcp_app(memory=memory)
-            app.mount(mcp_mount_path, mcp_app)
-            logger.info(f"MCP server enabled at {mcp_mount_path}/{{bank_id}}/sse")
+            # Returns (app, lifespan) - lifespan must be passed to parent FastAPI
+            # because Starlette's Mount doesn't forward lifespan events
+            mcp_app, mcp_lifespan = create_mcp_app(memory=memory)
         except ImportError as e:
             logger.error(f"MCP server requested but dependencies not available: {e}")
             logger.error("Install with: pip install hindsight-api[mcp]")
             raise
 
+    # Import and create HTTP API if enabled
+    if http_api_enabled:
+        from .http import create_app as create_http_app
+
+        app = create_http_app(
+            memory=memory,
+            initialize_memory=initialize_memory,
+            mcp_lifespan=mcp_lifespan,
+        )
+        logger.info("HTTP REST API enabled")
+    else:
+        # Create minimal FastAPI app with MCP lifespan if needed
+        if mcp_lifespan:
+            app = FastAPI(title="Hindsight API", version="0.0.7", lifespan=mcp_lifespan)
+        else:
+            app = FastAPI(title="Hindsight API", version="0.0.7")
+        logger.info("HTTP REST API disabled")
+
+    # Mount MCP server if enabled
+    if mcp_app:
+        app.mount(mcp_mount_path, mcp_app)
+        logger.info(f"MCP server enabled at {mcp_mount_path}/{{bank_id}}/mcp")
+
     return app
 
 

hindsight-api/hindsight_api/api/http.py

@@ -706,7 +706,7 @@ class DeleteResponse(BaseModel):
     deleted_count: int | None = None
 
 
-def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
+def create_app(memory: MemoryEngine, initialize_memory: bool = True, mcp_lifespan=None) -> FastAPI:
     """
     Create and configure the FastAPI application.
 
@@ -714,6 +714,7 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
         memory: MemoryEngine instance (already initialized with required parameters).
                 Migrations are controlled by the MemoryEngine's run_migrations parameter.
         initialize_memory: Whether to initialize memory system on startup (default: True)
+        mcp_lifespan: Optional MCP lifespan context manager to run (for Streamable HTTP)
 
     Returns:
         Configured FastAPI application
@@ -746,7 +747,14 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
             await memory.initialize()
             logging.info("Memory system initialized")
 
-        yield
+        # Run MCP lifespan if provided (required for Streamable HTTP transport)
+        # This must be run on the parent app because Starlette's Mount doesn't
+        # forward lifespan events to mounted apps
+        if mcp_lifespan:
+            async with mcp_lifespan(app):
+                yield
+        else:
+            yield
 
         # Shutdown: Cleanup memory system
         await memory.close()

hindsight-api/hindsight_api/api/mcp.py

@@ -115,66 +115,72 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
     return mcp
 
 
-class MCPMiddleware:
-    """ASGI middleware that extracts bank_id from path and sets context."""
+class MCPRouterMiddleware:
+    """ASGI middleware that extracts bank_id from path and routes to MCP server.
 
-    def __init__(self, app, memory: MemoryEngine):
-        self.app = app
-        self.memory = memory
-        self.mcp_server = create_mcp_server(memory)
-        self.mcp_app = self.mcp_server.http_app()
+    This middleware wraps the FastMCP http_app and:
+    1. Extracts bank_id from the URL path (pattern: /{bank_id}/mcp)
+    2. Sets the bank_id in a context variable for tools to access
+    3. Forwards requests to the underlying MCP app with the correct path (/mcp)
+
+    The middleware also handles lifespan events to ensure the MCP server's
+    session manager is properly initialized.
+    """
+
+    def __init__(self, mcp_http_app):
+        self.mcp_http_app = mcp_http_app
+        self._lifespan_started = False
 
     async def __call__(self, scope, receive, send):
+        logger.debug(f"MCPRouterMiddleware: type={scope['type']}, path={scope.get('path', 'N/A')}")
+
+        # Handle lifespan events - forward to MCP app
+        if scope["type"] == "lifespan":
+            logger.debug("MCPRouterMiddleware: handling lifespan")
+            await self.mcp_http_app(scope, receive, send)
+            return
+
+        # Only handle HTTP requests
         if scope["type"] != "http":
-            await self.mcp_app(scope, receive, send)
+            await self.mcp_http_app(scope, receive, send)
             return
 
         path = scope.get("path", "")
 
-        # Strip any mount prefix (e.g., /mcp) that FastAPI might not have stripped
+        # Strip any mount prefix (root_path) from the path
         root_path = scope.get("root_path", "")
         if root_path and path.startswith(root_path):
             path = path[len(root_path) :] or "/"
 
-        # Also handle case where mount path wasn't stripped (e.g., /mcp/...)
-        if path.startswith("/mcp/"):
-            path = path[4:]  # Remove /mcp prefix
-
-        # Extract bank_id from path: /{bank_id}/ or /{bank_id}
-        # http_app expects requests at /
+        # Path should now be like /{bank_id}/mcp or /{bank_id}
+        # Extract bank_id from first path segment
         if not path.startswith("/") or len(path) <= 1:
-            # No bank_id in path - return error
-            await self._send_error(send, 400, "bank_id required in path: /mcp/{bank_id}/")
+            await self._send_error(send, 400, "bank_id required in path: /{bank_id}/mcp")
             return
 
-        # Extract bank_id from first path segment
         parts = path[1:].split("/", 1)
         if not parts[0]:
-            await self._send_error(send, 400, "bank_id required in path: /mcp/{bank_id}/")
+            await self._send_error(send, 400, "bank_id required in path: /{bank_id}/mcp")
             return
 
         bank_id = parts[0]
-        new_path = "/" + parts[1] if len(parts) > 1 else "/"
+        # The remainder of the path after bank_id, or /mcp if nothing follows
+        remaining_path = "/" + parts[1] if len(parts) > 1 else "/mcp"
 
-        # Set bank_id context
+        logger.debug(f"MCPRouterMiddleware: bank_id={bank_id}, remaining_path={remaining_path}")
+
+        # Set bank_id context for this request
         token = _current_bank_id.set(bank_id)
         try:
+            # Create new scope with the path that FastMCP expects (/mcp)
             new_scope = scope.copy()
-            new_scope["path"] = new_path
+            new_scope["path"] = remaining_path
+            new_scope["raw_path"] = remaining_path.encode()
+            # Clear root_path since we've already handled the routing
+            new_scope["root_path"] = ""
 
-            # Wrap send to rewrite the SSE endpoint URL to include bank_id
-            # The SSE app sends "event: endpoint\ndata: /messages\n" but we need
-            # the client to POST to /{bank_id}/messages instead
-            async def send_wrapper(message):
-                if message["type"] == "http.response.body":
-                    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)
+            logger.debug(f"MCPRouterMiddleware: forwarding to MCP app with path={remaining_path}")
+            await self.mcp_http_app(new_scope, receive, send)
         finally:
             _current_bank_id.reset(token)
 
@@ -198,16 +204,29 @@ class MCPMiddleware:
 
 def create_mcp_app(memory: MemoryEngine):
     """
-    Create an ASGI app that handles MCP requests.
+    Create an ASGI app that handles MCP requests using Streamable HTTP transport.
 
-    URL pattern: /mcp/{bank_id}/
+    URL pattern: {mount_path}/{bank_id}/mcp
 
     The bank_id is extracted from the URL path and made available to tools.
+    Uses Streamable HTTP transport (recommended by MCP spec).
 
     Args:
         memory: MemoryEngine instance
 
     Returns:
-        ASGI application
+        Tuple of (ASGI application, lifespan context manager).
+        The lifespan MUST be passed to the parent FastAPI app for Streamable HTTP to work.
+        This is required because Starlette's Mount doesn't forward lifespan events.
     """
-    return MCPMiddleware(None, memory)
+    mcp_server = create_mcp_server(memory)
+    # Use Streamable HTTP transport (recommended)
+    mcp_http_app = mcp_server.http_app(transport="streamable-http")
+
+    # Wrap with router middleware for bank_id extraction
+    router_app = MCPRouterMiddleware(mcp_http_app)
+
+    # Return both the app and the lifespan
+    # The lifespan must be passed to the parent FastAPI app because
+    # Starlette's Mount doesn't forward lifespan events to mounted apps
+    return router_app, mcp_http_app.lifespan

hindsight-api/tests/test_mcp_api_integration.py

@@ -1,176 +1,204 @@
 """
 Integration test for the MCP (Model Context Protocol) server.
 
-Tests MCP endpoints by starting a FastAPI server with MCP enabled and using the MCP client.
+Tests MCP endpoints by starting a real FastAPI server with MCP enabled and using the MCP client.
+Uses Streamable HTTP transport, which is the recommended MCP transport that provides:
+- Single HTTP endpoint for all communication
+- Streaming responses via Server-Sent Events
+- Proper session management
+- Better performance than legacy SSE transport
 
-Note: MCP server is integrated with the web server. These tests require HINDSIGHT_API_MCP_ENABLED=true.
+These tests verify the full integration flow including:
+- Server startup with proper lifespan management
+- MCP client connection and session initialization
+- Tool listing and execution
+- Multi-tenant bank_id routing
+- Concurrent request handling
 """
 import asyncio
+import socket
 import pytest
 import pytest_asyncio
-import httpx
+import uvicorn
 from mcp import ClientSession
-from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 from hindsight_api.api import create_app
 
 
+def get_free_port():
+    """Get a free port on localhost."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        s.listen(1)
+        return s.getsockname()[1]
+
+
+class UvicornServer:
+    """Helper class to run uvicorn in background."""
+
+    def __init__(self, app, host: str, port: int):
+        self.config = uvicorn.Config(app=app, host=host, port=port, log_level="warning")
+        self.server = uvicorn.Server(self.config)
+        self._task = None
+
+    async def start(self):
+        """Start the server in background."""
+        self._task = asyncio.create_task(self.server.serve())
+        # Wait for server to be ready
+        while not self.server.started:
+            await asyncio.sleep(0.01)
+
+    async def stop(self):
+        """Stop the server."""
+        self.server.should_exit = True
+        if self._task:
+            await self._task
+
+
 @pytest_asyncio.fixture
 async def mcp_server(memory):
-    """Start the FastAPI app with MCP enabled and return the SSE URL."""
+    """Start a real FastAPI server with MCP enabled and return the MCP URL."""
+    import uuid
+
     # Memory is already initialized by the conftest fixture (with migrations)
     app = create_app(
         memory,
         initialize_memory=False,
-        mcp_api_enabled=True
+        mcp_api_enabled=True,
+        mcp_mount_path="/mcp"
     )
 
-    # Use httpx to create a test server
-    transport = httpx.ASGITransport(app=app)
-    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
-        # The MCP SSE endpoint is at /mcp/sse
-        # We need to yield the base URL for sse_client to connect
-        # However, sse_client expects a real URL, not a test client
-        # So we'll start a real server on a random port
-        pass
+    port = get_free_port()
+    server = UvicornServer(app, "127.0.0.1", port)
+    await server.start()
 
-    # For now, skip these tests as they require a real server
-    # The sse_client doesn't work with ASGI test transport
-    pytest.skip("MCP tests require a real running server. Run: HINDSIGHT_API_MCP_ENABLED=true uvicorn hindsight_api.api:app")
+    # Use a unique bank_id for tests - bank will be auto-created on first retain
+    bank_id = f"mcp-test-{uuid.uuid4().hex[:8]}"
+
+    # Return the Streamable HTTP URL: /mcp/{bank_id}/mcp
+    mcp_url = f"http://127.0.0.1:{port}/mcp/{bank_id}/mcp"
+
+    yield mcp_url
+
+    # Cleanup
+    await server.stop()
+    # Delete test bank data
+    await memory.delete_bank(bank_id)
 
 
 @pytest.mark.asyncio
-async def test_mcp_server_tools_via_sse(mcp_server):
-    """Test MCP server tools via SSE transport using proper MCP client."""
-    sse_url = mcp_server
+async def test_mcp_list_tools(mcp_server):
+    """Test that MCP server exposes the expected tools."""
+    mcp_url = mcp_server
 
-    async with sse_client(sse_url) as (read, write):
+    async with streamablehttp_client(mcp_url) as (read, write, _):
         async with ClientSession(read, write) as session:
             await session.initialize()
 
-            # Test 1: List tools
             tools_list = await session.list_tools()
-            print(f"Tools: {tools_list}")
             tool_names = [t.name for t in tools_list.tools]
-            assert "hindsight_search" in tool_names
-            assert "hindsight_put" in tool_names
 
-            # Test 2: Call hindsight_put
+            assert "retain" in tool_names
+            assert "recall" in tool_names
+
+
+@pytest.mark.asyncio
+async def test_mcp_retain_and_recall(mcp_server):
+    """Test retain and recall flow via MCP."""
+    mcp_url = mcp_server
+
+    async with streamablehttp_client(mcp_url) as (read, write, _):
+        async with ClientSession(read, write) as session:
+            await session.initialize()
+
+            # Store a memory via retain
             put_result = await session.call_tool(
-                "hindsight_put",
+                "retain",
                 arguments={
-                    "content": "User loves Python programming",
+                    "content": "User loves Python programming and prefers functional style",
                     "context": "programming_preferences",
-                    "explanation": "Storing user's programming language preference"
                 }
             )
-            print(f"Put result: {put_result}")
             assert put_result is not None
 
-            # Wait a bit for indexing
+            # Wait for indexing
             await asyncio.sleep(1)
 
-            # Test 3: Call hindsight_search
+            # Search for it via recall
             search_result = await session.call_tool(
-                "hindsight_search",
+                "recall",
                 arguments={
                     "query": "What programming languages does the user like?",
-                    "max_tokens": 4096,
-                    "explanation": "Searching for programming preferences"
                 }
             )
-            print(f"Search result: {search_result}")
             assert search_result is not None
 
 
 @pytest.mark.asyncio
-async def test_multiple_concurrent_requests(mcp_server):
+async def test_mcp_multiple_concurrent_requests(mcp_server):
     """Test multiple concurrent requests from a single session."""
-    sse_url = mcp_server
+    mcp_url = mcp_server
 
-    async with sse_client(sse_url) as (read, write):
+    async with streamablehttp_client(mcp_url) as (read, write, _):
         async with ClientSession(read, write) as session:
             await session.initialize()
 
-            # Fire off 10 concurrent search requests from same session
+            # Fire off concurrent recall requests
             async def make_search(idx):
                 try:
                     result = await session.call_tool(
-                        "hindsight_search",
+                        "recall",
                         arguments={
                             "query": f"test query {idx}",
-                            "explanation": f"Concurrent test {idx}"
                         }
                     )
                     return idx, "success", result
                 except Exception as e:
                     return idx, "error", str(e)
 
-            tasks = [make_search(i) for i in range(10)]
+            tasks = [make_search(i) for i in range(5)]
             results = await asyncio.gather(*tasks, return_exceptions=True)
 
-            # Check results
-            successes = 0
-            failures = 0
+            # Count successes
+            successes = sum(
+                1 for r in results
+                if not isinstance(r, Exception) and r[1] == "success"
+            )
 
-            for result in results:
-                if isinstance(result, Exception):
-                    print(f"Request failed with exception: {result}")
-                    failures += 1
-                else:
-                    idx, status, data = result
-                    if status == "success":
-                        successes += 1
-                    else:
-                        print(f"Request {idx} failed: {data}")
-                        failures += 1
-
-            print(f"Successes: {successes}, Failures: {failures}")
-
-            # We expect all requests to succeed
-            assert successes >= 8, f"Too many failures: {failures}/10"
+            # Most requests should succeed
+            assert successes >= 3, f"Too many failures: only {successes}/5 succeeded"
 
 
 @pytest.mark.asyncio
-async def test_race_condition_with_rapid_requests(mcp_server):
-    """Test rapid-fire requests with multiple sessions to trigger race condition."""
-    sse_url = mcp_server
+async def test_mcp_rapid_sessions(mcp_server):
+    """Test rapid-fire requests with multiple sessions."""
+    mcp_url = mcp_server
 
     async def rapid_session_search(idx):
         """Create a new session and immediately make a request."""
         try:
-            async with sse_client(sse_url) as (read, write):
+            async with streamablehttp_client(mcp_url) as (read, write, _):
                 async with ClientSession(read, write) as session:
                     await session.initialize()
-
-                    # Make request immediately after initialization
                     result = await session.call_tool(
-                        "hindsight_search",
+                        "recall",
                         arguments={
                             "query": f"rapid query {idx}",
-                            "max_tokens": 2048
                         }
                     )
                     return idx, "success", result
         except Exception as e:
             return idx, "error", str(e)
 
-    # Fire 20 requests with minimal delay, each with its own session
-    tasks = [rapid_session_search(i) for i in range(20)]
+    # Fire requests with their own sessions
+    tasks = [rapid_session_search(i) for i in range(10)]
     results = await asyncio.gather(*tasks)
 
-    # Analyze results
-    errors = []
-    for idx, status, data in results:
-        if status == "error":
-            errors.append((idx, data))
-
-    if errors:
-        print(f"Found {len(errors)} errors:")
-        for idx, error_msg in errors:
-            print(f"  Request {idx}: {error_msg}")
+    # Count errors
+    errors = [(idx, data) for idx, status, data in results if status == "error"]
 
     # Most requests should succeed
-    assert len(errors) < 5, f"Too many errors: {len(errors)}/20"
+    assert len(errors) < 5, f"Too many errors: {len(errors)}/10"
 
 
 if __name__ == "__main__":

uv.lock

@@ -1141,7 +1141,7 @@ wheels = [
 
 [[package]]
 name = "hindsight-all"
-version = "0.1.6"
+version = "0.1.7"
 source = { editable = "hindsight" }
 dependencies = [
     { name = "hindsight-api" },
@@ -1165,7 +1165,7 @@ provides-extras = ["test"]
 
 [[package]]
 name = "hindsight-api"
-version = "0.1.6"
+version = "0.1.7"
 source = { editable = "hindsight-api" }
 dependencies = [
     { name = "alembic" },
@@ -1249,7 +1249,7 @@ requires-dist = [
     { name = "sentence-transformers", specifier = ">=3.0.0,<3.3.0" },
     { name = "sqlalchemy", specifier = ">=2.0.44" },
     { name = "tiktoken", specifier = ">=0.12.0" },
-    { name = "torch", specifier = ">=2.0.0,<2.6.0" },
+    { name = "torch", specifier = ">=2.0.0" },
     { name = "transformers", specifier = ">=4.30.0,<4.46.0" },
     { name = "uvicorn", specifier = ">=0.38.0" },
     { name = "wsproto", specifier = ">=1.0.0" },
@@ -1269,7 +1269,7 @@ dev = [
 
 [[package]]
 name = "hindsight-client"
-version = "0.1.6"
+version = "0.1.7"
 source = { editable = "hindsight-clients/python" }
 dependencies = [
     { name = "aiohttp" },
@@ -1301,7 +1301,7 @@ provides-extras = ["test"]
 
 [[package]]
 name = "hindsight-dev"
-version = "0.1.6"
+version = "0.1.7"
 source = { editable = "hindsight-dev" }
 dependencies = [
     { name = "hindsight-api" },