From d533b432fdb2d9dbfd4410c5f30e552611b137ec Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Fri, 27 Mar 2026 14:42:05 +0100 Subject: [PATCH] fix(proxy): enforce budget limits across multi-pod deployments via Redis-backed spend counters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Budget checks on API keys, teams, and team members were not enforced in multi-pod deployments because user_api_key_cache is intentionally in-memory-only. Each pod tracked spend independently, so with N pods the effective budget was N × max_budget. Introduces a separate spend_counter_cache (DualCache wired to redis_usage_cache) with atomic increment/read helpers: - increment_spend_counters(): awaited in cost callback (not create_task) to update both in-memory and Redis before the next auth check - get_current_spend(): reads Redis first (cross-pod authoritative), falls back to in-memory, then to cached object .spend from DB Budget check functions (_virtual_key_max_budget_check, _team_max_budget_check, _check_team_member_budget) now read spend via get_current_spend() instead of cached object .spend fields. When Redis is not configured, falls back to in-memory-only counters (same as current single-instance behavior). Fixes #23714 --- litellm/proxy/auth/auth_checks.py | 77 +++++--- litellm/proxy/auth/user_api_key_auth.py | 16 +- .../proxy/common_utils/reset_budget_job.py | 84 +++++++- .../proxy/hooks/proxy_track_cost_callback.py | 18 +- litellm/proxy/proxy_server.py | 132 +++++++++++++ .../proxy/auth/test_auth_checks.py | 151 +++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 180 ++++++++++++++++++ 7 files changed, 618 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 815393467de..0efc3dab638 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2876,7 +2876,15 @@ async def _virtual_key_max_budget_check( Triggers a budget alert if the token is over it's max budget. """ - if valid_token.spend is not None and valid_token.max_budget is not None: + if valid_token.max_budget is not None: + from litellm.proxy.proxy_server import get_current_spend + + # Read spend from cross-pod counter (Redis-first) or cached object (fallback) + spend = await get_current_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + ) + #################################### # collect information for alerting # #################################### @@ -2888,7 +2896,7 @@ async def _virtual_key_max_budget_check( call_info = CallInfo( token=valid_token.token, - spend=valid_token.spend, + spend=spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, @@ -2909,9 +2917,9 @@ async def _virtual_key_max_budget_check( # collect information for alerting # #################################### - if valid_token.spend >= valid_token.max_budget: + if spend >= valid_token.max_budget: raise litellm.BudgetExceededError( - current_cost=valid_token.spend, + current_cost=spend, max_budget=valid_token.max_budget, ) @@ -3042,6 +3050,14 @@ async def _check_team_member_budget( team_member_budget = team_membership.litellm_budget_table.max_budget team_member_spend = team_membership.spend or 0.0 + # Read from cross-pod counter (Redis-first) if available + from litellm.proxy.proxy_server import get_current_spend + + team_member_spend = await get_current_spend( + counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + fallback_spend=team_member_spend, + ) + if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, @@ -3065,33 +3081,40 @@ async def _team_max_budget_check( if ( team_object is not None and team_object.max_budget is not None - and team_object.spend is not None - and team_object.spend > team_object.max_budget ): - if valid_token: - call_info = CallInfo( - token=valid_token.token, - spend=team_object.spend, - max_budget=team_object.max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.TEAM, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="team_budget", - user_info=call_info, - ) - ) + from litellm.proxy.proxy_server import get_current_spend - raise litellm.BudgetExceededError( - current_cost=team_object.spend, - max_budget=team_object.max_budget, - message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}", + # Read spend from cross-pod counter (Redis-first) or cached object (fallback) + spend = await get_current_spend( + counter_key=f"spend:team:{team_object.team_id}", + fallback_spend=team_object.spend or 0.0, ) + if spend > team_object.max_budget: + if valid_token: + call_info = CallInfo( + token=valid_token.token, + spend=spend, + max_budget=team_object.max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.TEAM, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=spend, + max_budget=team_object.max_budget, + message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", + ) + async def _team_soft_budget_check( team_object: Optional[LiteLLM_TeamTable], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index eba787c63b3..9dd2ab18c8a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1303,9 +1303,21 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_member_info.litellm_budget_table.max_budget ) if team_member_budget is not None and team_member_budget > 0: - if valid_token.team_member_spend > team_member_budget: + # Read from cross-pod counter (Redis-first) if available + from litellm.proxy.proxy_server import get_current_spend + + team_member_spend = valid_token.team_member_spend + if ( + valid_token.user_id is not None + and valid_token.team_id is not None + ): + team_member_spend = await get_current_spend( + counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", + fallback_spend=team_member_spend, + ) + if team_member_spend > team_member_budget: raise litellm.BudgetExceededError( - current_cost=valid_token.team_member_spend, + current_cost=team_member_spend, max_budget=team_member_budget, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 674214b19e5..5db919659ed 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -54,16 +54,49 @@ class ResetBudgetJob: """ Resets the budget for all LiteLLM Team Members if their budget has expired """ + budget_ids = [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + + # Reset spend counters for affected team members. + # Reset Redis directly so a transient failure doesn't leave stale + # counters that get_current_spend would read as authoritative. + try: + from litellm.proxy.proxy_server import spend_counter_cache + + memberships = ( + await self.prisma_client.db.litellm_teammembership.find_many( + where={"budget_id": {"in": budget_ids}} + ) + ) + for m in memberships: + counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" + # Always reset in-memory + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0 + ) + # Explicitly reset Redis with warning on failure + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset team member spend counter in Redis %s: %s. " + "Budget may be over-enforced until counter expires.", + counter_key, + redis_err, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to reset team member spend counters: %s", e + ) + return await self.prisma_client.db.litellm_teammembership.update_many( - where={ - "budget_id": { - "in": [ - budget.budget_id - for budget in budgets_to_reset - if budget.budget_id is not None - ] - } - }, + where={"budget_id": {"in": budget_ids}}, data={ "spend": 0, }, @@ -531,6 +564,39 @@ class ResetBudgetJob: """ try: item.spend = 0.0 + + # Reset the cross-pod spend counter. + # Reset Redis directly (not via DualCache) so a Redis failure + # doesn't silently leave a stale counter that get_current_spend + # would read as authoritative, permanently blocking the user. + from litellm.proxy.proxy_server import spend_counter_cache + + counter_key = None + if item_type == "key" and hasattr(item, "token") and item.token is not None: + counter_key = f"spend:key:{item.token}" + elif item_type == "team" and hasattr(item, "team_id") and item.team_id is not None: + counter_key = f"spend:team:{item.team_id}" + + if counter_key is not None: + # Always reset in-memory (local fallback) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0 + ) + # Explicitly reset Redis with warning on failure + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset spend counter in Redis for %s key=%s: %s. " + "Budget may be over-enforced until counter expires.", + item_type, + counter_key, + redis_err, + ) + if hasattr(item, "budget_duration") and item.budget_duration is not None: # Get standardized reset time based on budget duration from litellm.proxy.common_utils.timezone_utils import ( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 220c5066a6d..948d6dd33af 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -135,7 +135,11 @@ class _ProxyDBLogger(CustomLogger): start_time=None, end_time=None, # start/end time for completion ): - from litellm.proxy.proxy_server import proxy_logging_obj, update_cache + from litellm.proxy.proxy_server import ( + increment_spend_counters, + proxy_logging_obj, + update_cache, + ) verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") try: @@ -194,7 +198,17 @@ class _ProxyDBLogger(CustomLogger): org_id=org_id, ) - # update cache + # Atomically update spend counters (in-memory + Redis) + # for cross-pod budget enforcement. + await increment_spend_counters( + token=user_api_key, + team_id=team_id, + user_id=user_id, + response_cost=response_cost, + ) + + # update cache (fire-and-forget for backward compat: + # cached object fields, soft budget alerts, etc.) asyncio.create_task( update_cache( token=user_api_key, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d3d2ceb533..974f21d4f74 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1530,6 +1530,9 @@ shared_aiohttp_session: Optional[ user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) +spend_counter_cache = DualCache( + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value +) model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) @@ -1694,6 +1697,134 @@ def cost_tracking(): ) +async def get_current_spend(counter_key: str, fallback_spend: float) -> float: + """ + Read current spend from the cross-pod spend counter. + + Reads Redis FIRST (authoritative cross-pod value), not DualCache's + async_get_cache which returns in-memory first. This is critical: + DualCache.async_get_cache returns stale per-pod values because each + pod's in-memory cache is only updated by that pod's own increments. + + Fallback chain: + 1. Redis counter (cross-pod, authoritative) + 2. In-memory counter (single-instance or Redis failure) + 3. Cached object's .spend from DB (cold start, no counter yet) + """ + # 1. Try Redis first (cross-pod authoritative) + if spend_counter_cache.redis_cache is not None: + try: + val = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + if val is not None: + return float(val) + except Exception as e: + verbose_proxy_logger.debug( + "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", + counter_key, + e, + ) + + # 2. Fall back to in-memory counter (single-instance or Redis failure) + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) + + # 3. Final fallback: cached object's spend from DB + return fallback_spend + + +async def increment_spend_counters( + token: Optional[str], + team_id: Optional[str], + user_id: Optional[str], + response_cost: Optional[float], +): + """ + Atomically increment spend counters for budget enforcement. + + Uses spend_counter_cache (DualCache with Redis backend when available) + so counters are shared across all pods. Budget check functions read + from these counters via get_current_spend() (Redis-first). + + Awaited (not create_task) in the cost callback, so the counter is + updated before the next request's auth check runs. + """ + if response_cost is None or response_cost == 0: + return + + if token is not None: + # token arrives pre-hashed from metadata["user_api_key"] (auth flow + # hashes raw "sk-..." keys before they reach the callback). The + # startswith("sk-") check is a safety net matching update_cache — + # if a raw key somehow arrives, hash it; otherwise use as-is to + # avoid double-hashing (budget checks read valid_token.token which + # is single-hashed). + hashed_token = ( + hash_token(token=token) + if isinstance(token, str) and token.startswith("sk-") + else token + ) + await _init_and_increment_spend_counter( + counter_key=f"spend:key:{hashed_token}", + source_cache_key=hashed_token, + increment=response_cost, + ) + + if team_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:team:{team_id}", + source_cache_key=f"team_id:{team_id}", + increment=response_cost, + ) + + if user_id is not None and team_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:team_member:{user_id}:{team_id}", + source_cache_key=f"team_membership:{user_id}:{team_id}", + increment=response_cost, + ) + + +async def _init_and_increment_spend_counter( + counter_key: str, + source_cache_key: str, + increment: float, +): + """ + Initialize counter from cached object's DB-loaded spend if not yet set, + then atomically increment in both in-memory and Redis. + + On first access per pod: + 1. Check spend_counter_cache (in-memory -> Redis via DualCache for init check) + 2. If not found anywhere, read base spend from user_api_key_cache (DB-loaded object) + 3. Seed counter via async_increment_cache (not async_set_cache) to avoid a + check-then-set race: if two pods cold-start simultaneously, both may see + the counter as absent and seed it. Using increment instead of set means + the worst case is over-counting (conservative — blocks slightly early) + rather than under-counting (would allow overspend). + 4. Increment atomically (both in-memory + Redis) + """ + current = await spend_counter_cache.async_get_cache(key=counter_key) + if current is None: + source = await user_api_key_cache.async_get_cache(key=source_cache_key) + base_spend = 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 + if base_spend > 0: + await spend_counter_cache.async_increment_cache( + key=counter_key, value=base_spend + ) + + await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment + ) + + async def update_cache( # noqa: PLR0915 token: Optional[str], user_id: Optional[str], @@ -2528,6 +2659,7 @@ class ProxyConfig: ): ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache + spend_counter_cache.redis_cache = redis_usage_cache # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 69188fd200e..bd659ed518f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -29,10 +29,13 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, + _check_team_member_budget, _get_fuzzy_user_object, _get_team_db_check, _log_budget_lookup_failure, + _team_max_budget_check, _virtual_key_max_budget_alert_check, + _virtual_key_max_budget_check, _virtual_key_soft_budget_check, get_key_object, get_user_object, @@ -1629,3 +1632,151 @@ async def test_custom_auth_common_checks_opt_in(): parent_otel_span=None, ) mock_common.assert_called_once() + + +# ===================================================================== +# Spend counter budget check tests (v2 — Redis-backed spend counters) +# ===================================================================== + + +@pytest.mark.asyncio +async def test_virtual_key_budget_check_reads_from_spend_counter(): + """Budget check should use get_current_spend when counter exists, + even if cached object shows lower spend.""" + from litellm.proxy.utils import ProxyLogging + + valid_token = UserAPIKeyAuth( + token="test-hashed-token", + spend=0.0, # stale — counter has 1.5 + max_budget=1.0, + user_id="test-user", + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.budget_alerts = AsyncMock() + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:key:test-hashed-token": + return 1.5 + return fallback_spend + + with patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 1.5 + assert exc_info.value.max_budget == 1.0 + + +@pytest.mark.asyncio +async def test_virtual_key_budget_check_fallback_no_counter(): + """When counter doesn't exist, budget check should fall back + to cached object's spend via fallback_spend.""" + from litellm.proxy.utils import ProxyLogging + + valid_token = UserAPIKeyAuth( + token="test-hashed-token", + spend=15.0, + max_budget=10.0, + user_id="test-user", + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.budget_alerts = AsyncMock() + + # get_current_spend returns fallback_spend when no counter exists + async def mock_get_current_spend(counter_key, fallback_spend): + return fallback_spend + + with patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_max_budget_check( + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 15.0 + + +@pytest.mark.asyncio +async def test_team_budget_check_reads_from_spend_counter(): + """Team budget check should use get_current_spend when counter exists.""" + from litellm.proxy.utils import ProxyLogging + + team_object = LiteLLM_TeamTable( + team_id="test-team", + spend=0.0, # stale + max_budget=1.0, + ) + valid_token = UserAPIKeyAuth(token="test-token", team_id="test-team") + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.budget_alerts = AsyncMock() + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:team:test-team": + return 1.5 + return fallback_spend + + with patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _team_max_budget_check( + team_object=team_object, + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 1.5 + + +@pytest.mark.asyncio +async def test_team_member_budget_check_reads_from_spend_counter(): + """Team member budget check should use get_current_spend when counter exists.""" + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + team_object = LiteLLM_TeamTable(team_id="test-team") + user_object = LiteLLM_UserTable(user_id="test-user") + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user", + team_id="test-team", + ) + + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, # stale + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:team_member:test-user:test-team": + return 1.5 + return fallback_spend + + with patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 1.5 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index bd6162f225a..11435fd7520 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4352,3 +4352,183 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch): # add_deployment should NOT have been called since store_model_in_db is False mock_proxy_config.add_deployment.assert_not_called() + + +# ===================================================================== +# Spend counter tests (v2 — Redis-backed spend counters) +# ===================================================================== + + +@pytest.mark.asyncio +async def test_get_current_spend_reads_redis_first(): + """get_current_spend should prefer Redis over in-memory.""" + from litellm.caching.dual_cache import DualCache + + counter_cache = DualCache() + + # In-memory has stale value + counter_cache.in_memory_cache.set_cache(key="spend:key:test", value=0.30) + + # Mock Redis with cross-pod authoritative value + mock_redis = AsyncMock() + mock_redis.async_get_cache = AsyncMock(return_value=0.90) + counter_cache.redis_cache = mock_redis + + import litellm.proxy.proxy_server as ps + + original = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + + try: + from litellm.proxy.proxy_server import get_current_spend + + result = await get_current_spend( + counter_key="spend:key:test", + fallback_spend=0.0, + ) + # Should return Redis value (0.90), not in-memory (0.30) + assert result == 0.90 + mock_redis.async_get_cache.assert_called_once_with(key="spend:key:test") + finally: + ps.spend_counter_cache = original + + +@pytest.mark.asyncio +async def test_get_current_spend_fallback_to_in_memory(): + """When Redis is not configured, get_current_spend uses in-memory.""" + from litellm.caching.dual_cache import DualCache + + counter_cache = DualCache() # no redis_cache + counter_cache.in_memory_cache.set_cache(key="spend:key:test", value=0.50) + + import litellm.proxy.proxy_server as ps + + original = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + + try: + from litellm.proxy.proxy_server import get_current_spend + + result = await get_current_spend( + counter_key="spend:key:test", + fallback_spend=0.0, + ) + assert result == 0.50 + finally: + ps.spend_counter_cache = original + + +@pytest.mark.asyncio +async def test_increment_spend_counters_initializes_and_increments(): + """Counter should initialize from cached object spend, then increment. + + Uses a pre-hashed token to match production: metadata["user_api_key"] + is always hashed by the auth flow before reaching the cost callback. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_VerificationTokenView, hash_token + + key_cache = DualCache() + counter_cache = DualCache() + + # In production, the auth flow hashes the raw key before it reaches + # the cost callback. Simulate that by passing the hashed token. + hashed_token = hash_token("sk-test-token-for-counter") + + # Simulate a cached key object with existing spend from DB + cached_key = LiteLLM_VerificationTokenView( + token=hashed_token, + spend=5.0, + max_budget=10.0, + ) + key_cache.in_memory_cache.set_cache(key=hashed_token, value=cached_key) + + import litellm.proxy.proxy_server as ps + + original_key_cache = ps.user_api_key_cache + original_counter_cache = ps.spend_counter_cache + ps.user_api_key_cache = key_cache + ps.spend_counter_cache = counter_cache + + try: + from litellm.proxy.proxy_server import increment_spend_counters + + # Pass pre-hashed token (as the cost callback would in production) + await increment_spend_counters( + token=hashed_token, + team_id=None, + user_id=None, + response_cost=0.50, + ) + + # Counter should be: base(5.0) + increment(0.50) = 5.50 + counter = counter_cache.in_memory_cache.get_cache( + key=f"spend:key:{hashed_token}" + ) + assert counter == 5.50 + + # Second increment — counter already exists, just increment + await increment_spend_counters( + token=hashed_token, + team_id=None, + user_id=None, + response_cost=0.25, + ) + + counter = counter_cache.in_memory_cache.get_cache( + key=f"spend:key:{hashed_token}" + ) + assert counter == 5.75 + finally: + ps.user_api_key_cache = original_key_cache + ps.spend_counter_cache = original_counter_cache + + +@pytest.mark.asyncio +async def test_increment_spend_counters_team_and_member(): + """Counter should track team and team member spend separately.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamTable + + key_cache = DualCache() + counter_cache = DualCache() + + # Cached team object + team_obj = LiteLLM_TeamTable(team_id="team-1", spend=2.0) + key_cache.in_memory_cache.set_cache(key="team_id:team-1", value=team_obj) + + # Cached team membership + key_cache.in_memory_cache.set_cache( + key="team_membership:user-1:team-1", + value={"user_id": "user-1", "team_id": "team-1", "spend": 1.0}, + ) + + import litellm.proxy.proxy_server as ps + + original_key_cache = ps.user_api_key_cache + original_counter_cache = ps.spend_counter_cache + ps.user_api_key_cache = key_cache + ps.spend_counter_cache = counter_cache + + try: + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id="team-1", + user_id="user-1", + response_cost=0.30, + ) + + team_counter = counter_cache.in_memory_cache.get_cache( + key="spend:team:team-1" + ) + assert team_counter == 2.30 + + member_counter = counter_cache.in_memory_cache.get_cache( + key="spend:team_member:user-1:team-1" + ) + assert member_counter == 1.30 + finally: + ps.user_api_key_cache = original_key_cache + ps.spend_counter_cache = original_counter_cache