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

This commit is contained in:
Max Isbey 2026-03-13 14:43:54 +00:00 • committed by GitHub
commit e1fd62e0f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 117 additions and 12 deletions

View file

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

View file

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

View file

@ -577,3 +577,5 @@ async def streamable_http_client(
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
await read_stream.aclose()
await write_stream_reader.aclose()

View file

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

View file

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

View file

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

2
uv.lock generated
View file

@ -847,7 +847,7 @@ docs = [
[package.metadata]
requires-dist = [
{ name = "anyio", specifier = ">=4.5" },
{ name = "anyio", specifier = ">=4.9" },
{ name = "httpx", specifier = ">=0.27.1" },
{ name = "httpx-sse", specifier = ">=0.4" },
{ name = "jsonschema", specifier = ">=4.20.0" },