From 3831e66d2bcd4e458ba4a5dd6cfa5095636e4c7f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:32:13 +0000 Subject: [PATCH 1/3] fix(budget_reservation): don't reserve budget on token counting routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/budget_reservation.py | 12 ++++- .../proxy/test_budget_reservation.py | 47 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 58a85171cc7..7b62fd44d09 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -144,6 +144,16 @@ async def _apply_over_budget_reservation_policy( ) +_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset({"/models", "/v1/models", "/utils/token_counter"}) +_UNBILLED_ROUTE_SUFFIXES: Final[tuple[str, ...]] = ("/v1/messages/count_tokens", ":countTokens") + + +def _is_unbilled_route(route: str) -> bool: + """Routes that never emit a cost-tracking callback. Reserving budget for them + is a permanent leak: nothing ever reconciles or releases the reservation.""" + return route in _UNBILLED_ROUTES or route.endswith(_UNBILLED_ROUTE_SUFFIXES) + + async def reserve_budget_for_request( request_body: dict, route: str, @@ -161,7 +171,7 @@ async def reserve_budget_for_request( ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None - if route in {"/models", "/v1/models", "/utils/token_counter"}: + if _is_unbilled_route(route): return None if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 34adb4d2091..7bdc73abf32 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2583,3 +2583,50 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat assert received == [{"content": "hi"}] streaming_logging_obj.async_post_call_streaming_hook.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [ + "/v1/messages/count_tokens", + "/anthropic/v1/messages/count_tokens", + "/v1beta/models/gemini-2.5-pro:countTokens", + "/models/gemini-2.5-pro:countTokens", + ], +) +async def test_token_counting_routes_never_reserve_budget(spend_counter_state, route): + """Token counting is free and never fires a cost callback, so a reservation + there is never reconciled and permanently bricks the key's spend counter.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-count-tokens", + spend=0.0, + max_budget=0.01, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.01, + ): + for _ in range(2): + assert ( + await reserve_budget_for_request( + request_body=_request_body(), + route=route, + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + is None + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-count-tokens") is None + + # a real completion on the same key is still budget enforced + assert await _reserve(valid_token, 0.01, key_cache, proxy_logging_obj) is not None From fb7d06da4b75f04ab6487e8dabb3ac0f0a54407a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:45:53 -0700 Subject: [PATCH 2/3] test(budget_reservation): type the tiny-budget reservation helper --- .../proxy/spend_tracking/test_budget_reservation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 9935dbceb7d..5e88268c283 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -53,7 +53,7 @@ async def test_non_exempt_llm_route_still_reserves_budget(): ANTHROPIC_MESSAGES: Final = [{"role": "user", "content": "hello!!!"}] -COUNT_TOKENS_REQUESTS: Final = ( +COUNT_TOKENS_REQUESTS: Final[tuple[tuple[str, dict[str, object]], ...]] = ( ("/v1/messages/count_tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), ("/v1beta/models/gemini-3.8-flash:countTokens", {"contents": [{"role": "user", "parts": [{"text": "hello!!!"}]}]}), ) @@ -68,7 +68,7 @@ def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache: return cache -async def _reserve_for_tiny_budget_key(route: str, request_body: dict) -> dict | None: +async def _reserve_for_tiny_budget_key(route: str, request_body: dict[str, object]) -> dict[str, object] | None: return await reserve_budget_for_request( request_body=request_body, route=route, @@ -85,7 +85,7 @@ async def _reserve_for_tiny_budget_key(route: str, request_body: dict) -> dict | @pytest.mark.asyncio @pytest.mark.parametrize(("route", "request_body"), COUNT_TOKENS_REQUESTS) async def test_repeated_token_counting_never_touches_a_tiny_budget( - spend_counter_cache: DualCache, route: str, request_body: dict + spend_counter_cache: DualCache, route: str, request_body: dict[str, object] ): counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}" @@ -97,8 +97,10 @@ async def test_repeated_token_counting_never_touches_a_tiny_budget( "/v1/messages", {"model": "claude-sonnet-5", "max_tokens": 16, "messages": ANTHROPIC_MESSAGES} ) assert completion is not None - assert completion["reserved_cost"] > 0 - assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(completion["reserved_cost"]) + reserved_cost: Final = completion["reserved_cost"] + assert isinstance(reserved_cost, float) + assert reserved_cost > 0 + assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reserved_cost) BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6" From 9c980b96d6f32fb2b1568e1e26659e8a8912b37e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:44:04 -0700 Subject: [PATCH 3/3] fix(budget_reservation): exempt vertex and bedrock count-tokens routes from budget reservation --- litellm/proxy/spend_tracking/budget_reservation.py | 10 ++++++++-- .../proxy/spend_tracking/test_budget_reservation.py | 9 +++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 31d7d236657..d985075464f 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -183,11 +183,17 @@ _UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( "/openai/v1/responses/input_tokens", } ) -_UNBILLED_ROUTE_SUFFIXES: Final[tuple[str, ...]] = ("/v1/messages/count_tokens", ":countTokens") +_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) +_TOKEN_COUNTING_ACTION: Final = "countTokens" + + +def _is_token_counting_route(route: str) -> bool: + resource, _, action = route.rsplit("/", 1)[-1].partition(":") + return resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION def _is_unbilled_route(route: str) -> bool: - return route in _UNBILLED_ROUTES or route.endswith(_UNBILLED_ROUTE_SUFFIXES) + return route in _UNBILLED_ROUTES or _is_token_counting_route(route) async def reserve_budget_for_request( diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 5e88268c283..de6c7c2a40a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -17,6 +17,10 @@ TOKEN_COUNTING_ROUTES: Final = ( "/v1/messages/count_tokens", "/v1beta/models/gemini-3.8-flash:countTokens", "/models/gemini-3.8-flash:countTokens", + "/bedrock/v1/messages/count-tokens", + "/bedrock/model/us.anthropic.claude-sonnet-4-6/count-tokens", + "/vertex_ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", + "/vertex-ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", ) @@ -56,6 +60,11 @@ ANTHROPIC_MESSAGES: Final = [{"role": "user", "content": "hello!!!"}] COUNT_TOKENS_REQUESTS: Final[tuple[tuple[str, dict[str, object]], ...]] = ( ("/v1/messages/count_tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), ("/v1beta/models/gemini-3.8-flash:countTokens", {"contents": [{"role": "user", "parts": [{"text": "hello!!!"}]}]}), + ( + "/vertex_ai/v1/projects/p/locations/us-east5/publishers/anthropic/models/count-tokens:rawPredict", + {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}, + ), + ("/bedrock/v1/messages/count-tokens", {"model": "claude-sonnet-5", "messages": ANTHROPIC_MESSAGES}), ) TINY_BUDGET_KEY_TOKEN: Final = "hashed-count-tokens-key"