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
1 changed files with 161 additions and 153 deletions
Download patch file Download diff file Expand all files Collapse all files

70
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,36 +540,40 @@ 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}")
logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
async with contextlib.AsyncExitStack() as stack:
# Only manage client lifecycle if we created it
if not client_provided:
await stack.enter_async_context(client)
async with contextlib.AsyncExitStack() as stack:
# Only manage client lifecycle if we created it
if not client_provided:
await stack.enter_async_context(client)
def start_get_stream() -> None:
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
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)
tg.start_soon(
transport.post_writer,
client,
write_stream_reader,
read_stream_writer,
write_stream,
start_get_stream,
tg,
)
async with (
read_stream_writer,
read_stream,
write_stream,
write_stream_reader,
anyio.create_task_group() as tg,
):
try:
yield read_stream, write_stream
finally:
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()
def start_get_stream() -> None:
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
tg.start_soon(
transport.post_writer,
client,
write_stream_reader,
read_stream_writer,
write_stream,
start_get_stream,
tg,
)
try:
yield read_stream, write_stream
finally:
if transport.session_id and terminate_on_close:
await transport.terminate_session(client)
tg.cancel_scope.cancel()