From 911ddf8ead1d3b906f89b683ee286511a188c526 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 10:44:54 -0700 Subject: [PATCH] fix(redis-store): keep main's polling bytes and per-failure outcomes --- .../proxy/response_polling/polling_handler.py | 48 ++++++++++++----- litellm/redis/stores/response_polling.py | 19 +++++-- .../test_response_polling_handler.py | 53 ++++++++++++++++--- 3 files changed, 96 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 9bf4874ac80..679c650969f 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -6,12 +6,12 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Final -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter 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.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.redis.stores import response_polling from litellm.redis.stores.response_polling import ResponsePollingStore, TransportResponsePollingStore from litellm.redis.transport.contract import RedisFailure @@ -86,9 +86,8 @@ class ResponsePollingHandler: ) 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) + _raise_if_breaker_open(await self._store.create(polling_id, response, ttl=self._ttl())) + verbose_proxy_logger.debug("Created initial polling state for %s with TTL=%ss", polling_id, self.ttl) return response @@ -148,9 +147,7 @@ class ResponsePollingHandler: if self._store is None: return - cached_state: Final = await self._store.get(polling_id) - if isinstance(cached_state, RedisFailure): - return + cached_state: Final = _unless_raising(await self._store.get(polling_id)) if cached_state is None: verbose_proxy_logger.warning("No cached state found for polling_id: %s", polling_id) return @@ -181,9 +178,7 @@ class ResponsePollingHandler: **{key: value for key, value in updates.items() if value is not None}, } - written: Final = await self._store.put(polling_id, state, ttl=self._ttl()) - if written is not None: - return + _raise_if_breaker_open(await self._store.update(polling_id, state, ttl=self._ttl())) output_count: Final = len(_JSON_LIST.validate_python(state.get("output", []))) verbose_proxy_logger.debug( @@ -196,8 +191,7 @@ class ResponsePollingHandler: if self._store is None: return None - cached_state: Final = await self._store.get(polling_id) - return None if isinstance(cached_state, RedisFailure) else cached_state + return _unless_raising(await self._store.get(polling_id)) async def cancel_polling(self, polling_id: str) -> bool: """ @@ -217,7 +211,33 @@ class ResponsePollingHandler: if self._store is None: return False - return await self._store.delete(polling_id) is None + deleted: Final = await self._store.delete(polling_id) + if isinstance(deleted, RedisFailure): + raise _cache_error(deleted) + return True + + +def _cache_error(failure: RedisFailure) -> Exception: + match failure.kind: + case "breaker_open": + return RedisCircuitBreakerOpenError(failure.message) + case "decode": + return ValueError(failure.message) + case "transport" | "server": + return RuntimeError(failure.message) + + +def _raise_if_breaker_open(written: RedisFailure | None) -> None: + if isinstance(written, RedisFailure) and written.kind == "breaker_open": + raise _cache_error(written) + + +def _unless_raising(read: Mapping[str, JsonValue] | None | RedisFailure) -> Mapping[str, JsonValue] | None: + if not isinstance(read, RedisFailure): + return read + if read.kind in ("breaker_open", "decode"): + raise _cache_error(read) + return None def should_use_polling_for_request( diff --git a/litellm/redis/stores/response_polling.py b/litellm/redis/stores/response_polling.py index df1f6d78736..01bb510c982 100644 --- a/litellm/redis/stores/response_polling.py +++ b/litellm/redis/stores/response_polling.py @@ -7,6 +7,7 @@ from typing import Final, Protocol from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.redis.transport.contract import RedisFailure, RedisTransport +from litellm.types.llms.openai import ResponsesAPIResponse KEY_PREFIX: Final = "litellm:polling:response:" @@ -19,7 +20,11 @@ def state_key(polling_id: str) -> str: class ResponsePollingStore(Protocol): - async def put(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: ... + async def create( + self, polling_id: str, response: ResponsesAPIResponse, *, ttl: timedelta + ) -> None | RedisFailure: ... + + async def update(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: ... async def get(self, polling_id: str) -> Mapping[str, JsonValue] | None | RedisFailure: ... @@ -30,9 +35,11 @@ class ResponsePollingStore(Protocol): 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 create(self, polling_id: str, response: ResponsesAPIResponse, *, ttl: timedelta) -> None | RedisFailure: + return await self._write(polling_id, json.dumps(response.model_dump_json()), ttl) + + async def update(self, polling_id: str, state: Mapping[str, object], *, ttl: timedelta) -> None | RedisFailure: + return await self._write(polling_id, json.dumps(json.dumps(state)), ttl) async def get(self, polling_id: str) -> Mapping[str, JsonValue] | None | RedisFailure: raw: Final = await self.transport.get(state_key(polling_id)) @@ -46,3 +53,7 @@ class TransportResponsePollingStore: 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 + + async def _write(self, polling_id: str, document: str, ttl: timedelta) -> None | RedisFailure: + written: Final = await self.transport.set(state_key(polling_id), document, ttl=ttl) + return written if isinstance(written, 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 d7b9fe44885..a85a3ae92b8 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -25,7 +25,8 @@ import pytest from fastapi import Request from pydantic import ValidationError -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError +from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from tests._support.fake_redis import fake_redis_cache @@ -39,6 +40,7 @@ async def _seed(cache: RedisCache, state: dict[str, object]) -> None: 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 @@ -578,18 +580,21 @@ class TestResponsePollingHandler: 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): + @pytest.mark.parametrize("stored", [b'{"id": "single-encoded"}', b'"[1, 2]"']) + async def test_a_state_that_does_not_decode_raises_from_reads_and_updates(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 + with pytest.raises((TypeError, ValueError)): + await handler.get_state("litellm_poll_x") + with pytest.raises((TypeError, ValueError)): + await handler.update_state("litellm_poll_x", status="completed") @pytest.mark.asyncio - async def test_an_unreachable_redis_reads_as_missing_and_fails_the_delete(self): + async def test_an_unreachable_redis_reads_as_missing_and_raises_from_the_delete(self): server = fakeredis.FakeServer() cache = fake_redis_cache(server) await _seed(cache, {"id": "litellm_poll_test", "status": "queued"}) @@ -597,10 +602,46 @@ class TestResponsePollingHandler: server.connected = False assert await handler.get_state("litellm_poll_test") is None - assert await handler.delete_polling("litellm_poll_test") is False + await handler.update_state("litellm_poll_test", status="completed") + with pytest.raises(RuntimeError): + await handler.delete_polling("litellm_poll_test") queued = await handler.create_initial_state(polling_id="litellm_poll_new", request_data={}) assert queued.status == "queued" + @pytest.mark.asyncio + async def test_an_open_breaker_raises_from_every_operation(self): + server = fakeredis.FakeServer() + cache = fake_redis_cache(server) + handler = ResponsePollingHandler(redis_cache=cache) + server.connected = False + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert await handler.get_state("litellm_poll_test") is None + + with pytest.raises(RedisCircuitBreakerOpenError): + await handler.get_state("litellm_poll_test") + with pytest.raises(RedisCircuitBreakerOpenError): + await handler.update_state("litellm_poll_test", status="completed") + with pytest.raises(RedisCircuitBreakerOpenError): + await handler.create_initial_state(polling_id="litellm_poll_new", request_data={}) + with pytest.raises(RedisCircuitBreakerOpenError): + await handler.delete_polling("litellm_poll_test") + + @pytest.mark.asyncio + async def test_initial_state_and_updates_keep_the_legacy_bytes_in_redis(self): + server = fakeredis.FakeServer() + raw = fakeredis.FakeAsyncRedis(server=server) + handler = ResponsePollingHandler(redis_cache=fake_redis_cache(server)) + key = ResponsePollingHandler.get_cache_key("litellm_poll_bytes") + + queued = await handler.create_initial_state( + polling_id="litellm_poll_bytes", request_data={"metadata": {"note": "caf\u00e9"}} + ) + assert await raw.get(key) == json.dumps(queued.model_dump_json()).encode() + + await handler.update_state("litellm_poll_bytes", status="completed") + updated = {**json.loads(queued.model_dump_json()), "status": "completed"} + assert await raw.get(key) == json.dumps(json.dumps(updated)).encode() + @pytest.mark.asyncio async def test_delete_polling_returns_false_without_redis(self): """Test that delete_polling returns False when Redis is not configured"""