Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 24 additions & 1 deletion src/mcp/client/streamable_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -337,6 +337,7 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
"server answered a request with 202 Accepted",
code=INVALID_REQUEST,
)
await self._drain_response(response)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1: A server that keeps an SSE response stream open after sending the final JSON-RPC reply can now leave this POST task blocked indefinitely, because draining waits for EOF even though stream termination is only a protocol SHOULD. This is particularly problematic for long-lived response streams and can retain the connection/task until a timeout or client shutdown. Consider retaining the previous close-on-completion behavior for streams without a known finite body, and only drain responses whose framing guarantees EOF is immediately available.

Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At src/mcp/client/streamable_http.py, line 340:

<comment>A server that keeps an SSE response stream open after sending the final JSON-RPC reply can now leave this POST task blocked indefinitely, because draining waits for EOF even though stream termination is only a protocol SHOULD. This is particularly problematic for long-lived response streams and can retain the connection/task until a timeout or client shutdown. Consider retaining the previous close-on-completion behavior for streams without a known finite body, and only drain responses whose framing guarantees EOF is immediately available.</comment>

<file context>
@@ -337,6 +337,7 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                         "server answered a request with 202 Accepted",
                         code=INVALID_REQUEST,
                     )
+                await self._drain_response(response)
                 return
 
</file context>

return

if response.status_code >= 400:
Expand Down Expand Up @@ -388,6 +389,10 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
error_data = ErrorData(code=INVALID_REQUEST, message=f"Unexpected content type: {content_type}")
error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=message.id, error=error_data))
await ctx.read_stream_writer.send(error_msg)
else:
# A notification POST has no response body; drain it so the
# connection returns to the pool instead of being discarded.
await self._drain_response(response)

async def _handle_json_response(
self,
Expand All @@ -408,6 +413,24 @@ async def _handle_json_response(
error_msg = SessionMessage(JSONRPCError(jsonrpc="2.0", id=request_id, error=error_data))
await read_stream_writer.send(error_msg)

async def _drain_response(self, response: httpx2.Response) -> None:
"""Consume a response body to EOF so httpx can return the TCP connection
to its pool instead of discarding it. In streamable-HTTP legacy mode
every POST owns a response stream; abandoning it mid-body (the previous
``aclose()`` call) costs one TCP connection per JSON-RPC exchange. A 202
or notification body drains instantly; an SSE body the EventSource
already iterated is drained via the raw stream because ``aiter_raw``
raises ``StreamConsumed`` once iteration has started.
"""
try:
await response.aread()
except httpx2.StreamConsumed:
try:
async for _ in response.stream: # type: ignore[attr-defined]
pass
except Exception: # pragma: lax no cover
logger.debug("failed to drain response stream", exc_info=True)

async def _handle_sse_response(
self,
response: httpx2.Response,
Expand Down Expand Up @@ -442,7 +465,7 @@ async def _handle_sse_response(
# If the SSE event indicates completion, like returning response/error
# break the loop
if is_complete:
await response.aclose()
await self._drain_response(response)
return # Normal completion, no reconnect needed
except Exception:
logger.debug("SSE stream ended", exc_info=True) # pragma: lax no cover
Expand Down
88 changes: 88 additions & 0 deletions tests/client/test_streamable_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -748,3 +748,91 @@ async def test_resolving_an_abandoned_request_after_the_reader_closed_is_contain
_abandoned_request_context(http, send), "evt-7", None, MAX_RECONNECTION_ATTEMPTS
)
send.close()


@pytest.mark.anyio
async def test_legacy_mode_reuses_tcp_connections_across_exchanges() -> None:
"""Regression test for #3281: in streamable-HTTP legacy mode the client must
drain each POST's response body to EOF so httpx returns the TCP connection
to its pool, instead of `aclose()`-ing an unread stream and opening one
connection per JSON-RPC exchange.

Without the drain, every exchange (initialize, initialized notification,
tools/list, DELETE) opens a fresh connection. With it, at least one POST
reuses the previous exchange's connection, so distinct connections < posts.
"""
import json
import socket
import threading
import time

import uvicorn

from mcp.client.client import Client
from mcp.server.mcpserver import MCPServer

server = MCPServer(name="conn-reuse", version="1.0.0")

@server.tool()
def echo(text: str) -> str:
"""Echo a message back verbatim."""
return text

with socket.socket() as s:
s.bind(("127.0.0.1", 0))
port = s.getsockname()[1]

uvicorn_srv = uvicorn.Server(
uvicorn.Config(server.streamable_http_app(), host="127.0.0.1", port=port, log_level="error")
)
thread = threading.Thread(target=uvicorn_srv.run, daemon=True)
thread.start()
try:
# Wait for the server to accept connections.
deadline = time.monotonic() + 10
while time.monotonic() < deadline:
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
break
except OSError:
time.sleep(0.05)

class _TrackingTransport(httpx2.AsyncBaseTransport):
def __init__(self) -> None:
self.inner = httpx2.AsyncHTTPTransport()
self.log: list[tuple[str, int]] = []

async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
resp = await self.inner.handle_async_request(request)
try:
method = json.loads(request.content).get("method", request.method)
except Exception:
method = request.method
stream = resp.extensions.get("network_stream")
self.log.append((method, id(stream) if stream is not None else -1))
return resp

async def aclose(self) -> None:
await self.inner.aclose()

transport = _TrackingTransport()
async with httpx2.AsyncClient(transport=transport, timeout=30) as http:
async with Client(
streamable_http_client(f"http://127.0.0.1:{port}/mcp", http_client=http),
mode="legacy",
) as client:
await client.list_tools()

exchanges = transport.log
distinct = len({conn_id for _, conn_id in exchanges})
# The POST exchanges (initialize, notifications/initialized, tools/list,
# DELETE) must share fewer TCP connections than the number of exchanges.
# A long-lived GET resumption stream is expected to hold its own
# connection, so we only require strict sharing overall.
assert distinct < len(exchanges), (
f"legacy mode opened one connection per exchange ({distinct} distinct "
f"for {len(exchanges)} exchanges): response bodies are not being drained"
)
finally:
uvicorn_srv.should_exit = True
thread.join(timeout=3)
Loading