From 59b04e1c9fd6604fe750b9783bb3ad9eed597039 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 10:09:52 -0700 Subject: [PATCH] refactor(redis-store): move response polling state behind a store Background response polling state now goes through the shared Redis transport behind TransportResponsePollingStore, which owns the keys, codec and expiry. The wire format is kept byte for byte so old and new pods read each other's state --- .../proxy/response_polling/polling_handler.py | 168 ++++----- litellm/redis/stores/__init__.py | 0 litellm/redis/stores/response_polling.py | 48 +++ .../test_response_polling_handler.py | 341 +++++++++--------- 4 files changed, 287 insertions(+), 270 deletions(-) create mode 100644 litellm/redis/stores/__init__.py create mode 100644 litellm/redis/stores/response_polling.py diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index fe224a52cfa..9bf4874ac80 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 collections.abc import Mapping, Sequence +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 @@ -97,13 +97,13 @@ class ResponsePollingHandler: self, polling_id: str, status: ResponsesAPIStatus | None = None, - usage: dict | None = None, - error: dict | None = None, - incomplete_details: dict | None = None, - reasoning: dict | None = None, + usage: Mapping[str, object] | None = None, + error: Mapping[str, object] | None = None, + incomplete_details: Mapping[str, object] | None = None, + reasoning: Mapping[str, object] | None = None, tool_choice: object | None = None, - tools: list | None = None, - output: list | None = None, + tools: Sequence[object] | None = None, + output: Sequence[object] | None = None, # Additional ResponsesAPIResponse fields model: str | None = None, instructions: str | None = None, @@ -111,7 +111,7 @@ class ResponsePollingHandler: top_p: float | None = None, max_output_tokens: int | None = None, previous_response_id: str | None = None, - text: dict | None = None, + text: Mapping[str, object] | None = None, truncation: str | None = None, parallel_tool_calls: bool | None = None, user: str | None = None, @@ -145,79 +145,45 @@ 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)) + updates: Final[Mapping[str, object]] = { + "status": "failed" if error else status or None, + "output": output, + "usage": usage or None, + "error": error or None, + "incomplete_details": incomplete_details or None, + "reasoning": reasoning, + "tool_choice": tool_choice, + "tools": tools, + "model": model, + "instructions": instructions, + "temperature": temperature, + "top_p": top_p, + "max_output_tokens": max_output_tokens, + "previous_response_id": previous_response_id, + "text": text, + "truncation": truncation, + "parallel_tool_calls": parallel_tool_calls, + "user": user, + "store": store, + } + state: Final[Mapping[str, object]] = { + **cached_state, + **{key: value for key, value in updates.items() if value is not None}, + } - # Update status (using OpenAI native status values) - if status: - state["status"] = status - - # Replace full output list if provided - if output is not None: - state["output"] = output - - # Update usage - if usage: - state["usage"] = usage - - # Handle error (sets status to OpenAI's "failed") - if error: - state["status"] = "failed" - state["error"] = error # Use OpenAI's 'error' field - - # Handle incomplete details - if incomplete_details: - state["incomplete_details"] = incomplete_details - - # Update reasoning, tool_choice, tools from response.completed - if reasoning is not None: - state["reasoning"] = reasoning - if tool_choice is not None: - state["tool_choice"] = tool_choice - if tools is not None: - state["tools"] = tools - - # Update additional ResponsesAPIResponse fields - if model is not None: - state["model"] = model - if instructions is not None: - state["instructions"] = instructions - if temperature is not None: - state["temperature"] = temperature - if top_p is not None: - state["top_p"] = top_p - if max_output_tokens is not None: - state["max_output_tokens"] = max_output_tokens - if previous_response_id is not None: - state["previous_response_id"] = previous_response_id - if text is not None: - state["text"] = text - if truncation is not None: - state["truncation"] = truncation - if parallel_tool_calls is not None: - state["parallel_tool_calls"] = parallel_tool_calls - if user is not None: - state["user"] = user - 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 +193,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 +214,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/__init__.py b/litellm/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/redis/stores/response_polling.py b/litellm/redis/stores/response_polling.py new file mode 100644 index 00000000000..df1f6d78736 --- /dev/null +++ b/litellm/redis/stores/response_polling.py @@ -0,0 +1,48 @@ +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 RedisFailure, RedisTransport + +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) -> Mapping[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.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) -> Mapping[str, JsonValue] | None | RedisFailure: + raw: Final = await self.transport.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.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..d7b9fe44885 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", @@ -407,30 +403,62 @@ class TestResponsePollingHandler: await handler.update_state( polling_id="litellm_poll_test", + status="completed", 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 + @pytest.mark.asyncio + async def test_update_state_preserves_false_empty_and_zero_values(self): + cache: Final = fake_redis_cache(fakeredis.FakeServer()) + await _seed(cache, {"id": "litellm_poll_test", "status": "in_progress", "usage": {"total_tokens": 3}}) + handler: Final = ResponsePollingHandler(redis_cache=cache) + + await handler.update_state( + polling_id="litellm_poll_test", + usage={}, + output=[], + tools=[], + reasoning={}, + instructions="", + temperature=0, + parallel_tool_calls=False, + store=False, + ) + + assert await _legacy_read(cache, "litellm_poll_test") == { + "id": "litellm_poll_test", + "status": "in_progress", + "usage": {"total_tokens": 3}, + "output": [], + "tools": [], + "reasoning": {}, + "instructions": "", + "temperature": 0, + "parallel_tool_calls": False, + "store": False, + } + @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 +468,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 +487,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 +511,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 +522,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 +542,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 +568,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 +627,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 +724,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 +922,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 +1316,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 +1378,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 +1506,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 +1543,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 +1584,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 +1619,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 +1652,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 +1717,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 +1741,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 +1777,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 +1850,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 +1862,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"] == []