Streamable HTTP - improve usability, fast mcp and auth (#641)

examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py

@@ -1,58 +1,22 @@
 import contextlib
 import logging
-from http import HTTPStatus
-from uuid import uuid4
+from collections.abc import AsyncIterator
 
 import anyio
 import click
 import mcp.types as types
 from mcp.server.lowlevel import Server
-from mcp.server.streamable_http import (
-    MCP_SESSION_ID_HEADER,
-    StreamableHTTPServerTransport,
-)
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from pydantic import AnyUrl
 from starlette.applications import Starlette
-from starlette.requests import Request
-from starlette.responses import Response
 from starlette.routing import Mount
+from starlette.types import Receive, Scope, Send
 
 from .event_store import InMemoryEventStore
 
 # Configure logging
 logger = logging.getLogger(__name__)
 
-# Global task group that will be initialized in the lifespan
-task_group = None
-
-# Event store for resumability
-# The InMemoryEventStore enables resumability support for StreamableHTTP transport.
-# It stores SSE events with unique IDs, allowing clients to:
-#   1. Receive event IDs for each SSE message
-#   2. Resume streams by sending Last-Event-ID in GET requests
-#   3. Replay missed events after reconnection
-# Note: This in-memory implementation is for demonstration ONLY.
-# For production, use a persistent storage solution.
-event_store = InMemoryEventStore()
-
-
-@contextlib.asynccontextmanager
-async def lifespan(app):
-    """Application lifespan context manager for managing task group."""
-    global task_group
-
-    async with anyio.create_task_group() as tg:
-        task_group = tg
-        logger.info("Application started, task group initialized!")
-        try:
-            yield
-        finally:
-            logger.info("Application shutting down, cleaning up resources...")
-            if task_group:
-                tg.cancel_scope.cancel()
-                task_group = None
-            logger.info("Resources cleaned up successfully.")
-
 
 @click.command()
 @click.option("--port", default=3000, help="Port to listen on for HTTP")
@@ -156,60 +120,38 @@ def main(
             )
         ]
 
-    # We need to store the server instances between requests
-    server_instances = {}
-    # Lock to prevent race conditions when creating new sessions
-    session_creation_lock = anyio.Lock()
+    # Create event store for resumability
+    # The InMemoryEventStore enables resumability support for StreamableHTTP transport.
+    # It stores SSE events with unique IDs, allowing clients to:
+    #   1. Receive event IDs for each SSE message
+    #   2. Resume streams by sending Last-Event-ID in GET requests
+    #   3. Replay missed events after reconnection
+    # Note: This in-memory implementation is for demonstration ONLY.
+    # For production, use a persistent storage solution.
+    event_store = InMemoryEventStore()
+
+    # Create the session manager with our app and event store
+    session_manager = StreamableHTTPSessionManager(
+        app=app,
+        event_store=event_store,  # Enable resumability
+        json_response=json_response,
+    )
 
     # ASGI handler for streamable HTTP connections
-    async def handle_streamable_http(scope, receive, send):
-        request = Request(scope, receive)
-        request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
-        if (
-            request_mcp_session_id is not None
-            and request_mcp_session_id in server_instances
-        ):
-            transport = server_instances[request_mcp_session_id]
-            logger.debug("Session already exists, handling request directly")
-            await transport.handle_request(scope, receive, send)
-        elif request_mcp_session_id is None:
-            # try to establish new session
-            logger.debug("Creating new transport")
-            # Use lock to prevent race conditions when creating new sessions
-            async with session_creation_lock:
-                new_session_id = uuid4().hex
-                http_transport = StreamableHTTPServerTransport(
-                    mcp_session_id=new_session_id,
-                    is_json_response_enabled=json_response,
-                    event_store=event_store,  # Enable resumability
-                )
-                server_instances[http_transport.mcp_session_id] = http_transport
-                logger.info(f"Created new transport with session ID: {new_session_id}")
+    async def handle_streamable_http(
+        scope: Scope, receive: Receive, send: Send
+    ) -> None:
+        await session_manager.handle_request(scope, receive, send)
 
-                async def run_server(task_status=None):
-                    async with http_transport.connect() as streams:
-                        read_stream, write_stream = streams
-                        if task_status:
-                            task_status.started()
-                        await app.run(
-                            read_stream,
-                            write_stream,
-                            app.create_initialization_options(),
-                        )
-
-                if not task_group:
-                    raise RuntimeError("Task group is not initialized")
-
-                await task_group.start(run_server)
-
-                # Handle the HTTP request and return the response
-                await http_transport.handle_request(scope, receive, send)
-        else:
-            response = Response(
-                "Bad Request: No valid session ID provided",
-                status_code=HTTPStatus.BAD_REQUEST,
-            )
-            await response(scope, receive, send)
+    @contextlib.asynccontextmanager
+    async def lifespan(app: Starlette) -> AsyncIterator[None]:
+        """Context manager for managing session manager lifecycle."""
+        async with session_manager.run():
+            logger.info("Application started with StreamableHTTP session manager!")
+            try:
+                yield
+            finally:
+                logger.info("Application shutting down...")
 
     # Create an ASGI application using the transport
     starlette_app = Starlette(