fix(proxy): give spend_counter_cache its own 10000-entry in-memory cache

litellm/constants.py

@@ -1879,6 +1879,7 @@
 SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30
 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000
+SPEND_COUNTER_CACHE_MAX_SIZE: Final = 10000
 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000
 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated
 # callers from forcing a DB query per request for unknown names, while bounding

litellm/proxy/proxy_server.py

@@ -256,7 +256,7 @@ def generate_feedback_box():
 import litellm._redis
 from litellm import Router
 from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger
-from litellm.caching.caching import DualCache, RedisCache
+from litellm.caching.caching import DualCache, InMemoryCache, RedisCache
 from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
 from litellm.caching.redis_cluster_cache import RedisClusterCache
 from litellm.constants import (
@@ -285,6 +285,7 @@ def generate_feedback_box():
     REALTIME_SESSION_FAILURE_LOGGED_KEY,
     REALTIME_SESSION_SUCCESS_LOGGED_KEY,
     ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
+    SPEND_COUNTER_CACHE_MAX_SIZE,
     USER_SPEND_ALERTS_JOB_ID,
     WEEKLY_SPEND_REPORT_JOB_ID,
 )
@@ -2506,7 +2507,10 @@ def general_settings_view() -> Mapping[str, object]:
 user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
     default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
 )
-spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value)
+spend_counter_cache: Final = DualCache(
+    in_memory_cache=InMemoryCache(max_size_in_memory=SPEND_COUNTER_CACHE_MAX_SIZE),
+    default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
+)
 cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS)
 model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache)
 litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)

tests/integration/contracts.json

@@ -107,6 +107,9 @@
     "tests/integration/spend/test_cache_and_quota.py::test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores": [
       "quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores"
     ],
+    "tests/integration/spend/test_spend_counter_scope_churn.py::test_over_budget_key_stays_blocked_while_hundreds_of_end_users_take_traffic": [
+      "quota_management.budget.spend_counter.survives_scope_churn"
+    ],
     "tests/integration/spend/test_cache_and_quota.py::test_different_system_messages_do_not_share_a_cached_response": [
       "quota_management.response_cache.system_messages_partition_cache_identity"
     ],

tests/integration/spend/test_spend_counter_scope_churn.py

@@ -0,0 +1,70 @@
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import yaml
+from integration._support.client import Gateway
+from integration._support.process import owned_proxy
+from pydantic import TypeAdapter
+
+GENERIC_CACHE_SIZE: Final = 200
+CHURN_SCOPES: Final = 250
+
+
+def _chat(candidate: Gateway, model: str, key: str, user: str) -> httpx.Response:
+    return candidate.request(
+        "POST",
+        "/v1/chat/completions",
+        {"model": model, "messages": [{"role": "user", "content": f"scope churn {uuid.uuid4().hex}"}], "user": user},
+        key=key,
+    )
+
+
+@pytest.mark.covers("quota_management.budget.spend_counter.survives_scope_churn")
+@pytest.mark.timeout(180)
+def test_over_budget_key_stays_blocked_while_hundreds_of_end_users_take_traffic(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    shared: Final = TypeAdapter(dict[str, object]).validate_python(
+        yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    )
+    general_settings: Final = TypeAdapter(dict[str, object]).validate_python(shared["general_settings"])
+    path: Final = tmp_path / "in-memory-counters.yaml"
+    path.write_text(
+        yaml.safe_dump(
+            {
+                "model_list": [],
+                "general_settings": {**general_settings, "proxy_batch_write_at": 600},
+                "router_settings": shared["router_settings"],
+            }
+        )
+    )
+    run: Final = uuid.uuid4().hex
+    with (
+        owned_proxy(gateway, tmp_path, {}, config=path, remove_environment=("REDIS_HOST", "REDIS_PORT")) as candidate,
+        candidate.scenario() as scenario,
+        httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as upstream,
+    ):
+        model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
+        probe: Final = scenario.key(models=[model], max_budget=0.06)
+        churn: Final = scenario.key(models=[model])
+
+        def _churn_status(index: int) -> int:
+            return _chat(candidate, model, churn, f"churn-{run}-{index}").status_code
+
+        first: Final = _chat(candidate, model, probe, f"probe-{run}")
+        assert first.status_code == 200 and first.json()["usage"]["total_tokens"] == 40, first.text
+        blocked: Final = _chat(candidate, model, probe, f"probe-{run}")
+        assert blocked.status_code == 422 and blocked.json()["error"]["type"] == "budget_exceeded", blocked.text
+        with ThreadPoolExecutor(max_workers=8) as pool:
+            statuses: Final = tuple(pool.map(_churn_status, range(CHURN_SCOPES)))
+        assert statuses.count(200) > GENERIC_CACHE_SIZE, statuses
+        upstream.get("/__observations").raise_for_status()
+        after: Final = _chat(candidate, model, probe, f"probe-{run}")
+        assert after.status_code == 422 and after.json()["error"]["type"] == "budget_exceeded", (
+            f"over-budget key admitted after {statuses.count(200)} other scopes took traffic: {after.status_code} {after.text}"
+        )
+        assert upstream.get("/__observations").json()["requests"] == []

tests/test_litellm/proxy/test_proxy_server.py

@@ -8553,6 +8553,31 @@ async def test_get_current_spend_fallback_to_in_memory():
         ps.spend_counter_cache = original
 
 
+@pytest.mark.asyncio
+async def test_over_budget_key_spend_survives_hundreds_of_other_scopes_incrementing():
+    from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY
+    from litellm.proxy.proxy_server import get_current_spend, increment_spend_counters, spend_counter_cache
+
+    token = "unit-40221-probe"
+    other_scopes = DEFAULT_MAX_SIZE_IN_MEMORY + 100
+    end_users = [f"unit-40221-{index}" for index in range(other_scopes)]
+    try:
+        await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.25)
+        for end_user in end_users:
+            await increment_spend_counters(
+                token=None, team_id=None, user_id=None, response_cost=0.01, end_user_id=end_user
+            )
+        spend = await get_current_spend(counter_key=f"spend:key:{token}", fallback_spend=0.0)
+    finally:
+        for counter_key in [f"spend:key:{token}", *(f"spend:end_user:{end_user}" for end_user in end_users)]:
+            spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
+
+    assert spend == 0.25, (
+        f"key spend read back as {spend} after {other_scopes} other scopes were charged; "
+        "the counter was evicted and the budget check fell back to the stale value (#40221)"
+    )
+
+
 @pytest.mark.asyncio
 async def test_increment_spend_counters_initializes_and_increments():
     """Counter should initialize from cached object spend, then increment.