1093 lines
48 KiB
Python
1093 lines
48 KiB
Python
"""StreamableHTTP Server Transport Module
|
|
|
|
This module implements an HTTP transport layer with Streamable HTTP.
|
|
|
|
The transport handles bidirectional communication using HTTP requests and
|
|
responses, with streaming support for long-running operations.
|
|
"""
|
|
|
|
import logging
|
|
import re
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import AsyncGenerator, Awaitable, Callable
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
from http import HTTPStatus
|
|
from typing import Any, Final
|
|
|
|
import anyio
|
|
import pydantic_core
|
|
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
|
from mcp_types import (
|
|
DEFAULT_NEGOTIATED_VERSION,
|
|
INTERNAL_ERROR,
|
|
INVALID_PARAMS,
|
|
INVALID_REQUEST,
|
|
PARSE_ERROR,
|
|
ErrorData,
|
|
JSONRPCError,
|
|
JSONRPCMessage,
|
|
JSONRPCRequest,
|
|
JSONRPCResponse,
|
|
RequestId,
|
|
jsonrpc_message_adapter,
|
|
)
|
|
from mcp_types.version import is_version_at_least
|
|
from pydantic import ValidationError
|
|
from sse_starlette import EventSourceResponse
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
from starlette.types import Receive, Scope, Send
|
|
|
|
from mcp.server.transport_security import TransportSecurityMiddleware, TransportSecuritySettings
|
|
from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams
|
|
from mcp.shared._stream_protocols import ReadStream, WriteStream
|
|
from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER
|
|
from mcp.shared.message import CloseSSEStreamCallback, ServerMessageMetadata, SessionMessage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Header names
|
|
MCP_SESSION_ID_HEADER = "mcp-session-id"
|
|
LAST_EVENT_ID_HEADER = "last-event-id"
|
|
|
|
# Content types
|
|
CONTENT_TYPE_JSON = "application/json"
|
|
CONTENT_TYPE_SSE = "text/event-stream"
|
|
|
|
# Special key for the standalone GET stream
|
|
GET_STREAM_KEY = "_GET_stream"
|
|
|
|
# Buffer for the per-request `_request_streams` so the serial `message_router`
|
|
# can deposit a response and move on instead of head-of-line blocking the
|
|
# whole session on a lazily-started `sse_writer`. See #1764.
|
|
REQUEST_STREAM_BUFFER_SIZE: Final = 16
|
|
|
|
# Error code answering a request that settled without a response (e.g. it was
|
|
# cancelled) on this 2025-era wire, which ends a request's stream only with a
|
|
# response. Mirrors LSP's RequestCancelled; not sent by the 2026 transports, where
|
|
# the spec forbids answering a cancelled request. See
|
|
# `StreamableHTTPServerTransport._terminate_unanswered_request`.
|
|
REQUEST_CANCELLED: Final = -32800
|
|
|
|
# Session ID validation pattern (visible ASCII characters ranging from 0x21 to 0x7E)
|
|
# Pattern ensures entire string contains only valid characters by using ^ and $ anchors
|
|
SESSION_ID_PATTERN = re.compile(r"^[\x21-\x7E]+$")
|
|
|
|
# Type aliases
|
|
StreamId = str
|
|
EventId = str
|
|
# An SSE event-dict as accepted by sse-starlette (`event`, `data`, `id`, `retry`).
|
|
SSEEvent = dict[str, Any]
|
|
|
|
|
|
def check_accept_headers(request: Request) -> tuple[bool, bool]:
|
|
"""Return (has_json, has_sse) for the request's Accept header, with RFC 7231 wildcard handling.
|
|
|
|
Supports wildcard media types per RFC 7231, section 5.3.2:
|
|
- */* matches any media type
|
|
- application/* matches any application/ subtype
|
|
- text/* matches any text/ subtype
|
|
"""
|
|
accept_header = request.headers.get("accept", "")
|
|
accept_types = [media_type.strip().split(";")[0].strip().lower() for media_type in accept_header.split(",")]
|
|
|
|
has_wildcard = "*/*" in accept_types
|
|
has_json = has_wildcard or any(t in (CONTENT_TYPE_JSON, "application/*") for t in accept_types)
|
|
has_sse = has_wildcard or any(t in (CONTENT_TYPE_SSE, "text/*") for t in accept_types)
|
|
|
|
return has_json, has_sse
|
|
|
|
|
|
@dataclass
|
|
class EventMessage:
|
|
"""A JSONRPCMessage with an optional event ID for stream resumability."""
|
|
|
|
message: JSONRPCMessage
|
|
event_id: str | None = None
|
|
|
|
|
|
EventCallback = Callable[[EventMessage], Awaitable[None]]
|
|
|
|
|
|
class EventStore(ABC):
|
|
"""Interface for resumability support via event storage."""
|
|
|
|
@abstractmethod
|
|
async def store_event(self, stream_id: StreamId, message: JSONRPCMessage | None) -> EventId:
|
|
"""Stores an event for later retrieval.
|
|
|
|
Args:
|
|
stream_id: ID of the stream the event belongs to
|
|
message: The JSON-RPC message to store, or None for priming events
|
|
|
|
Returns:
|
|
The generated event ID for the stored event.
|
|
"""
|
|
pass # pragma: no cover
|
|
|
|
@abstractmethod
|
|
async def replay_events_after(
|
|
self,
|
|
last_event_id: EventId,
|
|
send_callback: EventCallback,
|
|
) -> StreamId | None:
|
|
"""Replays events that occurred after the specified event ID.
|
|
|
|
Args:
|
|
last_event_id: The ID of the last event the client received
|
|
send_callback: A callback function to send events to the client
|
|
|
|
Returns:
|
|
The stream ID of the replayed events, or None if no events were found.
|
|
"""
|
|
pass # pragma: no cover
|
|
|
|
|
|
class StreamableHTTPServerTransport:
|
|
"""HTTP server transport with event streaming support for MCP.
|
|
|
|
Handles JSON-RPC messages in HTTP POST requests with SSE streaming.
|
|
Supports optional JSON responses and session management.
|
|
"""
|
|
|
|
# Server notification streams for POST requests as well as standalone SSE stream
|
|
_read_stream_writer: ContextSendStream[SessionMessage | Exception] | None = None
|
|
_read_stream: ContextReceiveStream[SessionMessage | Exception] | None = None
|
|
_write_stream: ContextSendStream[SessionMessage] | None = None
|
|
_write_stream_reader: ContextReceiveStream[SessionMessage] | None = None
|
|
_security: TransportSecurityMiddleware
|
|
|
|
def __init__(
|
|
self,
|
|
mcp_session_id: str | None,
|
|
is_json_response_enabled: bool = False,
|
|
event_store: EventStore | None = None,
|
|
security_settings: TransportSecuritySettings | None = None,
|
|
retry_interval: int | None = None,
|
|
) -> None:
|
|
"""Initialize a new StreamableHTTP server transport.
|
|
|
|
Args:
|
|
mcp_session_id: Optional session identifier for this connection.
|
|
Must contain only visible ASCII characters (0x21-0x7E).
|
|
is_json_response_enabled: If True, answer each request POST with a single
|
|
JSON body instead of an SSE stream, which removes
|
|
the request-scoped back-channel: a server-initiated
|
|
request tied to the call raises `NoBackChannelError`
|
|
and its notifications are dropped (see
|
|
`TransportContext.can_send_request`). Default is False.
|
|
event_store: Event store for resumability support. If provided,
|
|
resumability will be enabled, allowing clients to
|
|
reconnect and resume messages.
|
|
security_settings: Optional security settings for DNS rebinding protection.
|
|
retry_interval: Retry interval in milliseconds to suggest to clients in SSE
|
|
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.
|
|
|
|
Raises:
|
|
ValueError: If the session ID contains invalid characters.
|
|
"""
|
|
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)")
|
|
|
|
self.mcp_session_id = mcp_session_id
|
|
self.is_json_response_enabled = is_json_response_enabled
|
|
self._event_store = event_store
|
|
self._security = TransportSecurityMiddleware(security_settings)
|
|
self._retry_interval = retry_interval
|
|
self._request_streams: dict[
|
|
RequestId,
|
|
tuple[
|
|
MemoryObjectSendStream[EventMessage],
|
|
MemoryObjectReceiveStream[EventMessage],
|
|
],
|
|
] = {}
|
|
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[SSEEvent]] = {}
|
|
self._terminated = False
|
|
# Idle timeout cancel scope; managed by the session manager.
|
|
self.idle_scope: anyio.CancelScope | None = None
|
|
|
|
@property
|
|
def is_terminated(self) -> bool:
|
|
"""Check if this transport has been explicitly terminated."""
|
|
return self._terminated
|
|
|
|
def _message_metadata(
|
|
self,
|
|
request: Request,
|
|
*,
|
|
close_sse_stream: CloseSSEStreamCallback | None = None,
|
|
close_standalone_sse_stream: CloseSSEStreamCallback | None = None,
|
|
on_request_unanswered: Callable[[], Awaitable[None]] | None = None,
|
|
) -> ServerMessageMetadata:
|
|
"""The metadata this transport frames every inbound message with.
|
|
|
|
The one place `can_send_request` is stamped, so no construction site can
|
|
forget it: a JSON body carries only the response, so in JSON-response mode
|
|
the request-scoped channel cannot carry a server-initiated request (see
|
|
`TransportContext.can_send_request`).
|
|
"""
|
|
return ServerMessageMetadata(
|
|
request_context=request,
|
|
close_sse_stream=close_sse_stream,
|
|
close_standalone_sse_stream=close_standalone_sse_stream,
|
|
on_request_unanswered=on_request_unanswered,
|
|
can_send_request=not self.is_json_response_enabled,
|
|
)
|
|
|
|
def close_sse_stream(self, request_id: RequestId) -> None:
|
|
"""Close SSE connection for a specific request without terminating the stream.
|
|
|
|
This method closes the HTTP connection for the specified request, triggering
|
|
client reconnection. Events continue to be stored in the event store and will
|
|
be replayed when the client reconnects with Last-Event-ID.
|
|
|
|
Use this to implement polling behavior during long-running operations -
|
|
the client will reconnect after the retry interval specified in the priming event.
|
|
|
|
Args:
|
|
request_id: The request ID whose SSE stream should be closed.
|
|
|
|
Note:
|
|
This is a no-op if there is no active stream for the request ID.
|
|
Requires event_store to be configured for events to be stored during
|
|
the disconnect.
|
|
"""
|
|
writer = self._sse_stream_writers.pop(request_id, None)
|
|
if writer: # pragma: no branch
|
|
writer.close()
|
|
|
|
# Also close and remove request streams
|
|
if request_id in self._request_streams: # pragma: no branch
|
|
send_stream, receive_stream = self._request_streams.pop(request_id)
|
|
send_stream.close()
|
|
receive_stream.close()
|
|
|
|
def close_standalone_sse_stream(self) -> None:
|
|
"""Close the standalone GET SSE stream, triggering client reconnection.
|
|
|
|
This method closes the HTTP connection for the standalone GET stream used
|
|
for unsolicited server-to-client notifications. The client SHOULD reconnect
|
|
with Last-Event-ID to resume receiving notifications.
|
|
|
|
Use this to implement polling behavior for the notification stream -
|
|
the client will reconnect after the retry interval specified in the priming event.
|
|
|
|
Note:
|
|
This is a no-op if there is no active standalone SSE stream.
|
|
Requires event_store to be configured for events to be stored during
|
|
the disconnect.
|
|
"""
|
|
self.close_sse_stream(GET_STREAM_KEY)
|
|
|
|
def _create_session_message(
|
|
self,
|
|
message: JSONRPCRequest,
|
|
request: Request,
|
|
request_id: RequestId,
|
|
protocol_version: str,
|
|
) -> SessionMessage:
|
|
"""Create a session message with metadata including close_sse_stream callback.
|
|
|
|
The close_sse_stream callbacks are only provided when the client supports
|
|
resumability (protocol version >= 2025-11-25). Old clients can't resume if
|
|
the stream is closed early because they didn't receive a priming event.
|
|
Every request carries `on_request_unanswered`, so a request that settles
|
|
without a response is still terminated on this era's wire.
|
|
"""
|
|
end_stream = partial(self._terminate_unanswered_request, message.id)
|
|
# Only provide close callbacks when client supports resumability
|
|
if self._event_store and is_version_at_least(protocol_version, "2025-11-25"):
|
|
|
|
async def close_stream_callback() -> None:
|
|
self.close_sse_stream(request_id)
|
|
|
|
async def close_standalone_stream_callback() -> None:
|
|
self.close_standalone_sse_stream()
|
|
|
|
metadata = self._message_metadata(
|
|
request,
|
|
close_sse_stream=close_stream_callback,
|
|
close_standalone_sse_stream=close_standalone_stream_callback,
|
|
on_request_unanswered=end_stream,
|
|
)
|
|
else:
|
|
metadata = self._message_metadata(request, on_request_unanswered=end_stream)
|
|
|
|
return SessionMessage(message, metadata=metadata)
|
|
|
|
async def _mint_priming_event(self, stream_id: StreamId, protocol_version: str) -> SSEEvent | None:
|
|
"""Store the priming cursor for `stream_id` and return its SSE wire form.
|
|
|
|
Called before the request is dispatched so the priming row precedes
|
|
anything `message_router` can store for this stream. Returns `None`
|
|
when no event store is configured or the client predates 2025-11-25
|
|
(older clients cannot parse the empty-data event).
|
|
"""
|
|
if not self._event_store:
|
|
return None
|
|
if not is_version_at_least(protocol_version, "2025-11-25"):
|
|
return None
|
|
priming_event_id = await self._event_store.store_event(stream_id, None)
|
|
priming_event: SSEEvent = {"id": priming_event_id, "data": ""}
|
|
if self._retry_interval is not None:
|
|
priming_event["retry"] = self._retry_interval
|
|
return priming_event
|
|
|
|
async def _run_sse_writer(
|
|
self,
|
|
request_id: RequestId,
|
|
sse_stream_writer: MemoryObjectSendStream[SSEEvent],
|
|
request_stream_reader: MemoryObjectReceiveStream[EventMessage],
|
|
priming_event: SSEEvent | None,
|
|
) -> None:
|
|
"""Forward `_request_streams[request_id]` onto the SSE wire for one POST."""
|
|
try:
|
|
async with sse_stream_writer, request_stream_reader:
|
|
if priming_event is not None:
|
|
await sse_stream_writer.send(priming_event)
|
|
async for event_message in request_stream_reader:
|
|
await sse_stream_writer.send(self._create_event_data(event_message))
|
|
if isinstance(event_message.message, JSONRPCResponse | JSONRPCError):
|
|
break
|
|
except anyio.ClosedResourceError: # pragma: lax no cover
|
|
logger.debug("SSE stream closed by close_sse_stream()")
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Error in SSE writer")
|
|
finally:
|
|
logger.debug("Closing SSE writer")
|
|
self._sse_stream_writers.pop(request_id, None)
|
|
await self._clean_up_memory_streams(request_id)
|
|
|
|
def _create_error_response(
|
|
self,
|
|
error_message: str,
|
|
status_code: HTTPStatus,
|
|
error_code: int = INVALID_REQUEST,
|
|
headers: dict[str, str] | None = None,
|
|
) -> Response:
|
|
"""Create an error response with a simple string message."""
|
|
response_headers = {"Content-Type": CONTENT_TYPE_JSON}
|
|
if headers:
|
|
response_headers.update(headers)
|
|
|
|
if self.mcp_session_id:
|
|
response_headers[MCP_SESSION_ID_HEADER] = self.mcp_session_id
|
|
|
|
# Return a properly formatted JSON error response
|
|
error_response = JSONRPCError(
|
|
jsonrpc="2.0",
|
|
id=None,
|
|
error=ErrorData(code=error_code, message=error_message),
|
|
)
|
|
|
|
return Response(
|
|
error_response.model_dump_json(by_alias=True, exclude_unset=True),
|
|
status_code=status_code,
|
|
headers=response_headers,
|
|
)
|
|
|
|
def _create_json_response(
|
|
self,
|
|
response_message: JSONRPCMessage | None,
|
|
status_code: HTTPStatus = HTTPStatus.OK,
|
|
headers: dict[str, str] | None = None,
|
|
) -> Response:
|
|
"""Create a JSON response from a JSONRPCMessage."""
|
|
response_headers = {"Content-Type": CONTENT_TYPE_JSON}
|
|
if headers:
|
|
response_headers.update(headers) # pragma: no cover
|
|
|
|
if self.mcp_session_id:
|
|
response_headers[MCP_SESSION_ID_HEADER] = self.mcp_session_id
|
|
|
|
return Response(
|
|
response_message.model_dump_json(by_alias=True, exclude_unset=True) if response_message else None,
|
|
status_code=status_code,
|
|
headers=response_headers,
|
|
)
|
|
|
|
def _get_session_id(self, request: Request) -> str | None:
|
|
"""Extract the session ID from request headers."""
|
|
return request.headers.get(MCP_SESSION_ID_HEADER)
|
|
|
|
def _create_event_data(self, event_message: EventMessage) -> SSEEvent:
|
|
"""Create event data dictionary from an EventMessage."""
|
|
event_data = {
|
|
"event": "message",
|
|
"data": event_message.message.model_dump_json(by_alias=True, exclude_unset=True),
|
|
}
|
|
|
|
# If an event ID was provided, include it
|
|
if event_message.event_id:
|
|
event_data["id"] = event_message.event_id
|
|
|
|
return event_data
|
|
|
|
async def _terminate_unanswered_request(self, request_id: RequestId) -> None:
|
|
"""Terminate a request that settled without a response (e.g. cancelled).
|
|
|
|
The 2025-era wire ends a request's stream only with a response for its
|
|
id - and stores that response so a resuming client's replay terminates
|
|
too - so this era answers a cancelled request with `REQUEST_CANCELLED`
|
|
where the dispatcher itself stays silent (the 2026 transports MUST NOT
|
|
answer). It is written through the same ordered channel as the request's
|
|
other messages, so it cannot overtake anything already queued for it.
|
|
"""
|
|
assert self._write_stream is not None # a dispatched request implies connect() ran
|
|
error = ErrorData(code=REQUEST_CANCELLED, message="Request cancelled")
|
|
await self._write_stream.send(SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error)))
|
|
|
|
async def _clean_up_memory_streams(self, request_id: RequestId) -> None:
|
|
"""Clean up memory streams for a given request ID."""
|
|
if request_id in self._request_streams: # pragma: no branch
|
|
try:
|
|
# Close the request stream
|
|
await self._request_streams[request_id][0].aclose()
|
|
await self._request_streams[request_id][1].aclose()
|
|
except Exception: # pragma: no cover
|
|
# During cleanup, we catch all exceptions since streams might be in various states
|
|
logger.debug("Error closing memory streams - may already be closed")
|
|
finally:
|
|
# Remove the request stream from the mapping
|
|
self._request_streams.pop(request_id, None)
|
|
|
|
async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
"""Application entry point that handles all HTTP requests."""
|
|
request = Request(scope, receive)
|
|
|
|
# Validate request headers for DNS rebinding protection
|
|
is_post = request.method == "POST"
|
|
error_response = await self._security.validate_request(request, is_post=is_post)
|
|
if error_response:
|
|
await error_response(scope, receive, send)
|
|
return
|
|
|
|
if self._terminated:
|
|
# If the session has been terminated, return 404 Not Found
|
|
response = self._create_error_response(
|
|
"Not Found: Session has been terminated",
|
|
HTTPStatus.NOT_FOUND,
|
|
)
|
|
await response(scope, receive, send)
|
|
return
|
|
|
|
if request.method == "POST":
|
|
await self._handle_post_request(scope, request, receive, send)
|
|
elif request.method == "GET":
|
|
await self._handle_get_request(request, send)
|
|
elif request.method == "DELETE":
|
|
await self._handle_delete_request(request, send)
|
|
else:
|
|
await self._handle_unsupported_request(request, send)
|
|
|
|
def _check_content_type(self, request: Request) -> bool:
|
|
"""Check if the request has the correct Content-Type."""
|
|
content_type = request.headers.get("content-type", "")
|
|
content_type_parts = [part.strip() for part in content_type.split(";")[0].split(",")]
|
|
|
|
return any(part == CONTENT_TYPE_JSON for part in content_type_parts)
|
|
|
|
async def _validate_accept_header(self, request: Request, scope: Scope, send: Send) -> bool:
|
|
"""Validate Accept header based on response mode. Returns True if valid."""
|
|
has_json, has_sse = check_accept_headers(request)
|
|
if self.is_json_response_enabled:
|
|
# For JSON-only responses, only require application/json
|
|
if not has_json:
|
|
response = self._create_error_response(
|
|
"Not Acceptable: Client must accept application/json",
|
|
HTTPStatus.NOT_ACCEPTABLE,
|
|
)
|
|
await response(scope, request.receive, send)
|
|
return False
|
|
# For SSE responses, require both content types
|
|
elif not (has_json and has_sse):
|
|
response = self._create_error_response(
|
|
"Not Acceptable: Client must accept both application/json and text/event-stream",
|
|
HTTPStatus.NOT_ACCEPTABLE,
|
|
)
|
|
await response(scope, request.receive, send)
|
|
return False
|
|
return True
|
|
|
|
async def _handle_post_request(self, scope: Scope, request: Request, receive: Receive, send: Send) -> None:
|
|
"""Handle POST requests containing JSON-RPC messages."""
|
|
writer = self._read_stream_writer
|
|
if writer is None: # pragma: no cover
|
|
raise ValueError("No read stream writer available. Ensure connect() is called first.")
|
|
try:
|
|
# Validate Accept header
|
|
if not await self._validate_accept_header(request, scope, send):
|
|
return
|
|
|
|
# Validate Content-Type
|
|
if not self._check_content_type(request): # pragma: no cover
|
|
response = self._create_error_response(
|
|
"Unsupported Media Type: Content-Type must be application/json",
|
|
HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
|
|
)
|
|
await response(scope, receive, send)
|
|
return
|
|
|
|
# Parse the body - only read it once
|
|
body = await request.body()
|
|
|
|
try:
|
|
raw_message = pydantic_core.from_json(body)
|
|
except ValueError as e:
|
|
response = self._create_error_response(f"Parse error: {str(e)}", HTTPStatus.BAD_REQUEST, PARSE_ERROR)
|
|
await response(scope, receive, send)
|
|
return
|
|
|
|
try:
|
|
message = jsonrpc_message_adapter.validate_python(raw_message, by_name=False)
|
|
except ValidationError as e:
|
|
response = self._create_error_response(
|
|
f"Validation error: {str(e)}",
|
|
HTTPStatus.BAD_REQUEST,
|
|
INVALID_PARAMS,
|
|
)
|
|
await response(scope, receive, send)
|
|
return
|
|
|
|
# Check if this is an initialization request
|
|
is_initialization_request = isinstance(message, JSONRPCRequest) and message.method == "initialize"
|
|
|
|
if is_initialization_request:
|
|
# Check if the server already has an established session
|
|
if self.mcp_session_id:
|
|
# Check if request has a session ID
|
|
request_session_id = self._get_session_id(request)
|
|
|
|
# If request has a session ID but doesn't match, return 404
|
|
if request_session_id and request_session_id != self.mcp_session_id: # pragma: no cover
|
|
response = self._create_error_response(
|
|
"Not Found: Invalid or expired session ID",
|
|
HTTPStatus.NOT_FOUND,
|
|
)
|
|
await response(scope, receive, send)
|
|
return
|
|
elif not await self._validate_request_headers(request, send):
|
|
return
|
|
|
|
# For notifications and responses only, return 202 Accepted
|
|
if not isinstance(message, JSONRPCRequest):
|
|
# Create response object and send it
|
|
response = self._create_json_response(
|
|
None,
|
|
HTTPStatus.ACCEPTED,
|
|
)
|
|
await response(scope, receive, send)
|
|
|
|
# Process the message after sending the response
|
|
session_message = SessionMessage(message, metadata=self._message_metadata(request))
|
|
await writer.send(session_message)
|
|
|
|
return
|
|
|
|
# Extract protocol version for priming event decision.
|
|
# For initialize requests, get from request params.
|
|
# For other requests, get from header (already validated).
|
|
protocol_version = (
|
|
str(message.params.get("protocolVersion", DEFAULT_NEGOTIATED_VERSION))
|
|
if is_initialization_request and message.params
|
|
else request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
|
|
)
|
|
|
|
request_id = str(message.id)
|
|
|
|
if self.is_json_response_enabled:
|
|
self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](
|
|
REQUEST_STREAM_BUFFER_SIZE
|
|
)
|
|
request_stream_reader = self._request_streams[request_id][1]
|
|
# Process the message
|
|
metadata = self._message_metadata(
|
|
request, on_request_unanswered=partial(self._terminate_unanswered_request, message.id)
|
|
)
|
|
session_message = SessionMessage(message, metadata=metadata)
|
|
await writer.send(session_message)
|
|
try:
|
|
# `message_router` deposits only this request's own response
|
|
# here: anything else scoped to the request has no wire in
|
|
# JSON-response mode.
|
|
event_message = await request_stream_reader.receive()
|
|
except (anyio.EndOfStream, anyio.ClosedResourceError):
|
|
# The stream closed with no response: the session was
|
|
# terminated while this request was in flight.
|
|
logger.debug(f"Session terminated with request {request_id} in flight; no response to send")
|
|
response = self._create_error_response(
|
|
"Session terminated before the request completed",
|
|
HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
INTERNAL_ERROR,
|
|
)
|
|
else:
|
|
response = self._create_json_response(event_message.message)
|
|
finally:
|
|
await self._clean_up_memory_streams(request_id)
|
|
await response(scope, receive, send)
|
|
else:
|
|
# Mint the priming event before any per-request state exists:
|
|
# `EventStore.store_event` is user code and may raise, in which
|
|
# case the outer handler returns a 500 with nothing to clean up.
|
|
# Still strictly precedes dispatch, so storage order == wire order.
|
|
priming_event = await self._mint_priming_event(request_id, protocol_version)
|
|
|
|
sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
|
|
self._sse_stream_writers[request_id] = sse_stream_writer
|
|
self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](
|
|
REQUEST_STREAM_BUFFER_SIZE
|
|
)
|
|
request_stream_reader = self._request_streams[request_id][1]
|
|
|
|
headers = {
|
|
"Cache-Control": "no-cache, no-transform",
|
|
"Connection": "keep-alive",
|
|
"Content-Type": CONTENT_TYPE_SSE,
|
|
**({MCP_SESSION_ID_HEADER: self.mcp_session_id} if self.mcp_session_id else {}),
|
|
}
|
|
response = EventSourceResponse(
|
|
content=sse_stream_reader,
|
|
data_sender_callable=partial(
|
|
self._run_sse_writer, request_id, sse_stream_writer, request_stream_reader, priming_event
|
|
),
|
|
headers=headers,
|
|
)
|
|
|
|
# Start the SSE response (this will send headers immediately)
|
|
try:
|
|
# First send the response to establish the SSE connection
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(response, scope, receive, send)
|
|
# Then send the message to be processed by the server
|
|
session_message = self._create_session_message(message, request, request_id, protocol_version)
|
|
await writer.send(session_message)
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("SSE response error")
|
|
await sse_stream_writer.aclose()
|
|
await self._clean_up_memory_streams(request_id)
|
|
finally:
|
|
await sse_stream_reader.aclose()
|
|
|
|
except Exception as err:
|
|
logger.exception("Error handling POST request")
|
|
response = self._create_error_response(
|
|
"Error handling POST request",
|
|
HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
INTERNAL_ERROR,
|
|
)
|
|
await response(scope, receive, send)
|
|
await writer.send(Exception(err))
|
|
return
|
|
|
|
async def _handle_get_request(self, request: Request, send: Send) -> None:
|
|
"""Handle GET request to establish SSE.
|
|
|
|
This allows the server to communicate to the client without the client
|
|
first sending data via HTTP POST. The server can send JSON-RPC requests
|
|
and notifications on this stream.
|
|
"""
|
|
writer = self._read_stream_writer
|
|
if writer is None: # pragma: no cover
|
|
raise ValueError("No read stream writer available. Ensure connect() is called first.")
|
|
|
|
# Validate Accept header - must include text/event-stream
|
|
_, has_sse = check_accept_headers(request)
|
|
|
|
if not has_sse:
|
|
response = self._create_error_response(
|
|
"Not Acceptable: Client must accept text/event-stream",
|
|
HTTPStatus.NOT_ACCEPTABLE,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
return
|
|
|
|
if not await self._validate_request_headers(request, send):
|
|
return
|
|
|
|
# Handle resumability: check for Last-Event-ID header
|
|
if last_event_id := request.headers.get(LAST_EVENT_ID_HEADER):
|
|
await self._replay_events(last_event_id, request, send)
|
|
return
|
|
|
|
headers = {
|
|
"Cache-Control": "no-cache, no-transform",
|
|
"Connection": "keep-alive",
|
|
"Content-Type": CONTENT_TYPE_SSE,
|
|
}
|
|
|
|
if self.mcp_session_id: # pragma: no branch
|
|
headers[MCP_SESSION_ID_HEADER] = self.mcp_session_id
|
|
|
|
# Check if we already have an active GET stream
|
|
if GET_STREAM_KEY in self._request_streams:
|
|
response = self._create_error_response(
|
|
"Conflict: Only one SSE stream is allowed per session",
|
|
HTTPStatus.CONFLICT,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
return
|
|
|
|
# Create SSE stream
|
|
sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
|
|
|
|
async def standalone_sse_writer():
|
|
try:
|
|
# Create a standalone message stream for server-initiated messages
|
|
|
|
self._request_streams[GET_STREAM_KEY] = anyio.create_memory_object_stream[EventMessage](
|
|
REQUEST_STREAM_BUFFER_SIZE
|
|
)
|
|
standalone_stream_reader = self._request_streams[GET_STREAM_KEY][1]
|
|
|
|
async with sse_stream_writer, standalone_stream_reader:
|
|
# Process messages from the standalone stream
|
|
async for event_message in standalone_stream_reader:
|
|
# For the standalone stream, we handle:
|
|
# - JSONRPCNotification (server sends notifications to client)
|
|
# - JSONRPCRequest (server sends requests to client)
|
|
# We should NOT receive JSONRPCResponse
|
|
|
|
# Send the message via SSE
|
|
event_data = self._create_event_data(event_message)
|
|
await sse_stream_writer.send(event_data)
|
|
except anyio.ClosedResourceError:
|
|
# Session teardown can close the stream while the writer is between dequeues.
|
|
pass
|
|
except Exception:
|
|
logger.exception("Error in standalone SSE writer") # pragma: no cover
|
|
finally:
|
|
logger.debug("Closing standalone SSE writer")
|
|
await self._clean_up_memory_streams(GET_STREAM_KEY)
|
|
|
|
# Create and start EventSourceResponse
|
|
response = EventSourceResponse(
|
|
content=sse_stream_reader,
|
|
data_sender_callable=standalone_sse_writer,
|
|
headers=headers,
|
|
)
|
|
|
|
try:
|
|
# This will send headers immediately and establish the SSE connection
|
|
await response(request.scope, request.receive, send)
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Error in standalone SSE response")
|
|
await self._clean_up_memory_streams(GET_STREAM_KEY)
|
|
finally:
|
|
await sse_stream_writer.aclose()
|
|
await sse_stream_reader.aclose()
|
|
|
|
async def _handle_delete_request(self, request: Request, send: Send) -> None:
|
|
"""Handle DELETE requests for explicit session termination."""
|
|
# Validate session ID
|
|
if not self.mcp_session_id: # pragma: no cover
|
|
# If no session ID set, return Method Not Allowed
|
|
response = self._create_error_response(
|
|
"Method Not Allowed: Session termination not supported",
|
|
HTTPStatus.METHOD_NOT_ALLOWED,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
return
|
|
|
|
if not await self._validate_request_headers(request, send): # pragma: no cover
|
|
return
|
|
|
|
await self.terminate()
|
|
|
|
response = self._create_json_response(
|
|
None,
|
|
HTTPStatus.OK,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
|
|
async def terminate(self) -> None:
|
|
"""Terminate the current session, closing all streams.
|
|
|
|
Once terminated, all requests with this session ID will receive 404 Not Found.
|
|
"""
|
|
|
|
self._terminated = True
|
|
logger.info(f"Terminating session: {self.mcp_session_id}")
|
|
|
|
# We need a copy of the keys to avoid modification during iteration
|
|
request_stream_keys = list(self._request_streams.keys())
|
|
|
|
# Close all request streams asynchronously
|
|
for key in request_stream_keys:
|
|
await self._clean_up_memory_streams(key)
|
|
|
|
# Clear the request streams dictionary immediately
|
|
self._request_streams.clear()
|
|
try:
|
|
if self._read_stream_writer is not None: # pragma: no branch
|
|
await self._read_stream_writer.aclose()
|
|
if self._read_stream is not None: # pragma: no branch
|
|
await self._read_stream.aclose()
|
|
if self._write_stream_reader is not None: # pragma: no branch
|
|
await self._write_stream_reader.aclose()
|
|
if self._write_stream is not None: # pragma: no branch
|
|
await self._write_stream.aclose()
|
|
except Exception as e: # pragma: no cover
|
|
# During cleanup, we catch all exceptions since streams might be in various states
|
|
logger.debug(f"Error closing streams: {e}")
|
|
|
|
async def _handle_unsupported_request(self, request: Request, send: Send) -> None:
|
|
"""Handle unsupported HTTP methods."""
|
|
headers = {
|
|
"Content-Type": CONTENT_TYPE_JSON,
|
|
"Allow": "GET, POST, DELETE",
|
|
}
|
|
if self.mcp_session_id: # pragma: no branch
|
|
headers[MCP_SESSION_ID_HEADER] = self.mcp_session_id
|
|
|
|
response = self._create_error_response(
|
|
"Method Not Allowed",
|
|
HTTPStatus.METHOD_NOT_ALLOWED,
|
|
headers=headers,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
|
|
async def _validate_request_headers(self, request: Request, send: Send) -> bool:
|
|
# Protocol-version validation lives in the manager's era-routing: only
|
|
# values in `HANDSHAKE_PROTOCOL_VERSIONS` (or no header at all) reach
|
|
# this transport, so the legacy version-gate is gone.
|
|
return await self._validate_session(request, send)
|
|
|
|
async def _validate_session(self, request: Request, send: Send) -> bool:
|
|
"""Validate the session ID in the request."""
|
|
if not self.mcp_session_id:
|
|
# If we're not using session IDs, return True
|
|
return True
|
|
|
|
# Get the session ID from the request headers
|
|
request_session_id = self._get_session_id(request)
|
|
|
|
# If no session ID provided but required, return error
|
|
if not request_session_id:
|
|
response = self._create_error_response(
|
|
"Bad Request: Missing session ID",
|
|
HTTPStatus.BAD_REQUEST,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
return False
|
|
|
|
# If session ID doesn't match, return error
|
|
if request_session_id != self.mcp_session_id: # pragma: no cover
|
|
response = self._create_error_response(
|
|
"Not Found: Invalid or expired session ID",
|
|
HTTPStatus.NOT_FOUND,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
return False
|
|
|
|
return True
|
|
|
|
async def _replay_events(self, last_event_id: str, request: Request, send: Send) -> None:
|
|
"""Replays events that would have been sent after the specified event ID.
|
|
|
|
Only used when resumability is enabled.
|
|
"""
|
|
event_store = self._event_store
|
|
if not event_store:
|
|
return # pragma: no cover
|
|
|
|
try:
|
|
headers = {
|
|
"Cache-Control": "no-cache, no-transform",
|
|
"Connection": "keep-alive",
|
|
"Content-Type": CONTENT_TYPE_SSE,
|
|
}
|
|
|
|
if self.mcp_session_id: # pragma: no branch
|
|
headers[MCP_SESSION_ID_HEADER] = self.mcp_session_id
|
|
|
|
# The manager only routes supported (or absent) header values to this transport
|
|
replay_protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER, DEFAULT_NEGOTIATED_VERSION)
|
|
|
|
# Create SSE stream for replay
|
|
sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[SSEEvent](0)
|
|
|
|
async def replay_sender():
|
|
try:
|
|
async with sse_stream_writer:
|
|
# Define an async callback for sending events
|
|
async def send_event(event_message: EventMessage) -> None:
|
|
event_data = self._create_event_data(event_message)
|
|
await sse_stream_writer.send(event_data)
|
|
|
|
# Replay past events and get the stream ID
|
|
stream_id = await event_store.replay_events_after(last_event_id, send_event)
|
|
|
|
# If stream ID not in mapping, create it
|
|
if stream_id and stream_id not in self._request_streams: # pragma: no branch
|
|
try:
|
|
# Register SSE writer so close_sse_stream() can close it
|
|
self._sse_stream_writers[stream_id] = sse_stream_writer
|
|
|
|
# Prime the resumed connection so the client sees the stream
|
|
# is re-registered. The replay→live-tail ordering window here
|
|
# is pre-existing and tracked separately.
|
|
priming_event = await self._mint_priming_event(stream_id, replay_protocol_version)
|
|
if priming_event is not None:
|
|
await sse_stream_writer.send(priming_event)
|
|
|
|
# Create new request streams for this connection
|
|
self._request_streams[stream_id] = anyio.create_memory_object_stream[EventMessage](
|
|
REQUEST_STREAM_BUFFER_SIZE
|
|
)
|
|
msg_reader = self._request_streams[stream_id][1]
|
|
|
|
# Forward messages to SSE
|
|
async with msg_reader:
|
|
async for event_message in msg_reader:
|
|
event_data = self._create_event_data(event_message)
|
|
|
|
await sse_stream_writer.send(event_data)
|
|
finally:
|
|
self._sse_stream_writers.pop(stream_id, None)
|
|
await self._clean_up_memory_streams(stream_id)
|
|
except anyio.ClosedResourceError: # pragma: lax no cover
|
|
# Expected when close_sse_stream() is called
|
|
logger.debug("Replay SSE stream closed by close_sse_stream()")
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Error in replay sender")
|
|
|
|
# Create and start EventSourceResponse
|
|
response = EventSourceResponse(
|
|
content=sse_stream_reader,
|
|
data_sender_callable=replay_sender,
|
|
headers=headers,
|
|
)
|
|
|
|
try:
|
|
await response(request.scope, request.receive, send)
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Error in replay response")
|
|
finally:
|
|
await sse_stream_writer.aclose()
|
|
await sse_stream_reader.aclose()
|
|
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Error replaying events")
|
|
response = self._create_error_response(
|
|
"Error replaying events",
|
|
HTTPStatus.INTERNAL_SERVER_ERROR,
|
|
INTERNAL_ERROR,
|
|
)
|
|
await response(request.scope, request.receive, send)
|
|
|
|
@asynccontextmanager
|
|
async def connect(
|
|
self,
|
|
) -> AsyncGenerator[
|
|
tuple[
|
|
ReadStream[SessionMessage | Exception],
|
|
WriteStream[SessionMessage],
|
|
],
|
|
None,
|
|
]:
|
|
"""Context manager that provides read and write streams for a connection.
|
|
|
|
Yields:
|
|
Tuple of (read_stream, write_stream) for bidirectional communication
|
|
"""
|
|
|
|
# Create the memory streams for this connection
|
|
|
|
read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0)
|
|
write_stream, write_stream_reader = create_context_streams[SessionMessage](0)
|
|
|
|
# Store the streams
|
|
self._read_stream_writer = read_stream_writer
|
|
self._read_stream = read_stream
|
|
self._write_stream_reader = write_stream_reader
|
|
self._write_stream = write_stream
|
|
|
|
# Start a task group for message routing
|
|
async with anyio.create_task_group() as tg:
|
|
# Create a message router that distributes messages to request streams
|
|
async def message_router():
|
|
try:
|
|
async for session_message in write_stream_reader: # pragma: no branch
|
|
# Determine which request stream(s) should receive this message
|
|
message = session_message.message
|
|
target_request_id = None
|
|
# Check if this is a response with a known request id.
|
|
# Null-id errors (e.g., parse errors) fall through to
|
|
# the GET stream since they can't be correlated.
|
|
if isinstance(message, JSONRPCResponse | JSONRPCError) and message.id is not None:
|
|
target_request_id = str(message.id)
|
|
# Extract related_request_id from meta if it exists
|
|
elif (
|
|
session_message.metadata is not None
|
|
and isinstance(
|
|
session_message.metadata,
|
|
ServerMessageMetadata,
|
|
)
|
|
and session_message.metadata.related_request_id is not None
|
|
):
|
|
related_request_id = session_message.metadata.related_request_id
|
|
if self.is_json_response_enabled:
|
|
# A JSON body carries only the response: this message
|
|
# has no wire form (nor a replay), so drop it before
|
|
# storing or queueing rather than park it (#1764).
|
|
logger.debug(f"Dropped message related to request {related_request_id} in JSON mode")
|
|
continue
|
|
target_request_id = str(related_request_id)
|
|
|
|
request_stream_id = target_request_id if target_request_id is not None else GET_STREAM_KEY
|
|
|
|
# Store the event if we have an event store,
|
|
# regardless of whether a client is connected
|
|
# messages will be replayed on the re-connect
|
|
event_id = None
|
|
if self._event_store:
|
|
event_id = await self._event_store.store_event(request_stream_id, message)
|
|
logger.debug(f"Stored {event_id} from {request_stream_id}")
|
|
|
|
if request_stream_id in self._request_streams:
|
|
try:
|
|
# Send both the message and the event ID
|
|
await self._request_streams[request_stream_id][0].send(EventMessage(message, event_id))
|
|
except (anyio.BrokenResourceError, anyio.ClosedResourceError): # pragma: no cover
|
|
# Stream might be closed, remove from registry
|
|
self._request_streams.pop(request_stream_id, None)
|
|
else:
|
|
logger.debug(
|
|
f"""Request stream {request_stream_id} not found
|
|
for message. Still processing message as the client
|
|
might reconnect and replay."""
|
|
)
|
|
except anyio.ClosedResourceError:
|
|
if self._terminated: # pragma: lax no cover
|
|
logger.debug("Read stream closed by client")
|
|
else:
|
|
logger.exception("Unexpected closure of read stream in message router")
|
|
except Exception: # pragma: lax no cover
|
|
logger.exception("Error in message router")
|
|
|
|
# Start the message router
|
|
tg.start_soon(message_router)
|
|
|
|
try:
|
|
# Yield the streams for the caller to use
|
|
yield read_stream, write_stream
|
|
finally:
|
|
for stream_id in list(self._request_streams.keys()):
|
|
await self._clean_up_memory_streams(stream_id)
|
|
self._request_streams.clear()
|
|
|
|
# Clean up the read and write streams
|
|
try:
|
|
await read_stream_writer.aclose()
|
|
await read_stream.aclose()
|
|
await write_stream_reader.aclose()
|
|
await write_stream.aclose()
|
|
except Exception as e: # pragma: no cover
|
|
# During cleanup, we catch all exceptions since streams might be in various states
|
|
logger.debug(f"Error closing streams: {e}")
|