"""An in-process, full-duplex HTTP transport for driving ASGI applications from httpx. `httpx.ASGITransport` runs the application to completion and only then hands the buffered response to the caller, so a server that streams its response — the streamable HTTP transport's SSE responses — can never converse with the client mid-request: a server-initiated request nested inside a still-open call deadlocks. `StreamingASGITransport` removes that limitation by running the application as a background task and forwarding every `http.response.body` chunk to the client the moment it is sent. Everything happens on the one event loop: no sockets, no threads, no sleeps, no extra dependencies. The behavioural contract, pinned by `test_bridge.py`: - The request body is buffered before the application is invoked (MCP requests are small JSON documents); the response streams chunk by chunk. - Closing the response — or the whole client — delivers `http.disconnect` to the application, exactly as a real server sees when its peer goes away. - An exception the application raises before sending `http.response.start` fails the originating request with that same exception. After the response has started, a failure is visible to the client only through the response itself (status code, truncated body) — the same signal a real server over a real socket would give. The transport owns an anyio task group for the application tasks; it is opened and closed by `httpx.AsyncClient`'s own context manager, so use the client as a context manager (the suite always does). Closing the transport cancels every running application task by default; set `cancel_on_close=False` to wait for the application's own disconnect handling instead. """ import math from collections.abc import AsyncIterator from types import TracebackType import anyio import anyio.abc import httpx from anyio.streams.memory import MemoryObjectReceiveStream from starlette.types import ASGIApp, Message, Scope from mcp.shared._compat import resync_tracer class _StreamingResponseBody(httpx.AsyncByteStream): """A response body that yields chunks as the application produces them. Closing it tells the application the client has gone away (`http.disconnect`), mirroring a peer that drops the connection mid-response. """ def __init__(self, chunks: MemoryObjectReceiveStream[bytes], client_disconnected: anyio.Event) -> None: self._chunks = chunks self._client_disconnected = client_disconnected async def __aiter__(self) -> AsyncIterator[bytes]: async for chunk in self._chunks: yield chunk async def aclose(self) -> None: self._client_disconnected.set() await self._chunks.aclose() class StreamingASGITransport(httpx.AsyncBaseTransport): """Drive an ASGI application in-process, streaming each response as it is produced. With `cancel_on_close` (the default), closing the transport cancels every application task still running so harness teardown can never hang. Setting it to False makes the transport wait for the application's own disconnect handling to complete instead, which is the path the legacy SSE server transport relies on for resource cleanup. """ _task_group: anyio.abc.TaskGroup def __init__(self, app: ASGIApp, *, cancel_on_close: bool = True) -> None: self._app = app self._cancel_on_close = cancel_on_close async def __aenter__(self) -> "StreamingASGITransport": self._task_group = anyio.create_task_group() await self._task_group.__aenter__() return self async def __aexit__( self, exc_type: type[BaseException] | None = None, exc_value: BaseException | None = None, traceback: TracebackType | None = None, ) -> None: # httpx closes every streamed response before closing the transport, so by now each # application task has been delivered `http.disconnect`. Either cancel immediately, or wait # for the application's own disconnect handling to unwind. if self._cancel_on_close: self._task_group.cancel_scope.cancel() await self._task_group.__aexit__(exc_type, exc_value, traceback) await resync_tracer() async def handle_async_request(self, request: httpx.Request) -> httpx.Response: assert isinstance(request.stream, httpx.AsyncByteStream) request_body = b"".join([chunk async for chunk in request.stream]) scope: Scope = { "type": "http", "asgi": {"version": "3.0"}, "http_version": "1.1", "method": request.method, "scheme": request.url.scheme, "path": request.url.path, "raw_path": request.url.raw_path.split(b"?", maxsplit=1)[0], "query_string": request.url.query, "root_path": "", "headers": [(name.lower(), value) for name, value in request.headers.raw], "server": (request.url.host, request.url.port), "client": ("127.0.0.1", 1234), } request_delivered = False client_disconnected = anyio.Event() response_started = anyio.Event() response_status = 0 response_headers: list[tuple[bytes, bytes]] = [] application_error: Exception | None = None chunk_writer, chunk_reader = anyio.create_memory_object_stream[bytes](math.inf) async def receive_request() -> Message: nonlocal request_delivered if not request_delivered: request_delivered = True return {"type": "http.request", "body": request_body, "more_body": False} await client_disconnected.wait() return {"type": "http.disconnect"} async def send_response(message: Message) -> None: nonlocal response_status, response_headers if message["type"] == "http.response.start": response_status = message["status"] response_headers = list(message.get("headers", [])) response_started.set() return assert message["type"] == "http.response.body" body: bytes = message.get("body", b"") if body: await chunk_writer.send(body) if not message.get("more_body", False): await chunk_writer.aclose() async def run_application() -> None: nonlocal application_error try: await self._app(scope, receive_request, send_response) except Exception as exc: # The bridge is the application's outermost boundary: a crash # must fail the originating request (or show up in the already-started response), # never tear down the task group shared with every other in-flight request. application_error = exc finally: response_started.set() await chunk_writer.aclose() self._task_group.start_soon(run_application) try: await response_started.wait() if application_error is not None: raise application_error except BaseException: # No response will be built, so close the reader the response body would have owned # and tell the application its peer has gone away. client_disconnected.set() await chunk_reader.aclose() raise return httpx.Response( status_code=response_status, headers=response_headers, stream=_StreamingResponseBody(chunk_reader, client_disconnected), request=request, )