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