fix: close all memory stream ends in client transport cleanup (#2266)

Author: maxisbey

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
     "Programming Language :: Python :: 3.14",
 ]
 dependencies = [
-    "anyio>=4.5",
+    "anyio>=4.9",
     "httpx>=0.27.1",
     "httpx-sse>=0.4",
     "pydantic>=2.12.0",

src/mcp/client/sse.py

@@ -160,3 +160,5 @@ async def post_writer(endpoint_url: str):
         finally:
             await read_stream_writer.aclose()
             await write_stream.aclose()
+            await read_stream.aclose()
+            await write_stream_reader.aclose()

src/mcp/client/streamable_http.py

@@ -577,3 +577,5 @@ def start_get_stream() -> None:
         finally:
             await read_stream_writer.aclose()
             await write_stream.aclose()
+            await read_stream.aclose()
+            await write_stream_reader.aclose()

src/mcp/client/websocket.py

@@ -38,11 +38,10 @@ async def websocket_client(
     write_stream: MemoryObjectSendStream[SessionMessage]
     write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
-    read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
-    write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
-
     # Connect using websockets, requesting the "mcp" subprotocol
     async with ws_connect(url, subprotocols=[Subprotocol("mcp")]) as ws:
+        read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
+        write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
         async def ws_reader():
             """Reads text messages from the WebSocket, parses them as JSON-RPC messages,
@@ -68,7 +67,13 @@ async def ws_writer():
                     msg_dict = session_message.message.model_dump(by_alias=True, mode="json", exclude_unset=True)
                     await ws.send(json.dumps(msg_dict))
 
-        async with anyio.create_task_group() as tg:
+        async with (
+            read_stream_writer,
+            read_stream,
+            write_stream,
+            write_stream_reader,
+            anyio.create_task_group() as tg,
+        ):
             # Start reader and writer tasks
             tg.start_soon(ws_reader)
             tg.start_soon(ws_writer)

tests/client/test_transport_stream_cleanup.py

@@ -0,0 +1,102 @@
+"""Regression tests for memory stream leaks in client transports.
+
+When a connection error occurs (404, 403, ConnectError), transport context
+managers must close ALL 4 memory stream ends they created. anyio memory streams
+are paired but independent — closing the writer does NOT close the reader.
+Unclosed stream ends emit ResourceWarning on GC, which pytest promotes to a
+test failure in whatever test happens to be running when GC triggers.
+
+These tests force GC after the transport context exits, so any leaked stream
+triggers a ResourceWarning immediately and deterministically here, rather than
+nondeterministically in an unrelated later test.
+"""
+
+import gc
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+import httpx
+import pytest
+
+from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamable_http_client
+from mcp.client.websocket import websocket_client
+
+
+@contextmanager
+def _assert_no_memory_stream_leak() -> Iterator[None]:
+    """Fail if any anyio MemoryObject stream emits ResourceWarning during the block.
+
+    Uses a custom sys.unraisablehook to capture ONLY MemoryObject stream leaks,
+    ignoring unrelated resources (e.g. PipeHandle from flaky stdio tests on the
+    same xdist worker). gc.collect() is forced after the block to make leaks
+    deterministic.
+    """
+    leaked: list[str] = []
+    old_hook = sys.unraisablehook
+
+    def hook(args: "sys.UnraisableHookArgs") -> None:  # pragma: no cover
+        # Only executes if a leak occurs (i.e. the bug is present).
+        # args.object is the __del__ function (not the stream instance) when
+        # unraisablehook fires from a finalizer, so check exc_value — the
+        # actual ResourceWarning("Unclosed <MemoryObjectSendStream at ...>").
+        # Non-MemoryObject unraisables (e.g. PipeHandle leaked by an earlier
+        # flaky test on the same xdist worker) are deliberately ignored —
+        # this test should not fail for another test's resource leak.
+        if "MemoryObject" in str(args.exc_value):
+            leaked.append(str(args.exc_value))
+
+    sys.unraisablehook = hook
+    try:
+        yield
+        gc.collect()
+        assert not leaked, f"Memory streams leaked: {leaked}"
+    finally:
+        sys.unraisablehook = old_hook
+
+
+@pytest.mark.anyio
+async def test_sse_client_closes_all_streams_on_connection_error(free_tcp_port: int) -> None:
+    """sse_client must close all 4 stream ends when the connection fails.
+
+    Before the fix, only read_stream_writer and write_stream were closed in
+    the finally block. read_stream and write_stream_reader were leaked.
+    """
+    with _assert_no_memory_stream_leak():
+        # sse_client enters a task group BEFORE connecting, so anyio wraps the
+        # ConnectError from aconnect_sse in an ExceptionGroup.
+        with pytest.raises(Exception) as exc_info:  # noqa: B017
+            async with sse_client(f"http://127.0.0.1:{free_tcp_port}/sse"):
+                pytest.fail("should not reach here")  # pragma: no cover
+
+        assert exc_info.group_contains(httpx.ConnectError)
+        # exc_info holds the traceback → holds frame locals → keeps leaked
+        # streams alive. Must drop it before gc.collect() can detect a leak.
+        del exc_info
+
+
+@pytest.mark.anyio
+async def test_streamable_http_client_closes_all_streams_on_exit() -> None:
+    """streamable_http_client must close all 4 stream ends on exit.
+
+    Before the fix, read_stream was never closed — not even on the happy path.
+    This test enters and exits the context without sending any messages, so no
+    network connection is ever attempted (streamable_http connects lazily).
+    """
+    with _assert_no_memory_stream_leak():
+        async with streamable_http_client("http://127.0.0.1:1/mcp"):
+            pass
+
+
+@pytest.mark.anyio
+async def test_websocket_client_closes_all_streams_on_connection_error(free_tcp_port: int) -> None:
+    """websocket_client must close all 4 stream ends when ws_connect fails.
+
+    Before the fix, there was no try/finally at all — if ws_connect raised,
+    all 4 streams were leaked.
+    """
+    with _assert_no_memory_stream_leak():
+        with pytest.raises(OSError):
+            async with websocket_client(f"ws://127.0.0.1:{free_tcp_port}/ws"):
+                pytest.fail("should not reach here")  # pragma: no cover

tests/shared/test_sse.py

@@ -544,12 +544,6 @@ def test_sse_server_transport_endpoint_validation(endpoint: str, expected_result
         assert sse._endpoint.startswith("/")
 
 
-# ResourceWarning filter: When mocking aconnect_sse, the sse_client's internal task
-# group doesn't receive proper cancellation signals, so the sse_reader task's finally
-# block (which closes read_stream_writer) doesn't execute. This is a test artifact -
-# the actual code path (`if not sse.data: continue`) IS exercised and works correctly.
-# Production code with real SSE connections cleans up properly.
-@pytest.mark.filterwarnings("ignore::ResourceWarning")
 @pytest.mark.anyio
 async def test_sse_client_handles_empty_keepalive_pings() -> None:
     """Test that SSE client properly handles empty data lines (keep-alive pings).