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

tests/shared/test_streamable_http.py

@@ -4,13 +4,10 @@ Tests for the StreamableHTTP server and client transport.
 Contains tests for both server and client sides of the StreamableHTTP transport.
 """
 
-import contextlib
 import multiprocessing
 import socket
 import time
 from collections.abc import Generator
-from http import HTTPStatus
-from uuid import uuid4
 
 import anyio
 import httpx
@@ -19,8 +16,6 @@ import requests
 import uvicorn
 from pydantic import AnyUrl
 from starlette.applications import Starlette
-from starlette.requests import Request
-from starlette.responses import Response
 from starlette.routing import Mount
 
 import mcp.types as types
@@ -37,6 +32,7 @@ from mcp.server.streamable_http import (
     StreamableHTTPServerTransport,
     StreamId,
 )
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import (
     ClientMessageMetadata,
@@ -184,7 +180,7 @@ class ServerTest(Server):
 def create_app(
     is_json_response_enabled=False, event_store: EventStore | None = None
 ) -> Starlette:
-    """Create a Starlette application for testing that matches the example server.
+    """Create a Starlette application for testing using the session manager.
 
     Args:
         is_json_response_enabled: If True, use JSON responses instead of SSE streams.
@@ -193,85 +189,20 @@ def create_app(
     # Create server instance
     server = ServerTest()
 
-    server_instances = {}
-    # Lock to prevent race conditions when creating new sessions
-    session_creation_lock = anyio.Lock()
-    task_group = None
+    # Create the session manager
+    session_manager = StreamableHTTPSessionManager(
+        app=server,
+        event_store=event_store,
+        json_response=is_json_response_enabled,
+    )
 
-    @contextlib.asynccontextmanager
-    async def lifespan(app):
-        """Application lifespan context manager for managing task group."""
-        nonlocal task_group
-
-        async with anyio.create_task_group() as tg:
-            task_group = tg
-            try:
-                yield
-            finally:
-                if task_group:
-                    tg.cancel_scope.cancel()
-                    task_group = None
-
-    async def handle_streamable_http(scope, receive, send):
-        request = Request(scope, receive)
-        request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
-
-        # Use existing transport if session ID matches
-        if (
-            request_mcp_session_id is not None
-            and request_mcp_session_id in server_instances
-        ):
-            transport = server_instances[request_mcp_session_id]
-
-            await transport.handle_request(scope, receive, send)
-        elif request_mcp_session_id is None:
-            async with session_creation_lock:
-                new_session_id = uuid4().hex
-
-                http_transport = StreamableHTTPServerTransport(
-                    mcp_session_id=new_session_id,
-                    is_json_response_enabled=is_json_response_enabled,
-                    event_store=event_store,
-                )
-
-                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 server.run(
-                            read_stream,
-                            write_stream,
-                            server.create_initialization_options(),
-                        )
-
-                if task_group is None:
-                    response = Response(
-                        "Internal Server Error: Task group is not initialized",
-                        status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
-                    )
-                    await response(scope, receive, send)
-                    return
-
-                # Store the instance before starting the task to prevent races
-                server_instances[http_transport.mcp_session_id] = http_transport
-                await task_group.start(run_server)
-
-                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)
-
-    # Create an ASGI application
+    # Create an ASGI application that uses the session manager
     app = Starlette(
         debug=True,
         routes=[
-            Mount("/mcp", app=handle_streamable_http),
+            Mount("/mcp", app=session_manager.handle_request),
         ],
-        lifespan=lifespan,
+        lifespan=lambda app: session_manager.run(),
     )
 
     return app