Skip to content

Commit ad1fcb9

Browse files
committed
refactor: fold post_writer signaling closes into async with
Both post_writer tasks (sse, streamable_http) used a try/finally to close signaling stream ends after the loop. Since the except handlers only log and don't touch those streams, the closes can move into the existing async with header — __aexit__ runs before except, which is fine when the handler doesn't send on the closed stream. sse_reader can't do the same: its except Exception handler sends on read_stream_writer (L120), so it must stay open until after the handler runs.
1 parent 473a678 commit ad1fcb9

File tree

2 files changed

+2
-7
lines changed

2 files changed

+2
-7
lines changed

src/mcp/client/sse.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
123123

124124
async def post_writer(endpoint_url: str):
125125
try:
126-
async with write_stream_reader:
126+
async with write_stream_reader, write_stream:
127127
async for session_message in write_stream_reader:
128128
logger.debug(f"Sending client message: {session_message}")
129129
response = await client.post(
@@ -138,8 +138,6 @@ async def post_writer(endpoint_url: str):
138138
logger.debug(f"Client message sent successfully: {response.status_code}")
139139
except Exception: # pragma: lax no cover
140140
logger.exception("Error in post_writer")
141-
finally:
142-
await write_stream.aclose()
143141

144142
# On Python 3.14, coverage.py reports a phantom branch arc on this
145143
# line (->yield) when nested two async-with levels deep. The branch

src/mcp/client/streamable_http.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -440,7 +440,7 @@ async def post_writer(
440440
) -> None:
441441
"""Handle writing requests to the server."""
442442
try:
443-
async with write_stream_reader:
443+
async with write_stream_reader, read_stream_writer, write_stream:
444444
async for session_message in write_stream_reader:
445445
message = session_message.message
446446
metadata = (
@@ -480,9 +480,6 @@ async def handle_request_async():
480480

481481
except Exception: # pragma: lax no cover
482482
logger.exception("Error in post_writer")
483-
finally:
484-
await read_stream_writer.aclose()
485-
await write_stream.aclose()
486483

487484
async def terminate_session(self, client: httpx.AsyncClient) -> None:
488485
"""Terminate the session by sending a DELETE request."""

0 commit comments

Comments
 (0)