Raise a distinct error when a Streamable HTTP session is lost (HTTP 404) #3799

"""A lost Streamable HTTP session reaches the client as an ordinary MCPError.

Requires: mcp (tested with 2.3.0) and uvicorn. Run: python repro.py

A small ASGI server stands in for an MCP server with sessions. It answers
`initialize` with an `Mcp-Session-Id`, and then each `tools/call` by tool name:

- gone-empty:      HTTP 404, empty body (the server lost the session)
- gone-python-sdk: HTTP 404 with the Python SDK server's body (-32600 "Session not found")
- gone-ts-sdk:     HTTP 404 with the TypeScript SDK server's body (-32001 "Session not found")
- invalid:         HTTP 200 with a JSON-RPC -32600 error (an ordinary invalid request)
- ok:              a normal result

It prints what the client raises for each. Exits 0 when nothing distinguishes a
lost session: every 404 case raises plain MCPError, and the empty-body 404 has
the same type and code as the ordinary invalid request.
"""

import asyncio
import json
import socket

import uvicorn
from mcp import ClientSession, MCPError
from mcp.client.streamable_http import streamable_http_client

SESSION_LOST_BODIES = {
    "gone-empty": None,
    "gone-python-sdk": {"jsonrpc": "2.0", "id": "server-error", "error": {"code": -32600, "message": "Session not found"}},
    "gone-ts-sdk": {"jsonrpc": "2.0", "id": None, "error": {"code": -32001, "message": "Session not found"}},
}


async def app(scope, receive, send):
    async def respond(status, body=None, headers=()):
        payload = b"" if body is None else json.dumps(body).encode()
        hdrs = [(b"content-type", b"application/json"), *headers] if body is not None else list(headers)
        await send({"type": "http.response.start", "status": status, "headers": hdrs})
        await send({"type": "http.response.body", "body": payload})

    if scope["type"] != "http":
        return
    if scope["method"] != "POST":
        return await respond(405 if scope["method"] == "GET" else 200)
    body, more = b"", True
    while more:
        message = await receive()
        body += message.get("body", b"")
        more = message.get("more_body", False)
    rpc = json.loads(body)
    method, rid = rpc.get("method"), rpc.get("id")
    if method == "initialize":
        result = {
            "protocolVersion": rpc["params"]["protocolVersion"],
            "capabilities": {"tools": {}},
            "serverInfo": {"name": "fake", "version": "0"},
        }
        return await respond(200, {"jsonrpc": "2.0", "id": rid, "result": result}, [(b"mcp-session-id", b"session-1")])
    if rid is None:  # notifications
        return await respond(202)
    if method == "tools/list":  # the client lists tools to check a call's result
        tools = [{"name": "ok", "inputSchema": {"type": "object"}}]
        return await respond(200, {"jsonrpc": "2.0", "id": rid, "result": {"tools": tools}})
    if method == "tools/call":
        name = rpc["params"]["name"]
        if name in SESSION_LOST_BODIES:
            return await respond(404, SESSION_LOST_BODIES[name])
        if name == "invalid":
            return await respond(200, {"jsonrpc": "2.0", "id": rid, "error": {"code": -32600, "message": "Invalid request"}})
        result = {"content": [{"type": "text", "text": "ok"}], "isError": False}
        return await respond(200, {"jsonrpc": "2.0", "id": rid, "result": result})
    return await respond(200, {"jsonrpc": "2.0", "id": rid, "error": {"code": -32601, "message": "Method not found"}})


async def call(url: str, tool: str):
    async with streamable_http_client(url) as (read, write), ClientSession(read, write) as session:
        await session.initialize()
        try:
            result = await session.call_tool(tool, {})
            return ("result", result.content[0].text)
        except MCPError as err:
            return (type(err).__name__, err.code, err.error.message)


async def main():
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        port = s.getsockname()[1]
    server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning"))
    task = asyncio.create_task(server.serve())
    while not server.started:
        await asyncio.sleep(0.05)
    url = f"http://127.0.0.1:{port}/mcp"
    seen = {}
    for tool in ["ok", "invalid", *SESSION_LOST_BODIES]:
        seen[tool] = await call(url, tool)
        print(f"{tool:16} -> {seen[tool]}")
    server.should_exit = True
    await task

    assert seen["ok"] == ("result", "ok")
    assert all(seen[t][0] == "MCPError" for t in SESSION_LOST_BODIES), "a lost session got its own type"
    assert seen["gone-empty"][:2] == seen["invalid"][:2], "the empty-body 404 differs from an invalid request"


if __name__ == "__main__":
    asyncio.run(main())