fix: close all memory stream ends in client transport cleanup (#2266)
This commit is contained in:
parent
2c73a2a881
commit
e1fd62e0f3
7 changed files with 117 additions and 12 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
102
tests/client/test_transport_stream_cleanup.py
Normal file
102
tests/client/test_transport_stream_cleanup.py
Normal 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
|
||||
|
|
@ -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
2
uv.lock
generated
|
|
@ -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" },
|
||||
|
|
|
|||
Loading…
Reference in a new issue