streamable_http (legacy mode): early response.aclose() discards a TCP connection per JSON-RPC exchange — deterministic reproduction for #2707 #3281

"""Repro: mcp 2.0.0 streamable_http legacy-mode client uses a fresh TCP
connection for every JSON-RPC exchange (early response.aclose() before the
SSE body is drained; follow-up to issue #2707 / PR #2712).

Run:  python repro_connection_reuse.py
Deps: pip install "mcp==2.0.0" uvicorn httpx
"""

import asyncio
import json
import threading
import time

import httpx
import uvicorn

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

PORT = 8971

server = MCPServer(name="repro", version="1.0.0")


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


class StreamIdTransport(httpx.AsyncBaseTransport):
    """Log the identity of the network stream used by each request.

    Two requests sharing a TCP connection see the same network_stream
    object; a distinct id per request means no connection reuse.
    """

    def __init__(self):
        self.inner = httpx.AsyncHTTPTransport()
        self.log = []

    async def handle_async_request(self, request):
        resp = await self.inner.handle_async_request(request)
        try:
            method = json.loads(request.content).get("method", request.method)
        except Exception:
            method = request.method
        self.log.append((method, id(resp.extensions.get("network_stream"))))
        return resp

    async def aclose(self):
        await self.inner.aclose()


async def run(mode: str) -> None:
    t = StreamIdTransport()
    async with httpx.AsyncClient(transport=t, timeout=30) as hc:
        async with Client(
            streamable_http_client(f"http://127.0.0.1:{PORT}/mcp", http_client=hc),
            mode=mode,
        ) as client:
            await client.list_tools()
    ids = [i for _, i in t.log]
    print(
        f"mode={mode!r}: {len(t.log)} requests "
        f"{[m for m, _ in t.log]} -> "
        f"{len(set(ids))} distinct TCP connection(s)"
    )


async def raw_httpx_control() -> None:
    """Same server, same three JSON-RPC exchanges, plain httpx: reuses."""
    t = StreamIdTransport()
    headers = {
        "Accept": "application/json, text/event-stream",
        "Content-Type": "application/json",
    }
    async with httpx.AsyncClient(transport=t, timeout=30) as hc:
        r = await hc.post(
            f"http://127.0.0.1:{PORT}/mcp",
            json={
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {},
                    "clientInfo": {"name": "raw", "version": "1"},
                },
            },
            headers=headers,
        )
        await r.aread()
        h2 = dict(headers, **{"Mcp-Session-Id": r.headers.get("mcp-session-id")})
        r = await hc.post(
            f"http://127.0.0.1:{PORT}/mcp",
            json={"jsonrpc": "2.0", "method": "notifications/initialized"},
            headers=h2,
        )
        await r.aread()
        r = await hc.post(
            f"http://127.0.0.1:{PORT}/mcp",
            json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
            headers=h2,
        )
        await r.aread()
    ids = [i for _, i in t.log]
    print(
        f"raw httpx control: {len(t.log)} requests -> "
        f"{len(set(ids))} distinct TCP connection(s)"
    )


def main() -> None:
    config = uvicorn.Config(
        server.streamable_http_app(), host="127.0.0.1", port=PORT, log_level="error"
    )
    srv = uvicorn.Server(config)
    th = threading.Thread(target=srv.run, daemon=True)
    th.start()
    time.sleep(1.5)
    asyncio.run(run("legacy"))
    asyncio.run(run("auto"))
    asyncio.run(raw_httpx_control())
    srv.should_exit = True
    th.join(timeout=3)


if __name__ == "__main__":
    main()