modelcontextprotocol/python-sdk
Watch
1
0
Fork
You've already forked python-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: connect-first stream lifecycle for sse and streamable_http (#2292)

Browse source 
Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
This commit is contained in:
Max Isbey 2026-03-16 23:30:20 +00:00 • committed by GitHub
commit 75a80b6f07
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 161 additions and 153 deletions
Download patch file Download diff file Expand all files Collapse all files

39
src/mcp/client/sse.py
View file

@ -57,23 +57,17 @@ async def sse_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)
async with anyio.create_task_group() as tg:
try:
logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
async with httpx_client_factory(
headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout)
) as client:
async with aconnect_sse(
client,
"GET",
url,
) as event_source:
async with aconnect_sse(client, "GET", url) as event_source:
event_source.response.raise_for_status()
logger.debug("SSE connection established")
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
try:
async for sse in event_source.aiter_sse(): # pragma: no branch
@ -107,9 +101,7 @@ async def sse_client(
if not sse.data:
continue
try:
message = types.jsonrpc_message_adapter.validate_json(
sse.data, by_name=False
)
message = types.jsonrpc_message_adapter.validate_json(sse.data, by_name=False)
logger.debug(f"Received server message: {message}")
except Exception as exc: # pragma: no cover
logger.exception("Error parsing server message") # pragma: no cover
@ -131,7 +123,7 @@ async def sse_client(
async def post_writer(endpoint_url: str):
try:
async with write_stream_reader:
async with write_stream_reader, write_stream:
async for session_message in write_stream_reader:
logger.debug(f"Sending client message: {session_message}")
response = await client.post(
@ -146,19 +138,20 @@ async def sse_client(
logger.debug(f"Client message sent successfully: {response.status_code}")
except Exception: # pragma: lax no cover
logger.exception("Error in post_writer")
finally:
await write_stream.aclose()
# On Python 3.14, coverage.py reports a phantom branch arc on this
# line (->yield) when nested two async-with levels deep. The branch
# is the unreachable "did __aexit__ suppress?" arm for memory streams.
async with ( # pragma: no branch
read_stream_writer,
read_stream,
write_stream,
write_stream_reader,
anyio.create_task_group() as tg,
):
endpoint_url = await tg.start(sse_reader)
logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}")
tg.start_soon(post_writer, endpoint_url)
try:
yield read_stream, write_stream
finally:
tg.cancel_scope.cancel()
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
await read_stream.aclose()
await write_stream_reader.aclose()

26
src/mcp/client/streamable_http.py
View file

@ -440,7 +440,7 @@ class StreamableHTTPTransport:
) -> None:
"""Handle writing requests to the server."""
try:
async with write_stream_reader:
async with write_stream_reader, read_stream_writer, write_stream:
async for session_message in write_stream_reader:
message = session_message.message
metadata = (
@ -480,9 +480,6 @@ class StreamableHTTPTransport:
except Exception: # pragma: lax no cover
logger.exception("Error in post_writer")
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
async def terminate_session(self, client: httpx.AsyncClient) -> None:
"""Terminate the session by sending a DELETE request."""
@ -533,9 +530,6 @@ async def streamable_http_client(
Example:
See examples/snippets/clients/ for usage patterns.
"""
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
# Determine if we need to create and manage the client
client_provided = http_client is not None
client = http_client
@ -546,8 +540,6 @@ async def streamable_http_client(
transport = StreamableHTTPTransport(url)
async with anyio.create_task_group() as tg:
try:
logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
async with contextlib.AsyncExitStack() as stack:
@ -555,6 +547,17 @@ async def streamable_http_client(
if not client_provided:
await stack.enter_async_context(client)
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
async with (
read_stream_writer,
read_stream,
write_stream,
write_stream_reader,
anyio.create_task_group() as tg,
):
def start_get_stream() -> None:
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
@ -574,8 +577,3 @@ async def streamable_http_client(
if transport.session_id and terminate_on_close:
await transport.terminate_session(client)
tg.cancel_scope.cancel()
finally:
await read_stream_writer.aclose()
await write_stream.aclose()
await read_stream.aclose()
await write_stream_reader.aclose()

37
tests/client/test_transport_stream_cleanup.py
View file

@ -58,22 +58,39 @@ def _assert_no_memory_stream_leak() -> Iterator[None]:
@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.
"""sse_client creates streams only after the SSE connection succeeds, so a
ConnectError propagates directly with nothing to leak.
Before the fix, only read_stream_writer and write_stream were closed in
the finally block. read_stream and write_stream_reader were leaked.
Before the fix, streams were created before connecting and only 2 of 4 were
closed in the finally block.
"""
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
with pytest.raises(httpx.ConnectError):
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_sse_client_closes_all_streams_on_http_error() -> None:
"""sse_client creates streams only after raise_for_status() passes, so an
HTTPStatusError from a 4xx/5xx response propagates bare (not wrapped in an
ExceptionGroup) with nothing to leak — the task group is never entered.
"""
def return_403(request: httpx.Request) -> httpx.Response:
return httpx.Response(403)
def mock_factory(
headers: dict[str, str] | None = None,
timeout: httpx.Timeout | None = None,
auth: httpx.Auth | None = None,
) -> httpx.AsyncClient:
return httpx.AsyncClient(transport=httpx.MockTransport(return_403))
with _assert_no_memory_stream_leak():
with pytest.raises(httpx.HTTPStatusError):
async with sse_client("http://test/sse", httpx_client_factory=mock_factory):
pytest.fail("should not reach here") # pragma: no cover
@pytest.mark.anyio