Merge pull request #40114 from BerriAI/litellm_fix_background_polling_disconnect_guard

fix(responses): keep background polling alive after the client disconnects

litellm/proxy/response_polling/background_streaming.py

@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias
 
 from fastapi import Request, Response
 from fastapi.responses import StreamingResponse
+from starlette.types import Message
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
@@ -74,6 +75,20 @@ class _StreamEventParser:
     parse: Callable[[str], _StreamEvent] = staticmethod(json.loads)
 
 
+async def _never_receive() -> Message:
+    await asyncio.Event().wait()
+    raise AssertionError("unreachable")
+
+
+def detach_request_from_client(request: Request) -> Request:
+    """Same scope (headers, parsed body, auth) but a receive() that never yields http.disconnect.
+
+    The polling client closes its connection right after getting the polling id, so the
+    upstream call must not be cancelled by the client-disconnect guards.
+    """
+    return Request(request.scope, _never_receive)
+
+
 async def background_streaming_task(
     polling_id: str,
     data: dict[str, object],
@@ -123,7 +138,7 @@ async def background_streaming_task(
         # Pre-call checks (rate limits, guardrails, budget) were already run
         # before polling ID creation, so skip them here to avoid double-counting.
         response: Final[StreamingResponse] = await processor.base_process_llm_request(
-            request=request,
+            request=detach_request_from_client(request),
             fastapi_response=fastapi_response,
             user_api_key_dict=user_api_key_dict,
             route_type="aresponses",

tests/proxy_unit_tests/test_response_polling_handler.py

@@ -14,13 +14,14 @@ These tests ensure the polling handler correctly manages response state
 following the OpenAI Response API format.
 """
 
+import asyncio
 import json
 from datetime import datetime, timezone
 from typing import Any, Dict, Optional
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-
+from fastapi import Request
 
 from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
 
@@ -1414,7 +1415,7 @@ def _make_background_streaming_kwargs(
         polling_id=polling_id,
         data={"model": "gpt-4o", "stream": False, "background": True},
         polling_handler=polling_handler,
-        request=Mock(),
+        request=Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}),
         fastapi_response=Mock(),
         user_api_key_dict=Mock(),
         general_settings={},
@@ -1663,6 +1664,63 @@ class TestBackgroundStreamingTerminalEvents:
         final_call = handler.update_state.call_args_list[-1]
         assert final_call.kwargs["status"] == "completed"
 
+    @pytest.mark.asyncio
+    async def test_polling_client_disconnect_does_not_cancel_upstream_call(self):
+        """The polling client hangs up right after getting its polling id. The detached task
+        must still stream the upstream response through the client-disconnect guards."""
+        from litellm.proxy.common_request_processing import create_response
+        from litellm.proxy.response_polling.background_streaming import (
+            background_streaming_task,
+        )
+
+        async def client_already_left():
+            return {"type": "http.disconnect"}
+
+        async def slow_upstream_stream():
+            await asyncio.sleep(0.05)
+            for event in (
+                {"type": "response.in_progress"},
+                {
+                    "type": "response.completed",
+                    "response": {
+                        "id": "resp_123",
+                        "status": "completed",
+                        "usage": {"input_tokens": 13, "output_tokens": 10},
+                        "model": "gpt-4o",
+                        "output": [{"id": "item_1", "type": "message"}],
+                    },
+                },
+            ):
+                yield f"data: {json.dumps(event)}\n\n"
+
+        async def upstream_call_behind_disconnect_guard(**kwargs):
+            return await create_response(
+                slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"]
+            )
+
+        handler = AsyncMock(spec=ResponsePollingHandler)
+        kwargs = _make_background_streaming_kwargs("poll_7", handler)
+        kwargs["request"] = Request(
+            {
+                "type": "http",
+                "method": "POST",
+                "path": "/v1/responses",
+                "headers": [(b"x-litellm-call-id", b"call-123")],
+                "query_string": b"",
+            },
+            client_already_left,
+        )
+
+        with patch(  # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests
+            "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing"
+        ) as MockProcessor:
+            MockProcessor.return_value.base_process_llm_request = upstream_call_behind_disconnect_guard
+            await background_streaming_task(**kwargs)
+
+        final_call = handler.update_state.call_args_list[-1]
+        assert final_call.kwargs["status"] == "completed"
+        assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10}
+
 
 class TestEdgeCases:
     """Test edge cases and error scenarios"""