"""Implements StreamableHTTP transport for MCP clients.""" from __future__ import annotations as _annotations import contextlib import logging from collections.abc import AsyncGenerator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass import anyio import httpx from anyio.abc import TaskGroup from httpx_sse import EventSource, ServerSentEvent, aconnect_sse from mcp_types import ( INTERNAL_ERROR, INVALID_REQUEST, METHOD_NOT_FOUND, PARSE_ERROR, ErrorData, JSONRPCError, JSONRPCMessage, JSONRPCNotification, JSONRPCRequest, JSONRPCResponse, RequestId, jsonrpc_message_adapter, ) from pydantic import ValidationError from mcp.client._transport import TransportStreams from mcp.shared._compat import resync_tracer from mcp.shared._context_streams import ContextReceiveStream, ContextSendStream, create_context_streams from mcp.shared._httpx_utils import create_mcp_http_client from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER from mcp.shared.message import ClientMessageMetadata, SessionMessage logger = logging.getLogger(__name__) # TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here. SessionMessageOrError = SessionMessage | Exception StreamWriter = ContextSendStream[SessionMessageOrError] StreamReader = ContextReceiveStream[SessionMessage] MCP_SESSION_ID = "mcp-session-id" LAST_EVENT_ID = "last-event-id" # Reconnection defaults DEFAULT_RECONNECTION_DELAY_MS = 1000 # 1 second fallback when server doesn't provide retry MAX_RECONNECTION_ATTEMPTS = 2 # Max retry attempts before giving up class StreamableHTTPError(Exception): """Base exception for StreamableHTTP transport errors.""" class ResumptionError(StreamableHTTPError): """Raised when resumption request is invalid.""" @dataclass class RequestContext: """Context for a request operation.""" client: httpx.AsyncClient session_id: str | None session_message: SessionMessage metadata: ClientMessageMetadata | None read_stream_writer: StreamWriter class StreamableHTTPTransport: """StreamableHTTP client transport implementation.""" def __init__(self, url: str) -> None: """Initialize the StreamableHTTP transport. Args: url: The endpoint URL. """ self.url = url self.session_id: str | None = None # Captured from each stamped POST's metadata. Reused on outbound HTTP that carries # no per-message header (transport-internal GET/DELETE, and dispatcher-written # response/error/cancel POSTs that bypass the session's stamp). Cleared when an # `initialize` POST goes out so a probe-stamped value cannot leak onto the handshake. self._protocol_version_header: str | None = None def _prepare_headers(self) -> dict[str, str]: """Build MCP-specific request headers for any outbound HTTP request. These are merged with the ``httpx.AsyncClient`` defaults (these take precedence). The cached ``MCP-Protocol-Version`` is included whenever present so messages that don't pass through the session's stamp — response/error/cancel POSTs, transport-internal GET/DELETE — still carry the negotiated version. Per-message headers are layered on top by the caller. """ headers: dict[str, str] = { "accept": "application/json, text/event-stream", "content-type": "application/json", } if self.session_id: headers[MCP_SESSION_ID] = self.session_id if self._protocol_version_header: headers[MCP_PROTOCOL_VERSION_HEADER] = self._protocol_version_header return headers def _is_initialization_request(self, message: JSONRPCMessage) -> bool: """Check if the message is an initialization request.""" return isinstance(message, JSONRPCRequest) and message.method == "initialize" def _is_initialized_notification(self, message: JSONRPCMessage) -> bool: """Check if the message is an initialized notification.""" return isinstance(message, JSONRPCNotification) and message.method == "notifications/initialized" def _maybe_extract_session_id_from_response(self, response: httpx.Response) -> None: """Extract and store session ID from response headers.""" new_session_id = response.headers.get(MCP_SESSION_ID) if new_session_id: self.session_id = new_session_id logger.info(f"Received session ID: {self.session_id}") async def _handle_sse_event( self, sse: ServerSentEvent, read_stream_writer: StreamWriter, original_request_id: RequestId | None = None, resumption_callback: Callable[[str], Awaitable[None]] | None = None, ) -> bool: """Handle an SSE event, returning True if the response is complete.""" if sse.event == "message": # Handle priming events (empty data with ID) for resumability if not sse.data: # Call resumption callback for priming events that have an ID if sse.id and resumption_callback: await resumption_callback(sse.id) return False try: message = jsonrpc_message_adapter.validate_json(sse.data, by_name=False) logger.debug(f"SSE message: {message}") # If this is a response and we have original_request_id, replace it if original_request_id is not None and isinstance(message, JSONRPCResponse | JSONRPCError): message.id = original_request_id session_message = SessionMessage(message) await read_stream_writer.send(session_message) # Call resumption token callback if we have an ID if sse.id and resumption_callback: await resumption_callback(sse.id) # If this is a response or error return True indicating completion # Otherwise, return False to continue listening return isinstance(message, JSONRPCResponse | JSONRPCError) # Forwarding to a closed read stream lands here when the caller cancels mid-SSE # (BrokenResourceError, not a parse failure); coverage is timing-dependent in the # streaming story's modern HTTP cancellation leg. except Exception as exc: # pragma: lax no cover logger.exception("Error parsing SSE message") if original_request_id is not None: error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse SSE message: {exc}") error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=original_request_id, error=error_data)) await read_stream_writer.send(error_msg) return True await read_stream_writer.send(exc) return False else: # pragma: no cover logger.warning(f"Unknown SSE event: {sse.event}") return False async def handle_get_stream(self, client: httpx.AsyncClient, read_stream_writer: StreamWriter) -> None: """Handle GET stream for server-initiated messages with auto-reconnect.""" last_event_id: str | None = None retry_interval_ms: int | None = None attempt: int = 0 while attempt < MAX_RECONNECTION_ATTEMPTS: # pragma: no branch try: if not self.session_id: return headers = self._prepare_headers() if last_event_id: headers[LAST_EVENT_ID] = last_event_id async with aconnect_sse(client, "GET", self.url, headers=headers) as event_source: event_source.response.raise_for_status() logger.debug("GET SSE connection established") async for sse in event_source.aiter_sse(): # Track last event ID for reconnection if sse.id: last_event_id = sse.id # Track retry interval from server if sse.retry is not None: retry_interval_ms = sse.retry await self._handle_sse_event(sse, read_stream_writer) # Stream ended normally (server closed) - reset attempt counter attempt = 0 except Exception: logger.debug("GET stream error", exc_info=True) attempt += 1 if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover logger.debug(f"GET stream max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") return # Wait before reconnecting delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS logger.info(f"GET stream disconnected, reconnecting in {delay_ms}ms...") await anyio.sleep(delay_ms / 1000.0) async def _handle_resumption_request(self, ctx: RequestContext) -> None: """Handle a resumption request using GET with SSE.""" headers = self._prepare_headers() if ctx.metadata and ctx.metadata.resumption_token: headers[LAST_EVENT_ID] = ctx.metadata.resumption_token else: raise ResumptionError("Resumption request requires a resumption token") # pragma: no cover # Extract original request ID to map responses original_request_id = None if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch original_request_id = ctx.session_message.message.id async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source: event_source.response.raise_for_status() logger.debug("Resumption GET SSE connection established") async for sse in event_source.aiter_sse(): # pragma: no branch is_complete = await self._handle_sse_event( sse, ctx.read_stream_writer, original_request_id, ctx.metadata.on_resumption_token_update if ctx.metadata else None, ) if is_complete: await event_source.response.aclose() break async def _handle_post_request(self, ctx: RequestContext) -> None: """Handle a POST request with response processing.""" message = ctx.session_message.message is_initialization = self._is_initialization_request(message) if is_initialization: # `initialize` is the negotiation, not a "subsequent request" — discard any # probe-stamped value so the discover→fallback path can't leak it onto the handshake. self._protocol_version_header = None headers = self._prepare_headers() if ctx.metadata is not None and ctx.metadata.headers is not None: headers.update(ctx.metadata.headers) if MCP_PROTOCOL_VERSION_HEADER in ctx.metadata.headers: self._protocol_version_header = ctx.metadata.headers[MCP_PROTOCOL_VERSION_HEADER] async with ctx.client.stream( "POST", self.url, json=message.model_dump(by_alias=True, mode="json", exclude_unset=True), headers=headers, ) as response: if response.status_code == 202: logger.debug("Received 202 Accepted") return if response.status_code >= 400: if isinstance(message, JSONRPCRequest): # A spec-correct server may return the JSON-RPC error in the # body at a non-2xx status (e.g. 400 for INVALID_PARAMS, 404 # for METHOD_NOT_FOUND). Surface that error rather than the # status-derived stand-in below. if response.headers.get("content-type", "").lower().startswith("application/json"): try: body = await response.aread() parsed = jsonrpc_message_adapter.validate_json(body, by_name=False) if isinstance(parsed, JSONRPCError): # The server may have set `id: null` (request rejected before its # id was parsed); use this request's id so correlation works. reply = JSONRPCError(jsonrpc="2.0", id=message.id, error=parsed.error) await ctx.read_stream_writer.send(SessionMessage(reply)) return except (httpx.StreamError, ValidationError): pass logger.debug("Non-2xx body was not a JSON-RPC error; using fallback") if response.status_code == 404: if self.session_id is None: # No session yet → 404 is the HTTP-level spelling of # METHOD_NOT_FOUND (gateway / legacy server doesn't know # this method); "Session terminated" would be a lie here. error_data = ErrorData(code=METHOD_NOT_FOUND, message="Not Found") else: error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated") else: error_data = ErrorData(code=INTERNAL_ERROR, message="Server returned an error response") session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)) await ctx.read_stream_writer.send(session_message) return if is_initialization: self._maybe_extract_session_id_from_response(response) # Per https://modelcontextprotocol.io/specification/2025-06-18/basic#notifications: # The server MUST NOT send a response to notifications. if isinstance(message, JSONRPCRequest): content_type = response.headers.get("content-type", "").lower() if content_type.startswith("application/json"): await self._handle_json_response(response, ctx.read_stream_writer, request_id=message.id) elif content_type.startswith("text/event-stream"): await self._handle_sse_response(response, ctx) else: logger.error(f"Unexpected content type: {content_type}") error_data = ErrorData(code=INVALID_REQUEST, message=f"Unexpected content type: {content_type}") error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data)) await ctx.read_stream_writer.send(error_msg) async def _handle_json_response( self, response: httpx.Response, read_stream_writer: StreamWriter, *, request_id: RequestId, ) -> None: """Handle JSON response from the server.""" try: content = await response.aread() message = jsonrpc_message_adapter.validate_json(content, by_name=False) session_message = SessionMessage(message) await read_stream_writer.send(session_message) except (httpx.StreamError, ValidationError) as exc: logger.exception("Error parsing JSON response") error_data = ErrorData(code=PARSE_ERROR, message=f"Failed to parse JSON response: {exc}") error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data)) await read_stream_writer.send(error_msg) async def _handle_sse_response( self, response: httpx.Response, ctx: RequestContext, ) -> None: """Handle SSE response from the server.""" last_event_id: str | None = None retry_interval_ms: int | None = None # The caller (_handle_post_request) only reaches here inside # isinstance(message, JSONRPCRequest), so this is always a JSONRPCRequest. assert isinstance(ctx.session_message.message, JSONRPCRequest) original_request_id = ctx.session_message.message.id try: event_source = EventSource(response) async for sse in event_source.aiter_sse(): # pragma: no branch # Track last event ID for potential reconnection if sse.id: last_event_id = sse.id # Track retry interval from server if sse.retry is not None: retry_interval_ms = sse.retry is_complete = await self._handle_sse_event( sse, ctx.read_stream_writer, original_request_id=original_request_id, resumption_callback=(ctx.metadata.on_resumption_token_update if ctx.metadata else None), ) # If the SSE event indicates completion, like returning response/error # break the loop if is_complete: await response.aclose() return # Normal completion, no reconnect needed except Exception: logger.debug("SSE stream ended", exc_info=True) # pragma: lax no cover # Stream ended without response - reconnect if we received an event with ID if last_event_id is not None: # pragma: no branch logger.info("SSE stream disconnected, reconnecting...") await self._handle_reconnection(ctx, last_event_id, retry_interval_ms) async def _handle_reconnection( self, ctx: RequestContext, last_event_id: str, retry_interval_ms: int | None = None, attempt: int = 0, ) -> None: """Reconnect with Last-Event-ID to resume stream after server disconnect.""" # Bail if max retries exceeded if attempt >= MAX_RECONNECTION_ATTEMPTS: # pragma: no cover logger.debug(f"Max reconnection attempts ({MAX_RECONNECTION_ATTEMPTS}) exceeded") return # Always wait - use server value or default delay_ms = retry_interval_ms if retry_interval_ms is not None else DEFAULT_RECONNECTION_DELAY_MS await anyio.sleep(delay_ms / 1000.0) headers = self._prepare_headers() headers[LAST_EVENT_ID] = last_event_id # Extract original request ID to map responses original_request_id = None if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch original_request_id = ctx.session_message.message.id try: async with aconnect_sse(ctx.client, "GET", self.url, headers=headers) as event_source: event_source.response.raise_for_status() logger.info("Reconnected to SSE stream") # Track for potential further reconnection reconnect_last_event_id: str = last_event_id reconnect_retry_ms = retry_interval_ms async for sse in event_source.aiter_sse(): if sse.id: # pragma: no branch reconnect_last_event_id = sse.id if sse.retry is not None: reconnect_retry_ms = sse.retry is_complete = await self._handle_sse_event( sse, ctx.read_stream_writer, original_request_id, ctx.metadata.on_resumption_token_update if ctx.metadata else None, ) if is_complete: await event_source.response.aclose() return # Stream ended again without response - reconnect again (reset attempt counter) logger.info("SSE stream disconnected, reconnecting...") await self._handle_reconnection(ctx, reconnect_last_event_id, reconnect_retry_ms, 0) except Exception as e: # pragma: no cover logger.debug(f"Reconnection failed: {e}") # Try to reconnect again if we still have an event ID await self._handle_reconnection(ctx, last_event_id, retry_interval_ms, attempt + 1) async def post_writer( self, client: httpx.AsyncClient, write_stream_reader: StreamReader, read_stream_writer: StreamWriter, write_stream: ContextSendStream[SessionMessage], start_get_stream: Callable[[], None], tg: TaskGroup, ) -> None: """Handle writing requests to the server.""" try: async with write_stream_reader, read_stream_writer, write_stream: async def _handle_message(session_message: SessionMessage) -> None: message = session_message.message metadata = ( session_message.metadata if isinstance(session_message.metadata, ClientMessageMetadata) else None ) # Check if this is a resumption request is_resumption = bool(metadata and metadata.resumption_token) logger.debug(f"Sending client message: {message}") # Handle initialized notification if self._is_initialized_notification(message): start_get_stream() ctx = RequestContext( client=client, session_id=self.session_id, session_message=session_message, metadata=metadata, read_stream_writer=read_stream_writer, ) async def handle_request_async(): if is_resumption: await self._handle_resumption_request(ctx) else: await self._handle_post_request(ctx) # If this is a request, start a new task to handle it if isinstance(message, JSONRPCRequest): tg.start_soon(handle_request_async) else: await handle_request_async() async for session_message in write_stream_reader: sender_ctx = write_stream_reader.last_context if sender_ctx is not None: async with anyio.create_task_group() as tg_local: sender_ctx.run(tg_local.start_soon, _handle_message, session_message) else: await _handle_message(session_message) # pragma: no cover except Exception: # pragma: lax no cover logger.exception("Error in post_writer") async def terminate_session(self, client: httpx.AsyncClient) -> None: """Terminate the session by sending a DELETE request.""" if not self.session_id: return # pragma: no cover try: headers = self._prepare_headers() response = await client.delete(self.url, headers=headers) if response.status_code == 405: logger.debug("Server does not allow session termination") elif response.status_code not in (200, 204): logger.warning(f"Session termination failed: {response.status_code}") # pragma: no cover except Exception as exc: # pragma: no cover logger.warning(f"Session termination failed: {exc}") # TODO(Marcelo): Check the TODO below, and cover this with tests if necessary. def get_session_id(self) -> str | None: """Get the current session ID.""" return self.session_id # pragma: no cover # TODO(Marcelo): I've dropped the `get_session_id` callback because it breaks the Transport protocol. Is that needed? # It's a completely wrong abstraction, so removal is a good idea. But if we need the client to find the session ID, # we should think about a better way to do it. I believe we can achieve it with other means. @asynccontextmanager async def streamable_http_client( url: str, *, http_client: httpx.AsyncClient | None = None, terminate_on_close: bool = True, ) -> AsyncGenerator[TransportStreams, None]: """Client transport for StreamableHTTP. Args: url: The MCP server endpoint URL. http_client: Optional pre-configured httpx.AsyncClient. If None, a default client with recommended MCP timeouts will be created. To configure headers, authentication, or other HTTP settings, create an httpx.AsyncClient and pass it here. terminate_on_close: If True, send a DELETE request to terminate the session when the context exits. Yields: Tuple containing: - read_stream: Stream for reading messages from the server - write_stream: Stream for sending messages to the server Example: See examples/snippets/clients/ for usage patterns. """ # Determine if we need to create and manage the client client_provided = http_client is not None client = http_client if client is None: # Create default client with recommended MCP timeouts client = create_mcp_http_client() transport = StreamableHTTPTransport(url) logger.debug(f"Connecting to StreamableHTTP endpoint: {url}") async with contextlib.AsyncExitStack() as stack: # Only manage client lifecycle if we created it if not client_provided: await stack.enter_async_context(client) read_stream_writer, read_stream = create_context_streams[SessionMessage | Exception](0) write_stream, write_stream_reader = create_context_streams[SessionMessage](0) async with ( read_stream_writer, read_stream, write_stream, write_stream_reader, anyio.create_task_group() as tg, ): def start_get_stream() -> None: tg.start_soon(transport.handle_get_stream, client, read_stream_writer) tg.start_soon( transport.post_writer, client, write_stream_reader, read_stream_writer, write_stream, start_get_stream, tg, ) try: yield read_stream, write_stream finally: if transport.session_id and terminate_on_close: await transport.terminate_session(client) tg.cancel_scope.cancel() await resync_tracer()