Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions docs/client/transports.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@ That is the whole production client. `Client` wraps the URL in `streamable_http_

Nothing was resolved, fetched or spawned when you wrote `Client("http://...")`. That line is free.

!!! note "Event-stream cleanup"
Streamable HTTP and legacy SSE close each event iterator when they stop reading, if it
exposes `aclose()`. This runs its cleanup on early return or cancellation instead of leaving
it to garbage collection. HTTPX2 remains responsible for nested response iterators.

### Bring your own `httpx2.AsyncClient`

The moment you need an `Authorization` header, a cookie, a proxy, mTLS, or a different timeout, build the `httpx2.AsyncClient` yourself and hand it to `streamable_http_client`:
Expand Down
86 changes: 44 additions & 42 deletions src/mcp/client/sse.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
McpHttpClientFactory,
create_mcp_http_client,
request_within_origin,
sse_events,
sse_within_origin,
)
from mcp.shared.message import SessionMessage
Expand Down Expand Up @@ -76,48 +77,49 @@ async def sse_client(

async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
try:
async for sse in event_source: # pragma: no branch
logger.debug(f"Received SSE event: {sse.event}")
match sse.event:
case "endpoint":
endpoint_url = urljoin(url, sse.data)
logger.debug(f"Received endpoint URL: {endpoint_url}")

url_parsed = urlparse(url)
endpoint_parsed = urlparse(endpoint_url)
if ( # pragma: no cover
url_parsed.netloc != endpoint_parsed.netloc
or url_parsed.scheme != endpoint_parsed.scheme
):
error_msg = ( # pragma: no cover
f"Endpoint origin does not match connection origin: {endpoint_url}"
)
logger.error(error_msg) # pragma: no cover
raise ValueError(error_msg) # pragma: no cover

if on_session_created:
session_id = _extract_session_id_from_endpoint(endpoint_url)
if session_id:
on_session_created(session_id)

task_status.started(endpoint_url)

case "message":
# Skip empty data (keep-alive pings)
if not sse.data:
continue
try:
message = types.jsonrpc_message_adapter.validate_json(sse.data, by_name=False)
logger.debug(f"Received server message: {message}")
except Exception as exc: # pragma: no cover
logger.exception("Error parsing server message") # pragma: no cover
await read_stream_writer.send(exc) # pragma: no cover
continue # pragma: no cover

session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
case _: # pragma: no cover
logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover
async with sse_events(event_source) as events:
async for sse in events: # pragma: no branch
logger.debug(f"Received SSE event: {sse.event}")
match sse.event:
case "endpoint":
endpoint_url = urljoin(url, sse.data)
logger.debug(f"Received endpoint URL: {endpoint_url}")

url_parsed = urlparse(url)
endpoint_parsed = urlparse(endpoint_url)
if ( # pragma: no cover
url_parsed.netloc != endpoint_parsed.netloc
or url_parsed.scheme != endpoint_parsed.scheme
):
error_msg = ( # pragma: no cover
f"Endpoint origin does not match connection origin: {endpoint_url}"
)
logger.error(error_msg) # pragma: no cover
raise ValueError(error_msg) # pragma: no cover

if on_session_created:
session_id = _extract_session_id_from_endpoint(endpoint_url)
if session_id:
on_session_created(session_id)

task_status.started(endpoint_url)

case "message":
# Skip empty data (keep-alive pings)
if not sse.data:
continue
try:
message = types.jsonrpc_message_adapter.validate_json(sse.data, by_name=False)
logger.debug(f"Received server message: {message}")
except Exception as exc: # pragma: no cover
logger.exception("Error parsing server message") # pragma: no cover
await read_stream_writer.send(exc) # pragma: no cover
continue # pragma: no cover

session_message = SessionMessage(message)
await read_stream_writer.send(session_message)
case _: # pragma: no cover
logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover
except SSEError as sse_exc: # pragma: lax no cover
logger.exception("Encountered SSE exception")
raise sse_exc
Expand Down
60 changes: 35 additions & 25 deletions src/mcp/client/streamable_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
create_mcp_http_client,
redirect_location,
request_within_origin,
sse_events,
sse_within_origin,
stream_within_origin,
)
Expand Down Expand Up @@ -231,15 +232,18 @@ async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer
if last_event_id:
headers[LAST_EVENT_ID] = last_event_id

async with sse_within_origin(client, self.url, headers=headers) as event_source:
async with (
sse_within_origin(client, self.url, headers=headers) as event_source,
sse_events(event_source) as events,
):
if (redirect := _unfollowed_redirect(event_source.response)) is not None:
# The same GET would be redirected again, so retrying cannot help.
logger.warning(f"GET stream not opened: {redirect}")
return
event_source.response.raise_for_status()
logger.debug("GET SSE connection established")

async for sse in event_source:
async for sse in events:
# Track last event ID for reconnection
if sse.id:
last_event_id = sse.id
Expand Down Expand Up @@ -278,7 +282,10 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
if isinstance(ctx.session_message.message, JSONRPCRequest): # pragma: no branch
original_request_id = ctx.session_message.message.id

async with sse_within_origin(ctx.client, self.url, headers=headers) as event_source:
async with (
sse_within_origin(ctx.client, self.url, headers=headers) as event_source,
sse_events(event_source) as events,
):
if (redirect := _unfollowed_redirect(event_source.response)) is not None:
logger.warning(redirect)
assert original_request_id is not None
Expand All @@ -289,7 +296,7 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
event_source.response.raise_for_status()
logger.debug("Resumption GET SSE connection established")

async for sse in event_source: # pragma: no branch
async for sse in events: # pragma: no branch
is_complete = await self._handle_sse_event(
sse,
ctx.read_stream_writer,
Expand Down Expand Up @@ -464,27 +471,27 @@ async def _handle_sse_response(
original_request_id = ctx.session_message.message.id

try:
event_source = EventSource(response)
async for sse in event_source: # pragma: no branch
# Track last event ID for potential reconnection
if sse.id:
last_event_id = sse.id
async with sse_events(EventSource(response)) as events:
async for sse in events: # 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
# 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
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

Expand Down Expand Up @@ -542,15 +549,18 @@ async def _handle_reconnection(
headers[LAST_EVENT_ID] = last_event_id

try:
async with sse_within_origin(ctx.client, self.url, headers=headers) as event_source:
async with (
sse_within_origin(ctx.client, self.url, headers=headers) as event_source,
sse_events(event_source) as events,
):
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:
async for sse in events:
if sse.id: # pragma: no branch
reconnect_last_event_id = sse.id
if sse.retry is not None:
Expand Down
21 changes: 19 additions & 2 deletions src/mcp/shared/_httpx_utils.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
"""Utilities for creating and using httpx2 AsyncClient instances in the MCP transports."""

from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, AsyncIterator
from contextlib import asynccontextmanager
from typing import Any, Protocol
from typing import Any

import httpx2
from typing_extensions import Protocol, runtime_checkable

__all__ = ["create_mcp_http_client", "MCP_DEFAULT_TIMEOUT", "MCP_DEFAULT_SSE_READ_TIMEOUT"]

Expand Down Expand Up @@ -165,6 +166,22 @@ async def sse_within_origin(
yield httpx2.EventSource(response)


@runtime_checkable
class _AsyncClosable(Protocol):
async def aclose(self) -> None: ...


@asynccontextmanager
async def sse_events(source: httpx2.EventSource) -> AsyncGenerator[AsyncIterator[httpx2.ServerSentEvent]]:
"""Close the outer EventSource iterator if supported; HTTPX2 owns its nested iterators."""
events = source.__aiter__()
try:
yield events
finally:
if isinstance(events, _AsyncClosable):
await events.aclose()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Shield iterator cleanup from task cancellation

When a reader is cancelled—such as during transport teardown or modern request cancellation—AnyIO's level cancellation remains active while this finally block runs, so an aclose() implementation that reaches an async checkpoint is cancelled before cleanup completes. The new tests miss this because their closers perform no checkpoint, while real iterator cleanup commonly awaits nested resources; repeated cancellations can therefore retain stream resources despite the newly documented guarantee. Run the close operation inside a shielded cancellation scope, ideally with an appropriate bound.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: When an SSE reader is cancelled, sse_events awaits aclose() inside the already-cancelled AnyIO scope, so asynchronous iterator cleanup can be cancelled and the response resources can leak. Run events.aclose() inside a shielded anyio.CancelScope.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/mcp/shared/_httpx_utils.py, line 182:

<comment>When an SSE reader is cancelled, `sse_events` awaits `aclose()` inside the already-cancelled AnyIO scope, so asynchronous iterator cleanup can be cancelled and the response resources can leak. Run `events.aclose()` inside a shielded `anyio.CancelScope`.</comment>

<file context>
@@ -165,6 +166,22 @@ async def sse_within_origin(
+        yield events
+    finally:
+        if isinstance(events, _AsyncClosable):
+            await events.aclose()
+
+
</file context>



def redirect_location(response: httpx2.Response) -> httpx2.URL | None:
"""Where `response` redirects to, for use in a message: without userinfo, query or fragment,
which can carry state that does not belong in an error or a log line. None if not a redirect."""
Expand Down
79 changes: 79 additions & 0 deletions tests/client/test_streamable_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,11 @@
from tests.shared.test_dispatcher import Recorder, echo_handlers


@pytest.fixture(autouse=True)
def _module_runner_lease() -> None:
"""Opt out of the shared runner because iterator cleanup parametrizes `anyio_backend`."""


@pytest.mark.parametrize(
("raw", "expected", "wrapped"),
[
Expand Down Expand Up @@ -917,6 +922,80 @@ def handler(request: httpx2.Request) -> httpx2.Response:
assert seen == [("GET http://test/mcp", "evt-41")]


@pytest.mark.anyio
@pytest.mark.parametrize("iterator_kind", ["generator", "closable", "plain"])
@pytest.mark.parametrize("anyio_backend", ["asyncio", "trio"])
async def test_resumed_response_accepts_async_iterators_and_closes_them_when_supported(
monkeypatch: pytest.MonkeyPatch, iterator_kind: str
) -> None:
"""SDK-defined: resumption accepts any EventSource async iterator and closes it when supported.

Substitute the public iterator boundary to isolate representation from HTTPX2's nested-generator cleanup.
"""
expected = JSONRPCResponse(jsonrpc="2.0", id="resume-1", result={"ok": True})
event = httpx2.ServerSentEvent(data=expected.model_dump_json(by_alias=True))
body = f"data: {event.data}\n\n"
closed: list[bool] = []

async def generate() -> AsyncIterator[httpx2.ServerSentEvent]:
try:
yield event
finally:
closed.append(True)

class EventIterator:
def __aiter__(self) -> AsyncIterator[httpx2.ServerSentEvent]:
return self

async def __anext__(self) -> httpx2.ServerSentEvent:
return event

class ClosingEventIterator:
def __aiter__(self) -> AsyncIterator[httpx2.ServerSentEvent]:
return self

async def __anext__(self) -> httpx2.ServerSentEvent:
return event

async def aclose(self) -> None:
closed.append(True)

iterators: dict[str, AsyncIterator[httpx2.ServerSentEvent]] = {
"generator": generate(),
"closable": ClosingEventIterator(),
"plain": EventIterator(),
}

def iterate(source: httpx2.EventSource) -> AsyncIterator[httpx2.ServerSentEvent]:
assert source.response.text == body
return iterators[iterator_kind]

monkeypatch.setattr(httpx2.EventSource, "__aiter__", iterate)
token = "evt-41"

def handler(request: httpx2.Request) -> httpx2.Response:
assert request.method == "GET"
assert request.headers["last-event-id"] == token
return httpx2.Response(200, headers={"content-type": "text/event-stream"}, text=body)

with anyio.fail_after(5):
async with (
httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as http,
streamable_http_client("http://test/mcp", http_client=http) as (read, write),
):
await write.send(
SessionMessage(
JSONRPCRequest(jsonrpc="2.0", id=expected.id, method="tools/call", params={}),
metadata=ClientMessageMetadata(resumption_token=token),
)
)
reply = await read.receive()

assert isinstance(reply, SessionMessage)
assert reply.message == expected
assert closed == ([] if iterator_kind == "plain" else [True])


async def _redirected_call_error(url: str, location: str) -> str:
"""Send one request through streamable_http_client to a server answering `url` with a 307 to
`location`, and return the message of the error that resolves it."""
Expand Down
9 changes: 7 additions & 2 deletions tests/interaction/transports/test_hosting_resume.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
"""

import json
from collections.abc import AsyncGenerator

import anyio
import httpx2
Expand Down Expand Up @@ -70,9 +71,13 @@ def _tools_call(request_id: int, name: str, arguments: dict[str, object]) -> str


async def _read_events(response: httpx2.Response, count: int) -> list[ServerSentEvent]:
"""Read exactly `count` SSE events from a streaming response without closing it."""
"""Read exactly `count` SSE events and close the iterator."""
source = aiter(EventSource(response))
return [await anext(source) for _ in range(count)]
try:
return [await anext(source) for _ in range(count)]
finally:
assert isinstance(source, AsyncGenerator)
await source.aclose()


@requirement("hosting:resume:event-ids")
Expand Down
8 changes: 6 additions & 2 deletions tests/shared/test_sse.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,12 @@ async def test_raw_sse_connection() -> None:
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

lines = response.aiter_lines()
assert await anext(lines) == "event: endpoint"
assert (await anext(lines)).startswith("data: /messages/?session_id=")
try:
assert await anext(lines) == "event: endpoint"
assert (await anext(lines)).startswith("data: /messages/?session_id=")
finally:
assert isinstance(lines, AsyncGenerator)
await lines.aclose()


@pytest.mark.anyio
Expand Down
Loading