fix(redis-store): keep main's polling bytes and per-failure outcomes

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(

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

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"""