427 lines
21 KiB
Python
427 lines
21 KiB
Python
"""StreamableHTTP Session Manager for MCP servers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
import math
|
|
from collections.abc import AsyncIterator
|
|
from typing import TYPE_CHECKING, Any, Final
|
|
from uuid import uuid4
|
|
|
|
import anyio
|
|
from anyio.abc import TaskStatus
|
|
from mcp_types import DEFAULT_NEGOTIATED_VERSION, INTERNAL_ERROR, INVALID_REQUEST, ErrorData, JSONRPCError
|
|
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
from mcp.server._streamable_http_modern import handle_modern_request
|
|
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context
|
|
from mcp.server.connection import Connection
|
|
from mcp.server.runner import serve_connection, serve_loop
|
|
from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, EventStore, StreamableHTTPServerTransport
|
|
from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE as DEFAULT_MAX_REQUEST_BODY_SIZE
|
|
from mcp.server.transport_security import RequestBodyLimitMiddleware as RequestBodyLimitMiddleware
|
|
from mcp.server.transport_security import TransportSecuritySettings
|
|
from mcp.shared._compat import resync_tracer
|
|
from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER
|
|
from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
|
|
from mcp.shared.transport_context import TransportContext
|
|
|
|
if TYPE_CHECKING:
|
|
from mcp.server.lowlevel.server import Server
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_SESSION_IDLE_TIMEOUT: Final = 30 * 60
|
|
"""Default idle period in seconds after which a stateful Streamable HTTP session is closed (30 minutes)."""
|
|
|
|
DEFAULT_MAX_SESSIONS: Final = 10_000
|
|
"""Default maximum number of concurrent stateful Streamable HTTP sessions per session manager."""
|
|
|
|
|
|
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
|
|
5. Idle session cleanup
|
|
|
|
Important: Only one StreamableHTTPSessionManager instance should be created
|
|
per application. The instance cannot be reused after its run() context has
|
|
completed. If you need to restart the manager, create a new instance.
|
|
|
|
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.
|
|
security_settings: Optional transport security settings.
|
|
retry_interval: Retry interval in milliseconds to suggest to clients in SSE retry field. Used for SSE
|
|
polling behavior.
|
|
session_idle_timeout: Idle timeout in seconds for stateful sessions. A session that has had no HTTP
|
|
request in flight for this long (no request being served, no open GET stream) is terminated and
|
|
removed; its ID then answers 404 and the client has to initialize a new session. When retry_interval
|
|
is also configured, ensure the idle timeout comfortably exceeds the retry interval to avoid reaping
|
|
sessions during normal SSE polling gaps. Defaults to 1800 (30 minutes); None disables the timeout so
|
|
sessions live until the client deletes them or the manager shuts down. Unused in stateless mode.
|
|
max_request_body_size: Maximum size in bytes for Streamable HTTP request bodies. Requests that
|
|
exceed this limit receive a 413 response before parsing or session creation. Defaults to 4 MiB.
|
|
max_sessions: Maximum number of concurrent stateful sessions. While that many sessions are open, a
|
|
request that would open another one receives a 503 response; existing sessions are unaffected and
|
|
room frees up as they end or expire. Defaults to 10 000; None removes the limit. Unused in stateless
|
|
mode.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
app: Server[Any],
|
|
event_store: EventStore | None = None,
|
|
json_response: bool = False,
|
|
stateless: bool = False,
|
|
security_settings: TransportSecuritySettings | None = None,
|
|
retry_interval: int | None = None,
|
|
session_idle_timeout: float | None = DEFAULT_SESSION_IDLE_TIMEOUT,
|
|
max_request_body_size: int = DEFAULT_MAX_REQUEST_BODY_SIZE,
|
|
max_sessions: int | None = DEFAULT_MAX_SESSIONS,
|
|
):
|
|
if session_idle_timeout is not None and not (math.isfinite(session_idle_timeout) and session_idle_timeout > 0):
|
|
raise ValueError("session_idle_timeout must be a positive, finite number of seconds")
|
|
if max_request_body_size <= 0:
|
|
raise ValueError("max_request_body_size must be a positive number of bytes")
|
|
if max_sessions is not None and max_sessions <= 0:
|
|
raise ValueError("max_sessions must be a positive number of sessions or None")
|
|
|
|
self.app = app
|
|
self.event_store = event_store
|
|
self.json_response = json_response
|
|
self.stateless = stateless
|
|
self.security_settings = security_settings
|
|
self.retry_interval = retry_interval
|
|
self.session_idle_timeout = session_idle_timeout
|
|
self.max_request_body_size = max_request_body_size
|
|
self.max_sessions = max_sessions
|
|
self.asgi_app = RequestBodyLimitMiddleware(self._handle_request, max_request_body_size)
|
|
|
|
# Session tracking (only used if not stateless)
|
|
self._session_creation_lock = anyio.Lock()
|
|
self._server_instances: dict[str, StreamableHTTPServerTransport] = {}
|
|
# Identity of the credential that created each session; requests for a
|
|
# session must present the same credential.
|
|
self._session_owners: dict[str, AuthorizationContext] = {}
|
|
|
|
# The task group and lifespan state are set during run()
|
|
self._task_group = None
|
|
self._lifespan_state: Any = None
|
|
# Thread-safe tracking of run() calls
|
|
self._run_lock = anyio.Lock()
|
|
self._has_started = False
|
|
|
|
@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.
|
|
|
|
Important: This method can only be called once per instance. The same
|
|
StreamableHTTPSessionManager instance cannot be reused after this
|
|
context manager exits. Create a new instance if you need to restart.
|
|
|
|
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
|
|
"""
|
|
# Thread-safe check to ensure run() is only called once
|
|
async with self._run_lock:
|
|
if self._has_started:
|
|
raise RuntimeError(
|
|
"StreamableHTTPSessionManager .run() can only be called "
|
|
"once per instance. Create a new instance if you need to run again."
|
|
)
|
|
self._has_started = True
|
|
|
|
async with self.app.lifespan(self.app) as lifespan_state, anyio.create_task_group() as tg:
|
|
# Store for handle_request: lifespan is entered once for the
|
|
# manager's lifetime, not per request (per-connection cleanup
|
|
# belongs on `connection.exit_stack`).
|
|
self._lifespan_state = lifespan_state
|
|
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
|
|
self._lifespan_state = None
|
|
# Clear any remaining server instances
|
|
self._server_instances.clear()
|
|
self._session_owners.clear()
|
|
await resync_tracer()
|
|
|
|
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.
|
|
"""
|
|
await self.asgi_app(scope, receive, send)
|
|
|
|
async def _handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if self._task_group is None:
|
|
raise RuntimeError("Task group is not initialized. Make sure to use run().")
|
|
|
|
# TODO(L49): header-only era-routing for now; body-primary classification
|
|
# is a follow-up. The legacy paths below own only the known
|
|
# initialize-handshake versions; anything else (including unknown
|
|
# values) goes to the modern entry so the classifier can validate it
|
|
# and return a structured rejection. 2025 paths below remain unchanged.
|
|
header = MCP_PROTOCOL_VERSION_HEADER.encode("ascii")
|
|
pv = next((v.decode("latin-1") for k, v in scope["headers"] if k == header), None)
|
|
if pv is not None and pv not in HANDSHAKE_PROTOCOL_VERSIONS:
|
|
await handle_modern_request(
|
|
self.app, self.security_settings, self.json_response, self._lifespan_state, scope, receive, send
|
|
)
|
|
return
|
|
|
|
# Dispatch to the appropriate handler
|
|
if self.stateless:
|
|
await self._handle_stateless_request(pv, scope, receive, send)
|
|
else:
|
|
await self._handle_stateful_request(scope, receive, send)
|
|
|
|
async def _handle_stateless_request(
|
|
self, protocol_version_hint: str | None, scope: Scope, receive: Receive, send: Send
|
|
) -> None:
|
|
"""Process request in stateless mode - creating a new transport for each request."""
|
|
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
|
|
security_settings=self.security_settings,
|
|
)
|
|
|
|
# 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()
|
|
dispatcher: JSONRPCDispatcher[TransportContext] = JSONRPCDispatcher(
|
|
read_stream,
|
|
write_stream,
|
|
inline_methods=frozenset({"initialize"}),
|
|
# No session ID means a server-to-client request can be
|
|
# written to this POST's response stream, but the client's
|
|
# reply has nowhere to land — `can_send_request=False`
|
|
# makes the per-request channel raise `NoBackChannelError`
|
|
# for requests while still allowing notifications.
|
|
transport_builder=lambda _md: TransportContext(kind="streamable-http", can_send_request=False),
|
|
)
|
|
# Born-ready, no standalone channel: the legacy stateless path
|
|
# never opens a GET stream and need not see `initialize`. The
|
|
# header (or the spec's default-absent value) seeds
|
|
# `ctx.protocol_version`.
|
|
connection = Connection.from_envelope(
|
|
protocol_version_hint if protocol_version_hint is not None else DEFAULT_NEGOTIATED_VERSION,
|
|
None,
|
|
None,
|
|
)
|
|
try:
|
|
await serve_connection(
|
|
self.app, dispatcher, connection=connection, lifespan_state=self._lifespan_state
|
|
)
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Stateless session crashed")
|
|
|
|
# The per-request server task only ends once the transport is
|
|
# terminated, so terminate it even if the request was cancelled.
|
|
assert self._task_group is not None
|
|
try:
|
|
await self._task_group.start(run_stateless_server)
|
|
await http_transport.handle_request(scope, receive, send)
|
|
finally:
|
|
with anyio.CancelScope(shield=True):
|
|
await http_transport.terminate()
|
|
|
|
async def _handle_stateful_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
"""Process request in stateful mode - maintaining session state between requests."""
|
|
request = Request(scope, receive)
|
|
request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
|
|
|
|
user = scope.get("user")
|
|
requestor = authorization_context(user) if isinstance(user, AuthenticatedUser) else None
|
|
|
|
# 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]
|
|
if requestor != self._session_owners.get(request_mcp_session_id):
|
|
# A session can only be used with the credential that created
|
|
# it. Respond exactly as if the session did not exist.
|
|
logger.warning(
|
|
"Rejecting request for session %s: credential does not match the one that created the session",
|
|
request_mcp_session_id[:64],
|
|
)
|
|
await _error_response("Session not found", 404)(scope, receive, send)
|
|
return
|
|
logger.debug("Session already exists, handling request directly")
|
|
await transport.handle_request(scope, receive, send)
|
|
if transport.is_terminated:
|
|
# The client ended the session (DELETE): forget it now rather
|
|
# than when its server task winds down.
|
|
await self._discard_session(request_mcp_session_id, transport)
|
|
return
|
|
|
|
if request_mcp_session_id is None:
|
|
# New session case. Admission (the session limit and registration)
|
|
# is decided under the lock; the request itself is served outside
|
|
# it, so one client that is slow to send its opening request does
|
|
# not hold up the others.
|
|
async with self._session_creation_lock:
|
|
http_transport = self._admit_session(requestor)
|
|
if http_transport is None:
|
|
logger.warning("Refusing to open a new session: %d sessions are already open", self.max_sessions)
|
|
await _error_response("Too many open sessions", 503, INTERNAL_ERROR)(scope, receive, send)
|
|
return
|
|
await self._serve_opening_request(http_transport, scope, receive, send)
|
|
else:
|
|
# Unknown or expired session ID - return 404 per MCP spec
|
|
# TODO(L62): Align error code once spec clarifies
|
|
# See: https://github.com/modelcontextprotocol/python-sdk/issues/1821
|
|
logger.info(f"Rejected request with unknown or expired session ID: {request_mcp_session_id[:64]}")
|
|
await _error_response("Session not found", 404)(scope, receive, send)
|
|
|
|
def _admit_session(self, requestor: AuthorizationContext | None) -> StreamableHTTPServerTransport | None:
|
|
"""Register a new session for `requestor` and return its transport, or None at the session limit."""
|
|
if self.max_sessions is not None and len(self._server_instances) >= self.max_sessions:
|
|
return None
|
|
http_transport = StreamableHTTPServerTransport(
|
|
mcp_session_id=uuid4().hex,
|
|
is_json_response_enabled=self.json_response,
|
|
event_store=self.event_store, # May be None (no resumability)
|
|
security_settings=self.security_settings,
|
|
retry_interval=self.retry_interval,
|
|
idle_timeout=self.session_idle_timeout,
|
|
)
|
|
session_id = http_transport.mcp_session_id
|
|
assert session_id is not None
|
|
if requestor is not None:
|
|
self._session_owners[session_id] = requestor
|
|
self._server_instances[session_id] = http_transport
|
|
logger.info(f"Created new transport with session ID: {session_id}")
|
|
return http_transport
|
|
|
|
async def _serve_opening_request(
|
|
self, http_transport: StreamableHTTPServerTransport, scope: Scope, receive: Receive, send: Send
|
|
) -> None:
|
|
"""Start the session's server task and let its transport answer the request that opens it.
|
|
|
|
Without a session ID only an initialize request can succeed, so if this
|
|
one is refused, fails or is cancelled (or the session's server task
|
|
cannot even be started) nothing was established: the session is
|
|
discarded again rather than kept (with its server task) around.
|
|
"""
|
|
session_id = http_transport.mcp_session_id
|
|
assert session_id is not None
|
|
|
|
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()
|
|
try:
|
|
# The transport cancels its idle scope once no request
|
|
# has been in flight for `session_idle_timeout`; that
|
|
# ends the loop and execution continues after the
|
|
# `with` block. Without a timeout there is nothing to fire.
|
|
idle_scope = http_transport.idle_scope
|
|
if idle_scope is None:
|
|
idle_scope = anyio.CancelScope()
|
|
with idle_scope:
|
|
# Drive via `serve_loop` (not `Server.run()`) so the
|
|
# manager's already-entered lifespan is reused
|
|
# rather than re-entered per session.
|
|
await serve_loop(
|
|
self.app,
|
|
read_stream,
|
|
write_stream,
|
|
lifespan_state=self._lifespan_state,
|
|
session_id=session_id,
|
|
)
|
|
|
|
if idle_scope.cancelled_caught:
|
|
logger.info(f"Session {session_id} idle timeout")
|
|
except Exception:
|
|
logger.exception(f"Session {session_id} crashed")
|
|
finally:
|
|
# However the session ended (client DELETE, idle
|
|
# timeout, crash), discard it.
|
|
await self._discard_session(session_id, http_transport)
|
|
|
|
established = False
|
|
try:
|
|
assert self._task_group is not None
|
|
await self._task_group.start(run_server)
|
|
status = await _send_and_report_status(http_transport.handle_request, scope, receive, send)
|
|
established = status is not None and status < 400
|
|
finally:
|
|
if not established: # pragma: no branch
|
|
await self._discard_session(session_id, http_transport)
|
|
|
|
async def _discard_session(self, session_id: str, transport: StreamableHTTPServerTransport) -> None:
|
|
"""Stop tracking the session and make sure its transport refuses anything that still reaches it.
|
|
|
|
The session is forgotten first, before any await, so its ID answers 404
|
|
from the moment this is called; terminating the transport is shielded so
|
|
it completes even while the caller is being cancelled.
|
|
"""
|
|
self._server_instances.pop(session_id, None)
|
|
self._session_owners.pop(session_id, None)
|
|
if not transport.is_terminated:
|
|
with anyio.CancelScope(shield=True):
|
|
await transport.terminate()
|
|
|
|
|
|
def _error_response(message: str, status_code: int, code: int = INVALID_REQUEST) -> Response:
|
|
"""A JSON-RPC error body (no request id) with the given HTTP status."""
|
|
body = JSONRPCError(jsonrpc="2.0", id=None, error=ErrorData(code=code, message=message))
|
|
return Response(
|
|
body.model_dump_json(by_alias=True, exclude_unset=True), status_code=status_code, media_type="application/json"
|
|
)
|
|
|
|
|
|
async def _send_and_report_status(app: ASGIApp, scope: Scope, receive: Receive, send: Send) -> int | None:
|
|
"""Run `app` for one request and return the HTTP status it answered with (None if it sent no response)."""
|
|
status: int | None = None
|
|
|
|
async def watch_status(message: Message) -> None:
|
|
nonlocal status
|
|
if message["type"] == "http.response.start":
|
|
status = message["status"]
|
|
await send(message)
|
|
|
|
await app(scope, receive, watch_status)
|
|
return status
|
|
|
|
|
|
class StreamableHTTPASGIApp:
|
|
"""ASGI application for Streamable HTTP server transport."""
|
|
|
|
def __init__(self, session_manager: StreamableHTTPSessionManager):
|
|
self.session_manager = session_manager
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
await self.session_manager.asgi_app(scope, receive, send)
|