From dcfde1b89952d5417d8f9a4af7f7a8e1bac8d20d Mon Sep 17 00:00:00 2001
From: user <70670632+stuxf@users.noreply.github.com>
Date: Thu, 30 Apr 2026 20:04:16 -0700
Subject: [PATCH] fallback to plain org cache for spend counters

---
 litellm/proxy/proxy_server.py                 | 34 +++++++++++++------
 .../proxy/test_budget_reservation.py          | 31 +++++++++++++++++
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 3eb536d0850..927d0ae8e8a 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -2048,7 +2048,7 @@ async def _increment_org_spend_counter(
 
     await _init_and_increment_unreserved_spend_counter(
         counter_key=f"spend:org:{org_id}",
-        source_cache_key=f"org_id:{org_id}:with_budget",
+        source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
         increment=response_cost,
         reserved_counter_keys=reserved_counter_keys,
     )
@@ -2056,7 +2056,7 @@ async def _increment_org_spend_counter(
 
 async def _init_and_increment_unreserved_spend_counter(
     counter_key: str,
-    source_cache_key: str,
+    source_cache_key: Union[str, List[str]],
     increment: float,
     reserved_counter_keys: Set[str],
 ) -> None:
@@ -2072,7 +2072,7 @@ async def _init_and_increment_unreserved_spend_counter(
 
 async def _init_and_increment_spend_counter(
     counter_key: str,
-    source_cache_key: str,
+    source_cache_key: Union[str, List[str]],
     increment: float,
 ):
     """
@@ -2117,7 +2117,7 @@ async def _init_and_increment_window_spend_counter(
 
 async def _ensure_spend_counter_initialized(
     counter_key: str,
-    source_cache_key: str,
+    source_cache_key: Union[str, List[str]],
 ):
     is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key)
     if is_warm is False:
@@ -2130,19 +2130,31 @@ async def _ensure_spend_counter_initialized(
         )
         if db_spend is None:
             # DB unavailable - fall back to in-process cache (may be stale).
-            source = await user_api_key_cache.async_get_cache(key=source_cache_key)
-            base_spend: float = 0.0
-            if source is not None:
-                if isinstance(source, dict):
-                    base_spend = source.get("spend", 0.0) or 0.0
-                else:
-                    base_spend = getattr(source, "spend", 0.0) or 0.0
+            base_spend = await _get_source_cache_base_spend(
+                source_cache_key=source_cache_key
+            )
             if base_spend > 0:
                 await _increment_spend_counter_cache(
                     counter_key=counter_key, increment=base_spend
                 )
 
 
+async def _get_source_cache_base_spend(
+    source_cache_key: Union[str, List[str]],
+) -> float:
+    source_cache_keys = (
+        [source_cache_key] if isinstance(source_cache_key, str) else source_cache_key
+    )
+    for cache_key in source_cache_keys:
+        source = await user_api_key_cache.async_get_cache(key=cache_key)
+        if source is None:
+            continue
+        if isinstance(source, dict):
+            return float(source.get("spend", 0.0) or 0.0)
+        return float(getattr(source, "spend", 0.0) or 0.0)
+    return 0.0
+
+
 async def _ensure_window_spend_counter_initialized(
     counter_key: str,
     entity_type: str,
diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py
index 72f434ede46..d1684b7ef1f 100644
--- a/tests/test_litellm/proxy/test_budget_reservation.py
+++ b/tests/test_litellm/proxy/test_budget_reservation.py
@@ -509,6 +509,37 @@ async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_stat
     ) == pytest.approx(2.25)
 
 
+@pytest.mark.asyncio
+async def test_should_seed_org_counter_from_plain_org_cache(spend_counter_state):
+    counter_cache, key_cache = spend_counter_state
+    await key_cache.async_set_cache(
+        key="org_id:org-counter-plain",
+        value=LiteLLM_OrganizationTable(
+            organization_id="org-counter-plain",
+            organization_alias="shared-org",
+            budget_id="org-budget-id",
+            spend=2.0,
+            models=[],
+            created_by="test",
+            updated_by="test",
+        ).model_dump(),
+    )
+
+    from litellm.proxy.proxy_server import increment_spend_counters
+
+    await increment_spend_counters(
+        token=None,
+        team_id=None,
+        user_id=None,
+        org_id="org-counter-plain",
+        response_cost=0.25,
+    )
+
+    assert counter_cache.in_memory_cache.get_cache(
+        key="spend:org:org-counter-plain"
+    ) == pytest.approx(2.25)
+
+
 @pytest.mark.asyncio
 async def test_should_cap_known_estimate_to_remaining_budget(
     spend_counter_state,