From dd6df01166e5d448837a6bf6fda749070e01ae3c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:48:29 +0000 Subject: [PATCH] refactor(redis): move the tool discovery buffer behind a typed store --- .../redis_update_buffer.py | 60 ++++++++----------- litellm/redis/stores/tool_discovery_buffer.py | 36 +++++++++++ .../test_redis_update_buffer.py | 15 +++-- .../stores/test_tool_discovery_buffer.py | 50 ++++++++++++++++ 4 files changed, 117 insertions(+), 44 deletions(-) create mode 100644 litellm/redis/stores/tool_discovery_buffer.py create mode 100644 tests/unit/redis/stores/test_tool_discovery_buffer.py diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 2ba570efacb..9076c2dc681 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -20,10 +20,8 @@ from litellm.caching import RedisCache from litellm.constants import ( MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, - REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, ) -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -59,12 +57,9 @@ from litellm.redis.stores.spend_update_buffer import ( TransportSpendUpdateBufferStore, spend_buffer_key, ) +from litellm.redis.stores.tool_discovery_buffer import ToolDiscoveryBufferStore, TransportToolDiscoveryBufferStore from litellm.redis.transport.contract import RedisFailure, raise_redis_failure from litellm.secret_managers.main import str_to_bool -from litellm.types.caching import ( - RedisPipelineLpopOperation, - RedisPipelineRpushOperation, -) from litellm.types.services import ServiceTypes SPEND_QUEUE_TARGET: Final = "spend_queue" @@ -125,6 +120,10 @@ def _parked_spend_logs(redis_cache: RedisCache) -> SpendLogsParkingStore: return TransportSpendLogsParkingStore(redis_cache.transport()) +def _tool_discoveries(redis_cache: RedisCache) -> ToolDiscoveryBufferStore: + return TransportToolDiscoveryBufferStore(redis_cache.transport()) + + def _zero_for_non_finite_constant(constant: str) -> float: verbose_proxy_logger.warning( "Redis spend buffer - reading a parked %s spend as 0, the spend it held is lost", constant @@ -160,9 +159,6 @@ def _merged_entity_transactions( ) -_TOOL_DISCOVERY_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem]) - - class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database @@ -575,46 +571,38 @@ class RedisUpdateBuffer: items: Final = tool_discovery_queue.flush() if not items: return - try: - _ = await self.redis_cache.async_rpush_pipeline( - rpush_list=[ - RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)]) - ] - ) - except (RedisError, OSError, asyncio.TimeoutError) as e: - verbose_proxy_logger.error( - "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s", - len(items), - str(e), - ) - for item in items: - tool_discovery_queue.add_update(item) + pushed: Final = await _tool_discoveries(self.redis_cache).push(items) + if not isinstance(pushed, RedisFailure): + return + if pushed.kind not in ("transport", "server"): + raise_redis_failure(pushed) + verbose_proxy_logger.error( + "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s", + len(items), + pushed.message, + ) + for item in items: + tool_discovery_queue.add_update(item) @with_service_target(SPEND_QUEUE_TARGET) async def restore_tool_discoveries_to_redis(self, items: Sequence[ToolDiscoveryQueueItem]) -> None: """Put drained tool rows back for the next drain after a write that provably did not apply.""" if self.redis_cache is None or not items: return - _ = await self.redis_cache.async_rpush_pipeline( - rpush_list=[ - RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(list(items))]) - ] - ) + restored: Final = await _tool_discoveries(self.redis_cache).push(items) + if isinstance(restored, RedisFailure): + raise_redis_failure(restored) @with_service_target(SPEND_QUEUE_TARGET) async def get_all_tool_discoveries_from_redis_buffer(self) -> tuple[ToolDiscoveryQueueItem, ...]: """Drain every worker's tool rows and fold them to one row per tool_name.""" if self.redis_cache is None: return () - popped: Final = await self.redis_cache.async_lpop_pipeline( - lpop_list=[ - RedisPipelineLpopOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT) - ] - ) - entries: Final = popped[0] if popped else None - if not entries: + batches: Final = await _tool_discoveries(self.redis_cache).drain(count=MAX_REDIS_BUFFER_DEQUEUE_COUNT) + if isinstance(batches, RedisFailure): + raise_redis_failure(batches) + if not batches: return () - batches: Final = tuple(_TOOL_DISCOVERY_ITEMS.validate_json(entry) for entry in entries) return aggregate_tool_discoveries(chain.from_iterable(batches)) @with_service_target(SPEND_QUEUE_TARGET) diff --git a/litellm/redis/stores/tool_discovery_buffer.py b/litellm/redis/stores/tool_discovery_buffer.py new file mode 100644 index 00000000000..9ebb7f4a5c5 --- /dev/null +++ b/litellm/redis/stores/tool_discovery_buffer.py @@ -0,0 +1,36 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.constants import REDIS_TOOL_DISCOVERY_BUFFER_KEY +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem]) + + +class ToolDiscoveryBufferStore(Protocol): + async def push(self, items: Sequence[ToolDiscoveryQueueItem]) -> int | RedisFailure: ... + + async def drain(self, *, count: int) -> tuple[tuple[ToolDiscoveryQueueItem, ...], ...] | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportToolDiscoveryBufferStore: + transport: RedisTransport + + async def push(self, items: Sequence[ToolDiscoveryQueueItem]) -> int | RedisFailure: + return await self.transport.rpush(REDIS_TOOL_DISCOVERY_BUFFER_KEY, safe_dumps(list(items))) + + async def drain(self, *, count: int) -> tuple[tuple[ToolDiscoveryQueueItem, ...], ...] | RedisFailure: + popped: Final = await self.transport.lpop(REDIS_TOOL_DISCOVERY_BUFFER_KEY, count) + if isinstance(popped, RedisFailure): + return popped + entries: Final = () if popped is None else (popped,) if isinstance(popped, bytes) else popped + try: + return tuple(tuple(_ITEMS.validate_json(entry)) for entry in entries) + except ValidationError as error: + return RedisFailure(kind="decode", message="Invalid tool discovery buffer entry", exception=error) diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index cd1289e7189..969792e467c 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -242,11 +242,10 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache): +async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache, raw_redis): """ When all queues are empty, pipeline should never be called. """ - mock_redis_cache.async_rpush_pipeline = AsyncMock() # All queues return empty empty_queue = AsyncMock() @@ -263,7 +262,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat daily_agent_spend_update_queue=empty_daily_queue, ) - mock_redis_cache.async_rpush_pipeline.assert_not_called() + assert await raw_redis.keys("*") == [] @pytest.mark.asyncio @@ -540,11 +539,10 @@ async def test_restored_window_spend_transactions_drain_back_unchanged( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache): +async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache, raw_redis): """Nothing to restore -> no Redis call.""" - mock_redis_cache.async_rpush_pipeline = AsyncMock() await redis_update_buffer.restore_transactions_to_redis() - mock_redis_cache.async_rpush_pipeline.assert_not_called() + assert await raw_redis.keys("*") == [] @pytest.mark.asyncio @@ -865,8 +863,9 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(): from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue - redis_cache = MagicMock() - redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + server: Final = fakeredis.FakeServer() + redis_cache: Final = fake_redis_cache(server) + server.connected = False queue = ToolDiscoveryQueue() queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) diff --git a/tests/unit/redis/stores/test_tool_discovery_buffer.py b/tests/unit/redis/stores/test_tool_discovery_buffer.py new file mode 100644 index 00000000000..9d460e2b5b9 --- /dev/null +++ b/tests/unit/redis/stores/test_tool_discovery_buffer.py @@ -0,0 +1,50 @@ +from typing import Final + +import fakeredis +import pytest + +from litellm.constants import REDIS_TOOL_DISCOVERY_BUFFER_KEY +from litellm.redis.stores.tool_discovery_buffer import TransportToolDiscoveryBufferStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + +KEY: Final = f"ns:{REDIS_TOOL_DISCOVERY_BUFFER_KEY}" + + +def _store(server: fakeredis.FakeServer) -> TransportToolDiscoveryBufferStore: + return TransportToolDiscoveryBufferStore(fake_redis_cache(server, namespace="ns").transport()) + + +@pytest.mark.asyncio +async def test_each_push_is_one_json_list_entry_and_a_drain_returns_every_batch(server: fakeredis.FakeServer) -> None: + store: Final = _store(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + + assert await store.push(({"tool_name": "a", "call_count": 2},)) == 1 + assert await store.push(({"tool_name": "b", "origin": None, "call_count": 1},)) == 2 + + assert await raw.lrange(KEY, 0, -1) == [ + b'[{"tool_name": "a", "call_count": 2}]', + b'[{"tool_name": "b", "origin": null, "call_count": 1}]', + ] + assert await store.drain(count=10) == ( + ({"tool_name": "a", "call_count": 2},), + ({"tool_name": "b", "origin": None, "call_count": 1},), + ) + assert await store.drain(count=10) == () + + +@pytest.mark.asyncio +async def test_bad_entries_rejections_and_outages_are_failures(server: fakeredis.FakeServer) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + store: Final = _store(server) + await raw.rpush(KEY, b'{"not": "a list"}') + malformed: Final = await store.drain(count=10) + await raw.set(KEY, "not a list") + rejected: Final = await store.push(({"tool_name": "a", "call_count": 1},)) + server.connected = False + outage: Final = await store.drain(count=10) + + assert isinstance(malformed, RedisFailure) and malformed.kind == "decode" + assert isinstance(rejected, RedisFailure) and rejected.kind == "server" + assert isinstance(outage, RedisFailure) and outage.kind == "transport"