fix http transport to be stateful

Author: shahar4499

fastapi_mcp/transport/http.py

@@ -1,74 +1,98 @@
 import logging
+import asyncio
 
 from fastapi import Request, Response, HTTPException
 from mcp.server.lowlevel.server import Server
-from mcp.server.streamable_http import StreamableHTTPServerTransport
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager, EventStore
 from mcp.server.transport_security import TransportSecuritySettings
 
 logger = logging.getLogger(__name__)
 
 
-class FastApiStreamableHttpTransport(StreamableHTTPServerTransport):
+class FastApiHttpSessionManager:
+    """
+    FastAPI-native wrapper around StreamableHTTPSessionManager
+    """
+
     def __init__(
         self,
-        mcp_session_id: str | None = None,
-        is_json_response_enabled: bool = True,  # Default to JSON for HTTP transport
-        event_store=None,
+        mcp_server: Server,
+        event_store: EventStore | None = None,
+        json_response: bool = True,  # Default to JSON for HTTP transport
         security_settings: TransportSecuritySettings | None = None,
-        mcp_server: Server | None = None,
     ):
-        super().__init__(
-            mcp_session_id=mcp_session_id,
-            is_json_response_enabled=is_json_response_enabled,
-            event_store=event_store,
-            security_settings=security_settings,
-        )
-        logger.debug(f"FastApiStreamableHttpTransport initialized with session_id: {mcp_session_id}")
-        self._mcp_server = mcp_server
-        self._server_running = False
-
-    async def handle_fastapi_request(self, request: Request, mcp_server: Server | None = None) -> Response:
+        self.mcp_server = mcp_server
+        self.event_store = event_store
+        self.json_response = json_response
+        self.security_settings = security_settings
+        self._session_manager: StreamableHTTPSessionManager | None = None
+        self._manager_task: asyncio.Task | None = None
+        self._manager_started = False
+        self._startup_lock = asyncio.Lock()
+
+    async def _ensure_session_manager_started(self) -> None:
         """
-        The approach here is different from FastApiSseTransport.
-        In FastApiSseTransport, we reimplement the SSE transport logic to have a more FastAPI-native transport.
-        It proved to be less bug-prone since it avoids deconstructing and reconstructing raw ASGI objects.
-
-        But, we took a different approach here because StreamableHTTPServerTransport handles more complexity,
-        and multiple request methods (GET/POST/DELETE), so we want to leverage that logic and avoid reimplementing.
+        Ensure the session manager is started.
 
-        We still ensure it works natively with FastAPI by capturing the ASGI response from the SDK and converting
-        it to a FastAPI Response.
+        This is called lazily on the first request to start the session manager
+        if it hasn't been started yet.
         """
-        logger.debug(f"Handling FastAPI request: {request.method} {request.url.path}")
-
-        # Use the stored server if available, or the passed one
-        server = self._mcp_server or mcp_server
-        if not server:
-            raise HTTPException(status_code=500, detail="No MCP server available")
-
-        # Initialize the transport if not already done
-        if not self._server_running:
-            import anyio
+        if self._manager_started:
+            return
+
+        async with self._startup_lock:
+            if self._manager_started:
+                return
+
+            logger.debug("Starting StreamableHTTP session manager")
+
+            # Create the session manager
+            # Note: We don't use stateless=True because we want to support sessions
+            # but sessions are optional as per the MCP spec
+            self._session_manager = StreamableHTTPSessionManager(
+                app=self.mcp_server,
+                event_store=self.event_store,
+                json_response=self.json_response,
+                stateless=False,  # Always support sessions, but they're optional
+                security_settings=self.security_settings,
+            )
 
-            async def start_server():
-                self._server_running = True
-                async with self.connect() as (reader, writer):
-                    await server.run(
-                        reader,
-                        writer,
-                        server.create_initialization_options(notification_options=None, experimental_capabilities={}),
-                        raise_exceptions=False,
-                    )
+            # Start the session manager in a background task
+            async def run_session_manager():
+                try:
+                    async with self._session_manager.run():
+                        logger.info("StreamableHTTP session manager is running")
+                        # Keep running until cancelled
+                        await asyncio.Event().wait()
+                except asyncio.CancelledError:
+                    logger.info("StreamableHTTP session manager is shutting down")
+                    raise
+                except Exception:
+                    logger.exception("Error in StreamableHTTP session manager")
+                    raise
+
+            self._manager_task = asyncio.create_task(run_session_manager())
+            self._manager_started = True
+
+            # Give the session manager a moment to initialize
+            await asyncio.sleep(0.1)
+
+    async def handle_fastapi_request(self, request: Request) -> Response:
+        """
+        Handle a FastAPI request by delegating to the session manager.
 
-            # Start the server in a background task
-            import asyncio
+        This converts FastAPI's Request/Response to ASGI scope/receive/send
+        and then converts the result back to a FastAPI Response.
+        """
+        # Ensure session manager is started
+        await self._ensure_session_manager_started()
 
-            asyncio.create_task(start_server())
+        if not self._session_manager:
+            raise HTTPException(status_code=500, detail="Session manager not initialized")
 
-            # Give the server a moment to initialize
-            await anyio.sleep(0.1)
+        logger.debug(f"Handling FastAPI request: {request.method} {request.url.path}")
 
-        # Capture the response from the SDK's handle_request method
+        # Capture the response from the session manager
         response_started = False
         response_status = 200
         response_headers = []
@@ -85,8 +109,8 @@ async def send_callback(message):
                 response_body += message.get("body", b"")
 
         try:
-            # Delegate to the SDK's handle_request method with ASGI interface
-            await self.handle_request(request.scope, request.receive, send_callback)
+            # Delegate to the session manager's handle_request method
+            await self._session_manager.handle_request(request.scope, request.receive, send_callback)
 
             # Convert the captured ASGI response to a FastAPI Response
             headers_dict = {name.decode(): value.decode() for name, value in response_headers}
@@ -98,5 +122,15 @@ async def send_callback(message):
             )
 
         except Exception:
-            logger.exception("Error in StreamableHTTPServerTransport")
+            logger.exception("Error in StreamableHTTPSessionManager")
             raise HTTPException(status_code=500, detail="Internal server error")
+
+    async def shutdown(self) -> None:
+        """Clean up the session manager and background task."""
+        if self._manager_task and not self._manager_task.done():
+            self._manager_task.cancel()
+            try:
+                await self._manager_task
+            except asyncio.CancelledError:
+                pass
+        self._manager_started = False