From a9d3c7cf320fba90cb69844996442db65886998f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:59:03 -0700 Subject: [PATCH] fix(proxy): free the parallel slot when a keepalive stream is cancelled before the first token With SSE keepalive pings on, a reader that left during the pings cancelled the upstream call but never reported the failure, so the key's parallel slot stayed taken until its TTL. The ping loop now reports one client disconnect, or the producer's own error when it failed first. A configured polling TTL that is not a number now falls back to the cache default instead of failing every state write. --- litellm/proxy/common_request_processing.py | 13 ++ .../proxy/response_polling/polling_handler.py | 15 +- ...test_background_response_polling_cancel.py | 83 +++++++-- .../hooks/test_parallel_request_limiter_v3.py | 79 +++++++- .../test_response_polling_handler.py | 172 ++++++++++-------- .../proxy/test_common_request_processing.py | 28 +++ 6 files changed, 295 insertions(+), 95 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cf411bbc2e1..624515c0832 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1389,6 +1389,13 @@ async def _abandon_late_producer(produce_task: "asyncio.Future[_LateResponseT]") await aclose_streaming_response(produce_task.result()) +def _abandoned_producer_failure(produce_task: "asyncio.Future[_LateResponseT]") -> Exception: + failure: Final = None if produce_task.cancelled() else produce_task.exception() + if isinstance(failure, Exception): + return failure + return HTTPException(status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail=LITELLM_CLIENT_DISCONNECT_DETAIL) + + async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None]: """Replay a Response that was built after a keepalive had already opened the wire.""" if not isinstance(produced, StreamingResponse): @@ -1476,6 +1483,12 @@ async def open_sse_before_first_byte( while not produce_task.done(): yield SSE_COMMENT_PING_BYTES await asyncio.wait((produce_task,), timeout=interval) + except BaseException: + await _abandon_late_producer(produce_task) + with anyio.CancelScope(shield=True): + await _sanitized_late_failure(_abandoned_producer_failure(produce_task), on_late_failure) + raise + try: try: produced: Final = produce_task.result() except Exception as exc: # noqa: BLE001 # the status line is already sent; surface it as a frame diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index e4ea6b873da..946fd1f3f62 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -125,13 +125,13 @@ class ResponsePollingHandler: self, polling_id: str, status: ResponsesAPIStatus | None = None, - usage: dict | None = None, - error: dict | None = None, - incomplete_details: dict | None = None, - reasoning: dict | None = None, + usage: Mapping[str, object] | None = None, + error: Mapping[str, object] | None = None, + incomplete_details: Mapping[str, object] | None = None, + reasoning: Mapping[str, object] | None = None, tool_choice: object | None = None, tools: Sequence[object] | None = None, - output: list | None = None, + output: Sequence[object] | None = None, # Additional ResponsesAPIResponse fields model: str | None = None, instructions: str | None = None, @@ -139,7 +139,7 @@ class ResponsePollingHandler: top_p: float | None = None, max_output_tokens: int | None = None, previous_response_id: str | None = None, - text: dict | None = None, + text: Mapping[str, object] | None = None, truncation: str | None = None, parallel_tool_calls: bool | None = None, user: str | None = None, @@ -244,7 +244,8 @@ class ResponsePollingHandler: try: written: Final = _SCRIPT_REPLY.validate_python( await write_unless_cancelled( - keys=(cache_key,), args=(_encoded_state(state), int(self.ttl), _CANCELLED_PREFIX) + keys=(cache_key,), + args=(_encoded_state(state), self.redis_cache.get_ttl(ttl=self.ttl), _CANCELLED_PREFIX), ) ) except Exception as write_error: # noqa: BLE001 # a Redis fault must not end the background stream diff --git a/tests/integration/streaming/test_background_response_polling_cancel.py b/tests/integration/streaming/test_background_response_polling_cancel.py index ee59d0708d7..88d89ab8bd5 100644 --- a/tests/integration/streaming/test_background_response_polling_cancel.py +++ b/tests/integration/streaming/test_background_response_polling_cancel.py @@ -1,14 +1,15 @@ import threading import uuid -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import ExitStack, contextmanager from pathlib import Path +from types import MappingProxyType from typing import Final import pytest import yaml from integration._support.client import Gateway, eventually, held, object_value, string_value -from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.process import DB_PUSH, graceful_stop_seconds, owned_proxy from integration._support.redis_process import owned_redis from integration._support.responses_stream import OPENAI_TARGET, completed, created, delta, frame, response_object from integration._support.wire import Reply, Request, Wire, wire_server @@ -22,9 +23,11 @@ _STREAM_SECONDS: Final = _DELTAS * _DELTA_PAUSE_SECONDS _OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 240 _MODEL_DISCOVERY: Final = ("GET", "/v1/models") _NO_MODELS: Final = Reply(body=b'{"object":"list","data":[]}') +_POLLING_VIA_CACHE: Final = MappingProxyType({"responses": {"background_mode": {"polling_via_cache": "all"}}}) +_KEEPALIVE_SECONDS: Final = 0.05 -def _config(directory: Path, wire: Wire) -> Path: +def _config(directory: Path, wire: Wire, litellm_settings: Mapping[str, JsonValue] = _POLLING_VIA_CACHE) -> Path: path: Final = directory / f"background-polling-cancel-{uuid.uuid4().hex}.yaml" path.write_text( yaml.safe_dump( @@ -43,7 +46,7 @@ def _config(directory: Path, wire: Wire) -> Path: "master_key": "os.environ/LITELLM_MASTER_KEY", "coordination_redis": {"host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, }, - "litellm_settings": {"responses": {"background_mode": {"polling_via_cache": "all"}}}, + "litellm_settings": dict(litellm_settings), } ) ) @@ -52,15 +55,14 @@ def _config(directory: Path, wire: Wire) -> Path: @contextmanager def _proxies_sharing_one_redis( - gateway: Gateway, directory: Path, wire: Wire, count: int + gateway: Gateway, directory: Path, config: Path, count: int, *, with_database: bool = False ) -> Iterator[tuple[Gateway, ...]]: - config: Final = _config(directory, wire) with owned_redis(directory / "redis") as cache, ExitStack() as proxies: overrides: Final = { "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "STORE_MODEL_IN_DB": "False", - "DISABLE_SCHEMA_UPDATE": "true", + "DISABLE_SCHEMA_UPDATE": str(not with_database).lower(), } yield tuple( proxies.enter_context( @@ -69,8 +71,8 @@ def _proxies_sharing_one_redis( directory / f"proxy-{index}", overrides, config=config, - database_setup=(), - remove_environment=_NO_DATABASE, + database_setup=DB_PUSH if with_database else (), + remove_environment=() if with_database else _NO_DATABASE, ) ) for index in range(count) @@ -88,9 +90,9 @@ def _serving(reply: Reply, requested: threading.Event) -> Callable[[Request], Re return respond -def _create_background_response(proxy: Gateway) -> str: +def _create_background_response(proxy: Gateway, key: str | None = None) -> str: created_response: Final = proxy.request( - "POST", "/v1/responses", {"model": _MODEL, "input": uuid.uuid4().hex, "background": True} + "POST", "/v1/responses", {"model": _MODEL, "input": uuid.uuid4().hex, "background": True}, key=key ) assert created_response.status_code == 200, created_response.text body: Final = object_value(created_response.json()) @@ -127,7 +129,7 @@ def test_a_cancel_sent_to_another_proxy_stops_the_background_stream_and_the_resp ) with ( wire_server(_serving(slow_stream, requested)) as wire, - _proxies_sharing_one_redis(gateway, tmp_path, wire, 2) as (streaming, cancelling), + _proxies_sharing_one_redis(gateway, tmp_path, _config(tmp_path, wire), 2) as (streaming, cancelling), ): polling_id: Final = _create_background_response(streaming) eventually(requested.is_set, lambda seen: seen, seconds=30) @@ -160,7 +162,7 @@ def test_a_background_response_cancelled_just_before_its_stream_finishes_stays_c ) with ( wire_server(_serving(held_stream, requested)) as wire, - _proxies_sharing_one_redis(gateway, tmp_path, wire, 1) as (proxy,), + _proxies_sharing_one_redis(gateway, tmp_path, _config(tmp_path, wire), 1) as (proxy,), ): polling_id: Final = _create_background_response(proxy) eventually(requested.is_set, lambda seen: seen, seconds=30) @@ -186,7 +188,7 @@ def test_a_delete_sent_to_another_proxy_stops_the_background_stream_and_the_resp ) with ( wire_server(_serving(slow_stream, requested)) as wire, - _proxies_sharing_one_redis(gateway, tmp_path, wire, 2) as (streaming, deleting), + _proxies_sharing_one_redis(gateway, tmp_path, _config(tmp_path, wire), 2) as (streaming, deleting), ): polling_id: Final = _create_background_response(streaming) eventually(requested.is_set, lambda seen: seen, seconds=30) @@ -199,3 +201,56 @@ def test_a_delete_sent_to_another_proxy_stops_the_background_stream_and_the_resp holding=_STREAM_SECONDS, seconds=2 * _STREAM_SECONDS, ) + + +def _streamed_response_status(proxy: Gateway, key: str) -> int: + return proxy.request( + "POST", "/v1/responses", {"model": _MODEL, "input": uuid.uuid4().hex, "stream": True}, key=key + ).status_code + + +@pytest.mark.timeout(_OWNED_PROXY_SECONDS + 90) +def test_a_background_response_cancelled_during_keepalive_pings_frees_its_parallel_slot( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + answer: Final = threading.Event() + quick_stream: Final = Reply( + content_type="text/event-stream", + chunks=(frame(created(identity)), frame(delta(identity, "done")), frame(completed(identity, "done"))), + ) + + def respond(request: Request) -> Reply: + if (request.method, request.target) == _MODEL_DISCOVERY: + return _NO_MODELS + assert request.method == "POST" and request.target.startswith(OPENAI_TARGET), request.target + if not requested.is_set(): + requested.set() + answer.wait() + return quick_stream + + with ( + wire_server(respond) as wire, + _proxies_sharing_one_redis( + gateway, + tmp_path, + _config(tmp_path, wire, {**_POLLING_VIA_CACHE, "sse_keepalive_ping_interval_seconds": _KEEPALIVE_SECONDS}), + 1, + with_database=True, + ) as (proxy,), + proxy.scenario() as scenario, + ): + try: + key: Final = scenario.key(max_parallel_requests=1) + polling_id: Final = _create_background_response(proxy, key) + eventually(requested.is_set, lambda seen: seen, seconds=30) + held(lambda: _status(proxy, polling_id), lambda status: status == "in_progress", holding=1) + assert _streamed_response_status(proxy, key) == 429 + _cancel(proxy, polling_id) + eventually( + lambda: _streamed_response_status(proxy, key), lambda status_code: status_code == 200, seconds=20 + ) + assert _status(proxy, polling_id) == "cancelled" + finally: + answer.set() diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index e056e723a04..593e55c7336 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -22,7 +22,7 @@ import litellm from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.litellm_core_utils import ptu_pricing from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY @@ -4837,6 +4837,83 @@ async def test_sse_keepalive_releases_parallel_capacity( assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("leaves_by", ["close", "cancel"]) +@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"]) +async def test_sse_keepalive_releases_parallel_capacity_when_the_reader_leaves_during_pings( + leaves_by: str, preadmitted: bool +) -> None: + from fastapi.responses import StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES + + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-keepalive-reader-leaves") + auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + + async def request(call_id: str) -> None: + data: Final = {"model": "test-model", "litellm_call_id": call_id} + admitted: Final = asyncio.Event() + never_answers: Final = asyncio.Event() + upstream_cancelled: Final = asyncio.Event() + reported: Final = asyncio.Queue[Exception]() + + async def admit() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion" + ) + + async def failure(exc: Exception) -> None: + await reported.put(exc) + await handler.async_post_call_failure_hook( + request_data=data, original_exception=exc, user_api_key_dict=auth + ) + + async def produce() -> StreamingResponse: + if not preadmitted: + await admit() + admitted.set() + try: + await never_answers.wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + return StreamingResponse(iter(())) + + if preadmitted: + await admit() + response: Final = await open_sse_before_first_byte(produce(), 0.001, on_late_failure=failure) + assert isinstance(response, StreamingResponse) + assert await anext(response.body_iterator) == SSE_COMMENT_PING_BYTES + await admitted.wait() + + async def read_until_cancelled() -> None: + async for _ in response.body_iterator: + continue + + if leaves_by == "close": + body: Final = response.body_iterator + assert isinstance(body, AsyncGenerator) + await body.aclose() + else: + reader: Final = asyncio.ensure_future(read_until_cancelled()) + await asyncio.sleep(0.01) + reader.cancel() + with pytest.raises(asyncio.CancelledError): + await reader + assert upstream_cancelled.is_set() + left: Final = reported.get_nowait() + assert isinstance(left, HTTPException) and left.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED + assert reported.empty() + + for index in range(3): + await contextvars.Context().run(asyncio.create_task, request(f"keepalive-left-{index}")) + assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0 + + @pytest.mark.asyncio async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None: from fastapi.responses import StreamingResponse diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 9195447816c..28806ff2acd 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -16,56 +16,62 @@ following the OpenAI Response API format. import asyncio import json +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Dict, Final, Optional +from typing import Dict, Final, NoReturn, Optional from unittest.mock import AsyncMock, Mock, patch import pytest from fastapi import Request from pydantic import ValidationError +from litellm.caching.base_cache import BaseCache from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler -class _RedisStandIn: +class _RedisStandIn(BaseCache): def __init__(self) -> None: + super().__init__() self.stored: Dict[str, str] = {} self.ttls: Dict[str, Optional[int]] = {} @classmethod - def holding(cls, state: Dict[str, Any], polling_id: str = "litellm_poll_test") -> "_RedisStandIn": - cache = cls() + def holding(cls, state: Mapping[str, object], polling_id: str = "litellm_poll_test") -> "_RedisStandIn": + cache: Final = cls() cache.stored[ResponsePollingHandler.get_cache_key(polling_id)] = json.dumps(json.dumps(state)) return cache - async def async_get_cache(self, key: str) -> Any: - raw = self.stored.get(key) + async def async_get_cache(self, key: str) -> object: + raw: Final = self.stored.get(key) return None if raw is None else json.loads(raw) - async def async_set_cache(self, key: str, value: Any, ttl: Optional[int] = None) -> None: + async def async_set_cache(self, key: str, value: object, ttl: Optional[int] = None) -> None: self.stored[key] = json.dumps(value) self.ttls[key] = ttl async def async_delete_cache(self, key: str) -> None: self.stored.pop(key, None) - def async_register_script(self, script: str): - async def write_unless_cancelled(keys, args): + async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> NoReturn: + raise NotImplementedError + + def async_register_script(self, script: str) -> Callable[..., Awaitable[int]]: + async def write_unless_cancelled(keys: Sequence[str], args: Sequence[object]) -> int: if keys[0] not in self.stored or self.stored[keys[0]].startswith(args[2]): return 0 self.stored[keys[0]] = args[0] self.ttls[keys[0]] = int(args[1]) return 1 - async def cancelled_or_gone(keys, args): + async def cancelled_or_gone(keys: Sequence[str], args: Sequence[object]) -> int: return int(self.stored.get(keys[0], "")[: len(args[0])] in ("", args[0])) return write_unless_cancelled if "'SET'" in script else cancelled_or_gone class _UnreachableRedisStandIn(_RedisStandIn): - def async_register_script(self, script: str): - async def write_fails(keys, args): + def async_register_script(self, script: str) -> Callable[..., Awaitable[int]]: + async def write_fails(keys: Sequence[str], args: Sequence[object]) -> NoReturn: raise ConnectionError("redis is unreachable") return write_fails @@ -540,8 +546,8 @@ class TestResponsePollingHandler: assert result is None @pytest.mark.asyncio - async def test_a_response_never_reads_as_stopped_without_redis(self): - handler = ResponsePollingHandler(redis_cache=None) + async def test_a_response_never_reads_as_stopped_without_redis(self) -> None: + handler: Final = ResponsePollingHandler(redis_cache=None) assert await handler.is_cancelled_or_gone("litellm_poll_test") is False @@ -570,8 +576,8 @@ class TestResponsePollingHandler: assert stored["status"] == "cancelled" @pytest.mark.asyncio - async def test_an_update_after_a_cancel_leaves_the_response_cancelled(self): - mock_redis = _RedisStandIn.holding( + async def test_an_update_after_a_cancel_leaves_the_response_cancelled(self) -> None: + mock_redis: Final = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -580,7 +586,7 @@ class TestResponsePollingHandler: "created_at": 1234567890, } ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler: Final = ResponsePollingHandler(redis_cache=mock_redis) await handler.cancel_polling("litellm_poll_test") await handler.update_state( @@ -606,8 +612,8 @@ class TestResponsePollingHandler: ("delete", True), ], ) - async def test_a_cancelled_or_deleted_response_reads_as_stopped(self, stop, expected): - handler = ResponsePollingHandler( + async def test_a_cancelled_or_deleted_response_reads_as_stopped(self, stop: str | None, expected: bool) -> None: + handler: Final = ResponsePollingHandler( redis_cache=_RedisStandIn.holding( {"id": "litellm_poll_test", "object": "response", "output": [], "status": "in_progress"} ) @@ -623,25 +629,25 @@ class TestResponsePollingHandler: assert await handler.is_cancelled_or_gone("litellm_poll_never_created") is True @pytest.mark.asyncio - async def test_a_failed_state_write_keeps_the_stored_state_and_does_not_raise(self): - state = { + async def test_a_failed_state_write_keeps_the_stored_state_and_does_not_raise(self) -> None: + state: Final = { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, } - handler = ResponsePollingHandler(redis_cache=_UnreachableRedisStandIn.holding(state)) + handler: Final = ResponsePollingHandler(redis_cache=_UnreachableRedisStandIn.holding(state)) await handler.update_state(polling_id="litellm_poll_test", status="completed") assert await handler.get_state("litellm_poll_test") == state @pytest.mark.asyncio - async def test_a_queued_response_reads_back_after_an_update(self): - handler = ResponsePollingHandler(redis_cache=_RedisStandIn()) + async def test_a_queued_response_reads_back_after_an_update(self) -> None: + handler: Final = ResponsePollingHandler(redis_cache=_RedisStandIn()) - queued = await handler.create_initial_state( + queued: Final = await handler.create_initial_state( polling_id="litellm_poll_test", request_data={"model": "gpt-4o", "metadata": {"trace": "t-1"}} ) await handler.update_state(polling_id="litellm_poll_test", status="in_progress", model="gpt-4o") @@ -724,6 +730,20 @@ class TestResponsePollingHandler: assert mock_redis.ttls == {ResponsePollingHandler.get_cache_key("litellm_poll_test"): 1800} + @pytest.mark.asyncio + @pytest.mark.parametrize("configured_ttl", [None, "1h"]) + async def test_an_unusable_configured_ttl_still_writes_the_state_with_the_cache_default( + self, configured_ttl: object + ) -> None: + cache: Final = _RedisStandIn.holding({"id": "litellm_poll_test", "status": "queued", "output": []}) + handler: Final = ResponsePollingHandler(redis_cache=cache, ttl=configured_ttl) + + await handler.update_state(polling_id="litellm_poll_test", status="in_progress") + + key: Final = ResponsePollingHandler.get_cache_key("litellm_poll_test") + assert json.loads(json.loads(cache.stored[key]))["status"] == "in_progress" + assert cache.ttls == {key: cache.default_ttl} + class TestStreamingEventProcessing: """ @@ -1544,7 +1564,7 @@ def _make_background_streaming_kwargs( def _streaming_polling_handler(status: str = "in_progress") -> AsyncMock: - handler = AsyncMock(spec=ResponsePollingHandler) + handler: Final = AsyncMock(spec=ResponsePollingHandler) handler.get_state.return_value = {"status": status} handler.is_cancelled_or_gone.return_value = status == "cancelled" return handler @@ -1932,11 +1952,11 @@ async def _let_the_loop_run() -> None: @pytest.mark.xdist_group("heavy_imports") class TestStreamUntilCancelled: @pytest.mark.asyncio - async def test_a_cancelled_state_stops_the_stream(self): + async def test_a_cancelled_state_stops_the_stream(self) -> None: from litellm.proxy.response_polling.background_streaming import stream_until_cancelled - state = _CancelledAfterReads(reads_before_cancel=2) - stopped = asyncio.Event() + state: Final = _CancelledAfterReads(reads_before_cancel=2) + stopped: Final = asyncio.Event() async def endless_stream() -> None: try: @@ -1953,11 +1973,11 @@ class TestStreamUntilCancelled: assert state.reads == 3 @pytest.mark.asyncio - async def test_failing_state_reads_leave_the_stream_running_to_its_end(self): + async def test_failing_state_reads_leave_the_stream_running_to_its_end(self) -> None: from litellm.proxy.response_polling.background_streaming import stream_until_cancelled - state = _UnreadableState() - finished = asyncio.Event() + state: Final = _UnreadableState() + finished: Final = asyncio.Event() async def stream_outliving_three_reads() -> None: while state.reads < 3: @@ -1972,10 +1992,10 @@ class TestStreamUntilCancelled: assert finished.is_set() @pytest.mark.asyncio - async def test_a_stream_that_finishes_first_leaves_no_watcher_reading(self): + async def test_a_stream_that_finishes_first_leaves_no_watcher_reading(self) -> None: from litellm.proxy.response_polling.background_streaming import stream_until_cancelled - state = _CancelledAfterReads(reads_before_cancel=1_000_000) + state: Final = _CancelledAfterReads(reads_before_cancel=1_000_000) async def short_stream() -> None: await _let_the_loop_run() @@ -1983,14 +2003,14 @@ class TestStreamUntilCancelled: await asyncio.wait_for( stream_until_cancelled("litellm_poll_test", short_stream(), state, poll_seconds=0), timeout=5 ) - reads_when_done = state.reads + reads_when_done: Final = state.reads await _let_the_loop_run() assert reads_when_done > 0 assert state.reads == reads_when_done @pytest.mark.asyncio - async def test_a_stream_error_reaches_the_caller(self): + async def test_a_stream_error_reaches_the_caller(self) -> None: from litellm.proxy.response_polling.background_streaming import stream_until_cancelled async def broken_stream() -> None: @@ -2002,12 +2022,12 @@ class TestStreamUntilCancelled: ) @pytest.mark.asyncio - async def test_cancelling_the_caller_stops_both_the_stream_and_the_watcher(self): + async def test_cancelling_the_caller_stops_both_the_stream_and_the_watcher(self) -> None: from litellm.proxy.response_polling.background_streaming import stream_until_cancelled - state = _CancelledAfterReads(reads_before_cancel=1_000_000) - started = asyncio.Event() - stopped = asyncio.Event() + state: Final = _CancelledAfterReads(reads_before_cancel=1_000_000) + started: Final = asyncio.Event() + stopped: Final = asyncio.Event() async def endless_stream() -> None: started.set() @@ -2017,26 +2037,26 @@ class TestStreamUntilCancelled: stopped.set() raise - caller = asyncio.create_task( + caller: Final = asyncio.create_task( stream_until_cancelled("litellm_poll_test", endless_stream(), state, poll_seconds=0) ) await started.wait() caller.cancel() with pytest.raises(asyncio.CancelledError): await caller - reads_when_cancelled = state.reads + reads_when_cancelled: Final = state.reads await _let_the_loop_run() assert stopped.is_set() assert state.reads == reads_when_cancelled @pytest.mark.asyncio - async def test_a_cancelled_background_stream_writes_no_state_after_it_stops(self): + async def test_a_cancelled_background_stream_writes_no_state_after_it_stops(self) -> None: from litellm.proxy.response_polling.background_streaming import background_streaming_task - stopped = asyncio.Event() + stopped: Final = asyncio.Event() - async def endless_body(): + async def endless_body() -> AsyncIterator[str]: yield f"data: {json.dumps({'type': 'response.in_progress'})}" try: await asyncio.Event().wait() @@ -2044,10 +2064,10 @@ class TestStreamUntilCancelled: stopped.set() raise - endless_response = Mock() + endless_response: Final = Mock() endless_response.body_iterator = endless_body() - handler = _streaming_polling_handler(status="cancelled") - kwargs = _make_background_streaming_kwargs("poll_cancelled", handler) + handler: Final = _streaming_polling_handler(status="cancelled") + kwargs: Final = _make_background_streaming_kwargs("poll_cancelled", handler) 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" @@ -2060,12 +2080,14 @@ class TestStreamUntilCancelled: kwargs["proxy_logging_obj"].post_call_failure_hook.assert_not_awaited() @pytest.mark.asyncio - async def test_a_cancel_while_the_stream_waits_for_its_first_event_leaves_the_failure_report_to_the_stream(self): + async def test_a_cancel_while_the_stream_waits_for_its_first_event_leaves_the_failure_report_to_the_stream( + self, + ) -> None: from litellm.proxy.response_polling.background_streaming import background_streaming_task - stopped = asyncio.Event() + stopped: Final = asyncio.Event() - async def body_before_its_first_event(): + async def body_before_its_first_event() -> AsyncIterator[str]: try: await asyncio.Event().wait() except asyncio.CancelledError: @@ -2073,10 +2095,10 @@ class TestStreamUntilCancelled: raise yield "data: [DONE]" - silent_response = Mock() + silent_response: Final = Mock() silent_response.body_iterator = body_before_its_first_event() - handler = _streaming_polling_handler(status="cancelled") - kwargs = _make_background_streaming_kwargs("poll_cancelled_before_first_event", handler) + handler: Final = _streaming_polling_handler(status="cancelled") + kwargs: Final = _make_background_streaming_kwargs("poll_cancelled_before_first_event", handler) 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" @@ -2088,12 +2110,12 @@ class TestStreamUntilCancelled: kwargs["proxy_logging_obj"].post_call_failure_hook.assert_not_awaited() @pytest.mark.asyncio - async def test_a_cancel_before_the_upstream_answers_reports_one_client_disconnect(self): + async def test_a_cancel_before_the_upstream_answers_reports_one_client_disconnect(self) -> None: from fastapi import HTTPException from litellm.proxy.response_polling.background_streaming import background_streaming_task - upstream_cancelled = asyncio.Event() + upstream_cancelled: Final = asyncio.Event() async def upstream_that_never_answers(**_: object) -> None: try: @@ -2102,8 +2124,8 @@ class TestStreamUntilCancelled: upstream_cancelled.set() raise - handler = _streaming_polling_handler(status="cancelled") - kwargs = _make_background_streaming_kwargs("poll_cancelled_before_upstream", handler) + handler: Final = _streaming_polling_handler(status="cancelled") + kwargs: Final = _make_background_streaming_kwargs("poll_cancelled_before_upstream", handler) 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" @@ -2111,9 +2133,11 @@ class TestStreamUntilCancelled: MockProcessor.return_value.base_process_llm_request = upstream_that_never_answers await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) - report = kwargs["proxy_logging_obj"].post_call_failure_hook + report: Final = kwargs["proxy_logging_obj"].post_call_failure_hook + assert isinstance(report, AsyncMock) report.assert_awaited_once() - reported = report.await_args.kwargs + assert report.await_args is not None + reported: Final = report.await_args.kwargs assert upstream_cancelled.is_set() assert isinstance(reported["original_exception"], HTTPException) assert reported["original_exception"].status_code == 499 @@ -2122,14 +2146,14 @@ class TestStreamUntilCancelled: assert reported["dispatched_call_type"] == "aresponses" @pytest.mark.asyncio - async def test_a_cancel_during_a_blocked_state_write_closes_the_upstream_stream(self): + async def test_a_cancel_during_a_blocked_state_write_closes_the_upstream_stream(self) -> None: from starlette.responses import StreamingResponse from litellm.proxy.response_polling.background_streaming import background_streaming_task - closed = asyncio.Event() + closed: Final = asyncio.Event() - async def body(): + async def body() -> AsyncIterator[str]: try: yield f"data: {json.dumps({'type': 'response.output_item.added', 'item': {'id': 'item_1'}})}" yield "data: [DONE]" @@ -2141,9 +2165,9 @@ class TestStreamUntilCancelled: if "output" in fields: await asyncio.Event().wait() - handler = _streaming_polling_handler(status="cancelled") + handler: Final = _streaming_polling_handler(status="cancelled") handler.update_state.side_effect = output_writes_block - kwargs = _make_background_streaming_kwargs("poll_cancelled_mid_write", handler) + kwargs: Final = _make_background_streaming_kwargs("poll_cancelled_mid_write", handler) 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" @@ -2155,7 +2179,6 @@ class TestStreamUntilCancelled: assert {call.kwargs.get("status") for call in handler.update_state.await_args_list} <= {None, "in_progress"} - _MESSAGE_ADDED: Final = { "type": "response.output_item.added", "item": {"id": "msg_1", "type": "message", "role": "assistant", "content": []}, @@ -2199,19 +2222,21 @@ _MESSAGE_DONE: Final = { ), ], ) -async def test_a_background_stream_writes_the_message_its_events_build(events, expected_content): +async def test_a_background_stream_writes_the_message_its_events_build( + events: tuple[Mapping[str, object], ...], expected_content: list[dict[str, str]] +) -> None: from litellm.proxy.response_polling.background_streaming import background_streaming_task - async def body(): + async def body() -> AsyncIterator[str]: yield "data: {not json" for event in events: yield f"data: {json.dumps(event)}" yield "data: [DONE]" - streamed = Mock() + streamed: Final = Mock() streamed.body_iterator = body() - handler = _streaming_polling_handler() - kwargs = _make_background_streaming_kwargs("poll_message", handler) + handler: Final = _streaming_polling_handler() + kwargs: Final = _make_background_streaming_kwargs("poll_message", handler) 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" @@ -2219,11 +2244,12 @@ async def test_a_background_stream_writes_the_message_its_events_build(events, e MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=streamed) await background_streaming_task(**kwargs) - writes = handler.update_state.await_args_list - written_outputs = [call.kwargs["output"] for call in writes if "output" in call.kwargs] + writes: Final = handler.update_state.await_args_list + written_outputs: Final = [call.kwargs["output"] for call in writes if "output" in call.kwargs] assert [list(item["content"]) for item in written_outputs[-1]] == [expected_content] assert handler.update_state.await_args_list[-1].kwargs["status"] == "completed" + class TestEdgeCases: """Test edge cases and error scenarios""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 078990c7d98..1e81a57feb7 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1,4 +1,5 @@ import asyncio +import collections.abc import contextvars import copy import datetime @@ -10784,6 +10785,33 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): assert collected[-1] == b"data: [DONE]\n\n" +@pytest.mark.asyncio +async def test_a_failure_that_lands_as_the_reader_leaves_is_reported_instead_of_a_disconnect(): + may_fail: Final = asyncio.Event() + audited: Final = asyncio.Queue[Exception]() + + async def fails_between_pings() -> StreamingResponse: + await may_fail.wait() + raise HTTPException(status_code=429, detail="rate limited") + + async def record(exc: Exception) -> None: + await audited.put(exc) + + response: Final = await open_sse_before_first_byte( + fails_between_pings(), ping_interval_seconds=0.05, on_late_failure=record + ) + body: Final = response.body_iterator + assert isinstance(body, collections.abc.AsyncGenerator) + assert await body.__anext__() == TTFT_PING + may_fail.set() + await asyncio.sleep(0.01) + await body.aclose() + + reported: Final = audited.get_nowait() + assert isinstance(reported, HTTPException) and reported.status_code == 429 + assert audited.empty() + + @pytest.mark.asyncio async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def slow_failure():