-
Notifications
You must be signed in to change notification settings - Fork 3.2k
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
Merged
+117
−12
Merged
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
ee46896
fix: close all memory stream ends in client transport cleanup
maxisbey ce00a76
test: isolate stream leak tests from other tests' resource leaks
maxisbey 6ed6261
test: remove ResourceWarning filter that was masking the stream leak bug
maxisbey 1e83583
refactor: use async with instead of try/finally for stream cleanup
maxisbey 906bcfc
fix: merge stream and task-group contexts to satisfy 3.11 branch cove…
maxisbey d297fcf
revert: use try/finally for sse and streamable_http stream cleanup
maxisbey d375f3e
test: use anyio's free_tcp_port fixture instead of hand-rolled helper
maxisbey 3332cfb
chore: bump anyio minimum to 4.9 for free_tcp_port fixture
maxisbey d532923
test: check exc_value instead of args.object in leak hook
maxisbey File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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.
Show resolved
Hide resolved
|
||
| """sse_client must close all 4 stream ends when the connection fails. | ||
|
|
||
maxisbey marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| 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 | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.