Expire idle Streamable HTTP sessions by default and cap concurrent sessions (#3395)
This commit is contained in:
parent
6705402e24
commit
d2290ca343
4 changed files with 915 additions and 234 deletions
|
|
@ -65,7 +65,12 @@ from mcp.server.context import HandlerResult, ServerMiddleware, ServerRequestCon
|
|||
from mcp.server.models import InitializationOptions
|
||||
from mcp.server.runner import serve_dual_era_loop
|
||||
from mcp.server.streamable_http import EventStore
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
|
||||
from mcp.server.streamable_http_manager import (
|
||||
DEFAULT_MAX_SESSIONS,
|
||||
DEFAULT_SESSION_IDLE_TIMEOUT,
|
||||
StreamableHTTPASGIApp,
|
||||
StreamableHTTPSessionManager,
|
||||
)
|
||||
from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE, TransportSecuritySettings
|
||||
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
||||
from mcp.shared.exceptions import MCPDeprecationWarning
|
||||
|
|
@ -722,6 +727,8 @@ class Server(Generic[LifespanResultT]):
|
|||
event_store: EventStore | None = None,
|
||||
retry_interval: int | None = None,
|
||||
max_request_body_size: int = DEFAULT_MAX_REQUEST_BODY_SIZE,
|
||||
session_idle_timeout: float | None = DEFAULT_SESSION_IDLE_TIMEOUT,
|
||||
max_sessions: int | None = DEFAULT_MAX_SESSIONS,
|
||||
transport_security: TransportSecuritySettings | None = None,
|
||||
host: str = "127.0.0.1",
|
||||
auth: AuthSettings | None = None,
|
||||
|
|
@ -747,6 +754,8 @@ class Server(Generic[LifespanResultT]):
|
|||
stateless=stateless_http,
|
||||
security_settings=transport_security,
|
||||
max_request_body_size=max_request_body_size,
|
||||
session_idle_timeout=session_idle_timeout,
|
||||
max_sessions=max_sessions,
|
||||
)
|
||||
self._session_manager = session_manager
|
||||
|
||||
|
|
|
|||
|
|
@ -93,7 +93,11 @@ from mcp.server.request_state import RequestStateBoundary, RequestStateSecurity
|
|||
from mcp.server.sse import SseServerTransport
|
||||
from mcp.server.stdio import stdio_server
|
||||
from mcp.server.streamable_http import EventStore
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
||||
from mcp.server.streamable_http_manager import (
|
||||
DEFAULT_MAX_SESSIONS,
|
||||
DEFAULT_SESSION_IDLE_TIMEOUT,
|
||||
StreamableHTTPSessionManager,
|
||||
)
|
||||
from mcp.server.subscriptions import InMemorySubscriptionBus, ListenHandler, SubscriptionBus
|
||||
from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE, TransportSecuritySettings
|
||||
from mcp.shared.exceptions import MCPError
|
||||
|
|
@ -388,6 +392,8 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
event_store: EventStore | None = ...,
|
||||
retry_interval: int | None = ...,
|
||||
max_request_body_size: int = ...,
|
||||
session_idle_timeout: float | None = ...,
|
||||
max_sessions: int | None = ...,
|
||||
transport_security: TransportSecuritySettings | None = ...,
|
||||
) -> None: ...
|
||||
|
||||
|
|
@ -1106,6 +1112,8 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
event_store: EventStore | None = None,
|
||||
retry_interval: int | None = None,
|
||||
max_request_body_size: int = DEFAULT_MAX_REQUEST_BODY_SIZE,
|
||||
session_idle_timeout: float | None = DEFAULT_SESSION_IDLE_TIMEOUT,
|
||||
max_sessions: int | None = DEFAULT_MAX_SESSIONS,
|
||||
transport_security: TransportSecuritySettings | None = None,
|
||||
) -> None:
|
||||
"""Run the server using StreamableHTTP transport."""
|
||||
|
|
@ -1118,6 +1126,8 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
event_store=event_store,
|
||||
retry_interval=retry_interval,
|
||||
max_request_body_size=max_request_body_size,
|
||||
session_idle_timeout=session_idle_timeout,
|
||||
max_sessions=max_sessions,
|
||||
transport_security=transport_security,
|
||||
host=host,
|
||||
)
|
||||
|
|
@ -1270,6 +1280,8 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
event_store: EventStore | None = None,
|
||||
retry_interval: int | None = None,
|
||||
max_request_body_size: int = DEFAULT_MAX_REQUEST_BODY_SIZE,
|
||||
session_idle_timeout: float | None = DEFAULT_SESSION_IDLE_TIMEOUT,
|
||||
max_sessions: int | None = DEFAULT_MAX_SESSIONS,
|
||||
transport_security: TransportSecuritySettings | None = None,
|
||||
host: str = "127.0.0.1",
|
||||
) -> Starlette:
|
||||
|
|
@ -1281,6 +1293,8 @@ class MCPServer(Generic[LifespanResultT]):
|
|||
event_store=event_store,
|
||||
retry_interval=retry_interval,
|
||||
max_request_body_size=max_request_body_size,
|
||||
session_idle_timeout=session_idle_timeout,
|
||||
max_sessions=max_sessions,
|
||||
transport_security=transport_security,
|
||||
host=host,
|
||||
auth=self.settings.auth,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ responses, with streaming support for long-running operations.
|
|||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable
|
||||
|
|
@ -167,6 +168,7 @@ class StreamableHTTPServerTransport:
|
|||
event_store: EventStore | None = None,
|
||||
security_settings: TransportSecuritySettings | None = None,
|
||||
retry_interval: int | None = None,
|
||||
idle_timeout: float | None = None,
|
||||
) -> None:
|
||||
"""Initialize a new StreamableHTTP server transport.
|
||||
|
||||
|
|
@ -187,12 +189,22 @@ class StreamableHTTPServerTransport:
|
|||
retry field. When set, the server will send a retry field in
|
||||
SSE priming events to control client reconnection timing for
|
||||
polling behavior. Only used when event_store is provided.
|
||||
idle_timeout: Seconds the session may go without any request in flight before
|
||||
`idle_scope` is cancelled. A request being served or an open GET
|
||||
stream holds the session open; the countdown starts each time the
|
||||
last in-flight request completes. The host enters `idle_scope`
|
||||
(available once `connect()` has been entered) around the session's
|
||||
message loop to end the session when it fires. Default is None: no
|
||||
`idle_scope`, the session never expires.
|
||||
|
||||
Raises:
|
||||
ValueError: If the session ID contains invalid characters.
|
||||
ValueError: If the session ID contains invalid characters, or if `idle_timeout`
|
||||
is not a positive, finite number.
|
||||
"""
|
||||
if mcp_session_id is not None and not SESSION_ID_PATTERN.fullmatch(mcp_session_id):
|
||||
raise ValueError("Session ID must only contain visible ASCII characters (0x21-0x7E)")
|
||||
if idle_timeout is not None and not (math.isfinite(idle_timeout) and idle_timeout > 0):
|
||||
raise ValueError("idle_timeout must be a positive, finite number of seconds")
|
||||
|
||||
self.mcp_session_id = mcp_session_id
|
||||
self.is_json_response_enabled = is_json_response_enabled
|
||||
|
|
@ -208,8 +220,11 @@ class StreamableHTTPServerTransport:
|
|||
] = {}
|
||||
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[SSEEvent]] = {}
|
||||
self._terminated = False
|
||||
# Idle timeout cancel scope; managed by the session manager.
|
||||
self._idle_timeout = idle_timeout
|
||||
self._requests_in_flight = 0
|
||||
self.idle_scope: anyio.CancelScope | None = None
|
||||
"""Created when `connect()` is entered if `idle_timeout` is set; cancelled once no request has been in
|
||||
flight for `idle_timeout` seconds."""
|
||||
|
||||
@property
|
||||
def is_terminated(self) -> bool:
|
||||
|
|
@ -458,6 +473,32 @@ class StreamableHTTPServerTransport:
|
|||
|
||||
async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
"""Application entry point that handles all HTTP requests."""
|
||||
if self.idle_scope is None or self._idle_timeout is None:
|
||||
await self._handle_request(scope, receive, send)
|
||||
return
|
||||
|
||||
if self.idle_scope.cancel_called:
|
||||
# The idle period already ran out and the host is ending this
|
||||
# session: answer as terminated rather than dispatch into a
|
||||
# message loop that is going away.
|
||||
if not self._terminated:
|
||||
await self.terminate()
|
||||
await self._handle_request(scope, receive, send)
|
||||
return
|
||||
|
||||
# A request in flight (an open GET stream included) holds the session:
|
||||
# the idle countdown is suspended while any is being served and
|
||||
# restarts when the last one completes.
|
||||
self._requests_in_flight += 1
|
||||
self.idle_scope.deadline = math.inf
|
||||
try:
|
||||
await self._handle_request(scope, receive, send)
|
||||
finally:
|
||||
self._requests_in_flight -= 1
|
||||
if not self._requests_in_flight:
|
||||
self.idle_scope.deadline = anyio.current_time() + self._idle_timeout
|
||||
|
||||
async def _handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
request = Request(scope, receive)
|
||||
|
||||
# Validate request headers for DNS rebinding protection
|
||||
|
|
@ -793,7 +834,7 @@ class StreamableHTTPServerTransport:
|
|||
await response(request.scope, request.receive, send)
|
||||
return
|
||||
|
||||
if not await self._validate_request_headers(request, send): # pragma: no cover
|
||||
if not await self._validate_request_headers(request, send):
|
||||
return
|
||||
|
||||
await self.terminate()
|
||||
|
|
@ -995,6 +1036,8 @@ class StreamableHTTPServerTransport:
|
|||
Yields:
|
||||
Tuple of (read_stream, write_stream) for bidirectional communication
|
||||
"""
|
||||
if self._idle_timeout is not None:
|
||||
self.idle_scope = anyio.CancelScope()
|
||||
|
||||
# Create the memory streams for this connection
|
||||
|
||||
|
|
|
|||
|
|
@ -4,17 +4,18 @@ from __future__ import annotations
|
|||
|
||||
import contextlib
|
||||
import logging
|
||||
import math
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import TYPE_CHECKING, Any
|
||||
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, INVALID_REQUEST, ErrorData, JSONRPCError
|
||||
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 Receive, Scope, Send
|
||||
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
|
||||
|
|
@ -34,6 +35,12 @@ if TYPE_CHECKING:
|
|||
|
||||
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.
|
||||
|
|
@ -45,7 +52,7 @@ class StreamableHTTPSessionManager:
|
|||
2. Resumability via an optional event store
|
||||
3. Connection management and lifecycle
|
||||
4. Request handling and transport setup
|
||||
5. Idle session cleanup via optional timeout
|
||||
5. Idle session cleanup
|
||||
|
||||
Important: Only one StreamableHTTPSessionManager instance should be created
|
||||
per application. The instance cannot be reused after its run() context has
|
||||
|
|
@ -62,13 +69,18 @@ class StreamableHTTPSessionManager:
|
|||
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: Optional idle timeout in seconds for stateful sessions. If set, sessions that
|
||||
receive no HTTP requests for this duration will be automatically terminated and removed. When
|
||||
retry_interval is also configured, ensure the idle timeout comfortably exceeds the retry interval to
|
||||
avoid reaping sessions during normal SSE polling gaps. Default is None (no timeout). A value of 1800
|
||||
(30 minutes) is recommended for most deployments.
|
||||
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__(
|
||||
|
|
@ -79,15 +91,16 @@ class StreamableHTTPSessionManager:
|
|||
stateless: bool = False,
|
||||
security_settings: TransportSecuritySettings | None = None,
|
||||
retry_interval: int | None = None,
|
||||
session_idle_timeout: float | 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 session_idle_timeout <= 0:
|
||||
raise ValueError("session_idle_timeout must be a positive number of seconds")
|
||||
if stateless and session_idle_timeout is not None:
|
||||
raise RuntimeError("session_idle_timeout is not supported in stateless mode")
|
||||
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
|
||||
|
|
@ -97,6 +110,7 @@ class StreamableHTTPSessionManager:
|
|||
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)
|
||||
|
|
@ -234,16 +248,15 @@ class StreamableHTTPSessionManager:
|
|||
except Exception: # pragma: lax no cover
|
||||
logger.exception("Stateless session crashed")
|
||||
|
||||
# Assert task group is not None for type checking
|
||||
# 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
|
||||
# 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)
|
||||
|
||||
# Terminate the transport after the request is handled
|
||||
await http_transport.terminate()
|
||||
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."""
|
||||
|
|
@ -263,109 +276,145 @@ class StreamableHTTPSessionManager:
|
|||
"Rejecting request for session %s: credential does not match the one that created the session",
|
||||
request_mcp_session_id[:64],
|
||||
)
|
||||
body = JSONRPCError(
|
||||
jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
|
||||
)
|
||||
response = Response(
|
||||
body.model_dump_json(by_alias=True, exclude_unset=True),
|
||||
status_code=404,
|
||||
media_type="application/json",
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
await _error_response("Session not found", 404)(scope, receive, send)
|
||||
return
|
||||
logger.debug("Session already exists, handling request directly")
|
||||
# Push back idle deadline on activity
|
||||
if transport.idle_scope is not None and self.session_idle_timeout is not None:
|
||||
transport.idle_scope.deadline = anyio.current_time() + self.session_idle_timeout # pragma: no cover
|
||||
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
|
||||
logger.debug("Creating new transport")
|
||||
# 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:
|
||||
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)
|
||||
security_settings=self.security_settings,
|
||||
retry_interval=self.retry_interval,
|
||||
)
|
||||
|
||||
assert http_transport.mcp_session_id is not None
|
||||
if requestor is not None:
|
||||
self._session_owners[http_transport.mcp_session_id] = requestor
|
||||
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()
|
||||
try:
|
||||
# Use a cancel scope for idle timeout — when the
|
||||
# deadline passes the scope cancels the loop and
|
||||
# execution continues after the ``with`` block.
|
||||
# Incoming requests push the deadline forward.
|
||||
idle_scope = anyio.CancelScope()
|
||||
if self.session_idle_timeout is not None:
|
||||
idle_scope.deadline = anyio.current_time() + self.session_idle_timeout
|
||||
http_transport.idle_scope = idle_scope
|
||||
|
||||
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=http_transport.mcp_session_id,
|
||||
)
|
||||
|
||||
if idle_scope.cancelled_caught:
|
||||
assert http_transport.mcp_session_id is not None
|
||||
logger.info(f"Session {http_transport.mcp_session_id} idle timeout")
|
||||
self._server_instances.pop(http_transport.mcp_session_id, None)
|
||||
self._session_owners.pop(http_transport.mcp_session_id, None)
|
||||
await http_transport.terminate()
|
||||
except Exception:
|
||||
logger.exception(f"Session {http_transport.mcp_session_id} crashed")
|
||||
finally:
|
||||
if ( # pragma: no branch
|
||||
http_transport.mcp_session_id
|
||||
and http_transport.mcp_session_id in self._server_instances
|
||||
and not http_transport.is_terminated
|
||||
):
|
||||
logger.info(
|
||||
"Cleaning up crashed session "
|
||||
f"{http_transport.mcp_session_id} from active instances."
|
||||
)
|
||||
del self._server_instances[http_transport.mcp_session_id]
|
||||
self._session_owners.pop(http_transport.mcp_session_id, None)
|
||||
|
||||
# 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)
|
||||
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]}")
|
||||
body = JSONRPCError(
|
||||
jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
|
||||
)
|
||||
response = Response(
|
||||
body.model_dump_json(by_alias=True, exclude_unset=True), status_code=404, media_type="application/json"
|
||||
)
|
||||
await response(scope, receive, send)
|
||||
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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue