StreamableHttp client transport (#573)
This commit is contained in:
parent
46523afe30
commit
9dfc925090
5 changed files with 498 additions and 15 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
258
src/mcp/client/streamable_http.py
Normal file
258
src/mcp/client/streamable_http.py
Normal 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()
|
||||
720
src/mcp/server/streamable_http.py
Normal file
720
src/mcp/server/streamable_http.py
Normal 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()
|
||||
855
tests/shared/test_streamable_http.py
Normal file
855
tests/shared/test_streamable_http.py
Normal 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
2
uv.lock
generated
|
|
@ -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 = [
|
||||
|
|
|
|||
Loading…
Reference in a new issue