refactor: use async with instead of try/finally for stream cleanup

Replace explicit try/finally + aclose() chains with async with on all
4 stream ends. Memory stream context managers are idempotent and have
no checkpoints in __aexit__, so this is semantically identical to the
try/finally form with the same teardown ordering — but the ownership
is stated once at creation time and can't drift.

For websocket_client, also move stream creation inside ws_connect so
a connection failure never creates streams in the first place.

Matches the existing pattern in shared/memory.py.

Author: maxisbey

src/mcp/client/sse.py

@@ -61,7 +61,7 @@ async def sse_client(
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
     async with anyio.create_task_group() as tg:
-        try:
+        async with read_stream_writer, read_stream, write_stream, write_stream_reader:
             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)
@@ -157,8 +157,3 @@ async def post_writer(endpoint_url: str):
                         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()

src/mcp/client/streamable_http.py

@@ -547,7 +547,7 @@ async def streamable_http_client(
     transport = StreamableHTTPTransport(url)
 
     async with anyio.create_task_group() as tg:
-        try:
+        async with read_stream_writer, read_stream, write_stream, write_stream_reader:
             logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
 
             async with contextlib.AsyncExitStack() as stack:
@@ -574,8 +574,3 @@ def start_get_stream() -> None:
                     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()

src/mcp/client/websocket.py

@@ -38,12 +38,12 @@ 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)
 
-    try:
-        # Connect using websockets, requesting the "mcp" subprotocol
-        async with ws_connect(url, subprotocols=[Subprotocol("mcp")]) as ws:
+        async with read_stream_writer, read_stream, write_stream, write_stream_reader:
 
             async def ws_reader():
                 """Reads text messages from the WebSocket, parses them as JSON-RPC messages,
@@ -79,8 +79,3 @@ async def ws_writer():
 
                 # Once the caller's 'async with' block exits, we shut down
                 tg.cancel_scope.cancel()
-    finally:
-        await read_stream_writer.aclose()
-        await write_stream.aclose()
-        await read_stream.aclose()
-        await write_stream_reader.aclose()