From 09516635621caccf7e3db4f537eb49ea49b8a453 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 12 Aug 2026 06:46:04 +0530 Subject: [PATCH] fix(proxy): close the upstream stream when a streaming body is never consumed Close unconsumed upstream streaming bodies. --- headroom/proxy/handlers/streaming.py | 20 ++++ tests/test_proxy_copilot_auth_hooks.py | 16 ++- .../test_proxy_streaming_ratelimit_headers.py | 99 +++++++++++++++++++ 3 files changed, 134 insertions(+), 1 deletion(-) diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index e265ba80c..ef9600381 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -1098,6 +1098,7 @@ class StreamingMixin: ) -> Response | StreamingResponse: """Actual streaming implementation, guarded by _stream_response's cleanup wrapper.""" from fastapi.responses import Response, StreamingResponse + from starlette.background import BackgroundTask from headroom.proxy.helpers import MAX_SSE_BUFFER_SIZE @@ -1656,10 +1657,29 @@ class StreamingMixin: ) yield f"event: headroom_pending_messages\ndata: {pending_event}\n\n".encode() + async def _release_upstream_stream() -> None: + # Guarantee the upstream HTTP/2 stream is released even when the + # body generator above is never iterated — the client disconnected + # before Starlette started sending the response body (routine when a + # harness like Claude Code cancels or supersedes an in-flight turn), + # so ``generate()`` never entered its own ``aclosing`` and nothing + # else closes ``upstream_response``. Each such request otherwise + # leaks one open h2 stream; they accumulate on the pooled upstream + # connection until it reaches SETTINGS_MAX_CONCURRENT_STREAMS (100) + # and no new stream can open ("Max outbound streams is 100, 100 + # open"), and the proxy goes unhealthy until restart (#2797). + # Starlette runs a response's ``background`` task after the body + # finishes *and* after an early client disconnect, so this fires in + # both cases. ``aclose()`` is idempotent, so on the normal path — + # where the generator already closed the stream — this is a no-op. + with contextlib.suppress(Exception): + await upstream_response.aclose() + return StreamingResponse( generate(), media_type="text/event-stream", headers=forwarded_headers, + background=BackgroundTask(_release_upstream_stream), ) async def _stream_response_bedrock( diff --git a/tests/test_proxy_copilot_auth_hooks.py b/tests/test_proxy_copilot_auth_hooks.py index eb0eecd4a..7a37a82c7 100644 --- a/tests/test_proxy_copilot_auth_hooks.py +++ b/tests/test_proxy_copilot_auth_hooks.py @@ -52,11 +52,22 @@ def _load_handler_module(monkeypatch: pytest.MonkeyPatch, module_name: str, rela responses_mod = types.ModuleType("fastapi.responses") class Response: - def __init__(self, content=None, status_code: int = 200, headers=None, media_type=None): + def __init__( + self, + content=None, + status_code: int = 200, + headers=None, + media_type=None, + background=None, + ): self.content = content self.status_code = status_code self.headers = headers or {} self.media_type = media_type + # The streaming forwarder attaches a background task that releases the + # upstream stream when the body is never consumed (#2882); the double + # must accept and store it so the real StreamingResponse call works. + self.background = background class StreamingResponse(Response): pass @@ -228,6 +239,9 @@ def test_streaming_response_applies_copilot_auth(monkeypatch: pytest.MonkeyPatch assert sent_headers["Authorization"] == "Bearer upstream-token" assert sent_headers["content-type"] == "application/json" assert response.status_code == 200 + # The Copilot auth hook and the #2882 upstream-stream cleanup coexist: the + # streaming response still carries its background release task. + assert response.background is not None def test_openai_chat_routes_copilot_requests_per_model(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_proxy_streaming_ratelimit_headers.py b/tests/test_proxy_streaming_ratelimit_headers.py index dfedc4915..5983dea54 100644 --- a/tests/test_proxy_streaming_ratelimit_headers.py +++ b/tests/test_proxy_streaming_ratelimit_headers.py @@ -463,6 +463,105 @@ class TestStreamingRatelimitHeaderForwarding: assert attempts["count"] == 2 assert chunks + @pytest.mark.asyncio + async def test_upstream_stream_closed_when_body_never_consumed(self): + """A never-iterated streaming body must still release the upstream stream (#2797). + + The upstream stream is opened before the body generator, and the + generator's own ``aclosing`` only runs if the body is iterated. When a + client disconnects before Starlette starts sending the body the + generator never runs, so the close must come from the response's + ``background`` task instead — otherwise every such request leaks an open + HTTP/2 stream and the pooled upstream connection eventually exhausts its + 100 concurrent streams ("Max outbound streams is 100, 100 open"). + """ + proxy = self._create_mock_proxy() + mock_response = self._create_mock_upstream_response() + + mock_request = MagicMock() + proxy.http_client.build_request = MagicMock(return_value=mock_request) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + result = await proxy._stream_response( + url="https://api.anthropic.com/v1/messages", + headers={"x-api-key": "sk-test"}, + body={ + "model": "claude-sonnet-4-20250514", + "max_tokens": 100, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + provider="anthropic", + model="claude-sonnet-4-20250514", + request_id="test-abandoned-stream", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + ) + + # Simulate the client disconnecting before the body is consumed: the + # generator is never iterated, so its aclosing never runs. + mock_response.aclose.assert_not_awaited() + + # Starlette runs the response's background task in exactly this case. + assert result.background is not None, "streaming response must carry a cleanup task" + await result.background() + + mock_response.aclose.assert_awaited() + + @pytest.mark.asyncio + async def test_upstream_stream_released_over_asgi_lifecycle_on_disconnect(self): + """Driving the real ASGI response through an early disconnect releases the stream. + + Rather than calling ``result.background()`` directly, this exercises the + Starlette response lifecycle with a client that disconnects immediately, + and asserts the upstream stream is closed by the end of it -- proving the + cleanup this PR attaches is actually invoked by Starlette, not merely + present on the response object. + """ + import asyncio + + proxy = self._create_mock_proxy() + mock_response = self._create_mock_upstream_response() + proxy.http_client.build_request = MagicMock(return_value=MagicMock()) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + result = await proxy._stream_response( + url="https://api.anthropic.com/v1/messages", + headers={"x-api-key": "sk-test"}, + body={ + "model": "claude-sonnet-4-20250514", + "max_tokens": 100, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + provider="anthropic", + model="claude-sonnet-4-20250514", + request_id="test-asgi-lifecycle", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + ) + + async def receive(): + # The client is already gone before the body is streamed. + return {"type": "http.disconnect"} + + async def send(_message): + return None + + scope = {"type": "http", "method": "POST", "headers": []} + await asyncio.wait_for(result(scope, receive, send), timeout=5.0) + + # By the end of the response lifecycle the upstream stream is released. + mock_response.aclose.assert_awaited() + @pytest.mark.asyncio async def test_codex_rate_limit_headers_captured_and_forwarded_in_streaming(self): """Codex x-codex-* headers must refresh /stats state AND reach the client.