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():