Merge pull request #26270 from BerriAI/litellm_/lucid-kowalevski-de832f

[Fix] Stabilize flaky spend accuracy tests + patch Redis buffer data-loss path

tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py

@@ -87,6 +87,89 @@ async def test_store_in_memory_spend_updates_uses_pipeline(
     assert len(rpush_list) == 3
 
 
+@pytest.mark.asyncio
+async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
+    redis_update_buffer, mock_redis_cache
+):
+    """
+    If async_rpush_pipeline raises, the already-drained transactions must be
+    put back into the in-memory queues so the next scheduler tick retries.
+    Without this, any transient Redis hiccup silently loses spend data.
+    """
+    from litellm.proxy._types import Litellm_EntityType
+    from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
+        DailySpendUpdateQueue,
+    )
+    from litellm.proxy.db.db_transaction_queue.spend_update_queue import (
+        SpendUpdateQueue,
+    )
+
+    mock_redis_cache.async_rpush_pipeline = AsyncMock(
+        side_effect=ConnectionError("redis went away")
+    )
+
+    spend_queue = SpendUpdateQueue()
+    daily_user_queue = DailySpendUpdateQueue()
+    daily_team_queue = DailySpendUpdateQueue()
+    daily_org_queue = DailySpendUpdateQueue()
+    daily_end_user_queue = DailySpendUpdateQueue()
+    daily_agent_queue = DailySpendUpdateQueue()
+
+    # Seed real queues with data so flush_and_get_aggregated returns it
+    await spend_queue.add_update(
+        {
+            "entity_type": Litellm_EntityType.KEY,
+            "entity_id": "key-abc",
+            "response_cost": 1.5,
+        }
+    )
+    await spend_queue.add_update(
+        {
+            "entity_type": Litellm_EntityType.TEAM,
+            "entity_id": "team-xyz",
+            "response_cost": 2.5,
+        }
+    )
+    await daily_user_queue.add_update(
+        {
+            "user1_day_model": {
+                "spend": 1.0,
+                "prompt_tokens": 10,
+                "completion_tokens": 20,
+            }
+        }
+    )
+
+    await redis_update_buffer.store_in_memory_spend_updates_in_redis(
+        spend_update_queue=spend_queue,
+        daily_spend_update_queue=daily_user_queue,
+        daily_team_spend_update_queue=daily_team_queue,
+        daily_org_spend_update_queue=daily_org_queue,
+        daily_end_user_spend_update_queue=daily_end_user_queue,
+        daily_agent_spend_update_queue=daily_agent_queue,
+    )
+
+    # After restore, the main spend queue should hold one item per
+    # (entity_type, entity_id) pair with the aggregated cost
+    restored_spend = (
+        await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
+    )
+    assert restored_spend["key_list_transactions"] == {"key-abc": 1.5}
+    assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5}
+
+    # Daily user queue should hold the same aggregated dict
+    restored_daily = (
+        await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions()
+    )
+    assert restored_daily == {
+        "user1_day_model": {
+            "spend": 1.0,
+            "prompt_tokens": 10,
+            "completion_tokens": 20,
+        }
+    }
+
+
 @pytest.mark.asyncio
 async def test_store_in_memory_spend_updates_all_empty_returns_early(
     redis_update_buffer, mock_redis_cache