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] 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"] == []