Merge pull request #36718 from BerriAI/litellm_fix_count_tokens_budget_reservation_leak

fix(budget_reservation): don't reserve budget on token counting routes

litellm/proxy/spend_tracking/budget_reservation.py

@@ -173,6 +173,29 @@ def _raise_counter_budget_exceeded(
     )
 
 
+_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset(
+    {
+        "/models",
+        "/v1/models",
+        "/utils/token_counter",
+        "/responses/input_tokens",
+        "/v1/responses/input_tokens",
+        "/openai/v1/responses/input_tokens",
+    }
+)
+_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 _is_token_counting_route(route)
+
+
 async def reserve_budget_for_request(
     request_body: dict,
     route: str,
@@ -190,14 +213,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",
-        "/responses/input_tokens",
-        "/v1/responses/input_tokens",
-        "/openai/v1/responses/input_tokens",
-    }:
+    if _is_unbilled_route(route):
         return None
     if get_model_from_request(request_body, route, llm_router=llm_router) is None:
         return None

tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py

@@ -2,6 +2,7 @@ from typing import Final
 
 import pytest
 
+import litellm.proxy.proxy_server as proxy_server
 from litellm.caching import DualCache
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -13,6 +14,13 @@ TOKEN_COUNTING_ROUTES: Final = (
     "/v1/responses/input_tokens",
     "/openai/v1/responses/input_tokens",
     "/utils/token_counter",
+    "/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",
 )
 
 
@@ -48,6 +56,62 @@ async def test_non_exempt_llm_route_still_reserves_budget():
     assert reservation["reserved_cost"] > 0
 
 
+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"
+
+
+@pytest.fixture
+def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache:
+    cache: Final = DualCache()
+    monkeypatch.setattr(proxy_server, "spend_counter_cache", cache)
+    monkeypatch.setattr(proxy_server, "prisma_client", None)
+    return cache
+
+
+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,
+        llm_router=None,
+        valid_token=UserAPIKeyAuth(token=TINY_BUDGET_KEY_TOKEN, max_budget=0.01, spend=0.0),
+        team_object=None,
+        user_object=None,
+        prisma_client=None,
+        user_api_key_cache=UserApiKeyCache(),
+        proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
+    )
+
+
+@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[str, object]
+):
+    counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
+
+    assert await _reserve_for_tiny_budget_key(route, request_body) is None
+    assert await _reserve_for_tiny_budget_key(route, request_body) is None
+    assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is None
+
+    completion: Final = await _reserve_for_tiny_budget_key(
+        "/v1/messages", {"model": "claude-sonnet-5", "max_tokens": 16, "messages": ANTHROPIC_MESSAGES}
+    )
+    assert completion is not None
+    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"
 CONVERSE_BODY: Final = {
     "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}],