diff --git a/src/fastmcp/server/http.py b/src/fastmcp/server/http.py index b37e58658..0a2c6e6e4 100644 --- a/src/fastmcp/server/http.py +++ b/src/fastmcp/server/http.py @@ -22,17 +22,10 @@ from starlette.responses import Response from starlette.routing import Mount, Route from starlette.types import Receive, Scope, Send +# This import is vendored until it is finalized in the upstream SDK +from fastmcp.server.streamable_http_manager import StreamableHTTPSessionManager from fastmcp.utilities.logging import get_logger -# Import these conditionally to handle case where they might not be available -try: - from mcp.server.streamable_http import EventStore - from mcp.server.streamable_http_manager import StreamableHTTPSessionManager - - STREAMABLE_HTTP_AVAILABLE = True -except ImportError: - STREAMABLE_HTTP_AVAILABLE = False - if TYPE_CHECKING: from fastmcp.server.server import FastMCP @@ -238,7 +231,7 @@ def create_sse_app( def create_streamable_http_app( server: FastMCP, streamable_http_path: str, - event_store: EventStore | None = None, + event_store: None = None, auth_server_provider: OAuthAuthorizationServerProvider | None = None, auth_settings: AuthSettings | None = None, json_response: bool = False, @@ -262,11 +255,6 @@ def create_streamable_http_app( Returns: A Starlette application with StreamableHTTP support """ - if not STREAMABLE_HTTP_AVAILABLE: - raise ImportError( - "StreamableHTTP transport is not available. Make sure your version of `mcp` is up-to-date." - ) - # Create session manager using the provided event store session_manager = StreamableHTTPSessionManager( app=server._mcp_server, diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 0fcbb15c5..7c11528a9 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -752,25 +752,19 @@ class FastMCP(Generic[LifespanResultT]): def streamable_http_app(self) -> Starlette: """Return an instance of the StreamableHTTP server app.""" - try: - from fastmcp.server.http import create_streamable_http_app + from fastmcp.server.http import create_streamable_http_app - return create_streamable_http_app( - server=self, - streamable_http_path=self.settings.streamable_http_path, - event_store=None, - auth_server_provider=self._auth_server_provider, - auth_settings=self.settings.auth, - json_response=self.settings.json_response, - stateless_http=self.settings.stateless_http, - debug=self.settings.debug, - additional_routes=self._additional_http_routes, - ) - except ImportError as e: - logger.error(f"Failed to create StreamableHTTP app: {e}") - raise ImportError( - "StreamableHTTP transport is not available. Make sure your version of `mcp` is up-to-date." - ) from e + return create_streamable_http_app( + server=self, + streamable_http_path=self.settings.streamable_http_path, + event_store=None, + auth_server_provider=self._auth_server_provider, + auth_settings=self.settings.auth, + json_response=self.settings.json_response, + stateless_http=self.settings.stateless_http, + debug=self.settings.debug, + additional_routes=self._additional_http_routes, + ) async def run_streamable_http_async( self, diff --git a/src/fastmcp/server/streamable_http_manager.py b/src/fastmcp/server/streamable_http_manager.py new file mode 100644 index 000000000..37cdc23ec --- /dev/null +++ b/src/fastmcp/server/streamable_http_manager.py @@ -0,0 +1,241 @@ +"""StreamableHTTP Session Manager for MCP servers.""" + +# follows https://github.com/modelcontextprotocol/python-sdk/blob/ihrpr/shttp/src/mcp/server/streamable_http_manager.py +# and can be removed once that spec is finalized + +from __future__ import annotations + +import contextlib +import logging +from collections.abc import AsyncIterator +from http import HTTPStatus +from typing import Any +from uuid import uuid4 + +import anyio +from anyio.abc import TaskStatus +from mcp.server.lowlevel.server import Server as MCPServer +from mcp.server.streamable_http import ( + MCP_SESSION_ID_HEADER, + EventStore, + StreamableHTTPServerTransport, +) +from starlette.requests import Request +from starlette.responses import Response +from starlette.types import Receive, Scope, Send + +logger = logging.getLogger(__name__) + + +class StreamableHTTPSessionManager: + """ + Manages StreamableHTTP sessions with optional resumability via event store. + + This class abstracts away the complexity of session management, event storage, + and request handling for StreamableHTTP transports. It handles: + + 1. Session tracking for clients + 2. Resumability via an optional event store + 3. Connection management and lifecycle + 4. Request handling and transport setup + + Args: + app: The MCP server instance + event_store: Optional event store for resumability support. + If provided, enables resumable connections where clients + can reconnect and receive missed events. + If None, sessions are still tracked but not resumable. + json_response: Whether to use JSON responses instead of SSE streams + stateless: If True, creates a completely fresh transport for each request + with no session tracking or state persistence between requests. + + """ + + def __init__( + self, + app: MCPServer[Any], + event_store: EventStore | None = None, + json_response: bool = False, + stateless: bool = False, + ): + self.app = app + self.event_store = event_store + self.json_response = json_response + self.stateless = stateless + + # Session tracking (only used if not stateless) + self._session_creation_lock = anyio.Lock() + self._server_instances: dict[str, StreamableHTTPServerTransport] = {} + + # The task group will be set during lifespan + self._task_group = None + + @contextlib.asynccontextmanager + async def run(self) -> AsyncIterator[None]: + """ + Run the session manager with proper lifecycle management. + + This creates and manages the task group for all session operations. + + Use this in the lifespan context manager of your Starlette app: + + @contextlib.asynccontextmanager + async def lifespan(app: Starlette) -> AsyncIterator[None]: + async with session_manager.run(): + yield + """ + async with anyio.create_task_group() as tg: + # Store the task group for later use + self._task_group = tg + logger.info("StreamableHTTP session manager started") + try: + yield # Let the application run + finally: + logger.info("StreamableHTTP session manager shutting down") + # Cancel task group to stop all spawned tasks + tg.cancel_scope.cancel() + self._task_group = None + # Clear any remaining server instances + self._server_instances.clear() + + async def handle_request( + self, + scope: Scope, + receive: Receive, + send: Send, + ) -> None: + """ + Process ASGI request with proper session handling and transport setup. + + Dispatches to the appropriate handler based on stateless mode. + + Args: + scope: ASGI scope + receive: ASGI receive function + send: ASGI send function + """ + if self._task_group is None: + raise RuntimeError( + "Task group is not initialized. Make sure to use the run()." + ) + + # Dispatch to the appropriate handler + if self.stateless: + await self._handle_stateless_request(scope, receive, send) + else: + await self._handle_stateful_request(scope, receive, send) + + async def _handle_stateless_request( + self, + scope: Scope, + receive: Receive, + send: Send, + ) -> None: + """ + Process request in stateless mode - creating a new transport for each request. + + Args: + scope: ASGI scope + receive: ASGI receive function + send: ASGI send function + """ + logger.debug("Stateless mode: Creating new transport for this request") + # No session ID needed in stateless mode + http_transport = StreamableHTTPServerTransport( + mcp_session_id=None, # No session tracking in stateless mode + is_json_response_enabled=self.json_response, + event_store=None, # No event store in stateless mode + ) + + # Start server in a new task + async def run_stateless_server( + *, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED + ): + async with http_transport.connect() as streams: + read_stream, write_stream = streams + task_status.started() + await self.app.run( + read_stream, + write_stream, + self.app.create_initialization_options(), + stateless=True, + ) + + # Assert task group is not None for type checking + assert self._task_group is not None + # Start the server task + await self._task_group.start(run_stateless_server) + + # Handle the HTTP request and return the response + await http_transport.handle_request(scope, receive, send) + + async def _handle_stateful_request( + self, + scope: Scope, + receive: Receive, + send: Send, + ) -> None: + """ + Process request in stateful mode - maintaining session state between requests. + + Args: + scope: ASGI scope + receive: ASGI receive function + send: ASGI send function + """ + request = Request(scope, receive) + request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER) + + # Existing session case + if ( + request_mcp_session_id is not None + and request_mcp_session_id in self._server_instances + ): + transport = self._server_instances[request_mcp_session_id] + logger.debug("Session already exists, handling request directly") + await transport.handle_request(scope, receive, send) + return + + if request_mcp_session_id is None: + # New session case + logger.debug("Creating new transport") + async with self._session_creation_lock: + new_session_id = uuid4().hex + http_transport = StreamableHTTPServerTransport( + mcp_session_id=new_session_id, + is_json_response_enabled=self.json_response, + event_store=self.event_store, # May be None (no resumability) + ) + + assert http_transport.mcp_session_id is not None + self._server_instances[http_transport.mcp_session_id] = http_transport + logger.info(f"Created new transport with session ID: {new_session_id}") + + # Define the server runner + async def run_server( + *, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED + ) -> None: + async with http_transport.connect() as streams: + read_stream, write_stream = streams + task_status.started() + await self.app.run( + read_stream, + write_stream, + self.app.create_initialization_options(), + stateless=False, # Stateful mode + ) + + # Assert task group is not None for type checking + assert self._task_group is not None + # Start the server task + await self._task_group.start(run_server) + + # Handle the HTTP request and return the response + await http_transport.handle_request(scope, receive, send) + else: + # Invalid session ID + response = Response( + "Bad Request: No valid session ID provided", + status_code=HTTPStatus.BAD_REQUEST, + ) + await response(scope, receive, send)