From a624324220fa0babbef25ef646e1cd63d2b832d1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 29 Oct 2026 18:36:06 -0700 Subject: [PATCH] refactor(redis-store): move response polling behind a store Co-Authored-By: Claude Opus 5.5 --- .../skills/redis-access/references/store.md | 7 +- .../proxy/response_polling/polling_handler.py | 75 ++--- litellm/redis/stores/response_polling.py | 50 +++ .../test_response_polling_handler.py | 309 ++++++++---------- 4 files changed, 231 insertions(+), 210 deletions(-) create mode 100644 litellm/redis/stores/response_polling.py diff --git a/.agents/skills/redis-access/references/store.md b/.agents/skills/redis-access/references/store.md index 3bd9b9850f2..75efc19bb8d 100644 --- a/.agents/skills/redis-access/references/store.md +++ b/.agents/skills/redis-access/references/store.md @@ -8,6 +8,8 @@ Keys, encoding, TTLs, Lua and reply decoding for each keyspace are spread across ## Shape +`litellm/redis/stores/response_polling.py` is the reference + - One module per keyspace in `litellm/redis/stores/.py` - The module defines a `Protocol` of domain operations, named for what the caller wants (`acquire`, `release`, `increment`, `get`), never for Redis commands - The Python implementation is a frozen dataclass that takes a `RedisTransport` and runs its own commands on it. It owns, and nothing outside it knows @@ -33,6 +35,8 @@ class TransportSessionBudgetStore: ## Callers +`ResponsePollingHandler` in `litellm/proxy/response_polling/polling_handler.py` is the reference caller + - Depend on the keyspace's `Protocol`, built once from `redis_cache.transport()` and injected where the caller can take it - `match` the result, ending in `assert_never` when it has more than two cases - Keep policy in the caller: which TTL applies, fail open or closed, and the `with_service_target` that labels its commands. The transport already logs failures and emits service events @@ -46,6 +50,7 @@ Add a store, not a new `RedisCache` method. A Python store is the finished state - Move the key builders, codec, TTLs, Lua and decoding verbatim. Keep policy in the caller - Keys and values already in Redis are a wire format. Old and new pods read each other's writes during a rolling deploy, so a move keeps them byte for byte + - polling state was written by `async_set_cache`, which runs `json.dumps` on a value that was already JSON text, so the store encodes with `json.dumps(json.dumps(state))` and decodes both layers - Changing a key, codec or TTL unit is never part of a move. It is its own change that writes a new versioned key and reads both for one release - Move one keyspace per change @@ -67,7 +72,7 @@ Response cache identity (normalization, hashing into a logical key) is backend-i - Test the store through its caller against `fake_redis_cache(server)` from `tests/_support/fake_redis.py`. That is a real `RedisCache` on a `fakeredis.FakeServer`, so namespace, breaker and failure mapping are the production ones - Cover each result the caller acts on, that the TTL lands (`pttl` on a raw `FakeAsyncRedis`), and an outage (`server.connected = False`) -- When a keyspace moves, seed with the old writer and read back with the old reader, so old and new pods read each other's writes +- When a keyspace moves, seed with the old writer and read back with the old reader, so old and new pods read each other's writes. See `_seed` and `_legacy_read` in `tests/unit/proxy/response_polling/test_response_polling_handler.py` - Assert returned values and Redis state, never the commands sent ## Porting a store to Rust diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index fe224a52cfa..e797858709f 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -2,9 +2,8 @@ Response Polling Handler for Background Responses with Cache """ -import json from collections.abc import Mapping -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Final from pydantic import TypeAdapter @@ -13,22 +12,26 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid4 from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores import response_polling +from litellm.redis.stores.response_polling import ResponsePollingStore, TransportResponsePollingStore +from litellm.redis.transport.contract import RedisFailure 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]) class ResponsePollingHandler: """Handles polling-based responses with Redis cache""" - CACHE_KEY_PREFIX = "litellm:polling:response:" + CACHE_KEY_PREFIX = response_polling.KEY_PREFIX POLLING_ID_PREFIX = "litellm_poll_" # Clear prefix to identify polling IDs def __init__(self, redis_cache: RedisCache | None = None, ttl: int = 3600): - self.redis_cache = redis_cache - self.ttl = ttl # Time-to-live for cache entries (default: 1 hour) + self.ttl = ttl + self._store: Final[ResponsePollingStore | None] = ( + TransportResponsePollingStore(redis_cache.transport()) if redis_cache is not None else None + ) @classmethod def generate_polling_id(cls) -> str: @@ -43,7 +46,10 @@ class ResponsePollingHandler: @classmethod def get_cache_key(cls, polling_id: str) -> str: """Get Redis cache key for a polling ID""" - return f"{cls.CACHE_KEY_PREFIX}{polling_id}" + return response_polling.state_key(polling_id) + + def _ttl(self) -> timedelta: + return timedelta(seconds=self.ttl) @with_service_target(_RESPONSE_POLLING_TARGET) async def create_initial_state( @@ -79,16 +85,10 @@ class ResponsePollingHandler: } ) - cache_key: Final = self.get_cache_key(polling_id) - - if self.redis_cache: - # Store ResponsesAPIResponse directly in Redis - await self.redis_cache.async_set_cache( - key=cache_key, - value=response.model_dump_json(), # Pydantic v2 method - ttl=self.ttl, - ) - verbose_proxy_logger.debug("Created initial polling state for %s with TTL=%ss", polling_id, self.ttl) + if self._store is not None: + written: Final = await self._store.put(polling_id, response.model_dump(mode="json"), ttl=self._ttl()) + if written is None: + verbose_proxy_logger.debug("Created initial polling state for %s with TTL=%ss", polling_id, self.ttl) return response @@ -145,19 +145,17 @@ class ResponsePollingHandler: user: User identifier store: Whether to store the response """ - if not self.redis_cache: + if self._store is None: return - cache_key: Final = self.get_cache_key(polling_id) - - # Get current state - cached_state: Final = await self.redis_cache.async_get_cache(cache_key) - if not cached_state: + cached_state: Final = await self._store.get(polling_id) + if isinstance(cached_state, RedisFailure): + return + if cached_state is None: verbose_proxy_logger.warning("No cached state found for polling_id: %s", polling_id) return - # Parse existing ResponsesAPIResponse from cache - state: Final = _JSON_OBJECT.validate_python(json.loads(cached_state)) + state: Final[dict[str, object]] = dict(cached_state) # Update status (using OpenAI native status values) if status: @@ -212,12 +210,9 @@ 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, - ) + written: Final = await self._store.put(polling_id, state, ttl=self._ttl()) + if written is not None: + return output_count: Final = len(_JSON_LIST.validate_python(state.get("output", []))) verbose_proxy_logger.debug( @@ -227,16 +222,11 @@ class ResponsePollingHandler: @with_service_target(_RESPONSE_POLLING_TARGET) async def get_state(self, polling_id: str) -> Mapping[str, object] | None: """Get current polling state from Redis""" - if not self.redis_cache: + if self._store is None: return None - cache_key: Final = self.get_cache_key(polling_id) - cached_state: Final = await self.redis_cache.async_get_cache(cache_key) - - if cached_state: - return _JSON_OBJECT.validate_python(json.loads(cached_state)) - - return None + cached_state: Final = await self._store.get(polling_id) + return None if isinstance(cached_state, RedisFailure) else cached_state async def cancel_polling(self, polling_id: str) -> bool: """ @@ -253,13 +243,10 @@ class ResponsePollingHandler: @with_service_target(_RESPONSE_POLLING_TARGET) async def delete_polling(self, polling_id: str) -> bool: """Delete a polling request from cache""" - if not self.redis_cache: + if self._store is None: return False - cache_key: Final = self.get_cache_key(polling_id) - # Use RedisCache's async_delete_cache method which handles Redis/RedisCluster - await self.redis_cache.async_delete_cache(cache_key) - return True + return await self._store.delete(polling_id) is None def should_use_polling_for_request( diff --git a/litellm/redis/stores/response_polling.py b/litellm/redis/stores/response_polling.py new file mode 100644 index 00000000000..0d4351edc0b --- /dev/null +++ b/litellm/redis/stores/response_polling.py @@ -0,0 +1,50 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import Delete, Get, RedisFailure, RedisTransport, Set + +KEY_PREFIX: Final = "litellm:polling:response:" + +_DOCUMENT: Final = TypeAdapter(str) +_STATE: Final = TypeAdapter(dict[str, JsonValue]) + + +def state_key(polling_id: str) -> str: + return f"{KEY_PREFIX}{polling_id}" + + +class ResponsePollingStore(Protocol): + async def put(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: ... + + async def get(self, polling_id: str) -> dict[str, JsonValue] | None | RedisFailure: ... + + async def delete(self, polling_id: str) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportResponsePollingStore: + transport: RedisTransport + + async def put(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: + written: Final = await self.transport.execute( + Set(state_key(polling_id), json.dumps(json.dumps(state)), ttl=ttl) + ) + return written if isinstance(written, RedisFailure) else None + + async def get(self, polling_id: str) -> dict[str, JsonValue] | None | RedisFailure: + raw: Final = await self.transport.execute(Get(state_key(polling_id))) + if raw is None or isinstance(raw, RedisFailure): + return raw + try: + return _STATE.validate_json(_DOCUMENT.validate_json(raw)) + except ValidationError as e: + return RedisFailure(kind="decode", message=f"polling state is not a JSON object: {e}") + + async def delete(self, polling_id: str) -> None | RedisFailure: + deleted: Final = await self.transport.execute(Delete(state_key(polling_id))) + return deleted if isinstance(deleted, RedisFailure) else None 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 087b273548a..9a3a128044a 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -20,11 +20,24 @@ from datetime import datetime, timezone from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, Mock, patch +import fakeredis import pytest from fastapi import Request from pydantic import ValidationError +from litellm.caching.redis_cache import RedisCache from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler +from tests._support.fake_redis import fake_redis_cache + + +async def _seed(cache: RedisCache, state: dict[str, object]) -> None: + await cache.async_set_cache( + key=ResponsePollingHandler.get_cache_key(str(state["id"])), value=json.dumps(state), ttl=3600 + ) + + +async def _legacy_read(cache: RedisCache, polling_id: str) -> dict[str, object]: + return json.loads(await cache.async_get_cache(ResponsePollingHandler.get_cache_key(polling_id))) from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @@ -58,15 +71,9 @@ class TestResponsePollingHandler: # OpenAI format assert ResponsePollingHandler.is_polling_id("resp_abc123") is False # Anthropic format - assert ( - ResponsePollingHandler.is_polling_id("msg_01XFDUDYJgAACzvnptvVoYEL") - is False - ) + assert ResponsePollingHandler.is_polling_id("msg_01XFDUDYJgAACzvnptvVoYEL") is False # Generic UUID - assert ( - ResponsePollingHandler.is_polling_id("550e8400-e29b-41d4-a716-446655440000") - is False - ) + assert ResponsePollingHandler.is_polling_id("550e8400-e29b-41d4-a716-446655440000") is False def test_get_cache_key_format(self): """Test that cache keys have the correct format""" @@ -80,8 +87,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_create_initial_state_returns_queued_status(self): """Test that create_initial_state returns response with queued status""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=3600) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer()), ttl=3600) polling_id = "litellm_poll_test123" request_data = { @@ -105,8 +111,9 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_create_initial_state_stores_in_redis(self): """Test that create_initial_state stores state in Redis with correct TTL""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=7200) + server = fakeredis.FakeServer() + cache = fake_redis_cache(server, namespace="ns") + handler = ResponsePollingHandler(redis_cache=cache, ttl=7200) polling_id = "litellm_poll_test123" request_data = {"model": "gpt-4o", "input": "Hello"} @@ -116,26 +123,18 @@ class TestResponsePollingHandler: request_data=request_data, ) - # Verify Redis was called with correct parameters - mock_redis.async_set_cache.assert_called_once() - call_args = mock_redis.async_set_cache.call_args + raw = fakeredis.FakeAsyncRedis(server=server) + assert await raw.keys("*") == [b"ns:litellm:polling:response:litellm_poll_test123"] + assert 7_190_000 < await raw.pttl("ns:litellm:polling:response:litellm_poll_test123") <= 7_200_000 - assert ( - call_args.kwargs["key"] == "litellm:polling:response:litellm_poll_test123" - ) - assert call_args.kwargs["ttl"] == 7200 - - # Verify the stored value is valid JSON - stored_value = call_args.kwargs["value"] - parsed = json.loads(stored_value) + parsed = await _legacy_read(cache, polling_id) assert parsed["id"] == polling_id assert parsed["status"] == "queued" @pytest.mark.asyncio async def test_create_initial_state_sets_created_at_timestamp(self): """Test that create_initial_state sets a valid created_at timestamp""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer())) before_time = int(datetime.now(timezone.utc).timestamp()) @@ -153,46 +152,45 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "queued", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=3600) + handler = ResponsePollingHandler(redis_cache=cache, ttl=3600) await handler.update_state( polling_id="litellm_poll_test", 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 _legacy_read(cache, "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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [{"id": "old_item", "type": "message"}], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=3600) + handler = ResponsePollingHandler(redis_cache=cache, ttl=3600) new_output = [ { @@ -212,8 +210,7 @@ class TestResponsePollingHandler: output=new_output, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert len(stored["output"]) == 2 assert stored["output"][0]["id"] == "item_1" @@ -222,18 +219,19 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) usage_data = {"input_tokens": 10, "output_tokens": 50, "total_tokens": 60} @@ -243,8 +241,7 @@ class TestResponsePollingHandler: usage=usage_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "completed" assert stored["usage"] == usage_data @@ -252,24 +249,23 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) reasoning_data = {"effort": "medium", "summary": "Step by step analysis"} tool_choice_data = {"type": "function", "function": {"name": "get_weather"}} - tools_data = [ - {"type": "function", "function": {"name": "get_weather", "parameters": {}}} - ] + tools_data = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] await handler.update_state( polling_id="litellm_poll_test", @@ -279,8 +275,7 @@ class TestResponsePollingHandler: tools=tools_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["reasoning"] == reasoning_data assert stored["tool_choice"] == tool_choice_data @@ -289,18 +284,19 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) # All ResponsesAPIResponse fields that can be updated await handler.update_state( @@ -324,8 +320,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 _legacy_read(cache, "litellm_poll_test") # Verify all fields are stored correctly assert stored["status"] == "completed" @@ -353,8 +348,9 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", @@ -363,10 +359,10 @@ class TestResponsePollingHandler: "created_at": 1234567890, "model": "gpt-4o", "temperature": 0.5, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) # Only update status await handler.update_state( @@ -374,8 +370,7 @@ class TestResponsePollingHandler: status="completed", ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") # Verify existing fields are preserved assert stored["status"] == "completed" @@ -386,18 +381,19 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) error_data = { "type": "internal_error", @@ -410,8 +406,7 @@ class TestResponsePollingHandler: error=error_data, ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "failed" assert stored["error"] == error_data @@ -419,18 +414,19 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) incomplete_details = {"reason": "max_output_tokens"} @@ -440,8 +436,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 _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "incomplete" assert stored["incomplete_details"] == incomplete_details @@ -460,27 +455,23 @@ 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 + server = fakeredis.FakeServer() + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(server)) - handler = ResponsePollingHandler(redis_cache=mock_redis) - - # Should not raise an exception await handler.update_state( polling_id="litellm_poll_test", status="in_progress", ) - # Should not try to set cache if nothing was found - mock_redis.async_set_cache.assert_not_called() + assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == [] # ==================== Get State Tests ==================== @pytest.mark.asyncio async def test_get_state_returns_cached_state(self): """Test that get_state returns the cached state""" - mock_redis = AsyncMock() - cached_state = { + cache = fake_redis_cache(fakeredis.FakeServer()) + cached_state: dict[str, object] = { "id": "litellm_poll_test", "object": "response", "status": "in_progress", @@ -488,9 +479,9 @@ class TestResponsePollingHandler: "created_at": 1234567890, "usage": {"input_tokens": 10, "output_tokens": 20}, } - mock_redis.async_get_cache.return_value = json.dumps(cached_state) + await _seed(cache, cached_state) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) result = await handler.get_state("litellm_poll_test") @@ -499,10 +490,7 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_get_state_returns_none_for_missing_state(self): """Test that get_state returns None when state doesn't exist""" - mock_redis = AsyncMock() - mock_redis.async_get_cache.return_value = None - - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer())) result = await handler.get_state("litellm_poll_nonexistent") @@ -522,25 +510,25 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) result = await handler.cancel_polling("litellm_poll_test") assert result is True - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["status"] == "cancelled" # ==================== Delete Polling Tests ==================== @@ -548,27 +536,38 @@ class TestResponsePollingHandler: @pytest.mark.asyncio async def test_delete_polling_removes_from_cache(self): """Test that delete_polling removes the entry from Redis""" - mock_redis = AsyncMock() - mock_async_client = AsyncMock() - mock_redis.redis_async_client = True # hasattr check - # init_async_client is a sync method that returns an async client - mock_redis.init_async_client = Mock(return_value=mock_async_client) - - # Mock async_delete_cache to actually call init_async_client and delete - async def mock_async_delete_cache(key): - client = mock_redis.init_async_client() - await client.delete(key) - - mock_redis.async_delete_cache = mock_async_delete_cache - - handler = ResponsePollingHandler(redis_cache=mock_redis) + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed(cache, {"id": "litellm_poll_test", "status": "queued"}) + handler = ResponsePollingHandler(redis_cache=cache) result = await handler.delete_polling("litellm_poll_test") assert result is True - mock_async_client.delete.assert_called_once_with( - "litellm:polling:response:litellm_poll_test" + assert await handler.get_state("litellm_poll_test") is None + + @pytest.mark.asyncio + @pytest.mark.parametrize("stored", [b"not json", b'{"id": "single-encoded"}', b'"[1, 2]"']) + async def test_a_state_that_does_not_decode_reads_as_missing(self, stored: bytes): + server = fakeredis.FakeServer() + await fakeredis.FakeAsyncRedis(server=server).set( + ResponsePollingHandler.get_cache_key("litellm_poll_x"), stored ) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(server)) + + assert await handler.get_state("litellm_poll_x") is None + + @pytest.mark.asyncio + async def test_an_unreachable_redis_reads_as_missing_and_fails_the_delete(self): + server = fakeredis.FakeServer() + cache = fake_redis_cache(server) + await _seed(cache, {"id": "litellm_poll_test", "status": "queued"}) + handler = ResponsePollingHandler(redis_cache=cache) + server.connected = False + + assert await handler.get_state("litellm_poll_test") is None + assert await handler.delete_polling("litellm_poll_test") is False + queued = await handler.create_initial_state(polling_id="litellm_poll_new", request_data={}) + assert queued.status == "queued" @pytest.mark.asyncio async def test_delete_polling_returns_false_without_redis(self): @@ -596,26 +595,30 @@ 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( + server = fakeredis.FakeServer() + cache = fake_redis_cache(server) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "queued", "output": [], "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis, ttl=1800) + handler = ResponsePollingHandler(redis_cache=cache, ttl=1800) await handler.update_state( polling_id="litellm_poll_test", status="in_progress", ) - call_args = mock_redis.async_set_cache.call_args - assert call_args.kwargs["ttl"] == 1800 + ttl_ms = await fakeredis.FakeAsyncRedis(server=server).pttl( + ResponsePollingHandler.get_cache_key("litellm_poll_test") + ) + assert 1_790_000 < ttl_ms <= 1_800_000 class TestStreamingEventProcessing: @@ -689,6 +692,7 @@ class TestBackgroundStreamingModule: def test_background_streaming_task_is_async(self): """Test that background_streaming_task is an async function""" import inspect + from litellm.proxy.response_polling.background_streaming import ( background_streaming_task, ) @@ -886,9 +890,7 @@ class TestProviderResolutionForPolling: model_name_to_deployment_indices = { "gpt-5": [0], } - model_list = [ - {"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}} - ] + model_list = [{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}}] model = "unknown-model" # Not in router polling_via_cache_enabled = ["openai"] @@ -1282,9 +1284,7 @@ class TestStreamingEventParsing: content_list = output_items[item_id]["content"] if content_index < len(content_list): if isinstance(content_list[content_index], dict): - content_list[content_index]["text"] = accumulated_text[ - key - ] + content_list[content_index]["text"] = accumulated_text[key] assert accumulated_text[("item_123", 0)] == "Hello World!" assert output_items["item_123"]["content"][0]["text"] == "Hello World!" @@ -1346,9 +1346,7 @@ class TestStreamingEventParsing: def test_parse_sse_format(self): """Test parsing Server-Sent Events format""" - raw_chunk = ( - b'data: {"type": "response.output_item.added", "item": {"id": "123"}}' - ) + raw_chunk = b'data: {"type": "response.output_item.added", "item": {"id": "123"}}' # Decode bytes to string if isinstance(raw_chunk, bytes): @@ -1476,9 +1474,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) # Find the final update_state call (last one) @@ -1515,9 +1511,7 @@ class TestBackgroundStreamingTerminalEvents: 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=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1558,17 +1552,13 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] assert final_call.kwargs["status"] == "incomplete" assert final_call.kwargs["error"] == error_payload - assert final_call.kwargs["incomplete_details"] == { - "reason": "max_output_tokens" - } + assert final_call.kwargs["incomplete_details"] == {"reason": "max_output_tokens"} assert final_call.kwargs["usage"] == {"input_tokens": 10, "output_tokens": 4096} @pytest.mark.asyncio @@ -1597,9 +1587,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1632,9 +1620,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1699,9 +1685,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1725,9 +1709,7 @@ class TestBackgroundStreamingTerminalEvents: with patch( "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" ) as MockProcessor: - MockProcessor.return_value.base_process_llm_request = AsyncMock( - return_value=mock_response - ) + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=mock_response) await background_streaming_task(**kwargs) final_call = handler.update_state.call_args_list[-1] @@ -1763,9 +1745,7 @@ class TestBackgroundStreamingTerminalEvents: yield f"data: {json.dumps(event)}\n\n" async def upstream_call_behind_disconnect_guard(**kwargs): - return await create_response( - slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"] - ) + return await create_response(slow_upstream_stream(), "text/event-stream", {}, request=kwargs["request"]) handler = AsyncMock(spec=ResponsePollingHandler) kwargs = _make_background_streaming_kwargs("poll_7", handler) @@ -1838,8 +1818,7 @@ class TestEdgeCases: @pytest.mark.asyncio async def test_create_initial_state_with_empty_metadata(self): """Test create_initial_state handles missing metadata gracefully""" - mock_redis = AsyncMock() - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(fakeredis.FakeServer())) response = await handler.create_initial_state( polling_id="litellm_poll_test", @@ -1851,26 +1830,26 @@ 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( + cache = fake_redis_cache(fakeredis.FakeServer()) + await _seed( + cache, { "id": "litellm_poll_test", "object": "response", "status": "in_progress", "output": [{"id": "item_1"}], # Has existing output "created_at": 1234567890, - } + }, ) - handler = ResponsePollingHandler(redis_cache=mock_redis) + handler = ResponsePollingHandler(redis_cache=cache) await handler.update_state( polling_id="litellm_poll_test", output=[], # Explicitly set empty ) - call_args = mock_redis.async_set_cache.call_args - stored = json.loads(call_args.kwargs["value"]) + stored = await _legacy_read(cache, "litellm_poll_test") assert stored["output"] == []