diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py index 226b0fecf9..780ae4ed3f 100644 --- a/src/mcp/client/streamable_http.py +++ b/src/mcp/client/streamable_http.py @@ -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 if response.status_code >= 400: @@ -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, @@ -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, @@ -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 diff --git a/tests/client/test_streamable_http.py b/tests/client/test_streamable_http.py index d21f520daf..d34c420c91 100644 --- a/tests/client/test_streamable_http.py +++ b/tests/client/test_streamable_http.py @@ -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)