From 06bebcc5bb614c4ac370346097709e86e9ab7e9b Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Tue, 20 Oct 2026 14:17:37 -0700 Subject: [PATCH] fix(proxy): restore daily tag spend to the in-memory queue when the Redis RPUSH fails (#47307) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../redis_update_buffer.py | 21 ++++-- .../test_redis_update_buffer.py | 68 +++++++++++++++++++ 2 files changed, 82 insertions(+), 7 deletions(-) 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 898aee493df..54d20bd4507 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -168,9 +168,11 @@ class RedisUpdateBuffer: @with_service_target(SPEND_QUEUE_TARGET) async def _store_transactions_in_redis( self, - transactions: Mapping[str, BaseDailySpendTransaction] | None, + transactions: dict[str, BaseDailySpendTransaction] | None, # mutable-ok: add_update requires dict redis_key: str, service_type: ServiceTypes, + *, + source_queue: DailySpendUpdateQueue, ) -> None: """ Helper method to store transactions in Redis and emit an event @@ -179,6 +181,7 @@ class RedisUpdateBuffer: transactions: The transactions to store redis_key: The Redis key to store under service_type: The service type for event emission + source_queue: The in-memory queue to restore transactions to if Redis push fails """ if transactions is None or len(transactions) == 0: return @@ -191,6 +194,15 @@ class RedisUpdateBuffer: key=redis_key, values=list_of_transactions, ) + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to push spend updates to Redis (redis_key=%s). " + "Restoring them to the in-memory queue for retry on next tick. Error: %s", + redis_key, + str(e), + ) + await source_queue.add_update(transactions) + else: verbose_proxy_logger.debug( "Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s", redis_key, @@ -200,12 +212,6 @@ class RedisUpdateBuffer: queue_size=current_redis_buffer_size, service=service_type, ) - except Exception as e: - verbose_proxy_logger.error( - "Spend tracking - failed to push spend updates to Redis (redis_key=%s). Error: %s", - redis_key, - str(e), - ) @with_service_target(SPEND_QUEUE_TARGET) async def store_in_memory_spend_updates_in_redis( @@ -793,6 +799,7 @@ class RedisUpdateBuffer: transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, + source_queue=daily_tag_spend_update_queue, ) @with_service_target(SPEND_QUEUE_TARGET) 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 5ecc8784384..f30606ac2f1 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 @@ -6,10 +6,78 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import DailyTagSpendTransaction +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.proxy_server import ProxyStartupEvent +def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction]: + return { + "research_2025-01-15_api-key-123_gpt-4o_openai_/chat/completions": DailyTagSpendTransaction( + tag="research", + request_id="request-123", + date="2025-01-15", + api_key="api-key-123", + model="gpt-4o", + model_group=None, + mcp_namespaced_tool_name=None, + custom_llm_provider="openai", + endpoint="/chat/completions", + prompt_tokens=12, + completion_tokens=7, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + compression_saved_tokens=0, + compression_savings_spend=0.0, + prompt_caching_savings_spend=0.0, + gateway_injected_caching_savings_spend=0.0, + spend=0.0042, + api_requests=1, + successful_requests=1, + failed_requests=0, + ) + } + + +@pytest.mark.asyncio +async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away")) + queue: Final = DailySpendUpdateQueue() + original_transactions: Final = _daily_tag_spend_update_transactions() + await queue.add_update(original_transactions) + + await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + + assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == original_transactions + + +@pytest.mark.asyncio +async def test_daily_tag_spend_restored_after_rpush_failure_is_pushed_once_on_next_tick( + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock +): + mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away")) + queue: Final = DailySpendUpdateQueue() + original_transactions: Final = _daily_tag_spend_update_transactions() + await queue.add_update(original_transactions) + + await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + mock_redis_cache.async_rpush = AsyncMock(return_value=1) + + await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) + + mock_redis_cache.async_rpush.assert_called_once() + redis_key: Final = mock_redis_cache.async_rpush.call_args.kwargs["key"] + values: Final = mock_redis_cache.async_rpush.call_args.kwargs["values"] + assert redis_key == REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY + assert json.loads(values[0]) == json.loads(safe_dumps(original_transactions)) + assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == {} + + @pytest.fixture def mock_redis_cache(): """Create a mock RedisCache instance"""