diff --git a/docs/client/transports.md b/docs/client/transports.md index 6d9d30f90d..728a1746e4 100644 --- a/docs/client/transports.md +++ b/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`: diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py index f72ff273a9..77c2ea58af 100644 --- a/src/mcp/client/sse.py +++ b/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 diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py index 82de50fd05..e6c08d9a0f 100644 --- a/src/mcp/client/streamable_http.py +++ b/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,7 +232,10 @@ 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}") @@ -239,7 +243,7 @@ async def handle_get_stream(self, client: httpx2.AsyncClient, read_stream_writer 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,7 +549,10 @@ 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") @@ -550,7 +560,7 @@ async def _handle_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: diff --git a/src/mcp/shared/_httpx_utils.py b/src/mcp/shared/_httpx_utils.py index 940b9f08cc..1800a61154 100644 --- a/src/mcp/shared/_httpx_utils.py +++ b/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.""" diff --git a/tests/client/test_streamable_http.py b/tests/client/test_streamable_http.py index c6e62ad94a..da68bb0ab9 100644 --- a/tests/client/test_streamable_http.py +++ b/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.""" diff --git a/tests/interaction/transports/test_hosting_resume.py b/tests/interaction/transports/test_hosting_resume.py index 2699ab32f1..df917513db 100644 --- a/tests/interaction/transports/test_hosting_resume.py +++ b/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") diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py index 77d1b28a0a..ef47ca774b 100644 --- a/tests/shared/test_sse.py +++ b/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