From 1d8650413d4237f9dfce6aec838035618a858d8b Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 11:51:35 +0000 Subject: [PATCH] fix(responses): keep background polling alive after the client disconnects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../response_polling/background_streaming.py | 17 +++++- .../test_response_polling_handler.py | 4 +- .../proxy/response_polling/__init__.py | 0 .../test_background_streaming.py | 52 +++++++++++++++++++ 4 files changed, 70 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/proxy/response_polling/__init__.py create mode 100644 tests/test_litellm/proxy/response_polling/test_background_streaming.py 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..7a4e0ca89c0 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -20,7 +20,7 @@ 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 +1414,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={}, diff --git a/tests/test_litellm/proxy/response_polling/__init__.py b/tests/test_litellm/proxy/response_polling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/response_polling/test_background_streaming.py b/tests/test_litellm/proxy/response_polling/test_background_streaming.py new file mode 100644 index 00000000000..98d914b1848 --- /dev/null +++ b/tests/test_litellm/proxy/response_polling/test_background_streaming.py @@ -0,0 +1,52 @@ +import asyncio +from collections.abc import AsyncGenerator +from typing import Final + +import pytest +from fastapi import Request +from fastapi.responses import JSONResponse, StreamingResponse +from starlette.types import Message + +from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED +from litellm.proxy.common_request_processing import create_response +from litellm.proxy.response_polling.background_streaming import detach_request_from_client + + +def _request_whose_client_already_left() -> Request: + async def receive() -> Message: + return {"type": "http.disconnect"} + + scope: Final = { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"x-litellm-call-id", b"call-123")], + "query_string": b"", + } + return Request(scope, receive) + + +async def _slow_first_chunk() -> AsyncGenerator[str, None]: + await asyncio.sleep(0.05) + yield 'data: {"type": "response.created"}\n\n' + + +@pytest.mark.asyncio +async def test_detached_request_survives_client_disconnect_before_first_chunk(): + original: Final = _request_whose_client_already_left() + + cancelled: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=original) + assert isinstance(cancelled, JSONResponse) + assert cancelled.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + + detached: Final = detach_request_from_client(original) + kept_alive: Final = await create_response(_slow_first_chunk(), "text/event-stream", {}, request=detached) + assert isinstance(kept_alive, StreamingResponse) + assert kept_alive.status_code == 200 + + +def test_detached_request_keeps_scope(): + original: Final = _request_whose_client_already_left() + detached: Final = detach_request_from_client(original) + assert detached.headers["x-litellm-call-id"] == "call-123" + assert detached.scope is original.scope