StreamableHttp client transport (#573)

This commit is contained in:
ihrpr 2025-05-02 13:59:02 +01:00 • committed by GitHub
commit 9dfc925090
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 498 additions and 15 deletions

View file

@ -7,7 +7,7 @@ import anyio
import click
import mcp.types as types
from mcp.server.lowlevel import Server
from mcp.server.streamableHttp import (
from mcp.server.streamable_http import (
MCP_SESSION_ID_HEADER,
StreamableHTTPServerTransport,
)

View file

@ -0,0 +1,258 @@
"""
StreamableHTTP Client Transport Module
This module implements the StreamableHTTP transport for MCP clients,
providing support for HTTP POST requests with optional SSE streaming responses
and session management.
"""
import logging
from contextlib import asynccontextmanager
from datetime import timedelta
from typing import Any
import anyio
import httpx
from httpx_sse import EventSource, aconnect_sse
from mcp.types import (
ErrorData,
JSONRPCError,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
)
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"
@asynccontextmanager
async def streamablehttp_client(
url: str,
headers: dict[str, Any] | None = None,
timeout: timedelta = timedelta(seconds=30),
sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
):
"""
Client transport for StreamableHTTP.
`sse_read_timeout` determines how long (in seconds) the client will wait for a new
event before disconnecting. All other HTTP operations are controlled by `timeout`.
Yields:
Tuple of (read_stream, write_stream, terminate_callback)
"""
read_stream_writer, read_stream = anyio.create_memory_object_stream[
JSONRPCMessage | Exception
](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[
JSONRPCMessage
](0)
async def get_stream():
"""
Optional GET stream for server-initiated messages
"""
nonlocal session_id
try:
# Only attempt GET if we have a session ID
if not session_id:
return
get_headers = request_headers.copy()
get_headers[MCP_SESSION_ID_HEADER] = session_id
async with aconnect_sse(
client,
"GET",
url,
headers=get_headers,
timeout=httpx.Timeout(timeout.seconds, read=sse_read_timeout.seconds),
) as event_source:
event_source.response.raise_for_status()
logger.debug("GET SSE connection established")
async for sse in event_source.aiter_sse():
if sse.event == "message":
try:
message = JSONRPCMessage.model_validate_json(sse.data)
logger.debug(f"GET message: {message}")
await read_stream_writer.send(message)
except Exception as exc:
logger.error(f"Error parsing GET message: {exc}")
await read_stream_writer.send(exc)
else:
logger.warning(f"Unknown SSE event from GET: {sse.event}")
except Exception as exc:
# GET stream is optional, so don't propagate errors
logger.debug(f"GET stream error (non-fatal): {exc}")
async def post_writer(client: httpx.AsyncClient):
nonlocal session_id
try:
async with write_stream_reader:
async for message in write_stream_reader:
# Add session ID to headers if we have one
post_headers = request_headers.copy()
if session_id:
post_headers[MCP_SESSION_ID_HEADER] = session_id
logger.debug(f"Sending client message: {message}")
# Handle initial initialization request
is_initialization = (
isinstance(message.root, JSONRPCRequest)
and message.root.method == "initialize"
)
if (
isinstance(message.root, JSONRPCNotification)
and message.root.method == "notifications/initialized"
):
tg.start_soon(get_stream)
async with client.stream(
"POST",
url,
json=message.model_dump(
by_alias=True, mode="json", exclude_none=True
),
headers=post_headers,
) as response:
if response.status_code == 202:
logger.debug("Received 202 Accepted")
continue
# Check for 404 (session expired/invalid)
if response.status_code == 404:
if isinstance(message.root, JSONRPCRequest):
jsonrpc_error = JSONRPCError(
jsonrpc="2.0",
id=message.root.id,
error=ErrorData(
code=32600,
message="Session terminated",
),
)
await read_stream_writer.send(
JSONRPCMessage(jsonrpc_error)
)
continue
response.raise_for_status()
# Extract session ID from response headers
if is_initialization:
new_session_id = response.headers.get(MCP_SESSION_ID_HEADER)
if new_session_id:
session_id = new_session_id
logger.info(f"Received session ID: {session_id}")
# Handle different response types
content_type = response.headers.get("content-type", "").lower()
if content_type.startswith(CONTENT_TYPE_JSON):
try:
content = await response.aread()
json_message = JSONRPCMessage.model_validate_json(
content
)
await read_stream_writer.send(json_message)
except Exception as exc:
logger.error(f"Error parsing JSON response: {exc}")
await read_stream_writer.send(exc)
elif content_type.startswith(CONTENT_TYPE_SSE):
# Parse SSE events from the response
try:
event_source = EventSource(response)
async for sse in event_source.aiter_sse():
if sse.event == "message":
try:
await read_stream_writer.send(
JSONRPCMessage.model_validate_json(
sse.data
)
)
except Exception as exc:
logger.exception("Error parsing message")
await read_stream_writer.send(exc)
else:
logger.warning(f"Unknown event: {sse.event}")
except Exception as e:
logger.exception("Error reading SSE stream:")
await read_stream_writer.send(e)
else:
# For 202 Accepted with no body
if response.status_code == 202:
logger.debug("Received 202 Accepted")
continue
error_msg = f"Unexpected content type: {content_type}"
logger.error(error_msg)
await read_stream_writer.send(ValueError(error_msg))
except Exception as exc:
logger.error(f"Error in post_writer: {exc}")
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
async def terminate_session():
"""
Terminate the session by sending a DELETE request.
"""
nonlocal session_id
if not session_id:
return # No session to terminate
try:
delete_headers = request_headers.copy()
delete_headers[MCP_SESSION_ID_HEADER] = session_id
response = await client.delete(
url,
headers=delete_headers,
)
if response.status_code == 405:
# Server doesn't allow client-initiated termination
logger.debug("Server does not allow session termination")
elif response.status_code != 200:
logger.warning(f"Session termination failed: {response.status_code}")
except Exception as exc:
logger.warning(f"Session termination failed: {exc}")
async with anyio.create_task_group() as tg:
try:
logger.info(f"Connecting to StreamableHTTP endpoint: {url}")
# Set up headers with required Accept header
request_headers = {
"Accept": f"{CONTENT_TYPE_JSON}, {CONTENT_TYPE_SSE}",
"Content-Type": CONTENT_TYPE_JSON,
**(headers or {}),
}
# Track session ID if provided by server
session_id: str | None = None
async with httpx.AsyncClient(
headers=request_headers,
timeout=httpx.Timeout(timeout.seconds, read=sse_read_timeout.seconds),
follow_redirects=True,
) as client:
tg.start_soon(post_writer, client)
try:
yield read_stream, write_stream, terminate_session
finally:
tg.cancel_scope.cancel()
finally:
await read_stream_writer.aclose()
await write_stream.aclose()

View file

@ -0,0 +1,720 @@
"""
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 json
import logging
import re
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from http import HTTPStatus
from typing import Any
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
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.types import (
INTERNAL_ERROR,
INVALID_PARAMS,
INVALID_REQUEST,
PARSE_ERROR,
ErrorData,
JSONRPCError,
JSONRPCMessage,
JSONRPCNotification,
JSONRPCRequest,
JSONRPCResponse,
RequestId,
)
logger = logging.getLogger(__name__)
# Maximum size for incoming messages
MAXIMUM_MESSAGE_SIZE = 4 * 1024 * 1024 # 4MB
# 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"
# 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]+$")
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: MemoryObjectSendStream[JSONRPCMessage | Exception] | None = (
None
)
_write_stream_reader: MemoryObjectReceiveStream[JSONRPCMessage] | None = None
def __init__(
self,
mcp_session_id: str | None,
is_json_response_enabled: bool = False,
) -> 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, return JSON responses for requests
instead of SSE streams. Default is False.
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._request_streams: dict[
RequestId, MemoryObjectSendStream[JSONRPCMessage]
] = {}
self._terminated = False
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="server-error", # We don't have a request ID for general errors
error=ErrorData(
code=error_code,
message=error_message,
),
)
return Response(
error_response.model_dump_json(by_alias=True, exclude_none=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)
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_none=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)
async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
"""Application entry point that handles all HTTP requests"""
request = Request(scope, receive)
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_accept_headers(self, request: Request) -> tuple[bool, bool]:
"""Check if the request accepts the required media types."""
accept_header = request.headers.get("accept", "")
accept_types = [media_type.strip() for media_type in accept_header.split(",")]
has_json = any(
media_type.startswith(CONTENT_TYPE_JSON) for media_type in accept_types
)
has_sse = any(
media_type.startswith(CONTENT_TYPE_SSE) for media_type in accept_types
)
return has_json, has_sse
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 _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:
raise ValueError(
"No read stream writer available. Ensure connect() is called first."
)
try:
# Check Accept headers
has_json, has_sse = self._check_accept_headers(request)
if 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, receive, send)
return
# Validate Content-Type
if not self._check_content_type(request):
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()
if len(body) > MAXIMUM_MESSAGE_SIZE:
response = self._create_error_response(
"Payload Too Large: Message exceeds maximum size",
HTTPStatus.REQUEST_ENTITY_TOO_LARGE,
)
await response(scope, receive, send)
return
try:
raw_message = json.loads(body)
except json.JSONDecodeError 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 = JSONRPCMessage.model_validate(raw_message)
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.root, JSONRPCRequest)
and message.root.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:
response = self._create_error_response(
"Not Found: Invalid or expired session ID",
HTTPStatus.NOT_FOUND,
)
await response(scope, receive, send)
return
# For non-initialization requests, validate the session
elif not await self._validate_session(request, send):
return
# For notifications and responses only, return 202 Accepted
if not isinstance(message.root, 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
await writer.send(message)
return
# Extract the request ID outside the try block for proper scope
request_id = str(message.root.id)
# Create promise stream for getting response
request_stream_writer, request_stream_reader = (
anyio.create_memory_object_stream[JSONRPCMessage](0)
)
# Register this stream for the request ID
self._request_streams[request_id] = request_stream_writer
if self.is_json_response_enabled:
# Process the message
await writer.send(message)
try:
# Process messages from the request-specific stream
# We need to collect all messages until we get a response
response_message = None
# Use similar approach to SSE writer for consistency
async for received_message in request_stream_reader:
# If it's a response, this is what we're waiting for
if isinstance(
received_message.root, JSONRPCResponse | JSONRPCError
):
response_message = received_message
break
# For notifications and request, keep waiting
else:
logger.debug(f"received: {received_message.root.method}")
# At this point we should have a response
if response_message:
# Create JSON response
response = self._create_json_response(response_message)
await response(scope, receive, send)
else:
# This shouldn't happen in normal operation
logger.error(
"No response message received before stream closed"
)
response = self._create_error_response(
"Error processing request: No response received",
HTTPStatus.INTERNAL_SERVER_ERROR,
)
await response(scope, receive, send)
except Exception as e:
logger.exception(f"Error processing JSON response: {e}")
response = self._create_error_response(
f"Error processing request: {str(e)}",
HTTPStatus.INTERNAL_SERVER_ERROR,
INTERNAL_ERROR,
)
await response(scope, receive, send)
finally:
# Clean up the request stream
if request_id in self._request_streams:
self._request_streams.pop(request_id, None)
await request_stream_reader.aclose()
await request_stream_writer.aclose()
else:
# Create SSE stream
sse_stream_writer, sse_stream_reader = (
anyio.create_memory_object_stream[dict[str, Any]](0)
)
async def sse_writer():
# Get the request ID from the incoming request message
try:
async with sse_stream_writer, request_stream_reader:
# Process messages from the request-specific stream
async for received_message in request_stream_reader:
# Build the event data
event_data = {
"event": "message",
"data": received_message.model_dump_json(
by_alias=True, exclude_none=True
),
}
await sse_stream_writer.send(event_data)
# If response, remove from pending streams and close
if isinstance(
received_message.root,
JSONRPCResponse | JSONRPCError,
):
if request_id:
self._request_streams.pop(request_id, None)
break
except Exception as e:
logger.exception(f"Error in SSE writer: {e}")
finally:
logger.debug("Closing SSE writer")
# Clean up the request-specific streams
if request_id and request_id in self._request_streams:
self._request_streams.pop(request_id, None)
# Create and start EventSourceResponse
# SSE stream mode (original behavior)
# Set up headers
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=sse_writer,
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
await writer.send(message)
except Exception:
logger.exception("SSE response error")
# Clean up the request stream if something goes wrong
if request_id and request_id in self._request_streams:
self._request_streams.pop(request_id, None)
except Exception as err:
logger.exception("Error handling POST request")
response = self._create_error_response(
f"Error handling POST request: {err}",
HTTPStatus.INTERNAL_SERVER_ERROR,
INTERNAL_ERROR,
)
await response(scope, receive, send)
if writer:
await writer.send(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:
raise ValueError(
"No read stream writer available. Ensure connect() is called first."
)
# Validate Accept header - must include text/event-stream
_, has_sse = self._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_session(request, send):
return
headers = {
"Cache-Control": "no-cache, no-transform",
"Connection": "keep-alive",
"Content-Type": CONTENT_TYPE_SSE,
}
if self.mcp_session_id:
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[
dict[str, Any]
](0)
async def standalone_sse_writer():
try:
# Create a standalone message stream for server-initiated messages
standalone_stream_writer, standalone_stream_reader = (
anyio.create_memory_object_stream[JSONRPCMessage](0)
)
# Register this stream using the special key
self._request_streams[GET_STREAM_KEY] = standalone_stream_writer
async with sse_stream_writer, standalone_stream_reader:
# Process messages from the standalone stream
async for received_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 = {
"event": "message",
"data": received_message.model_dump_json(
by_alias=True, exclude_none=True
),
}
await sse_stream_writer.send(event_data)
except Exception as e:
logger.exception(f"Error in standalone SSE writer: {e}")
finally:
logger.debug("Closing standalone SSE writer")
# Remove the stream from request_streams
self._request_streams.pop(GET_STREAM_KEY, None)
# 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 as e:
logger.exception(f"Error in standalone SSE response: {e}")
# Clean up the request stream
self._request_streams.pop(GET_STREAM_KEY, None)
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:
# 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_session(request, send):
return
self._terminate_session()
response = self._create_json_response(
None,
HTTPStatus.OK,
)
await response(request.scope, request.receive, send)
def _terminate_session(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 (synchronously)
for key in request_stream_keys:
try:
# Get the stream
stream = self._request_streams.get(key)
if stream:
# We must use close() here, not aclose() since this is a sync method
stream.close()
except Exception as e:
logger.debug(f"Error closing stream {key} during termination: {e}")
# Clear the request streams dictionary immediately
self._request_streams.clear()
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:
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_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:
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
@asynccontextmanager
async def connect(
self,
) -> AsyncGenerator[
tuple[
MemoryObjectReceiveStream[JSONRPCMessage | Exception],
MemoryObjectSendStream[JSONRPCMessage],
],
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 = anyio.create_memory_object_stream[
JSONRPCMessage | Exception
](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[
JSONRPCMessage
](0)
# Store the streams
self._read_stream_writer = read_stream_writer
self._write_stream_reader = write_stream_reader
# 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 message in write_stream_reader:
# Determine which request stream(s) should receive this message
target_request_id = None
if isinstance(
message.root, JSONRPCNotification | JSONRPCRequest
):
# Extract related_request_id from meta if it exists
if (
(params := getattr(message.root, "params", None))
and (meta := params.get("_meta"))
and (related_id := meta.get("related_request_id"))
is not None
):
target_request_id = str(related_id)
else:
target_request_id = str(message.root.id)
request_stream_id = target_request_id or GET_STREAM_KEY
if request_stream_id in self._request_streams:
try:
await self._request_streams[request_stream_id].send(
message
)
except (
anyio.BrokenResourceError,
anyio.ClosedResourceError,
):
# Stream might be closed, remove from registry
self._request_streams.pop(request_stream_id, None)
except Exception as e:
logger.exception(f"Error in message router: {e}")
# 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 in list(self._request_streams.values()):
try:
await stream.aclose()
except Exception:
pass
self._request_streams.clear()

View file

@ -0,0 +1,855 @@
"""
Tests for the StreamableHTTP server and client transport.
Contains tests for both server and client sides of the StreamableHTTP transport.
"""
import contextlib
import multiprocessing
import socket
import time
from collections.abc import Generator
from http import HTTPStatus
from uuid import uuid4
import anyio
import httpx
import pytest
import requests
import uvicorn
from pydantic import AnyUrl
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Mount
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from mcp.server import Server
from mcp.server.streamable_http import (
MCP_SESSION_ID_HEADER,
SESSION_ID_PATTERN,
StreamableHTTPServerTransport,
)
from mcp.shared.exceptions import McpError
from mcp.types import InitializeResult, TextContent, TextResourceContents, Tool
# Test constants
SERVER_NAME = "test_streamable_http_server"
TEST_SESSION_ID = "test-session-id-12345"
INIT_REQUEST = {
"jsonrpc": "2.0",
"method": "initialize",
"params": {
"clientInfo": {"name": "test-client", "version": "1.0"},
"protocolVersion": "2025-03-26",
"capabilities": {},
},
"id": "init-1",
}
# Test server implementation that follows MCP protocol
class ServerTest(Server):
def __init__(self):
super().__init__(SERVER_NAME)
@self.read_resource()
async def handle_read_resource(uri: AnyUrl) -> str | bytes:
if uri.scheme == "foobar":
return f"Read {uri.host}"
elif uri.scheme == "slow":
# Simulate a slow resource
await anyio.sleep(2.0)
return f"Slow response from {uri.host}"
raise ValueError(f"Unknown resource: {uri}")
@self.list_tools()
async def handle_list_tools() -> list[Tool]:
return [
Tool(
name="test_tool",
description="A test tool",
inputSchema={"type": "object", "properties": {}},
),
Tool(
name="test_tool_with_standalone_notification",
description="A test tool that sends a notification",
inputSchema={"type": "object", "properties": {}},
),
]
@self.call_tool()
async def handle_call_tool(name: str, args: dict) -> list[TextContent]:
# When the tool is called, send a notification to test GET stream
if name == "test_tool_with_standalone_notification":
ctx = self.request_context
await ctx.session.send_resource_updated(
uri=AnyUrl("http://test_resource")
)
return [TextContent(type="text", text=f"Called {name}")]
def create_app(is_json_response_enabled=False) -> Starlette:
"""Create a Starlette application for testing that matches the example server.
Args:
is_json_response_enabled: If True, use JSON responses instead of SSE streams.
"""
# Create server instance
server = ServerTest()
server_instances = {}
# Lock to prevent race conditions when creating new sessions
session_creation_lock = anyio.Lock()
task_group = None
@contextlib.asynccontextmanager
async def lifespan(app):
"""Application lifespan context manager for managing task group."""
nonlocal task_group
async with anyio.create_task_group() as tg:
task_group = tg
print("Application started, task group initialized!")
try:
yield
finally:
print("Application shutting down, cleaning up resources...")
if task_group:
tg.cancel_scope.cancel()
task_group = None
print("Resources cleaned up successfully.")
async def handle_streamable_http(scope, receive, send):
request = Request(scope, receive)
request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)
# Use existing transport if session ID matches
if (
request_mcp_session_id is not None
and request_mcp_session_id in server_instances
):
transport = server_instances[request_mcp_session_id]
await transport.handle_request(scope, receive, send)
elif request_mcp_session_id is None:
async with session_creation_lock:
new_session_id = uuid4().hex
http_transport = StreamableHTTPServerTransport(
mcp_session_id=new_session_id,
is_json_response_enabled=is_json_response_enabled,
)
async with http_transport.connect() as streams:
read_stream, write_stream = streams
async def run_server():
try:
await server.run(
read_stream,
write_stream,
server.create_initialization_options(),
)
except Exception as e:
print(f"Server exception: {e}")
if task_group is None:
response = Response(
"Internal Server Error: Task group is not initialized",
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
)
await response(scope, receive, send)
return
# Store the instance before starting the task to prevent races
server_instances[http_transport.mcp_session_id] = http_transport
task_group.start_soon(run_server)
await http_transport.handle_request(scope, receive, send)
else:
response = Response(
"Bad Request: No valid session ID provided",
status_code=HTTPStatus.BAD_REQUEST,
)
await response(scope, receive, send)
# Create an ASGI application
app = Starlette(
debug=True,
routes=[
Mount("/mcp", app=handle_streamable_http),
],
lifespan=lifespan,
)
return app
def run_server(port: int, is_json_response_enabled=False) -> None:
"""Run the test server.
Args:
port: Port to listen on.
is_json_response_enabled: If True, use JSON responses instead of SSE streams.
"""
print(
f"Starting test server on port {port} with "
f"json_enabled={is_json_response_enabled}"
)
app = create_app(is_json_response_enabled)
# Configure server
config = uvicorn.Config(
app=app,
host="127.0.0.1",
port=port,
log_level="info",
limit_concurrency=10,
timeout_keep_alive=5,
access_log=False,
)
# Start the server
server = uvicorn.Server(config=config)
# This is important to catch exceptions and prevent test hangs
try:
print("Server starting...")
server.run()
except Exception as e:
print(f"ERROR: Server failed to run: {e}")
import traceback
traceback.print_exc()
print("Server shutdown")
# Test fixtures - using same approach as SSE tests
@pytest.fixture
def basic_server_port() -> int:
"""Find an available port for the basic server."""
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
@pytest.fixture
def json_server_port() -> int:
"""Find an available port for the JSON response server."""
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
@pytest.fixture
def basic_server(basic_server_port: int) -> Generator[None, None, None]:
"""Start a basic server."""
proc = multiprocessing.Process(
target=run_server, kwargs={"port": basic_server_port}, daemon=True
)
proc.start()
# Wait for server to be running
max_attempts = 20
attempt = 0
while attempt < max_attempts:
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.connect(("127.0.0.1", basic_server_port))
break
except ConnectionRefusedError:
time.sleep(0.1)
attempt += 1
else:
raise RuntimeError(f"Server failed to start after {max_attempts} attempts")
yield
# Clean up
proc.kill()
proc.join(timeout=2)
if proc.is_alive():
print("server process failed to terminate")
@pytest.fixture
def json_response_server(json_server_port: int) -> Generator[None, None, None]:
"""Start a server with JSON response enabled."""
proc = multiprocessing.Process(
target=run_server,
kwargs={"port": json_server_port, "is_json_response_enabled": True},
daemon=True,
)
proc.start()
# Wait for server to be running
max_attempts = 20
attempt = 0
while attempt < max_attempts:
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.connect(("127.0.0.1", json_server_port))
break
except ConnectionRefusedError:
time.sleep(0.1)
attempt += 1
else:
raise RuntimeError(f"Server failed to start after {max_attempts} attempts")
yield
# Clean up
proc.kill()
proc.join(timeout=2)
if proc.is_alive():
print("server process failed to terminate")
@pytest.fixture
def basic_server_url(basic_server_port: int) -> str:
"""Get the URL for the basic test server."""
return f"http://127.0.0.1:{basic_server_port}"
@pytest.fixture
def json_server_url(json_server_port: int) -> str:
"""Get the URL for the JSON response test server."""
return f"http://127.0.0.1:{json_server_port}"
# Basic request validation tests
def test_accept_header_validation(basic_server, basic_server_url):
"""Test that Accept header is properly validated."""
# Test without Accept header
response = requests.post(
f"{basic_server_url}/mcp",
headers={"Content-Type": "application/json"},
json={"jsonrpc": "2.0", "method": "initialize", "id": 1},
)
assert response.status_code == 406
assert "Not Acceptable" in response.text
def test_content_type_validation(basic_server, basic_server_url):
"""Test that Content-Type header is properly validated."""
# Test with incorrect Content-Type
response = requests.post(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "text/plain",
},
data="This is not JSON",
)
assert response.status_code == 415
assert "Unsupported Media Type" in response.text
def test_json_validation(basic_server, basic_server_url):
"""Test that JSON content is properly validated."""
# Test with invalid JSON
response = requests.post(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
data="this is not valid json",
)
assert response.status_code == 400
assert "Parse error" in response.text
def test_json_parsing(basic_server, basic_server_url):
"""Test that JSON content is properly parse."""
# Test with valid JSON but invalid JSON-RPC
response = requests.post(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json={"foo": "bar"},
)
assert response.status_code == 400
assert "Validation error" in response.text
def test_method_not_allowed(basic_server, basic_server_url):
"""Test that unsupported HTTP methods are rejected."""
# Test with unsupported method (PUT)
response = requests.put(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json={"jsonrpc": "2.0", "method": "initialize", "id": 1},
)
assert response.status_code == 405
assert "Method Not Allowed" in response.text
def test_session_validation(basic_server, basic_server_url):
"""Test session ID validation."""
# session_id not used directly in this test
# Test without session ID
response = requests.post(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json={"jsonrpc": "2.0", "method": "list_tools", "id": 1},
)
assert response.status_code == 400
assert "Missing session ID" in response.text
def test_session_id_pattern():
"""Test that SESSION_ID_PATTERN correctly validates session IDs."""
# Valid session IDs (visible ASCII characters from 0x21 to 0x7E)
valid_session_ids = [
"test-session-id",
"1234567890",
"session!@#$%^&*()_+-=[]{}|;:,.<>?/",
"~`",
]
for session_id in valid_session_ids:
assert SESSION_ID_PATTERN.match(session_id) is not None
# Ensure fullmatch matches too (whole string)
assert SESSION_ID_PATTERN.fullmatch(session_id) is not None
# Invalid session IDs
invalid_session_ids = [
"", # Empty string
" test", # Space (0x20)
"test\t", # Tab
"test\n", # Newline
"test\r", # Carriage return
"test" + chr(0x7F), # DEL character
"test" + chr(0x80), # Extended ASCII
"test" + chr(0x00), # Null character
"test" + chr(0x20), # Space (0x20)
]
for session_id in invalid_session_ids:
# For invalid IDs, either match will fail or fullmatch will fail
if SESSION_ID_PATTERN.match(session_id) is not None:
# If match succeeds, fullmatch should fail (partial match case)
assert SESSION_ID_PATTERN.fullmatch(session_id) is None
def test_streamable_http_transport_init_validation():
"""Test that StreamableHTTPServerTransport validates session ID on init."""
# Valid session ID should initialize without errors
valid_transport = StreamableHTTPServerTransport(mcp_session_id="valid-id")
assert valid_transport.mcp_session_id == "valid-id"
# None should be accepted
none_transport = StreamableHTTPServerTransport(mcp_session_id=None)
assert none_transport.mcp_session_id is None
# Invalid session ID should raise ValueError
with pytest.raises(ValueError) as excinfo:
StreamableHTTPServerTransport(mcp_session_id="invalid id with space")
assert "Session ID must only contain visible ASCII characters" in str(excinfo.value)
# Test with control characters
with pytest.raises(ValueError):
StreamableHTTPServerTransport(mcp_session_id="test\nid")
with pytest.raises(ValueError):
StreamableHTTPServerTransport(mcp_session_id="test\n")
def test_session_termination(basic_server, basic_server_url):
"""Test session termination via DELETE and subsequent request handling."""
response = requests.post(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json=INIT_REQUEST,
)
assert response.status_code == 200
# Now terminate the session
session_id = response.headers.get(MCP_SESSION_ID_HEADER)
response = requests.delete(
f"{basic_server_url}/mcp",
headers={MCP_SESSION_ID_HEADER: session_id},
)
assert response.status_code == 200
# Try to use the terminated session
response = requests.post(
f"{basic_server_url}/mcp",
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
MCP_SESSION_ID_HEADER: session_id,
},
json={"jsonrpc": "2.0", "method": "ping", "id": 2},
)
assert response.status_code == 404
assert "Session has been terminated" in response.text
def test_response(basic_server, basic_server_url):
"""Test response handling for a valid request."""
mcp_url = f"{basic_server_url}/mcp"
response = requests.post(
mcp_url,
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json=INIT_REQUEST,
)
assert response.status_code == 200
# Now terminate the session
session_id = response.headers.get(MCP_SESSION_ID_HEADER)
# Try to use the terminated session
tools_response = requests.post(
mcp_url,
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
MCP_SESSION_ID_HEADER: session_id, # Use the session ID we got earlier
},
json={"jsonrpc": "2.0", "method": "tools/list", "id": "tools-1"},
stream=True,
)
assert tools_response.status_code == 200
assert tools_response.headers.get("Content-Type") == "text/event-stream"
def test_json_response(json_response_server, json_server_url):
"""Test response handling when is_json_response_enabled is True."""
mcp_url = f"{json_server_url}/mcp"
response = requests.post(
mcp_url,
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json=INIT_REQUEST,
)
assert response.status_code == 200
assert response.headers.get("Content-Type") == "application/json"
def test_get_sse_stream(basic_server, basic_server_url):
"""Test establishing an SSE stream via GET request."""
# First, we need to initialize a session
mcp_url = f"{basic_server_url}/mcp"
init_response = requests.post(
mcp_url,
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json=INIT_REQUEST,
)
assert init_response.status_code == 200
# Get the session ID
session_id = init_response.headers.get(MCP_SESSION_ID_HEADER)
assert session_id is not None
# Now attempt to establish an SSE stream via GET
get_response = requests.get(
mcp_url,
headers={
"Accept": "text/event-stream",
MCP_SESSION_ID_HEADER: session_id,
},
stream=True,
)
# Verify we got a successful response with the right content type
assert get_response.status_code == 200
assert get_response.headers.get("Content-Type") == "text/event-stream"
# Test that a second GET request gets rejected (only one stream allowed)
second_get = requests.get(
mcp_url,
headers={
"Accept": "text/event-stream",
MCP_SESSION_ID_HEADER: session_id,
},
stream=True,
)
# Should get CONFLICT (409) since there's already a stream
# Note: This might fail if the first stream fully closed before this runs,
# but generally it should work in the test environment where it runs quickly
assert second_get.status_code == 409
def test_get_validation(basic_server, basic_server_url):
"""Test validation for GET requests."""
# First, we need to initialize a session
mcp_url = f"{basic_server_url}/mcp"
init_response = requests.post(
mcp_url,
headers={
"Accept": "application/json, text/event-stream",
"Content-Type": "application/json",
},
json=INIT_REQUEST,
)
assert init_response.status_code == 200
# Get the session ID
session_id = init_response.headers.get(MCP_SESSION_ID_HEADER)
assert session_id is not None
# Test without Accept header
response = requests.get(
mcp_url,
headers={
MCP_SESSION_ID_HEADER: session_id,
},
stream=True,
)
assert response.status_code == 406
assert "Not Acceptable" in response.text
# Test with wrong Accept header
response = requests.get(
mcp_url,
headers={
"Accept": "application/json",
MCP_SESSION_ID_HEADER: session_id,
},
)
assert response.status_code == 406
assert "Not Acceptable" in response.text
# Client-specific fixtures
@pytest.fixture
async def http_client(basic_server, basic_server_url):
"""Create test client matching the SSE test pattern."""
async with httpx.AsyncClient(base_url=basic_server_url) as client:
yield client
@pytest.fixture
async def initialized_client_session(basic_server, basic_server_url):
"""Create initialized StreamableHTTP client session."""
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
read_stream,
write_stream,
_,
):
async with ClientSession(
read_stream,
write_stream,
) as session:
await session.initialize()
yield session
@pytest.mark.anyio
async def test_streamablehttp_client_basic_connection(basic_server, basic_server_url):
"""Test basic client connection with initialization."""
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
read_stream,
write_stream,
_,
):
async with ClientSession(
read_stream,
write_stream,
) as session:
# Test initialization
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.serverInfo.name == SERVER_NAME
@pytest.mark.anyio
async def test_streamablehttp_client_resource_read(initialized_client_session):
"""Test client resource read functionality."""
response = await initialized_client_session.read_resource(
uri=AnyUrl("foobar://test-resource")
)
assert len(response.contents) == 1
assert response.contents[0].uri == AnyUrl("foobar://test-resource")
assert response.contents[0].text == "Read test-resource"
@pytest.mark.anyio
async def test_streamablehttp_client_tool_invocation(initialized_client_session):
"""Test client tool invocation."""
# First list tools
tools = await initialized_client_session.list_tools()
assert len(tools.tools) == 2
assert tools.tools[0].name == "test_tool"
# Call the tool
result = await initialized_client_session.call_tool("test_tool", {})
assert len(result.content) == 1
assert result.content[0].type == "text"
assert result.content[0].text == "Called test_tool"
@pytest.mark.anyio
async def test_streamablehttp_client_error_handling(initialized_client_session):
"""Test error handling in client."""
with pytest.raises(McpError) as exc_info:
await initialized_client_session.read_resource(
uri=AnyUrl("unknown://test-error")
)
assert exc_info.value.error.code == 0
assert "Unknown resource: unknown://test-error" in exc_info.value.error.message
@pytest.mark.anyio
async def test_streamablehttp_client_session_persistence(
basic_server, basic_server_url
):
"""Test that session ID persists across requests."""
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
read_stream,
write_stream,
_,
):
async with ClientSession(
read_stream,
write_stream,
) as session:
# Initialize the session
result = await session.initialize()
assert isinstance(result, InitializeResult)
# Make multiple requests to verify session persistence
tools = await session.list_tools()
assert len(tools.tools) == 2
# Read a resource
resource = await session.read_resource(uri=AnyUrl("foobar://test-persist"))
assert isinstance(resource.contents[0], TextResourceContents) is True
content = resource.contents[0]
assert isinstance(content, TextResourceContents)
assert content.text == "Read test-persist"
@pytest.mark.anyio
async def test_streamablehttp_client_json_response(
json_response_server, json_server_url
):
"""Test client with JSON response mode."""
async with streamablehttp_client(f"{json_server_url}/mcp") as (
read_stream,
write_stream,
_,
):
async with ClientSession(
read_stream,
write_stream,
) as session:
# Initialize the session
result = await session.initialize()
assert isinstance(result, InitializeResult)
assert result.serverInfo.name == SERVER_NAME
# Check tool listing
tools = await session.list_tools()
assert len(tools.tools) == 2
# Call a tool and verify JSON response handling
result = await session.call_tool("test_tool", {})
assert len(result.content) == 1
assert result.content[0].type == "text"
assert result.content[0].text == "Called test_tool"
@pytest.mark.anyio
async def test_streamablehttp_client_get_stream(basic_server, basic_server_url):
"""Test GET stream functionality for server-initiated messages."""
import mcp.types as types
from mcp.shared.session import RequestResponder
notifications_received = []
# Define message handler to capture notifications
async def message_handler(
message: RequestResponder[types.ServerRequest, types.ClientResult]
| types.ServerNotification
| Exception,
) -> None:
if isinstance(message, types.ServerNotification):
notifications_received.append(message)
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
read_stream,
write_stream,
_,
):
async with ClientSession(
read_stream, write_stream, message_handler=message_handler
) as session:
# Initialize the session - this triggers the GET stream setup
result = await session.initialize()
assert isinstance(result, InitializeResult)
# Call the special tool that sends a notification
await session.call_tool("test_tool_with_standalone_notification", {})
# Verify we received the notification
assert len(notifications_received) > 0
# Verify the notification is a ResourceUpdatedNotification
resource_update_found = False
for notif in notifications_received:
if isinstance(notif.root, types.ResourceUpdatedNotification):
assert str(notif.root.params.uri) == "http://test_resource/"
resource_update_found = True
assert (
resource_update_found
), "ResourceUpdatedNotification not received via GET stream"
@pytest.mark.anyio
async def test_streamablehttp_client_session_termination(
basic_server, basic_server_url
):
"""Test client session termination functionality."""
# Create the streamablehttp_client with a custom httpx client to capture headers
async with streamablehttp_client(f"{basic_server_url}/mcp") as (
read_stream,
write_stream,
terminate_session,
):
async with ClientSession(read_stream, write_stream) as session:
# Initialize the session
result = await session.initialize()
assert isinstance(result, InitializeResult)
# Make a request to confirm session is working
tools = await session.list_tools()
assert len(tools.tools) == 2
# After exiting ClientSession context, explicitly terminate the session
await terminate_session()
with pytest.raises(
McpError,
match="Session terminated",
):
await session.list_tools()

2
uv.lock generated
View file

@ -1,4 +1,5 @@
version = 1
revision = 1
requires-python = ">=3.10"
[options]
@ -547,6 +548,7 @@ requires-dist = [
{ name = "uvicorn", marker = "sys_platform != 'emscripten'", specifier = ">=0.23.1" },
{ name = "websockets", marker = "extra == 'ws'", specifier = ">=15.0.1" },
]
provides-extras = ["cli", "rich", "ws"]
[package.metadata.requires-dev]
dev = [