Uh oh!
There was an error while loading. Please reload this page.
- Notifications
You must be signed in to change notification settings - Fork 3.8k
fix: close all memory stream ends in client transport cleanup#2266
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Uh oh!
There was an error while loading. Please reload this page.
Changes from all commits
ee46896ce00a766ed62611e83583906bcfcd297fcfd375f3e3332cfbd532923File filter
Filter by extension
Conversations
Uh oh!
There was an error while loading. Please reload this page.
Jump to
Uh oh!
There was an error while loading. Please reload this page.
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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: | ||
maxisbey marked this conversation as resolved.
Uh oh!There was an error while loading. Please reload this page. | ||
| """sse_client must close all 4 stream ends when the connection fails. | ||
maxisbey marked this conversation as resolved.
Uh oh!There was an error while loading. Please reload this page. | ||
| 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 | ||
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🟣 Not blocking — follow-up note. The same pattern this PR fixes exists in the server-side transports:
server/stdio.py,server/websocket.py, andserver/sse.pyeach create 4 stream ends but only close 2 (via innerasync within reader/writer tasks), with no outer try/finally —read_streamandwrite_streamare never closed by the transport itself. Ironically,server/stdio.pydoesn't follow theclient/stdio.pypattern this PR cites as the correct reference. This likely doesn't cause the current CI flakiness (server tests run in subprocesses, andtests/server/test_stdio.pymanually closes both streams viaasync with), but it's worth a follow-up issue to apply the same fix on the server side for consistency.Extended reasoning...
Summary
The server-side transport files have the exact same memory-stream leak pattern that this PR fixes on the client side. This is a pre-existing issue, not introduced by this PR, and does not need to block the PR — but it's directly relevant context that a reviewer should know about when approving a fix scoped to half of the affected code.
Affected files and the pattern
I verified each server transport against the code:
src/mcp/server/stdio.py(lines 52–83): Creates 4 stream ends via twocreate_memory_object_streamcalls. Thestdin_readertask usesasync with read_stream_writer:andstdout_writerusesasync with write_stream_reader:, so 2 ends are closed when the tasks complete. Butread_streamandwrite_streamare yielded to the caller and never closed by the transport. There is no outertry/finally.src/mcp/server/websocket.py(lines 28–58): Identical structure.ws_readerclosesread_stream_writer,ws_writercloseswrite_stream_reader.read_streamandwrite_streamare never closed. No outertry/finally.src/mcp/server/sse.py: Creates 6 stream ends (4 + an SSE pair). A subset is closed by inner task bodies;read_stream,write_stream, andsse_stream_readerare never explicitly closed.Comparison with the "correct" pattern
The PR description cites
client/stdio.pyas the reference correct pattern. Lines 208–211 ofclient/stdio.pyclose all 4 ends in thefinallyblock:Ironically,
server/stdio.py— the file most directly parallel toclient/stdio.py— does not follow this pattern.Why this doesn't cause current CI flakiness (addressing the counter-argument)
One might expect server-side leaks to cause the same knock-on
ResourceWarningfailures the PR describes. I investigated and this is largely not the case in the current test suite, for two reasons:Subprocess isolation:
tests/shared/test_ws.pyandtests/shared/test_sse.pyrun the server viamultiprocessing.Process(see theserverfixture intests/shared/test_sse.pyaround line 119). Leaked memory streams live in the subprocess's memory and die with the subprocess whenproc.kill()is called — they never trigger aResourceWarningin the pytest process.Manual close in test code:
tests/server/test_stdio.pyrunsstdio_serverin-process, but the test explicitly wraps the yielded streams:async with read_stream:(line 30) andasync with write_stream:(line 49). So the test compensates for the transport's missing cleanup.So the CI evidence cited in the PR description (client-side
test_basic_resources→test_tool_progresscascade) is genuinely a client-side problem, and this PR fully addresses the observed flakiness.Why it's still worth a follow-up
async withthe yielded streams, the leak surfaces.Recommended fix (for the follow-up, not this PR)
Same pattern as this PR — add the missing
aclose()calls in afinallyblock, or use the compoundasync withapproach fromclient/websocket.py. Since anyio'saclose()is idempotent, double-close is safe when the caller has already closed its end.