From 0f4ac2143bc3362620bedde2dec31ab8ff3169c3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:09:48 -0700 Subject: [PATCH 1/3] fix(responses): stop the background stream when a polling response is cancelled With polling via cache on, POST /v1/responses/{id}/cancel only wrote the cancelled status into Redis. The background task kept streaming from the provider, and its next state write merged over the cancel, so the response came back as completed with full output and full billing The background task now races the stream against a watcher that reads the stored state every second and stops the stream when it reads cancelled, on any worker or pod. Every state write goes through a Lua script that only sets the key when it exists and is not cancelled, so a write already in flight when the cancel lands can no longer overwrite it --- .../response_polling/background_streaming.py | 90 ++++- .../proxy/response_polling/polling_handler.py | 43 ++- ...test_background_response_polling_cancel.py | 169 +++++++++ .../test_response_polling_handler.py | 352 +++++++++++++++--- 4 files changed, 584 insertions(+), 70 deletions(-) create mode 100644 tests/integration/streaming/test_background_response_polling_cancel.py diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 1eb76f8f558..e9cad2dbe80 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,7 +10,7 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from collections.abc import AsyncIterable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, Callable, Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable from fastapi import Request, Response @@ -33,6 +33,8 @@ if TYPE_CHECKING: _JsonDict: TypeAlias = dict[str, object] _JsonList: TypeAlias = list[object] +CANCELLATION_POLL_SECONDS: Final = 1.0 + class _OutputItem(TypedDict, total=False): id: ReadOnly[str] @@ -81,6 +83,10 @@ class _StreamsBody(Protocol): def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ... +class PollingStateReader(Protocol): + async def get_state(self, polling_id: str) -> Mapping[str, object] | None: ... + + def _sse_frame_data(frame: str) -> str | None: return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None) @@ -99,6 +105,42 @@ def detach_request_from_client(request: Request) -> Request: return Request(request.scope, _never_receive) +async def _reads_cancelled(polling_id: str, polling_state: PollingStateReader) -> bool: + try: + state: Final = await polling_state.get_state(polling_id) + except Exception as error: # noqa: BLE001 # a failed read must never stop the upstream call + verbose_proxy_logger.warning( + "Could not read the polling state of %s, so it keeps streaming: %s", polling_id, error + ) + return False + return state is not None and state.get("status") == "cancelled" + + +async def _wait_until_cancelled(polling_id: str, polling_state: PollingStateReader, poll_seconds: float) -> None: + while not await _reads_cancelled(polling_id, polling_state): + await asyncio.sleep(poll_seconds) + + +async def stream_until_cancelled( + polling_id: str, + stream: Coroutine[object, object, None], + polling_state: PollingStateReader, + poll_seconds: float = CANCELLATION_POLL_SECONDS, +) -> None: + streaming: Final = asyncio.create_task(stream) + watching: Final = asyncio.create_task(_wait_until_cancelled(polling_id, polling_state, poll_seconds)) + try: + await asyncio.wait((streaming, watching), return_when=asyncio.FIRST_COMPLETED) + finally: + streaming.cancel() + watching.cancel() + await asyncio.gather(streaming, watching, return_exceptions=True) + if streaming.cancelled(): + verbose_proxy_logger.info("Stopped the background stream for %s because it was cancelled", polling_id) + return + streaming.result() + + async def background_streaming_task( polling_id: str, data: dict[str, object], @@ -117,7 +159,51 @@ async def background_streaming_task( user_max_tokens: int | None, user_api_base: str | None, version: str | None, -): +) -> None: + await stream_until_cancelled( + polling_id, + _stream_into_polling_state( + polling_id=polling_id, + data=data, + polling_handler=polling_handler, + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + general_settings=general_settings, + llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + select_data_generator=select_data_generator, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ), + polling_handler, + ) + + +async def _stream_into_polling_state( + polling_id: str, + data: dict[str, object], + polling_handler: ResponsePollingHandler, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth, + general_settings: dict[str, object], + llm_router: "Router | None", + proxy_config: "ProxyConfig", + proxy_logging_obj: "ProxyLogging", + select_data_generator: Callable[..., object] | None, + user_model: str | None, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, +) -> None: """ Background task to stream response and update cache diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index e50587b1a0e..9c8c8e27aed 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -3,7 +3,7 @@ Response Polling Handler for Background Responses with Cache """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timezone from typing import Final @@ -18,6 +18,24 @@ from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStatus _RESPONSE_POLLING_TARGET: Final = "response_polling" _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _JSON_LIST: Final = TypeAdapter(list[object]) +_SCRIPT_REPLY: Final = TypeAdapter(int) +_WRITTEN: Final = 1 + +_WRITE_UNLESS_CANCELLED_SCRIPT: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + return 0 +end +local decoded, stored = pcall(cjson.decode, current) +if decoded and type(stored) == 'string' then + decoded, stored = pcall(cjson.decode, stored) +end +if decoded and type(stored) == 'table' and stored['status'] == 'cancelled' then + return 0 +end +redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) +return 1 +""" class ResponsePollingHandler: @@ -102,7 +120,7 @@ class ResponsePollingHandler: incomplete_details: dict | None = None, reasoning: dict | None = None, tool_choice: object | None = None, - tools: list | None = None, + tools: Sequence[object] | None = None, output: list | None = None, # Additional ResponsesAPIResponse fields model: str | None = None, @@ -212,12 +230,19 @@ class ResponsePollingHandler: if store is not None: state["store"] = store - # Update cache with configured TTL - await self.redis_cache.async_set_cache( - key=cache_key, - value=json.dumps(state), - ttl=self.ttl, - ) + write_unless_cancelled: Final = self.redis_cache.async_register_script(_WRITE_UNLESS_CANCELLED_SCRIPT) + try: + written: Final = _SCRIPT_REPLY.validate_python( + await write_unless_cancelled(keys=(cache_key,), args=(json.dumps(json.dumps(state)), int(self.ttl))) + ) + except Exception as write_error: # noqa: BLE001 # a Redis fault must not end the background stream + verbose_proxy_logger.error("Failed to write the polling state for %s: %s", polling_id, write_error) + return + if written != _WRITTEN: + verbose_proxy_logger.debug( + "Skipped the polling state write for %s: the response was cancelled or has expired", polling_id + ) + return output_count: Final = len(_JSON_LIST.validate_python(state.get("output", []))) verbose_proxy_logger.debug( @@ -268,7 +293,7 @@ def should_use_polling_for_request( redis_cache, # RedisCache or None model: str, llm_router, # Router instance or None - native_background_mode: list[str] | None = None, # List of models that should use native background mode + native_background_mode: Sequence[str] | None = None, # List of models that should use native background mode ) -> bool: """ Determine if polling via cache should be used for a request. diff --git a/tests/integration/streaming/test_background_response_polling_cancel.py b/tests/integration/streaming/test_background_response_polling_cancel.py new file mode 100644 index 00000000000..b1c125005be --- /dev/null +++ b/tests/integration/streaming/test_background_response_polling_cancel.py @@ -0,0 +1,169 @@ +import threading +import uuid +from collections.abc import Callable, Iterator +from contextlib import ExitStack, contextmanager +from pathlib import Path +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.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 +from pydantic import JsonValue + +_MODEL: Final = "background-polling-cancel" +_NO_DATABASE: Final = ("DATABASE_URL", "DATABASE_URL_READ_REPLICA") +_DELTAS: Final = 40 +_DELTA_PAUSE_SECONDS: Final = 0.25 +_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":[]}') + + +def _config(directory: Path, wire: Wire) -> Path: + path: Final = directory / f"background-polling-cancel-{uuid.uuid4().hex}.yaml" + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": "openai/gpt-5.4-mini", + "api_base": wire.url, + "api_key": "synthetic-openai-key", + }, + } + ], + "general_settings": { + "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"}}}, + } + ) + ) + return path + + +@contextmanager +def _proxies_sharing_one_redis( + gateway: Gateway, directory: Path, wire: Wire, count: int +) -> 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", + } + yield tuple( + proxies.enter_context( + owned_proxy( + gateway, + directory / f"proxy-{index}", + overrides, + config=config, + database_setup=(), + remove_environment=_NO_DATABASE, + ) + ) + for index in range(count) + ) + + +def _serving(reply: Reply, requested: threading.Event) -> Callable[[Request], Reply]: + 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 + requested.set() + return reply + + return respond + + +def _create_background_response(proxy: Gateway) -> str: + created_response: Final = proxy.request( + "POST", "/v1/responses", {"model": _MODEL, "input": uuid.uuid4().hex, "background": True} + ) + assert created_response.status_code == 200, created_response.text + body: Final = object_value(created_response.json()) + assert body["status"] == "queued", body + return string_value(body["id"]) + + +def _cancel(proxy: Gateway, polling_id: str) -> None: + cancelled: Final = proxy.request("POST", f"/v1/responses/{polling_id}/cancel") + assert cancelled.status_code == 200, cancelled.text + assert object_value(cancelled.json())["status"] == "cancelled", cancelled.text + + +def _status(proxy: Gateway, polling_id: str) -> JsonValue: + polled: Final = proxy.request("GET", f"/v1/responses/{polling_id}") + assert polled.status_code == 200, polled.text + return object_value(polled.json())["status"] + + +@pytest.mark.timeout(2 * _OWNED_PROXY_SECONDS + 3 * _STREAM_SECONDS + 60) +def test_a_cancel_sent_to_another_proxy_stops_the_background_stream_and_the_response_stays_cancelled( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + slow_stream: Final = Reply( + content_type="text/event-stream", + chunks=( + frame(created(identity)), + *(frame(delta(identity, "word ")) for _ in range(_DELTAS)), + frame(completed(identity, "word " * _DELTAS)), + ), + pause_between_chunks=_DELTA_PAUSE_SECONDS, + ) + with ( + wire_server(_serving(slow_stream, requested)) as wire, + _proxies_sharing_one_redis(gateway, tmp_path, wire, 2) as (streaming, cancelling), + ): + polling_id: Final = _create_background_response(streaming) + eventually(requested.is_set, lambda seen: seen, seconds=30) + _cancel(cancelling, polling_id) + eventually(lambda: wire.disconnected.qsize(), lambda dropped: dropped == 1, seconds=_STREAM_SECONDS) + held( + lambda: _status(cancelling, polling_id), + lambda status: status == "cancelled", + holding=_STREAM_SECONDS, + seconds=2 * _STREAM_SECONDS, + ) + + +@pytest.mark.timeout(_OWNED_PROXY_SECONDS + 60) +def test_a_background_response_cancelled_just_before_its_stream_finishes_stays_cancelled( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + cancel_sent: Final = threading.Event() + held_stream: Final = Reply( + content_type="text/event-stream", + chunks=( + frame(created(identity)), + frame({"type": "response.in_progress", "response": response_object(identity, "in_progress")}), + frame(delta(identity, "done")), + frame(completed(identity, "done")), + ), + gate_after_first=cancel_sent, + ) + with ( + wire_server(_serving(held_stream, requested)) as wire, + _proxies_sharing_one_redis(gateway, tmp_path, wire, 1) as (proxy,), + ): + polling_id: Final = _create_background_response(proxy) + eventually(requested.is_set, lambda seen: seen, seconds=30) + _cancel(proxy, polling_id) + cancel_sent.set() + held(lambda: _status(proxy, polling_id), lambda status: status == "cancelled", holding=3) 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 42acad28488..8d33a7c2761 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -27,6 +27,48 @@ from pydantic import ValidationError from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler +class _RedisStandIn: + def __init__(self) -> None: + 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() + 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) + 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: + self.stored[key] = json.dumps(value) + self.ttls[key] = ttl + + def _status(self, key: str) -> Any: + stored = json.loads(self.stored[key]) + return json.loads(stored)["status"] if isinstance(stored, str) else stored["status"] + + def async_register_script(self, script: str): + async def write_unless_cancelled(keys, args): + if keys[0] not in self.stored or self._status(keys[0]) == "cancelled": + return 0 + self.stored[keys[0]] = args[0] + self.ttls[keys[0]] = int(args[1]) + return 1 + + return write_unless_cancelled + + +class _UnreachableRedisStandIn(_RedisStandIn): + def async_register_script(self, script: str): + async def write_fails(keys, args): + raise ConnectionError("redis is unreachable") + + return write_fails + + class TestResponsePollingHandler: """Test cases for ResponsePollingHandler""" @@ -151,8 +193,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_changes_status_to_in_progress(self): """Test that update_state can change status to in_progress""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -169,18 +210,14 @@ class TestResponsePollingHandler: status="in_progress", ) - # Verify the update was saved - mock_redis.async_set_cache.assert_called_once() - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "in_progress" @pytest.mark.asyncio async def test_update_state_replaces_full_output_list(self): """Test that update_state replaces the full output list""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -210,8 +247,7 @@ class TestResponsePollingHandler: output=new_output, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert len(stored["output"]) == 2 assert stored["output"][0]["id"] == "item_1" @@ -220,8 +256,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_usage(self): """Test that update_state correctly stores usage data""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -241,8 +276,7 @@ class TestResponsePollingHandler: usage=usage_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "completed" assert stored["usage"] == usage_data @@ -250,8 +284,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_reasoning_tools_tool_choice(self): """Test that update_state stores reasoning, tools, and tool_choice from response.completed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -277,8 +310,7 @@ class TestResponsePollingHandler: tools=tools_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["reasoning"] == reasoning_data assert stored["tool_choice"] == tool_choice_data @@ -287,8 +319,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_all_responses_api_fields(self): """Test that update_state stores all ResponsesAPIResponse fields from response.completed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -322,8 +353,7 @@ class TestResponsePollingHandler: incomplete_details={"reason": "max_output_tokens"}, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") # Verify all fields are stored correctly assert stored["status"] == "completed" @@ -351,8 +381,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_preserves_existing_fields(self): """Test that update_state preserves fields not being updated""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -372,8 +401,7 @@ class TestResponsePollingHandler: status="completed", ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") # Verify existing fields are preserved assert stored["status"] == "completed" @@ -384,8 +412,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_error_sets_failed_status(self): """Test that providing an error automatically sets status to failed""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -408,8 +435,7 @@ class TestResponsePollingHandler: error=error_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "failed" assert stored["error"] == error_data @@ -417,8 +443,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_with_incomplete_details(self): """Test that update_state stores incomplete_details""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -438,8 +463,7 @@ class TestResponsePollingHandler: incomplete_details=incomplete_details, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "incomplete" assert stored["incomplete_details"] == incomplete_details @@ -458,8 +482,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_handles_missing_cached_state(self): """Test that update_state handles case when cached state doesn't exist""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = None # Cache miss + mock_redis = _RedisStandIn() handler = ResponsePollingHandler(redis_cache=mock_redis) @@ -469,8 +492,7 @@ class TestResponsePollingHandler: status="in_progress", ) - # Should not try to set cache if nothing was found - mock_redis.async_set_cache.assert_not_called() + assert mock_redis.stored == {} # ==================== Get State Tests ==================== @@ -520,8 +542,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_cancel_polling_updates_status_to_cancelled(self): """Test that cancel_polling sets status to cancelled""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -537,10 +558,67 @@ class TestResponsePollingHandler: assert result is True - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["status"] == "cancelled" + @pytest.mark.asyncio + async def test_an_update_after_a_cancel_leaves_the_response_cancelled(self): + mock_redis = _RedisStandIn.holding( + { + "id": "litellm_poll_test", + "object": "response", + "status": "in_progress", + "output": [], + "created_at": 1234567890, + } + ) + handler = ResponsePollingHandler(redis_cache=mock_redis) + + await handler.cancel_polling("litellm_poll_test") + await handler.update_state( + polling_id="litellm_poll_test", + status="completed", + output=[{"id": "item_1", "type": "message"}], + ) + + assert await handler.get_state("litellm_poll_test") == { + "id": "litellm_poll_test", + "object": "response", + "status": "cancelled", + "output": [], + "created_at": 1234567890, + } + + @pytest.mark.asyncio + async def test_a_failed_state_write_keeps_the_stored_state_and_does_not_raise(self): + state = { + "id": "litellm_poll_test", + "object": "response", + "status": "in_progress", + "output": [], + "created_at": 1234567890, + } + handler = 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()) + + queued = 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") + + assert await handler.get_state("litellm_poll_test") == { + **queued.model_dump(), + "status": "in_progress", + "model": "gpt-4o", + } + # ==================== Delete Polling Tests ==================== @pytest.mark.asyncio @@ -594,8 +672,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_update_state_uses_configured_ttl(self): """Test that update_state uses the configured TTL""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -612,8 +689,7 @@ class TestResponsePollingHandler: status="in_progress", ) - call_args = mock_redis.async_set_cache.call_args - assert call_args.kwargs["ttl"] == 1800 + assert mock_redis.ttls == {ResponsePollingHandler.get_cache_key("litellm_poll_test"): 1800} class TestStreamingEventProcessing: @@ -1434,6 +1510,12 @@ def _make_background_streaming_kwargs( ) +def _streaming_polling_handler(status: str = "in_progress") -> AsyncMock: + handler = AsyncMock(spec=ResponsePollingHandler) + handler.get_state.return_value = {"status": status} + return handler + + @pytest.mark.xdist_group("heavy_imports") class TestBackgroundStreamingTerminalEvents: """ @@ -1468,7 +1550,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_1", handler) with patch( @@ -1507,7 +1589,7 @@ class TestBackgroundStreamingTerminalEvents: mock_response = Mock() mock_response.body_iterator = _body_iterator() - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_named_event", handler) with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests @@ -1550,7 +1632,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_2", handler) with patch( @@ -1589,7 +1671,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_3", handler) with patch( @@ -1624,7 +1706,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_4", handler) with patch( @@ -1653,7 +1735,7 @@ class TestBackgroundStreamingTerminalEvents: events: Final = [ {"type": "response.completed", "response": {"id": "resp_123", "status": "completed", "output": [item]}} ] - handler: Final = AsyncMock(spec=ResponsePollingHandler) + handler: Final = _streaming_polling_handler() with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" @@ -1691,7 +1773,7 @@ class TestBackgroundStreamingTerminalEvents: }, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_5", handler) with patch( @@ -1717,7 +1799,7 @@ class TestBackgroundStreamingTerminalEvents: {"type": "response.in_progress"}, ] mock_response = _make_sse_stream(events) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_6", handler) with patch( @@ -1765,7 +1847,7 @@ class TestBackgroundStreamingTerminalEvents: slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] ) - handler = AsyncMock(spec=ResponsePollingHandler) + handler = _streaming_polling_handler() kwargs = _make_background_streaming_kwargs("poll_7", handler) kwargs["request"] = Request( { @@ -1789,6 +1871,160 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["usage"] == {"input_tokens": 13, "output_tokens": 10} +class _CancelledAfterReads: + def __init__(self, reads_before_cancel: int) -> None: + self.reads = 0 + self.reads_before_cancel = reads_before_cancel + + async def get_state(self, polling_id: str) -> Dict[str, Any]: + self.reads += 1 + return {"id": polling_id, "status": "cancelled" if self.reads > self.reads_before_cancel else "in_progress"} + + +class _UnreadableState: + def __init__(self) -> None: + self.reads = 0 + + async def get_state(self, polling_id: str) -> Dict[str, Any]: + self.reads += 1 + raise ConnectionError("redis is unreachable") + + +async def _let_the_loop_run() -> None: + for _ in range(20): + await asyncio.sleep(0) + + +@pytest.mark.xdist_group("heavy_imports") +class TestStreamUntilCancelled: + @pytest.mark.asyncio + async def test_a_cancelled_state_stops_the_stream(self): + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state = _CancelledAfterReads(reads_before_cancel=2) + stopped = asyncio.Event() + + async def endless_stream() -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + + await asyncio.wait_for( + stream_until_cancelled("litellm_poll_test", endless_stream(), state, poll_seconds=0), timeout=5 + ) + + assert stopped.is_set() + assert state.reads == 3 + + @pytest.mark.asyncio + async def test_failing_state_reads_leave_the_stream_running_to_its_end(self): + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state = _UnreadableState() + finished = asyncio.Event() + + async def stream_outliving_three_reads() -> None: + while state.reads < 3: + await asyncio.sleep(0) + finished.set() + + await asyncio.wait_for( + stream_until_cancelled("litellm_poll_test", stream_outliving_three_reads(), state, poll_seconds=0), + timeout=5, + ) + + assert finished.is_set() + + @pytest.mark.asyncio + async def test_a_stream_that_finishes_first_leaves_no_watcher_reading(self): + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + state = _CancelledAfterReads(reads_before_cancel=1_000_000) + + async def short_stream() -> None: + await _let_the_loop_run() + + await asyncio.wait_for( + stream_until_cancelled("litellm_poll_test", short_stream(), state, poll_seconds=0), timeout=5 + ) + reads_when_done = 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): + from litellm.proxy.response_polling.background_streaming import stream_until_cancelled + + async def broken_stream() -> None: + raise ValueError("upstream broke") + + with pytest.raises(ValueError, match="upstream broke"): + await stream_until_cancelled( + "litellm_poll_test", broken_stream(), _CancelledAfterReads(1_000_000), poll_seconds=0 + ) + + @pytest.mark.asyncio + async def test_cancelling_the_caller_stops_both_the_stream_and_the_watcher(self): + 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() + + async def endless_stream() -> None: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + + caller = 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 + 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): + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + stopped = asyncio.Event() + + async def endless_body(): + yield f"data: {json.dumps({'type': 'response.in_progress'})}" + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + + endless_response = Mock() + endless_response.body_iterator = endless_body() + handler = _streaming_polling_handler(status="cancelled") + kwargs = _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" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=endless_response) + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + assert stopped.is_set() + assert {call.kwargs.get("status") for call in handler.update_state.await_args_list} <= {None, "in_progress"} + + class TestEdgeCases: """Test edge cases and error scenarios""" @@ -1849,8 +2085,7 @@ class TestEdgeCases: @pytest.mark.asyncio async def test_update_state_with_none_output_clears_output(self): """Test that output=[] explicitly sets empty output""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = json.dumps( + mock_redis = _RedisStandIn.holding( { "id": "litellm_poll_test", "object": "response", @@ -1867,8 +2102,7 @@ class TestEdgeCases: output=[], # Explicitly set empty ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await handler.get_state("litellm_poll_test") assert stored["output"] == [] From c4bcf3227477adb795f34021524c2e36efa4ee0b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:07:15 -0700 Subject: [PATCH 2/3] fix(responses): settle a cancelled background call and stop it on delete A cancel that lands before the provider answers now reports one client disconnect through the failure hooks, so the key's parallel slot and budget reservation come back and the spend logs get a failure row. The state guard and the watcher compare a status-first prefix with GETRANGE instead of decoding the stored response inside Redis, so each write and read costs the same whatever the response size. The watcher also stops when the response was deleted, and a cancelled stream closes its upstream response. The TTFT keepalive now cancels the in-flight call when its caller is cancelled before the first ping. --- litellm/proxy/common_request_processing.py | 38 +- .../response_polling/background_streaming.py | 367 +++++++++--------- .../proxy/response_polling/polling_handler.py | 36 +- ...test_background_response_polling_cancel.py | 32 ++ .../test_response_polling_handler.py | 217 ++++++++++- .../proxy/test_common_request_processing.py | 22 ++ 6 files changed, 506 insertions(+), 206 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9bb451495c8..cf411bbc2e1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1351,7 +1351,7 @@ def ttft_keepalive_interval(request_data: Mapping[str, object], llm_router: Rout return resolve_ttft_keepalive_interval(deployments, litellm.sse_keepalive_ping_interval_seconds) -async def _aclose_late_response(produced: Response) -> None: +async def aclose_streaming_response(produced: object) -> None: """Run the cleanup Starlette would have run, for a response it never called. Closing an already-closed async generator is a no-op, so this is safe to call @@ -1374,6 +1374,21 @@ async def _aclose_late_response(produced: Response) -> None: verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) +async def _abandon_late_producer(produce_task: "asyncio.Future[_LateResponseT]") -> None: + if not produce_task.done(): + produce_task.cancel() + with anyio.CancelScope(shield=True): + with contextlib.suppress(BaseException): + await produce_task + elif not produce_task.cancelled(): + # The upstream may have answered while nobody was draining this + # relay, e.g. the client vanished first. Nothing else holds that + # response, so its stream only gets closed here. + with anyio.CancelScope(shield=True): + with contextlib.suppress(BaseException): + await aclose_streaming_response(produce_task.result()) + + 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): @@ -1390,7 +1405,7 @@ async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None # Starlette never called this response, so the cleanup its __call__ would # have run has to happen here or the upstream LLM connection leaks. with anyio.CancelScope(shield=True): - await _aclose_late_response(produced) + await aclose_streaming_response(produced) async def _sanitized_late_failure( @@ -1446,7 +1461,11 @@ async def open_sse_before_first_byte( # Share admission state with stream cleanup before the producer copies this context. get_or_create_request_stash() produce_task: Final = asyncio.ensure_future(produce_response) - await asyncio.wait((produce_task,), timeout=interval) + try: + await asyncio.wait((produce_task,), timeout=interval) + except BaseException: + await _abandon_late_producer(produce_task) + raise if produce_task.done(): # Fast path: the upstream answered inside one interval, so nothing was # written early and this is byte-identical to not being wrapped at all. @@ -1474,18 +1493,7 @@ async def open_sse_before_first_byte( async for chunk in _relay_late_response(produced): yield chunk finally: - if not produce_task.done(): - produce_task.cancel() - with anyio.CancelScope(shield=True): - with contextlib.suppress(BaseException): - await produce_task - elif not produce_task.cancelled(): - # The upstream may have answered while nobody was draining this - # relay, e.g. the client vanished first. Nothing else holds that - # response, so its stream only gets closed here. - with anyio.CancelScope(shield=True): - with contextlib.suppress(BaseException): - await _aclose_late_response(produce_task.result()) + await _abandon_late_producer(produce_task) verbose_proxy_logger.info( "no upstream response after %ss, opening the SSE response early and sending keepalives", interval diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index e9cad2dbe80..fcdab710206 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -13,14 +13,16 @@ import json from collections.abc import AsyncIterable, Callable, Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable -from fastapi import Request, Response +import anyio +from fastapi import HTTPException, Request, Response from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_CLIENT_DISCONNECT_DETAIL, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing, aclose_streaming_response from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus @@ -84,7 +86,7 @@ class _StreamsBody(Protocol): class PollingStateReader(Protocol): - async def get_state(self, polling_id: str) -> Mapping[str, object] | None: ... + async def is_cancelled_or_gone(self, polling_id: str) -> bool: ... def _sse_frame_data(frame: str) -> str | None: @@ -105,19 +107,18 @@ def detach_request_from_client(request: Request) -> Request: return Request(request.scope, _never_receive) -async def _reads_cancelled(polling_id: str, polling_state: PollingStateReader) -> bool: +async def _reads_stopped(polling_id: str, polling_state: PollingStateReader) -> bool: try: - state: Final = await polling_state.get_state(polling_id) + return await polling_state.is_cancelled_or_gone(polling_id) except Exception as error: # noqa: BLE001 # a failed read must never stop the upstream call verbose_proxy_logger.warning( "Could not read the polling state of %s, so it keeps streaming: %s", polling_id, error ) return False - return state is not None and state.get("status") == "cancelled" -async def _wait_until_cancelled(polling_id: str, polling_state: PollingStateReader, poll_seconds: float) -> None: - while not await _reads_cancelled(polling_id, polling_state): +async def _wait_until_stopped(polling_id: str, polling_state: PollingStateReader, poll_seconds: float) -> None: + while not await _reads_stopped(polling_id, polling_state): await asyncio.sleep(poll_seconds) @@ -128,7 +129,7 @@ async def stream_until_cancelled( poll_seconds: float = CANCELLATION_POLL_SECONDS, ) -> None: streaming: Final = asyncio.create_task(stream) - watching: Final = asyncio.create_task(_wait_until_cancelled(polling_id, polling_state, poll_seconds)) + watching: Final = asyncio.create_task(_wait_until_stopped(polling_id, polling_state, poll_seconds)) try: await asyncio.wait((streaming, watching), return_when=asyncio.FIRST_COMPLETED) finally: @@ -136,7 +137,9 @@ async def stream_until_cancelled( watching.cancel() await asyncio.gather(streaming, watching, return_exceptions=True) if streaming.cancelled(): - verbose_proxy_logger.info("Stopped the background stream for %s because it was cancelled", polling_id) + verbose_proxy_logger.info( + "Stopped the background stream for %s because it was cancelled or deleted", polling_id + ) return streaming.result() @@ -217,41 +220,53 @@ async def _stream_into_polling_state( try: verbose_proxy_logger.info("Starting background streaming for %s", polling_id) - # Update status to in_progress (OpenAI format) - await polling_handler.update_state( - polling_id=polling_id, - status="in_progress", - ) + try: + # Update status to in_progress (OpenAI format) + await polling_handler.update_state( + polling_id=polling_id, + status="in_progress", + ) - # Force streaming mode and remove background flag - data["stream"] = True - data.pop("background", None) + # Force streaming mode and remove background flag + data["stream"] = True + data.pop("background", None) - # Create processor - processor: Final = ProxyBaseLLMRequestProcessing(data=data) + # Create processor + processor: Final = ProxyBaseLLMRequestProcessing(data=data) - # Make streaming request. - # Pre-call checks (rate limits, guardrails, budget) were already run - # before polling ID creation, so skip them here to avoid double-counting. - response: Final = await processor.base_process_llm_request( - request=detach_request_from_client(request), - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="aresponses", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - skip_pre_call_logic=True, - ) + # Make streaming request. + # Pre-call checks (rate limits, guardrails, budget) were already run + # before polling ID creation, so skip them here to avoid double-counting. + response: Final = await processor.base_process_llm_request( + request=detach_request_from_client(request), + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="aresponses", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + skip_pre_call_logic=True, + ) + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=HTTPException( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail=LITELLM_CLIENT_DISCONNECT_DETAIL + ), + request_data=data, + dispatched_call_type="aresponses", + ) + raise # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming @@ -315,158 +330,162 @@ async def _stream_into_polling_state( ) if isinstance(response, _StreamsBody): - async for chunk in response.body_iterator: - # Parse chunk - if isinstance(chunk, bytes): - chunk = chunk.decode("utf-8") + try: + async for chunk in response.body_iterator: + # Parse chunk + if isinstance(chunk, bytes): + chunk = chunk.decode("utf-8") - if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: - if chunk_data == "[DONE]": - break + if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: + if chunk_data == "[DONE]": + break - try: - event: _StreamEvent = _StreamEventParser.parse(chunk_data) - event_type = event.get("type", "") + try: + event: _StreamEvent = _StreamEventParser.parse(chunk_data) + event_type = event.get("type", "") - # Process different event types based on OpenAI streaming spec - if event_type == "response.output_item.added": - # New output item added - item = event.get("item", {}) - item_id = item.get("id") - if item_id: - output_items[item_id] = item - state_dirty = True + # Process different event types based on OpenAI streaming spec + if event_type == "response.output_item.added": + # New output item added + item = event.get("item", {}) + item_id = item.get("id") + if item_id: + output_items[item_id] = item + state_dirty = True - elif event_type == "response.content_part.added": - # Content part added to an output item - item_id = event.get("item_id") - content_part = event.get("part", {}) + elif event_type == "response.content_part.added": + # Content part added to an output item + item_id = event.get("item_id") + content_part = event.get("part", {}) - if item_id and item_id in output_items: - # Update the output item with new content - added_item = output_items[item_id] - output_items[item_id] = { - **added_item, - "content": (*added_item.get("content", ()), content_part), - } - state_dirty = True + if item_id and item_id in output_items: + # Update the output item with new content + added_item = output_items[item_id] + output_items[item_id] = { + **added_item, + "content": (*added_item.get("content", ()), content_part), + } + state_dirty = True - elif event_type == "response.output_text.delta": - # Text delta - accumulate text content - # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta - item_id = event.get("item_id") - content_index = event.get("content_index", 0) - delta = event.get("delta", "") + elif event_type == "response.output_text.delta": + # Text delta - accumulate text content + # https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta + item_id = event.get("item_id") + content_index = event.get("content_index", 0) + delta = event.get("delta", "") - if item_id and item_id in output_items: - # Accumulate text delta - key = (item_id, content_index) - if key not in accumulated_text: - accumulated_text[key] = "" - accumulated_text[key] += delta + if item_id and item_id in output_items: + # Accumulate text delta + key = (item_id, content_index) + if key not in accumulated_text: + accumulated_text[key] = "" + accumulated_text[key] += delta - # Update the content in output_items - delta_item = output_items[item_id] - if "content" in delta_item: - content_list = delta_item["content"] - if content_index < len(content_list): - content_entry = content_list[content_index] - if isinstance(content_entry, dict): - content_entry["text"] = accumulated_text[key] - state_dirty = True + # Update the content in output_items + delta_item = output_items[item_id] + if "content" in delta_item: + content_list = delta_item["content"] + if content_index < len(content_list): + content_entry = content_list[content_index] + if isinstance(content_entry, dict): + content_entry["text"] = accumulated_text[key] + state_dirty = True - elif event_type == "response.content_part.done": - # Content part completed - item_id = event.get("item_id") - content_part = event.get("part", {}) - content_index = event.get("content_index", 0) + elif event_type == "response.content_part.done": + # Content part completed + item_id = event.get("item_id") + content_part = event.get("part", {}) + content_index = event.get("content_index", 0) - if item_id and item_id in output_items: - # Update with final content from event - done_item = output_items[item_id] - if "content" in done_item: - content_list = done_item["content"] - if content_index < len(content_list): - output_items[item_id] = { - **done_item, - "content": tuple( - content_part if part_index == content_index else existing_part - for part_index, existing_part in enumerate(content_list) - ), - } - state_dirty = True + if item_id and item_id in output_items: + # Update with final content from event + done_item = output_items[item_id] + if "content" in done_item: + content_list = done_item["content"] + if content_index < len(content_list): + output_items[item_id] = { + **done_item, + "content": tuple( + content_part if part_index == content_index else existing_part + for part_index, existing_part in enumerate(content_list) + ), + } + state_dirty = True - elif event_type == "response.output_item.done": - # Output item completed - use final item data - item = event.get("item", {}) - item_id = item.get("id") - if item_id: - output_items[item_id] = item - state_dirty = True + elif event_type == "response.output_item.done": + # Output item completed - use final item data + item = event.get("item", {}) + item_id = item.get("id") + if item_id: + output_items[item_id] = item + state_dirty = True - elif event_type == "response.in_progress": - # Response is now in progress - # https://platform.openai.com/docs/api-reference/responses-streaming/response-in-progress - await polling_handler.update_state( - polling_id=polling_id, - status="in_progress", - ) + elif event_type == "response.in_progress": + # Response is now in progress + # https://platform.openai.com/docs/api-reference/responses-streaming/response-in-progress + await polling_handler.update_state( + polling_id=polling_id, + status="in_progress", + ) - elif event_type in ( - "response.completed", - "response.failed", - "response.incomplete", - "response.cancelled", - ): - # Terminal event - extract all ResponsesAPIResponse fields - # https://platform.openai.com/docs/api-reference/responses-streaming - response_data = event.get("response", {}) - terminal_status = response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ) + elif event_type in ( + "response.completed", + "response.failed", + "response.incomplete", + "response.cancelled", + ): + # Terminal event - extract all ResponsesAPIResponse fields + # https://platform.openai.com/docs/api-reference/responses-streaming + response_data = event.get("response", {}) + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), + ) - # Extract error for failed and incomplete responses - if event_type == "response.failed" or event_type == "response.incomplete": - terminal_error = response_data.get("error") + # Extract error for failed and incomplete responses + if event_type == "response.failed" or event_type == "response.incomplete": + terminal_error = response_data.get("error") - # Core response fields - usage_data = response_data.get("usage") - reasoning_data = response_data.get("reasoning") - tool_choice_data = response_data.get("tool_choice") - tools_data = response_data.get("tools") + # Core response fields + usage_data = response_data.get("usage") + reasoning_data = response_data.get("reasoning") + tool_choice_data = response_data.get("tool_choice") + tools_data = response_data.get("tools") - # Additional ResponsesAPIResponse fields - model_data = response_data.get("model") - instructions_data = response_data.get("instructions") - temperature_data = response_data.get("temperature") - top_p_data = response_data.get("top_p") - max_output_tokens_data = response_data.get("max_output_tokens") - previous_response_id_data = response_data.get("previous_response_id") - text_data = response_data.get("text") - truncation_data = response_data.get("truncation") - parallel_tool_calls_data = response_data.get("parallel_tool_calls") - user_data = response_data.get("user") - store_data = response_data.get("store") - incomplete_details_data = response_data.get("incomplete_details") + # Additional ResponsesAPIResponse fields + model_data = response_data.get("model") + instructions_data = response_data.get("instructions") + temperature_data = response_data.get("temperature") + top_p_data = response_data.get("top_p") + max_output_tokens_data = response_data.get("max_output_tokens") + previous_response_id_data = response_data.get("previous_response_id") + text_data = response_data.get("text") + truncation_data = response_data.get("truncation") + parallel_tool_calls_data = response_data.get("parallel_tool_calls") + user_data = response_data.get("user") + store_data = response_data.get("store") + incomplete_details_data = response_data.get("incomplete_details") - # Also update output from final response if available - if "output" in response_data: - final_output = response_data.get("output", []) - for item in final_output: - item_id = item.get("id") - if item_id: - output_items[item_id] = item - state_dirty = True + # Also update output from final response if available + if "output" in response_data: + final_output = response_data.get("output", []) + for item in final_output: + item_id = item.get("id") + if item_id: + output_items[item_id] = item + state_dirty = True - # Flush state to Redis if interval elapsed - await flush_state_if_needed() + # Flush state to Redis if interval elapsed + await flush_state_if_needed() - except json.JSONDecodeError as e: - verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) + except json.JSONDecodeError as e: + verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) - # Final flush to ensure all accumulated state is saved - await flush_state_if_needed(force=True) + # Final flush to ensure all accumulated state is saved + await flush_state_if_needed(force=True) + finally: + with anyio.CancelScope(shield=True): + await aclose_streaming_response(response) # Use the terminal status from the stream, default to "completed" final_status: Final = terminal_status or "completed" diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 9c8c8e27aed..e4ea6b873da 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -21,22 +21,32 @@ _JSON_LIST: Final = TypeAdapter(list[object]) _SCRIPT_REPLY: Final = TypeAdapter(int) _WRITTEN: Final = 1 + +def _encoded_state(state: Mapping[str, object]) -> str: + return json.dumps(json.dumps({"status": state["status"], **state})) + + +_CANCELLED_PREFIX: Final = _encoded_state({"status": "cancelled"}).removesuffix('}"') + _WRITE_UNLESS_CANCELLED_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then +if redis.call('EXISTS', KEYS[1]) == 0 then return 0 end -local decoded, stored = pcall(cjson.decode, current) -if decoded and type(stored) == 'string' then - decoded, stored = pcall(cjson.decode, stored) -end -if decoded and type(stored) == 'table' and stored['status'] == 'cancelled' then +if redis.call('GETRANGE', KEYS[1], 0, string.len(ARGV[3]) - 1) == ARGV[3] then return 0 end redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) return 1 """ +_CANCELLED_OR_GONE_SCRIPT: Final = """ +local head = redis.call('GETRANGE', KEYS[1], 0, string.len(ARGV[1]) - 1) +if head == '' or head == ARGV[1] then + return 1 +end +return 0 +""" + class ResponsePollingHandler: """Handles polling-based responses with Redis cache""" @@ -233,7 +243,9 @@ class ResponsePollingHandler: write_unless_cancelled: Final = self.redis_cache.async_register_script(_WRITE_UNLESS_CANCELLED_SCRIPT) try: written: Final = _SCRIPT_REPLY.validate_python( - await write_unless_cancelled(keys=(cache_key,), args=(json.dumps(json.dumps(state)), int(self.ttl))) + await write_unless_cancelled( + keys=(cache_key,), args=(_encoded_state(state), int(self.ttl), _CANCELLED_PREFIX) + ) ) except Exception as write_error: # noqa: BLE001 # a Redis fault must not end the background stream verbose_proxy_logger.error("Failed to write the polling state for %s: %s", polling_id, write_error) @@ -263,6 +275,14 @@ class ResponsePollingHandler: return None + @with_service_target(_RESPONSE_POLLING_TARGET) + async def is_cancelled_or_gone(self, polling_id: str) -> bool: + if not self.redis_cache: + return False + cancelled_or_gone: Final = self.redis_cache.async_register_script(_CANCELLED_OR_GONE_SCRIPT) + reply: Final = await cancelled_or_gone(keys=(self.get_cache_key(polling_id),), args=(_CANCELLED_PREFIX,)) + return _SCRIPT_REPLY.validate_python(reply) == 1 + async def cancel_polling(self, polling_id: str) -> bool: """ Cancel a polling request diff --git a/tests/integration/streaming/test_background_response_polling_cancel.py b/tests/integration/streaming/test_background_response_polling_cancel.py index b1c125005be..ee59d0708d7 100644 --- a/tests/integration/streaming/test_background_response_polling_cancel.py +++ b/tests/integration/streaming/test_background_response_polling_cancel.py @@ -167,3 +167,35 @@ def test_a_background_response_cancelled_just_before_its_stream_finishes_stays_c _cancel(proxy, polling_id) cancel_sent.set() held(lambda: _status(proxy, polling_id), lambda status: status == "cancelled", holding=3) + + +@pytest.mark.timeout(2 * _OWNED_PROXY_SECONDS + 3 * _STREAM_SECONDS + 60) +def test_a_delete_sent_to_another_proxy_stops_the_background_stream_and_the_response_stays_gone( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + requested: Final = threading.Event() + slow_stream: Final = Reply( + content_type="text/event-stream", + chunks=( + frame(created(identity)), + *(frame(delta(identity, "word ")) for _ in range(_DELTAS)), + frame(completed(identity, "word " * _DELTAS)), + ), + pause_between_chunks=_DELTA_PAUSE_SECONDS, + ) + with ( + wire_server(_serving(slow_stream, requested)) as wire, + _proxies_sharing_one_redis(gateway, tmp_path, wire, 2) as (streaming, deleting), + ): + polling_id: Final = _create_background_response(streaming) + eventually(requested.is_set, lambda seen: seen, seconds=30) + deleted: Final = deleting.request("DELETE", f"/v1/responses/{polling_id}") + assert deleted.status_code == 200, deleted.text + eventually(lambda: wire.disconnected.qsize(), lambda dropped: dropped == 1, seconds=_STREAM_SECONDS) + held( + lambda: streaming.request("GET", f"/v1/responses/{polling_id}").status_code, + lambda status_code: status_code == 404, + holding=_STREAM_SECONDS, + seconds=2 * _STREAM_SECONDS, + ) 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 8d33a7c2761..9195447816c 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -46,19 +46,21 @@ class _RedisStandIn: self.stored[key] = json.dumps(value) self.ttls[key] = ttl - def _status(self, key: str) -> Any: - stored = json.loads(self.stored[key]) - return json.loads(stored)["status"] if isinstance(stored, str) else stored["status"] + 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): - if keys[0] not in self.stored or self._status(keys[0]) == "cancelled": + 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 - return write_unless_cancelled + async def cancelled_or_gone(keys, args): + 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): @@ -537,6 +539,12 @@ 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) + + assert await handler.is_cancelled_or_gone("litellm_poll_test") is False + # ==================== Cancel Polling Tests ==================== @pytest.mark.asyncio @@ -589,6 +597,31 @@ class TestResponsePollingHandler: "created_at": 1234567890, } + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("stop", "expected"), + [ + (None, False), + ("cancel", True), + ("delete", True), + ], + ) + async def test_a_cancelled_or_deleted_response_reads_as_stopped(self, stop, expected): + handler = ResponsePollingHandler( + redis_cache=_RedisStandIn.holding( + {"id": "litellm_poll_test", "object": "response", "output": [], "status": "in_progress"} + ) + ) + await handler.update_state(polling_id="litellm_poll_test", output=[{"id": "item_1"}]) + + if stop == "cancel": + await handler.cancel_polling("litellm_poll_test") + if stop == "delete": + await handler.delete_polling("litellm_poll_test") + + assert await handler.is_cancelled_or_gone("litellm_poll_test") is expected + 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 = { @@ -1499,7 +1532,7 @@ def _make_background_streaming_kwargs( general_settings={}, llm_router=None, proxy_config=Mock(), - proxy_logging_obj=Mock(), + proxy_logging_obj=AsyncMock(), select_data_generator=Mock(), user_model=None, user_temperature=None, @@ -1513,6 +1546,7 @@ def _make_background_streaming_kwargs( def _streaming_polling_handler(status: str = "in_progress") -> AsyncMock: handler = AsyncMock(spec=ResponsePollingHandler) handler.get_state.return_value = {"status": status} + handler.is_cancelled_or_gone.return_value = status == "cancelled" return handler @@ -1876,16 +1910,16 @@ class _CancelledAfterReads: self.reads = 0 self.reads_before_cancel = reads_before_cancel - async def get_state(self, polling_id: str) -> Dict[str, Any]: + async def is_cancelled_or_gone(self, polling_id: str) -> bool: self.reads += 1 - return {"id": polling_id, "status": "cancelled" if self.reads > self.reads_before_cancel else "in_progress"} + return self.reads > self.reads_before_cancel class _UnreadableState: def __init__(self) -> None: self.reads = 0 - async def get_state(self, polling_id: str) -> Dict[str, Any]: + async def is_cancelled_or_gone(self, polling_id: str) -> bool: self.reads += 1 raise ConnectionError("redis is unreachable") @@ -2023,7 +2057,172 @@ class TestStreamUntilCancelled: assert stopped.is_set() assert {call.kwargs.get("status") for call in handler.update_state.await_args_list} <= {None, "in_progress"} + 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): + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + stopped = asyncio.Event() + + async def body_before_its_first_event(): + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + stopped.set() + raise + yield "data: [DONE]" + + silent_response = 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) + + 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 = AsyncMock(return_value=silent_response) + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + assert stopped.is_set() + 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): + from fastapi import HTTPException + + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + upstream_cancelled = asyncio.Event() + + async def upstream_that_never_answers(**_: object) -> None: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + handler = _streaming_polling_handler(status="cancelled") + kwargs = _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" + ) as MockProcessor: + 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.assert_awaited_once() + reported = report.await_args.kwargs + assert upstream_cancelled.is_set() + assert isinstance(reported["original_exception"], HTTPException) + assert reported["original_exception"].status_code == 499 + assert reported["request_data"] is kwargs["data"] + assert reported["user_api_key_dict"] is kwargs["user_api_key_dict"] + assert reported["dispatched_call_type"] == "aresponses" + + @pytest.mark.asyncio + async def test_a_cancel_during_a_blocked_state_write_closes_the_upstream_stream(self): + from starlette.responses import StreamingResponse + + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + closed = asyncio.Event() + + async def body(): + try: + yield f"data: {json.dumps({'type': 'response.output_item.added', 'item': {'id': 'item_1'}})}" + yield "data: [DONE]" + yield f"data: {json.dumps({'type': 'response.in_progress'})}" + finally: + closed.set() + + async def output_writes_block(**fields: object) -> None: + if "output" in fields: + await asyncio.Event().wait() + + handler = _streaming_polling_handler(status="cancelled") + handler.update_state.side_effect = output_writes_block + kwargs = _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" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=StreamingResponse(body())) + await asyncio.wait_for(background_streaming_task(**kwargs), timeout=5) + + assert closed.is_set() + 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": []}, +} +_PART_ADDED: Final = { + "type": "response.content_part.added", + "item_id": "msg_1", + "content_index": 0, + "part": {"type": "output_text", "text": ""}, +} +_DELTAS: Final = tuple( + {"type": "response.output_text.delta", "item_id": "msg_1", "content_index": 0, "delta": delta} + for delta in ("Hel", "lo") +) +_PART_DONE: Final = { + "type": "response.content_part.done", + "item_id": "msg_1", + "content_index": 0, + "part": {"type": "output_text", "text": "Hello!"}, +} +_MESSAGE_DONE: Final = { + "type": "response.output_item.done", + "item": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello, world"}], + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("events", "expected_content"), + [ + ((_MESSAGE_ADDED, _PART_ADDED, *_DELTAS), [{"type": "output_text", "text": "Hello"}]), + ((_MESSAGE_ADDED, _PART_ADDED, *_DELTAS, _PART_DONE), [{"type": "output_text", "text": "Hello!"}]), + ( + (_MESSAGE_ADDED, _PART_ADDED, *_DELTAS, _PART_DONE, _MESSAGE_DONE), + [{"type": "output_text", "text": "Hello, world"}], + ), + ], +) +async def test_a_background_stream_writes_the_message_its_events_build(events, expected_content): + from litellm.proxy.response_polling.background_streaming import background_streaming_task + + async def body(): + yield "data: {not json" + for event in events: + yield f"data: {json.dumps(event)}" + yield "data: [DONE]" + + streamed = Mock() + streamed.body_iterator = body() + handler = _streaming_polling_handler() + kwargs = _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" + ) as MockProcessor: + 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] + 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 fdcbcbfcd05..078990c7d98 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10587,6 +10587,28 @@ async def test_ttft_keepalive_cancels_the_in_flight_call_when_the_client_gives_u assert upstream_cancelled.is_set() +@pytest.mark.asyncio +async def test_ttft_keepalive_cancels_the_in_flight_call_when_its_caller_is_cancelled_before_the_first_ping(): + upstream_started = asyncio.Event() + upstream_cancelled = asyncio.Event() + + async def never_answers(): + upstream_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + caller = asyncio.create_task(open_sse_before_first_byte(never_answers(), ping_interval_seconds=60)) + await upstream_started.wait() + caller.cancel() + with pytest.raises(asyncio.CancelledError): + await caller + + assert upstream_cancelled.is_set() + + @pytest.mark.parametrize( "request_data, global_interval, expected", [ 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 3/3] 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():