From e589a0ed82023d6fa550e6ddfa005c5ebbf4951e Mon Sep 17 00:00:00 2001 From: jesus Date: Mon, 7 Sep 2026 12:07:31 +0000 Subject: [PATCH] test(responses): cover the polling disconnect regression through background_streaming_task Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 - Makefile | 2 +- .../test_response_polling_handler.py | 58 +++++++++++++++++++ .../proxy/response_polling/__init__.py | 0 .../test_background_streaming.py | 52 ----------------- 5 files changed, 59 insertions(+), 54 deletions(-) delete mode 100644 tests/test_litellm/proxy/response_polling/__init__.py delete mode 100644 tests/test_litellm/proxy/response_polling/test_background_streaming.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 761960145eb..33245ec5b5f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -155,7 +155,6 @@ jobs: tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints - tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/ocr_endpoints tests/test_litellm/proxy/vector_store_endpoints diff --git a/Makefile b/Makefile index b14aaaf896f..e17fdba3c85 100644 --- a/Makefile +++ b/Makefile @@ -298,7 +298,7 @@ test-unit-proxy-core: install-test-deps $(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20 test-unit-proxy-misc: install-test-deps - $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/response_polling tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20 test-unit-integrations: install-test-deps $(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20 diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 7a4e0ca89c0..f5335508ce4 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -14,6 +14,7 @@ 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 @@ -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( + "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""" diff --git a/tests/test_litellm/proxy/response_polling/__init__.py b/tests/test_litellm/proxy/response_polling/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/proxy/response_polling/test_background_streaming.py b/tests/test_litellm/proxy/response_polling/test_background_streaming.py deleted file mode 100644 index 98d914b1848..00000000000 --- a/tests/test_litellm/proxy/response_polling/test_background_streaming.py +++ /dev/null @@ -1,52 +0,0 @@ -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