Close client SSE iterators explicitly when supported

Author: Kludex

docs/client/transports.md

@@ -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`:

src/mcp/client/sse.py

@@ -16,6 +16,7 @@
     McpHttpClientFactory,
     create_mcp_http_client,
     request_within_origin,
+    sse_events,
     sse_within_origin,
 )
 from mcp.shared.message import SessionMessage
@@ -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

src/mcp/client/streamable_http.py

@@ -37,6 +37,7 @@
     create_mcp_http_client,
     redirect_location,
     request_within_origin,
+    sse_events,
     sse_within_origin,
     stream_within_origin,
 )
@@ -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
@@ -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
@@ -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,
@@ -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
 
@@ -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:

src/mcp/shared/_httpx_utils.py

@@ -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"]
 
@@ -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()
+
+
 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."""

tests/client/test_streamable_http.py

@@ -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"),
     [
@@ -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."""

tests/interaction/transports/test_hosting_resume.py

@@ -10,6 +10,7 @@
 """
 
 import json
+from collections.abc import AsyncGenerator
 
 import anyio
 import httpx2
@@ -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")

tests/shared/test_sse.py

@@ -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