"""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 anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream

from httpx_sse import EventSource, ServerSentEvent, aconnect_sse

from pydantic import ValidationError

from mcp.client._transport import TransportStreams

from mcp.shared._httpx_utils import create_mcp_http_client

from mcp.shared.message import ClientMessageMetadata, SessionMessage

from mcp.types import (

INTERNAL_ERROR,

INVALID_REQUEST,

PARSE_ERROR,

ErrorData,

InitializeResult,

JSONRPCError,

JSONRPCMessage,

JSONRPCNotification,

JSONRPCRequest,

JSONRPCResponse,

RequestId,

jsonrpc_message_adapter,

)

logger = logging.getLogger(__name__)

# TODO(Marcelo): Put the TransportStreams in a module under shared, so we can import here.

SessionMessageOrError = SessionMessage | Exception

StreamWriter = MemoryObjectSendStream[SessionMessageOrError]

StreamReader = MemoryObjectReceiveStream[SessionMessage]

MCP_SESSION_ID = "mcp-session-id"

MCP_PROTOCOL_VERSION = "mcp-protocol-version"

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

self.protocol_version: str | None = None

def _prepare_headers(self) -> dict[str, str]:

"""Build MCP-specific request headers.

These headers will be merged with the httpx.AsyncClient's default headers,

with these MCP-specific headers taking precedence.

"""

headers: dict[str, str] = {

"accept": "application/json, text/event-stream",

"content-type": "application/json",

}

# Add session headers if available

if self.session_id:

headers[MCP_SESSION_ID] = self.session_id

if self.protocol_version:

headers[MCP_PROTOCOL_VERSION] = self.protocol_version

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}")

def _maybe_extract_protocol_version_from_message(self, message: JSONRPCMessage) -> None:

"""Extract protocol version from initialization response message."""

if isinstance(message, JSONRPCResponse) and message.result: # pragma: no branch

try:

# Parse the result as InitializeResult for type safety

init_result = InitializeResult.model_validate(message.result, by_name=False)

self.protocol_version = str(init_result.protocol_version)

logger.info(f"Negotiated protocol version: {self.protocol_version}")

except Exception: # pragma: no cover

logger.warning("Failed to parse initialization response as InitializeResult", exc_info=True)

logger.warning(f"Raw result: {message.result}")

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,

is_initialization: bool = False,

) -> 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}")

# Extract protocol version from initialization response

if is_initialization:

self._maybe_extract_protocol_version_from_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)

except Exception as exc: # pragma: 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: # pragma: lax no cover

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."""

headers = self._prepare_headers()

message = ctx.session_message.message

is_initialization = self._is_initialization_request(message)

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 == 404: # pragma: no branch

if isinstance(message, JSONRPCRequest): # pragma: no branch

error_data = ErrorData(code=INVALID_REQUEST, message="Session terminated")

session_message = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data))

await ctx.read_stream_writer.send(session_message)

return

if response.status_code >= 400:

if isinstance(message, JSONRPCRequest):

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, is_initialization, request_id=message.id

)

elif content_type.startswith("text/event-stream"):

await self._handle_sse_response(response, ctx, is_initialization)

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,

is_initialization: bool = False,

*,

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)

# Extract protocol version from initialization response

if is_initialization:

self._maybe_extract_protocol_version_from_message(message)

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,

is_initialization: bool = False,

) -> 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),

is_initialization=is_initialization,

)

# 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: 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: MemoryObjectSendStream[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 for session_message in write_stream_reader:

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()

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: # pragma: lax no cover

return

try:

headers = self._prepare_headers()

response = await client.delete(self.url, headers=headers)

if response.status_code == 405: # pragma: lax no cover

logger.debug("Server does not allow session termination")

elif response.status_code not in (200, 204): # pragma: lax no cover

logger.warning(f"Session termination failed: {response.status_code}")

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 = anyio.create_memory_object_stream[SessionMessage | Exception](0)

write_stream, write_stream_reader = anyio.create_memory_object_stream[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()