From 1de5a82e04f836a0a581088d62e3c108ee2a7b05 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 04:44:54 +0000 Subject: [PATCH 1/4] fix(proxy): give spend_counter_cache its own 10000-entry in-memory cache The spend counter DualCache used the generic InMemoryCache default of 200 entries, so deployments with more than 200 active budget scopes evicted still-valid counters and reseeded them from lagging DB spend on the next request. A key the proxy had just rejected for being over budget could be admitted again until the batch writer flushed. Fixes #40221 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/proxy_server.py | 8 +- .../coverage_registry/quota_management.yaml | 1 + .../test_spend_counter_scope_churn_e2e.py | 78 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 15 ++++ 5 files changed, 101 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py diff --git a/litellm/constants.py b/litellm/constants.py index 1012ca831f2e..8d241886b51d 100644 --- a/litellm/constants.py +++ b/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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6fa7a7d761fd..cd5a5c70cefe 100644 --- a/litellm/proxy/proxy_server.py +++ b/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) diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index ad0914d455be..d7179b40ab0a 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -38,6 +38,7 @@ - {id: quota_management.budget.team_multi_window.blocks_then_resets, module: quota_management, tier: P1, behavior: budget, variant: team_multi_window, assertions: [blocks_then_resets], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Team budget_limits enforce and reset per window"} - {id: quota_management.budget.fallback.routes_to_fallback, module: quota_management, tier: P1, behavior: budget, variant: fallback, assertions: [routes_to_fallback], exercised_on: [messages], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "budget_fallbacks reroute to the fallback model once the primary's budget is exhausted"} - {id: quota_management.budget.spend_counter.reseed_matches_db, module: quota_management, tier: P2, behavior: budget, variant: spend_counter, assertions: [reseed_matches_db], exercised_on: [chat_completions], source: "proxy/spend_tracking/budget_reservation.py", rationale: "Concurrent cold-counter reseeds keep the enforcement counter equal to DB spend (#26829)"} +- {id: quota_management.budget.spend_counter.survives_scope_churn, module: quota_management, tier: P1, behavior: budget, variant: spend_counter, assertions: [survives_scope_churn], exercised_on: [chat_completions], source: "proxy/proxy_server.py:spend_counter_cache", rationale: "A key's in-memory spend counter is not evicted by traffic on hundreds of other budget scopes, so an over-budget key stays blocked instead of being re-admitted off lagging DB spend (#40221)"} - {id: quota_management.spend_tracking.chat_completions.logs_cost, module: quota_management, tier: P0, behavior: spend_tracking, variant: chat_completions, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A paid chat call writes a nonzero spend row"} - {id: quota_management.spend_tracking.stream.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Streaming responses aggregate token counts into a spend row"} - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} diff --git a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py new file mode 100644 index 000000000000..f640abf68a0f --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py @@ -0,0 +1,78 @@ +"""Live e2e: an over-budget key stays blocked while hundreds of other budget scopes take traffic (#40221). + +The proxy keeps one in-memory spend counter per budget scope (key, team, end user, tag, ...) +in ``spend_counter_cache``. Every request that names a new ``user`` creates a new end-user +counter. When that cache is sized like the generic in-memory cache (200 entries), the 201st +active scope evicts a still-valid counter, and the next request on the evicted key reseeds it +from the DB. The DB row lags the counter until ``proxy_batch_write_at`` flushes, so a key the +proxy had just rejected for being over budget is admitted again. + +The test blocks a tiny-budget key with a real call, then drives more than 200 distinct end users +on a second key, then re-probes the blocked key. It fails loudly when the probe key's DB spend +has already been flushed, because then a reseed would read the right value and the eviction +would be invisible. The proxy must run without Redis for the eviction to reach the DB; with +Redis the cross-pod counter answers first and this test only re-proves the plain budget block. +""" + +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MODEL = "claude-haiku-4-5" +TINY_BUDGET = 0.000001 +GENERIC_CACHE_SIZE = 200 +CHURN_SCOPES = 250 +CHURN_WORKERS = 24 + + +def _chat(client: BudgetClient, key: str, user: str) -> StreamingResponse: + return client.chat(key, MODEL, f"scope churn {unique_marker()}", max_tokens=1, user=user) + + +@pytest.mark.covers("quota_management.budget.spend_counter.survives_scope_churn") +def test_over_budget_key_stays_blocked_after_hundreds_of_other_scopes( + client: BudgetClient, resources: ResourceManager +) -> None: + run = unique_marker() + probe_key = client.generate_key(max_budget=TINY_BUDGET, models=[MODEL]) + resources.defer(lambda: client.delete_key(probe_key)) + churn_key = client.generate_key(models=[MODEL]) + resources.defer(lambda: client.delete_key(churn_key)) + churn_users = [f"churn-{run}-{index}" for index in range(CHURN_SCOPES)] + resources.defer(lambda: client.delete_customers(churn_users)) + + first = _chat(client, probe_key, user=f"probe-{run}") + assert first.ok, f"first call on a fresh key must be admitted: {first.status_code} {first.body[:300]}" + blocked = _chat(client, probe_key, user=f"probe-{run}") + assert is_budget_block(blocked), ( + f"precondition: the key must be over budget before the churn: {blocked.status_code} {blocked.body[:300]}" + ) + + def churn_call(user: str) -> StreamingResponse: + return _chat(client, churn_key, user=user) + + with ThreadPoolExecutor(max_workers=CHURN_WORKERS) as pool: + churn = list(pool.map(churn_call, churn_users)) + admitted = sum(1 for r in churn if r.ok) + assert admitted > GENERIC_CACHE_SIZE, ( + f"only {admitted}/{CHURN_SCOPES} churn calls succeeded; cannot exceed {GENERIC_CACHE_SIZE} live scopes" + ) + + db_spend = client.proxy.key_info(probe_key).spend or 0.0 + assert db_spend < TINY_BUDGET, ( + f"DB spend {db_spend} already flushed past the budget, so a reseed would block anyway and the " + "eviction path is not exercised; raise proxy_batch_write_at or run the churn faster" + ) + + after = _chat(client, probe_key, user=f"probe-{run}") + assert is_budget_block(after), ( + f"over-budget key was admitted after {admitted} other scopes took traffic: {after.status_code} " + f"{after.body[:300]}; its spend counter was evicted and reseeded from lagging DB spend (#40221)" + ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 26bfd5c52bd8..504a8ceea0e4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8553,6 +8553,21 @@ async def test_get_current_spend_fallback_to_in_memory(): ps.spend_counter_cache = original +def test_spend_counter_cache_keeps_key_counter_across_hundreds_of_other_scopes(): + from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY + from litellm.proxy.proxy_server import spend_counter_cache + + probe_key = "spend:key:unit-40221-probe" + spend_counter_cache.in_memory_cache.set_cache(key=probe_key, value=0.25) + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY + 100): + spend_counter_cache.in_memory_cache.set_cache(key=f"spend:end_user:unit-40221-{index}", value=0.01) + + assert spend_counter_cache.in_memory_cache.get_cache(key=probe_key) == 0.25, ( + f"key counter evicted after {DEFAULT_MAX_SIZE_IN_MEMORY + 100} other scopes; " + f"max_size_in_memory={spend_counter_cache.in_memory_cache.max_size_in_memory}" + ) + + @pytest.mark.asyncio async def test_increment_spend_counters_initializes_and_increments(): """Counter should initialize from cached object spend, then increment. From b0a174a7dfaf0428298d2e4888e605e0cb75b5f6 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 04:56:40 +0000 Subject: [PATCH 2/4] test(proxy): assert spend-counter retention through the public spend path and on any e2e stack Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_spend_counter_scope_churn_e2e.py | 20 ++++++------- tests/test_litellm/proxy/test_proxy_server.py | 28 +++++++++++++------ 2 files changed, 28 insertions(+), 20 deletions(-) diff --git a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py index f640abf68a0f..7b5b4d0fa12a 100644 --- a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py +++ b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py @@ -8,10 +8,12 @@ proxy had just rejected for being over budget is admitted again. The test blocks a tiny-budget key with a real call, then drives more than 200 distinct end users -on a second key, then re-probes the blocked key. It fails loudly when the probe key's DB spend -has already been flushed, because then a reseed would read the right value and the eviction -would be invisible. The proxy must run without Redis for the eviction to reach the DB; with -Redis the cross-pod counter answers first and this test only re-proves the plain budget block. +on a second key, then re-probes the blocked key. The invariant (blocked stays blocked) must hold +on every stack. The eviction itself is only reachable on a proxy without Redis whose DB row still +lags (``proxy_batch_write_at`` above the churn time); with Redis the cross-pod counter answers +first, and a flushed row reseeds correctly, so those stacks re-prove the plain budget block. +The probe key's DB spend at re-probe time is reported in the failure message so a reader can +tell which path was exercised. """ from concurrent.futures import ThreadPoolExecutor @@ -66,13 +68,9 @@ def churn_call(user: str) -> StreamingResponse: ) db_spend = client.proxy.key_info(probe_key).spend or 0.0 - assert db_spend < TINY_BUDGET, ( - f"DB spend {db_spend} already flushed past the budget, so a reseed would block anyway and the " - "eviction path is not exercised; raise proxy_batch_write_at or run the churn faster" - ) - after = _chat(client, probe_key, user=f"probe-{run}") assert is_budget_block(after), ( - f"over-budget key was admitted after {admitted} other scopes took traffic: {after.status_code} " - f"{after.body[:300]}; its spend counter was evicted and reseeded from lagging DB spend (#40221)" + f"over-budget key was admitted after {admitted} other scopes took traffic (DB spend {db_spend}): " + f"{after.status_code} {after.body[:300]}; its spend counter was evicted and reseeded from lagging " + "DB spend (#40221)" ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 504a8ceea0e4..647bcb082846 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -8553,18 +8553,28 @@ async def test_get_current_spend_fallback_to_in_memory(): ps.spend_counter_cache = original -def test_spend_counter_cache_keeps_key_counter_across_hundreds_of_other_scopes(): +@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 spend_counter_cache + from litellm.proxy.proxy_server import get_current_spend, increment_spend_counters, spend_counter_cache - probe_key = "spend:key:unit-40221-probe" - spend_counter_cache.in_memory_cache.set_cache(key=probe_key, value=0.25) - for index in range(DEFAULT_MAX_SIZE_IN_MEMORY + 100): - spend_counter_cache.in_memory_cache.set_cache(key=f"spend:end_user:unit-40221-{index}", value=0.01) + 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_counter_cache.in_memory_cache.get_cache(key=probe_key) == 0.25, ( - f"key counter evicted after {DEFAULT_MAX_SIZE_IN_MEMORY + 100} other scopes; " - f"max_size_in_memory={spend_counter_cache.in_memory_cache.max_size_in_memory}" + 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)" ) From 675de2f1f49b226343cb27b9342ddd9627b73224 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 05:23:28 +0000 Subject: [PATCH 3/4] test(e2e): keep the spend-counter scope churn case inside a feature class Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_spend_counter_scope_churn_e2e.py | 65 ++++++++++--------- 1 file changed, 33 insertions(+), 32 deletions(-) diff --git a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py index 7b5b4d0fa12a..c8a71285acc3 100644 --- a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py +++ b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py @@ -38,39 +38,40 @@ def _chat(client: BudgetClient, key: str, user: str) -> StreamingResponse: return client.chat(key, MODEL, f"scope churn {unique_marker()}", max_tokens=1, user=user) -@pytest.mark.covers("quota_management.budget.spend_counter.survives_scope_churn") -def test_over_budget_key_stays_blocked_after_hundreds_of_other_scopes( - client: BudgetClient, resources: ResourceManager -) -> None: - run = unique_marker() - probe_key = client.generate_key(max_budget=TINY_BUDGET, models=[MODEL]) - resources.defer(lambda: client.delete_key(probe_key)) - churn_key = client.generate_key(models=[MODEL]) - resources.defer(lambda: client.delete_key(churn_key)) - churn_users = [f"churn-{run}-{index}" for index in range(CHURN_SCOPES)] - resources.defer(lambda: client.delete_customers(churn_users)) +class TestSpendCounterScopeChurn: + @pytest.mark.covers("quota_management.budget.spend_counter.survives_scope_churn") + def test_over_budget_key_stays_blocked_after_hundreds_of_other_scopes( + self, client: BudgetClient, resources: ResourceManager + ) -> None: + run = unique_marker() + probe_key = client.generate_key(max_budget=TINY_BUDGET, models=[MODEL]) + resources.defer(lambda: client.delete_key(probe_key)) + churn_key = client.generate_key(models=[MODEL]) + resources.defer(lambda: client.delete_key(churn_key)) + churn_users = [f"churn-{run}-{index}" for index in range(CHURN_SCOPES)] + resources.defer(lambda: client.delete_customers(churn_users)) - first = _chat(client, probe_key, user=f"probe-{run}") - assert first.ok, f"first call on a fresh key must be admitted: {first.status_code} {first.body[:300]}" - blocked = _chat(client, probe_key, user=f"probe-{run}") - assert is_budget_block(blocked), ( - f"precondition: the key must be over budget before the churn: {blocked.status_code} {blocked.body[:300]}" - ) + first = _chat(client, probe_key, user=f"probe-{run}") + assert first.ok, f"first call on a fresh key must be admitted: {first.status_code} {first.body[:300]}" + blocked = _chat(client, probe_key, user=f"probe-{run}") + assert is_budget_block(blocked), ( + f"precondition: the key must be over budget before the churn: {blocked.status_code} {blocked.body[:300]}" + ) - def churn_call(user: str) -> StreamingResponse: - return _chat(client, churn_key, user=user) + def churn_call(user: str) -> StreamingResponse: + return _chat(client, churn_key, user=user) - with ThreadPoolExecutor(max_workers=CHURN_WORKERS) as pool: - churn = list(pool.map(churn_call, churn_users)) - admitted = sum(1 for r in churn if r.ok) - assert admitted > GENERIC_CACHE_SIZE, ( - f"only {admitted}/{CHURN_SCOPES} churn calls succeeded; cannot exceed {GENERIC_CACHE_SIZE} live scopes" - ) + with ThreadPoolExecutor(max_workers=CHURN_WORKERS) as pool: + churn = list(pool.map(churn_call, churn_users)) + admitted = sum(1 for r in churn if r.ok) + assert admitted > GENERIC_CACHE_SIZE, ( + f"only {admitted}/{CHURN_SCOPES} churn calls succeeded; cannot exceed {GENERIC_CACHE_SIZE} live scopes" + ) - db_spend = client.proxy.key_info(probe_key).spend or 0.0 - after = _chat(client, probe_key, user=f"probe-{run}") - assert is_budget_block(after), ( - f"over-budget key was admitted after {admitted} other scopes took traffic (DB spend {db_spend}): " - f"{after.status_code} {after.body[:300]}; its spend counter was evicted and reseeded from lagging " - "DB spend (#40221)" - ) + db_spend = client.proxy.key_info(probe_key).spend or 0.0 + after = _chat(client, probe_key, user=f"probe-{run}") + assert is_budget_block(after), ( + f"over-budget key was admitted after {admitted} other scopes took traffic (DB spend {db_spend}): " + f"{after.status_code} {after.body[:300]}; its spend counter was evicted and reseeded from lagging " + "DB spend (#40221)" + ) From ec97ed69c2390839a2a830093ea5b5652d3e6764 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 18:44:47 +0000 Subject: [PATCH 4/4] test(spend): move the spend-counter scope churn repro to tests/integration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../coverage_registry/quota_management.yaml | 1 - .../test_spend_counter_scope_churn_e2e.py | 77 ------------------- tests/integration/contracts.json | 3 + .../spend/test_spend_counter_scope_churn.py | 70 +++++++++++++++++ 4 files changed, 73 insertions(+), 78 deletions(-) delete mode 100644 tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py create mode 100644 tests/integration/spend/test_spend_counter_scope_churn.py diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index d7179b40ab0a..ad0914d455be 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -38,7 +38,6 @@ - {id: quota_management.budget.team_multi_window.blocks_then_resets, module: quota_management, tier: P1, behavior: budget, variant: team_multi_window, assertions: [blocks_then_resets], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Team budget_limits enforce and reset per window"} - {id: quota_management.budget.fallback.routes_to_fallback, module: quota_management, tier: P1, behavior: budget, variant: fallback, assertions: [routes_to_fallback], exercised_on: [messages], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "budget_fallbacks reroute to the fallback model once the primary's budget is exhausted"} - {id: quota_management.budget.spend_counter.reseed_matches_db, module: quota_management, tier: P2, behavior: budget, variant: spend_counter, assertions: [reseed_matches_db], exercised_on: [chat_completions], source: "proxy/spend_tracking/budget_reservation.py", rationale: "Concurrent cold-counter reseeds keep the enforcement counter equal to DB spend (#26829)"} -- {id: quota_management.budget.spend_counter.survives_scope_churn, module: quota_management, tier: P1, behavior: budget, variant: spend_counter, assertions: [survives_scope_churn], exercised_on: [chat_completions], source: "proxy/proxy_server.py:spend_counter_cache", rationale: "A key's in-memory spend counter is not evicted by traffic on hundreds of other budget scopes, so an over-budget key stays blocked instead of being re-admitted off lagging DB spend (#40221)"} - {id: quota_management.spend_tracking.chat_completions.logs_cost, module: quota_management, tier: P0, behavior: spend_tracking, variant: chat_completions, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A paid chat call writes a nonzero spend row"} - {id: quota_management.spend_tracking.stream.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Streaming responses aggregate token counts into a spend row"} - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} diff --git a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py b/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py deleted file mode 100644 index c8a71285acc3..000000000000 --- a/tests/e2e/quota_management/budgets/test_spend_counter_scope_churn_e2e.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Live e2e: an over-budget key stays blocked while hundreds of other budget scopes take traffic (#40221). - -The proxy keeps one in-memory spend counter per budget scope (key, team, end user, tag, ...) -in ``spend_counter_cache``. Every request that names a new ``user`` creates a new end-user -counter. When that cache is sized like the generic in-memory cache (200 entries), the 201st -active scope evicts a still-valid counter, and the next request on the evicted key reseeds it -from the DB. The DB row lags the counter until ``proxy_batch_write_at`` flushes, so a key the -proxy had just rejected for being over budget is admitted again. - -The test blocks a tiny-budget key with a real call, then drives more than 200 distinct end users -on a second key, then re-probes the blocked key. The invariant (blocked stays blocked) must hold -on every stack. The eviction itself is only reachable on a proxy without Redis whose DB row still -lags (``proxy_batch_write_at`` above the churn time); with Redis the cross-pod counter answers -first, and a flushed row reseeds correctly, so those stacks re-prove the plain budget block. -The probe key's DB spend at re-probe time is reported in the failure message so a reader can -tell which path was exercised. -""" - -from concurrent.futures import ThreadPoolExecutor - -import pytest - -from budget_client import BudgetClient, is_budget_block -from e2e_config import unique_marker -from e2e_http import StreamingResponse -from lifecycle import ResourceManager - -pytestmark = pytest.mark.e2e - -MODEL = "claude-haiku-4-5" -TINY_BUDGET = 0.000001 -GENERIC_CACHE_SIZE = 200 -CHURN_SCOPES = 250 -CHURN_WORKERS = 24 - - -def _chat(client: BudgetClient, key: str, user: str) -> StreamingResponse: - return client.chat(key, MODEL, f"scope churn {unique_marker()}", max_tokens=1, user=user) - - -class TestSpendCounterScopeChurn: - @pytest.mark.covers("quota_management.budget.spend_counter.survives_scope_churn") - def test_over_budget_key_stays_blocked_after_hundreds_of_other_scopes( - self, client: BudgetClient, resources: ResourceManager - ) -> None: - run = unique_marker() - probe_key = client.generate_key(max_budget=TINY_BUDGET, models=[MODEL]) - resources.defer(lambda: client.delete_key(probe_key)) - churn_key = client.generate_key(models=[MODEL]) - resources.defer(lambda: client.delete_key(churn_key)) - churn_users = [f"churn-{run}-{index}" for index in range(CHURN_SCOPES)] - resources.defer(lambda: client.delete_customers(churn_users)) - - first = _chat(client, probe_key, user=f"probe-{run}") - assert first.ok, f"first call on a fresh key must be admitted: {first.status_code} {first.body[:300]}" - blocked = _chat(client, probe_key, user=f"probe-{run}") - assert is_budget_block(blocked), ( - f"precondition: the key must be over budget before the churn: {blocked.status_code} {blocked.body[:300]}" - ) - - def churn_call(user: str) -> StreamingResponse: - return _chat(client, churn_key, user=user) - - with ThreadPoolExecutor(max_workers=CHURN_WORKERS) as pool: - churn = list(pool.map(churn_call, churn_users)) - admitted = sum(1 for r in churn if r.ok) - assert admitted > GENERIC_CACHE_SIZE, ( - f"only {admitted}/{CHURN_SCOPES} churn calls succeeded; cannot exceed {GENERIC_CACHE_SIZE} live scopes" - ) - - db_spend = client.proxy.key_info(probe_key).spend or 0.0 - after = _chat(client, probe_key, user=f"probe-{run}") - assert is_budget_block(after), ( - f"over-budget key was admitted after {admitted} other scopes took traffic (DB spend {db_spend}): " - f"{after.status_code} {after.body[:300]}; its spend counter was evicted and reseeded from lagging " - "DB spend (#40221)" - ) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index e1b5940935fb..d455c8b1da3f 100644 --- a/tests/integration/contracts.json +++ b/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" ], diff --git a/tests/integration/spend/test_spend_counter_scope_churn.py b/tests/integration/spend/test_spend_counter_scope_churn.py new file mode 100644 index 000000000000..1dd0577b4286 --- /dev/null +++ b/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"] == []