Vendor streamable http manager

This commit is contained in:
Jeremiah Lowin 2025-05-07 20:56:43 -04:00
commit ef6a01f286
3 changed files with 257 additions and 34 deletions

View file

@ -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,

View file

@ -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,

View file

@ -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)