diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 0fd242f2bc1..fac45d4391c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/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", diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 772d3622745..467c1332325 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/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"""