From 7aafa3e4bd888e1f3502d3e7256958b0e8202db5 Mon Sep 17 00:00:00 2001 From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:34:11 -0500 Subject: [PATCH] fix(proxy): harden team token budget enforcement Floor get_current_tokens_used at the recorded tokens_used column so a counter that reseeded low after a TTL expiry, or missed an increment when its logging payload failed, cannot read below real usage and let a team slip past the cap. This mirrors how get_current_spend reconciles a stale-low counter against recorded spend Also advance the team token counter for successful unpriced responses (response_cost None) that skip the spend pipeline, so an unpriced model no longer bypasses max_tokens_budget --- .../proxy/hooks/proxy_track_cost_callback.py | 38 ++++++++++----- litellm/proxy/proxy_server.py | 18 ++++--- .../unit/proxy/auth/test_team_token_budget.py | 47 +++++++++++++++++++ 3 files changed, 85 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 4ebc44630ed..33fe9e4d54b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -380,6 +380,16 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ) + # An unpriced model has no dollar cost and skips the spend pipeline below, but it + # still consumes tokens a team token budget must count, so advance the counter here. + if ( + response_cost is None + and team_id is not None + and ProxyUpdateSpend.disable_spend_updates() is not True + and not _is_unbilled_interaction_response(completion_response) + ): + await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object)) + if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) verbose_proxy_logger.debug( @@ -746,6 +756,22 @@ def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int: ) +async def _increment_team_token_counter(team_id: str, total_tokens: int) -> None: + """Advance the cross-pod ``tokens:team:{id}`` counter so the next request's + auth check sees this request's tokens before the batched ``tokens_used`` + flush lands. Best-effort, a counter hiccup must never fail a request that was + already served. + """ + if not total_tokens: + return + from litellm.proxy.proxy_server import increment_spend_counter_cache + + try: + await increment_spend_counter_cache(counter_key=f"tokens:team:{team_id}", increment=total_tokens) + except Exception: # noqa: BLE001 # best-effort cross-pod token counter; a Redis failure must not break the response or the spend write + verbose_proxy_logger.warning("Failed to increment team token-budget counter for team_id=%s", team_id) + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -898,17 +924,7 @@ async def _update_database_and_spend_counters_in_batch( "StandardLoggingPayload | None", kwargs.get("standard_logging_object"), # pyright: ignore[reportUnknownMemberType] # untyped callback kwargs ) - total_tokens: Final = _extract_total_tokens(sl_object) - if total_tokens: - from litellm.proxy.proxy_server import increment_spend_counter_cache - - try: - await increment_spend_counter_cache( - counter_key=f"tokens:team:{team_id}", - increment=total_tokens, - ) - except Exception: # noqa: BLE001 # best-effort cross-pod token counter; a Redis failure must not break the response or the spend write - verbose_proxy_logger.warning("Failed to increment team token-budget counter for team_id=%s", team_id) + await _increment_team_token_counter(team_id=team_id, total_tokens=_extract_total_tokens(sl_object)) return True diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 24fba4a5635..84ddc3d77b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3026,16 +3026,20 @@ async def get_current_tokens_used(counter_key: str, fallback_tokens: int) -> int Reads the Redis-backed ``tokens:team:{id}`` counter first (authoritative across pods), the same way :func:`get_current_spend` reads the spend counter, and falls back to the ``tokens_used`` column loaded in auth when - the counter is cold or Redis is unreachable. The counter is incremented - post-call in the cost callback and invalidated on the budget-reset tick, so - a read after a reset re-seeds from the freshly-zeroed column. Token budgets - tolerate the spend path's fallback-to-DB behavior for the same reason spend - does: the column is the authoritative value the reset job maintains. + the counter is cold or Redis is unreachable. The counter read is floored at + that column. An idle counter expires after its TTL and the next post-call + increment re-seeds it from only that request's delta, so a bare counter read + can sit below the recorded usage and let a team keep spending past the cap. + Flooring at the authoritative column closes that gap the same way + ``get_current_spend`` reconciles a stale-low counter against recorded spend, + and also covers a counter that missed an increment when its logging payload + failed to build. """ value, _ = await read_spend_counter_cache_value(counter_key=counter_key) + floor: Final = int(fallback_tokens or 0) if value is not None: - return int(value) - return int(fallback_tokens or 0) + return max(int(value), floor) + return floor @with_service_target(SPEND_COUNTERS_TARGET) diff --git a/tests/unit/proxy/auth/test_team_token_budget.py b/tests/unit/proxy/auth/test_team_token_budget.py index 361b42adebd..30080a31ee2 100644 --- a/tests/unit/proxy/auth/test_team_token_budget.py +++ b/tests/unit/proxy/auth/test_team_token_budget.py @@ -138,3 +138,50 @@ def test_extract_total_tokens_is_zero_without_a_payload(): from litellm.proxy.hooks.proxy_track_cost_callback import _extract_total_tokens assert _extract_total_tokens(None) == 0 + + +@pytest.mark.asyncio +async def test_get_current_tokens_used_floors_at_recorded_db_usage(): + """A counter that reseeded low after a TTL expiry, or missed an increment when its + logging payload failed, must not read below the tokens_used column loaded in auth, + otherwise a team slips past the cap once the counter sits under recorded usage.""" + from litellm.proxy.proxy_server import get_current_tokens_used + + async def _stale_low(counter_key): + return 200, None + + with patch("litellm.proxy.proxy_server.read_spend_counter_cache_value", _stale_low): + used = await get_current_tokens_used(counter_key="tokens:team:t1", fallback_tokens=1_100) + assert used == 1_100 + + +@pytest.mark.asyncio +async def test_unpriced_model_advances_the_team_token_counter(): # test-quality-ok: no in-process counter state + """An unpriced model returns response_cost=None and skips the spend pipeline, but its + tokens must still count against the team, otherwise it bypasses max_tokens_budget. The + cross-pod counter is Redis-backed, so the issued increment is the only in-process signal.""" + from datetime import datetime + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + + kwargs = { + "model": "custom-unpriced", + "call_type": "acompletion", + "litellm_params": {"metadata": {"user_api_key_team_id": "t1"}}, + "standard_logging_object": { + "response_cost": None, + "response_cost_failure_debug_info": "missing custom price", + "request_tags": None, + "total_tokens": 321, + }, + "stream": False, + } + + with patch("litellm.proxy.proxy_server.increment_spend_counter_cache", new_callable=AsyncMock) as inc: + await ProxyDBLogger()._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), + ) + inc.assert_awaited_once_with(counter_key="tokens:team:t1", increment=321)