fix(proxy): size a Gemini request's budget hold from its prompt and maxOutputTokens

A Gemini generateContent body's budget hold now counts its contents and
systemInstruction and reads generationConfig.maxOutputTokens. Before, it held
the model's whole context window, about $0.85 on gemini-3.8-flash, so a key
with a small budget got 422s on back-to-back calls it could afford

Also removes get_litellm_virtual_key, which nothing calls since the Vertex
routes stopped using it

litellm/proxy/pass_through_endpoints/common_utils.py

@@ -1,18 +0,0 @@
-from typing import Final
-
-from fastapi import Request
-
-
-def get_litellm_virtual_key(request: Request) -> str:
-    """
-    Extract and format API key from request headers.
-    Prioritizes x-litellm-api-key over Authorization header.
-
-
-    Vertex JS SDK uses `Authorization` header, we use `x-litellm-api-key` to pass litellm virtual key
-
-    """
-    litellm_api_key: Final = request.headers.get("x-litellm-api-key")
-    if litellm_api_key:
-        return f"Bearer {litellm_api_key}"
-    return request.headers.get("Authorization", "")

litellm/proxy/spend_tracking/budget_reservation.py

@@ -1654,9 +1654,11 @@ _OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_outpu
 
 def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
     inference_config: Final = request_body.get("inferenceConfig")
+    generation_config: Final = request_body.get("generationConfig")
     candidates: Final = (
         *(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
         inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None,
+        generation_config.get("maxOutputTokens") if isinstance(generation_config, Mapping) else None,
     )
     return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None)
 

litellm/proxy/spend_tracking/input_tokens.py

@@ -15,10 +15,13 @@ import asyncio
 import json
 from collections.abc import Mapping, Sequence
 from types import MappingProxyType
-from typing import Final
+from typing import Final, TypeAlias
+
+from pydantic import TypeAdapter
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
 from litellm.responses.litellm_completion_transformation.compaction_items import (
     billable_input_since_newest_compaction,
     holds_compaction_summary,
@@ -35,7 +38,20 @@ from litellm.rust_bridge.token_counter import (
 
 TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS: Final = 30_000
 
-_INPUT_SIZE_FIELDS: Final = ("messages", "prompt", "input", "query", "documents", "tools", "tool_choice")
+_GEMINI_SYSTEM_INSTRUCTION_FIELDS: Final = ("systemInstruction", "system_instruction")
+_GeminiContents: TypeAlias = list[dict[str, object]] | dict[str, object]
+_GEMINI_CONTENTS: Final = TypeAdapter[_GeminiContents](_GeminiContents)
+_INPUT_SIZE_FIELDS: Final = (
+    "messages",
+    "prompt",
+    "input",
+    "query",
+    "documents",
+    "contents",
+    *_GEMINI_SYSTEM_INSTRUCTION_FIELDS,
+    "tools",
+    "tool_choice",
+)
 
 
 def _approximate_input_size(request_body: Mapping[str, object]) -> int:
@@ -48,7 +64,7 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int:
 
 
 async def count_input_tokens(
-    request_body: dict,
+    request_body: dict[str, object],
     raw_body: bytes | None,
     models: Sequence[str],
 ) -> Mapping[str, int]:
@@ -74,7 +90,7 @@ async def count_input_tokens(
 
 
 async def _count_group(
-    request_body: dict,
+    request_body: dict[str, object],
     raw_body: bytes | None,
     tokenizer: RustTokenizer | None,
     models: tuple[str, ...],
@@ -118,7 +134,7 @@ async def _native_counts(
 
 
 def _count_input_tokens_for_models(
-    request_body: dict,
+    request_body: dict[str, object],
     models: Sequence[str],
 ) -> Mapping[str, int]:
     return MappingProxyType(
@@ -142,6 +158,8 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None:
                 )
             except ValueError:
                 return _count_text_tokens(model=model, text=request_body.get("messages"))
+        if "contents" in request_body:
+            return _count_gemini_input_tokens(request_body=request_body, model=model)
         if "prompt" in request_body:
             return _count_text_tokens(model=model, text=request_body.get("prompt"))
         if "input" in request_body:
@@ -160,6 +178,20 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None:
     return None
 
 
+def _count_gemini_input_tokens(request_body: Mapping[str, object], model: str) -> int | None:
+    system_instruction: Final = next(
+        (request_body[field] for field in _GEMINI_SYSTEM_INSTRUCTION_FIELDS if field in request_body), None
+    )
+    translate: Final = GoogleGenAIAdapter().translate_generate_content_to_completion  # pyright: ignore[reportUnknownMemberType]  # adapter's **kwargs is untyped
+    chat_request: Final = translate(
+        model=model,
+        contents=_GEMINI_CONTENTS.validate_python(request_body["contents"]),
+        systemInstruction=system_instruction,
+        tools=request_body.get("tools"),
+    )
+    return count_input_tokens_for_model(request_body=chat_request, model=model)
+
+
 def _count_text_tokens(model: str, text: object) -> int:
     if text is None:
         return 0

tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py

@@ -1,46 +1,6 @@
-import json
-import traceback
-from unittest import mock
-from unittest.mock import MagicMock, patch
-
-import httpx
-import pytest
-from fastapi import Request, Response
-from fastapi.testclient import TestClient
-
 from litellm.passthrough.utils import CommonUtils
 
 
-from unittest.mock import Mock
-
-from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
-
-
-@pytest.mark.asyncio
-async def test_get_litellm_virtual_key():
-    """
-    Test that the get_litellm_virtual_key function correctly handles the API key authentication
-    """
-    # Test with x-litellm-api-key
-    mock_request = Mock()
-    mock_request.headers = {"x-litellm-api-key": "test-key-123"}
-    result = get_litellm_virtual_key(mock_request)
-    assert result == "Bearer test-key-123"
-
-    # Test with Authorization header
-    mock_request.headers = {"Authorization": "Bearer auth-key-456"}
-    result = get_litellm_virtual_key(mock_request)
-    assert result == "Bearer auth-key-456"
-
-    # Test with both headers (x-litellm-api-key should take precedence)
-    mock_request.headers = {
-        "x-litellm-api-key": "test-key-123",
-        "Authorization": "Bearer auth-key-456",
-    }
-    result = get_litellm_virtual_key(mock_request)
-    assert result == "Bearer test-key-123"
-
-
 def test_encode_bedrock_runtime_modelid_arn():
     # Test application-inference-profile ARN
     endpoint = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"

tests/unit/proxy/spend_tracking/test_budget_reservation.py

@@ -176,6 +176,52 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window():
     assert invoke_cost < converse_cost < 2 * invoke_cost
 
 
+GEMINI_PROMPT: Final = "Say hello in five words"
+GEMINI_SYSTEM: Final = "Answer tersely"
+
+
+@pytest.mark.parametrize(
+    "route",
+    (
+        "/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
+        "/gemini/v1beta/models/gemini-3.8-flash:generateContent",
+        "/v1beta/models/gemini-3.8-flash:generateContent",
+    ),
+)
+@pytest.mark.parametrize(
+    ("native_output_config", "chat_output_limit"),
+    (({"generationConfig": {"maxOutputTokens": 5}}, {"max_tokens": 5}), ({}, {})),
+)
+def test_gemini_generate_content_body_reserves_the_prompt_not_the_context_window(
+    route: str, native_output_config: dict[str, object], chat_output_limit: dict[str, object]
+) -> None:
+    native_cost: Final = estimate_request_max_cost(
+        request_body={
+            "systemInstruction": {"parts": [{"text": GEMINI_SYSTEM}]},
+            "contents": [{"role": "user", "parts": [{"text": GEMINI_PROMPT}]}],
+            **native_output_config,
+        },
+        route=route,
+        llm_router=None,
+        input_token_counts={},
+    )
+    chat_cost: Final = estimate_request_max_cost(
+        request_body={
+            "model": "gemini-3.8-flash",
+            "messages": [
+                {"role": "system", "content": GEMINI_SYSTEM},
+                {"role": "user", "content": GEMINI_PROMPT},
+            ],
+            **chat_output_limit,
+        },
+        route="/chat/completions",
+        llm_router=None,
+        input_token_counts={},
+    )
+    assert native_cost is not None
+    assert native_cost == chat_cost
+
+
 def _tiered_deployment(input_cost_per_token: float) -> Deployment:
     return Deployment(
         model_name="tiered-group",