feat(mock): report admission-time input token count in mock_response usage (#40590)

* feat(mock): report admission-time input token count in mock_response usage

Mock completions always reported prompt_tokens=10, so spend tracking, TPM metrics, budgets and the tokens-per-minute autoscaling signal saw 10 tokens for a 100k-token request. Budget reservation now carries the admission-time input token count in the reservation record, and mock_completion reads it back so mock traffic exercises the same spend and TPM paths as real traffic without any extra tokenizer work.

* fix(mock): keep a zero admission input token count instead of falling back to 10

---------

Co-authored-by: yassin <yassin@berri.ai>

litellm/main.py

@@ -827,6 +827,29 @@ async def _sleep_for_timeout_async(timeout: float | str | httpx.Timeout):
         await asyncio.sleep(timeout.connect)
 
 
+class _AdmissionReservation(BaseModel):
+    input_tokens: int | None = None
+
+
+class _AdmissionMetadata(BaseModel):
+    user_api_key_budget_reservation: _AdmissionReservation | None = None
+
+
+def admission_input_tokens(kwargs: Mapping[str, object]) -> int | None:
+    reservations: Final = (
+        _AdmissionMetadata.model_validate(kwargs.get(key) or {}).user_api_key_budget_reservation
+        for key in ("litellm_metadata", "metadata")
+    )
+    return next(
+        (
+            reservation.input_tokens
+            for reservation in reservations
+            if reservation and reservation.input_tokens is not None
+        ),
+        None,
+    )
+
+
 def mock_completion(
     model: str,
     messages: list,
@@ -838,6 +861,7 @@ def mock_completion(
     logging=None,
     custom_llm_provider=None,
     timeout: float | str | httpx.Timeout | None = None,
+    prompt_tokens: int | None = None,
     **kwargs,
 ):
     """
@@ -953,13 +977,16 @@ def mock_completion(
                 ChatCompletionMessageToolCall(**tool_call) for tool_call in mock_tool_calls
             ]
 
+        usage_prompt_tokens: Final = (
+            prompt_tokens if prompt_tokens is not None else DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT
+        )
         setattr(
             model_response,
             "usage",
             Usage(
-                prompt_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
+                prompt_tokens=usage_prompt_tokens,
                 completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
-                total_tokens=DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
+                total_tokens=usage_prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
             ),
         )
 
@@ -5550,6 +5577,9 @@ def completion(
                 custom_llm_provider=custom_llm_provider,
                 mock_timeout=mock_timeout,
                 timeout=timeout,
+                prompt_tokens=admission_input_tokens(
+                    cast(Mapping[str, object], kwargs)  # cast-ok: completion's **kwargs is untyped
+                ),
             )
 
         ## RESPONSES API BRIDGE LOGIC ## - check if model has 'mode: responses' in litellm.model_cost map

litellm/proxy/spend_tracking/budget_reservation.py

@@ -315,6 +315,7 @@ async def reserve_budget_for_request(
         "entries": applied_entries,
         "finalized": False,
         "input_cost": min(float(input_cost or 0.0), reservation_cost),
+        "input_tokens": max(input_token_counts.values(), default=None),
     }
 
 

tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py

@@ -2,6 +2,7 @@ from typing import Final
 
 import pytest
 
+import litellm
 import litellm.proxy.proxy_server as proxy_server
 from litellm.caching import DualCache
 from litellm.proxy._types import UserAPIKeyAuth
@@ -56,6 +57,16 @@ async def test_non_exempt_llm_route_still_reserves_budget():
     assert reservation["reserved_cost"] > 0
 
 
+@pytest.mark.asyncio
+async def test_reservation_carries_the_admission_input_token_count():
+    reservation: Final = await _reserve("/v1/responses")
+    expected: Final = litellm.token_counter(model="gpt-4o", text="hello")
+
+    assert reservation is not None
+    assert expected > 0
+    assert reservation["input_tokens"] == expected
+
+
 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}),

tests/test_litellm/test_main.py

@@ -2394,6 +2394,33 @@ def test_image_edit_merges_headers_and_extra_headers():
     assert "extra_headers" not in handler_kwargs["image_edit_optional_request_params"]
 
 
+@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata"))
+@pytest.mark.parametrize("input_tokens", (51234, 0))
+def test_mock_completion_usage_reports_admission_input_tokens(metadata_key: str, input_tokens: int):
+    response = litellm.completion(
+        model="anthropic/claude-sonnet-5",
+        messages=[{"role": "user", "content": "hello"}],
+        mock_response="ok",
+        api_key="mock",
+        **{metadata_key: {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}},
+    )
+
+    assert response.usage.prompt_tokens == input_tokens
+    assert response.usage.total_tokens == input_tokens + response.usage.completion_tokens
+
+
+def test_mock_completion_usage_falls_back_to_default_without_admission_count():
+    response = litellm.completion(
+        model="anthropic/claude-sonnet-5",
+        messages=[{"role": "user", "content": "hello"}],
+        mock_response="ok",
+        api_key="mock",
+        metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}},
+    )
+
+    assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT
+
+
 def test_mock_completion_stream_with_model_response():
     """Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response."""
     from litellm import completion