mirror of
https://github.com/aljazceru/mcp-python-sdk.git
synced 2025-12-19 14:54:24 +01:00
StreamableHttp client transport (#573)
This commit is contained in:
ihrpr
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()
|
||||
Reference in New Issue
Block a user