Streamable HTTP notification sends a second response when the session writer is closed #3651

import json
import logging
import traceback
from functools import partial
from typing import NoReturn

import anyio
from mcp.server.streamable_http import StreamableHTTPServerTransport
from pydantic import BaseModel
from starlette.types import Message, Scope


class RequestState(BaseModel):
  body_delivered: bool


class ResponseState(BaseModel):
  started: bool
  completed: bool


async def receive_request(*, body: bytes, state: RequestState) -> Message:
  if not state.body_delivered:
    state.body_delivered = True
    return {"type": "http.request", "body": body, "more_body": False}

  return {"type": "http.disconnect"}


async def send_response(message: Message, *, state: ResponseState) -> None:
  if not state.started:
    if message["type"] != "http.response.start":
      raise RuntimeError(f"Expected ASGI message 'http.response.start', but got '{message['type']}'.")
    state.started = True
    return

  if not state.completed:
    if message["type"] != "http.response.body":
      raise RuntimeError(f"Expected ASGI message 'http.response.body', but got '{message['type']}'.")
    state.completed = not message.get("more_body", False)
    return

  raise RuntimeError(f"Unexpected ASGI message '{message['type']}' sent, after response already completed.")


def find_closed_resource_error(error: BaseException) -> anyio.ClosedResourceError | None:
  seen: set[int] = set()
  current: BaseException | None = error

  while current is not None and id(current) not in seen:
    if isinstance(current, anyio.ClosedResourceError):
      return current

    seen.add(id(current))
    current = current.__cause__ or current.__context__

  return None


def request_scope() -> Scope:
  return {
    "type": "http",
    "asgi": {"version": "3.0", "spec_version": "2.3"},
    "http_version": "1.1",
    "method": "POST",
    "scheme": "http",
    "path": "/mcp",
    "raw_path": b"/mcp",
    "query_string": b"",
    "root_path": "",
    "headers": [
      (b"accept", b"application/json, text/event-stream"),
      (b"content-type", b"application/json"),
      (b"mcp-protocol-version", b"2025-11-25"),
      (b"mcp-session-id", b"repro-session"),
    ],
    "client": ("127.0.0.1", 50000),
    "server": ("127.0.0.1", 8080),
  }


async def reproduce() -> None:
  logging.getLogger("mcp.server.streamable_http").setLevel(logging.CRITICAL)
  transport = StreamableHTTPServerTransport(mcp_session_id="repro-session")

  # A stateful session can remain addressable briefly after its central stream closes.
  async with transport.connect():
    pass

  body = json.dumps(
    {
      "jsonrpc": "2.0",
      "method": "notifications/initialized",
    }
  ).encode()
  request_state = RequestState(body_delivered=False)
  response_state = ResponseState(started=False, completed=False)

  try:
    await transport.handle_request(
      request_scope(),
      partial(receive_request, body=body, state=request_state),
      partial(send_response, state=response_state),
    )
  except RuntimeError as error:
    closed_resource_error = find_closed_resource_error(error)
    if "after response already completed" not in str(error) or closed_resource_error is None:
      raise

    traceback.print_exception(error)
    print("\nReproduced the ClosedResourceError -> duplicate http.response.start failure.")
    return

  raise AssertionError("The expected ASGI failure was not reproduced")


def main() -> NoReturn:
  anyio.run(reproduce)
  raise SystemExit(0)


if __name__ == "__main__":
  main()