fix(proxy): test_model skips the response cache, bounds the spend lookup, bills unsaved custom prices and checks spend rows against response_cost

litellm/proxy/health_endpoints/_health_endpoints.py

@@ -122,7 +122,11 @@ from litellm.router_api import deployment_usable_by_team
 from litellm.secret_managers.main import get_secret_bool
 from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields
 from litellm.types.router import Deployment, reject_server_owned_wif_fields
-from litellm.types.utils import AdminApiKeySecretReference, admin_api_key_secret_reference
+from litellm.types.utils import (
+    AdminApiKeySecretReference,
+    CustomPricingLiteLLMParams,
+    admin_api_key_secret_reference,
+)
 
 #### Health ENDPOINTS ####
 
@@ -2641,7 +2645,8 @@ _TEST_MODEL_TAG_PREFIX: Final = "litellm_test_model:"
 _SPEND_LOG_BY_TAG_SQL: Final = """
     SELECT tag, "spend"
     FROM "LiteLLM_SpendLogs" sl, jsonb_array_elements_text(sl."request_tags"::jsonb) AS tag
-    WHERE tag IN (SELECT jsonb_array_elements_text($1::jsonb))
+    WHERE sl."startTime" >= ($2::timestamptz AT TIME ZONE 'UTC')
+      AND tag IN (SELECT jsonb_array_elements_text($1::jsonb))
 """
 
 
@@ -2649,11 +2654,26 @@ def _test_model_tag(call_id: str) -> str:
     return f"{_TEST_MODEL_TAG_PREFIX}{call_id}"
 
 
-async def _spend_rows_once(run_query: _RawQuery, request_ids: frozenset[str]) -> Mapping[str, float]:
+def _custom_pricing_fields(model_info: Mapping[str, object] | None) -> Mapping[str, object]:
+    """The custom prices an unsaved model's ``model_info`` carries, sent on the call the way a saved deployment's
+    prices apply, so the check bills it as the proxy will once it is saved."""
+    return {
+        key: value
+        for key, value in (model_info or {}).items()
+        if key in CustomPricingLiteLLMParams.model_fields and value is not None
+    }
+
+
+async def _spend_rows_once(run_query: _RawQuery, request_ids: frozenset[str], since: datetime) -> Mapping[str, float]:
     """Spend per request id. The row's request_id is the provider's response id, so each turn is found by the
-    tag the check put on it, which the proxy writes to request_tags like any client tag."""
+    tag the check put on it, which the proxy writes to request_tags like any client tag. Only rows started at or
+    after ``since`` are scanned, through the ``startTime`` index, so the check never reads the whole table."""
     rows: Final = _SPEND_LOG_ROWS.validate_python(
-        await run_query(_SPEND_LOG_BY_TAG_SQL, json.dumps(sorted(_test_model_tag(call_id) for call_id in request_ids)))
+        await run_query(
+            _SPEND_LOG_BY_TAG_SQL,
+            json.dumps(sorted(_test_model_tag(call_id) for call_id in request_ids)),
+            since.isoformat(),
+        )
     )
     return MappingProxyType({row.tag.removeprefix(_TEST_MODEL_TAG_PREFIX): row.spend for row in rows})
 
@@ -2662,6 +2682,7 @@ async def _spend_log_rows(
     run_query: _RawQuery,
     request_ids: Sequence[str],
     *,
+    since: datetime,
     wait_seconds: float = TEST_MODEL_SPEND_LOG_WAIT_SECONDS,
     poll_seconds: float = 1.0,
 ) -> Mapping[str, float]:
@@ -2672,7 +2693,9 @@ async def _spend_log_rows(
     """
     deadline: Final = time.monotonic() + wait_seconds
     wanted: Final = frozenset(request_ids)
-    while not wanted <= (found := await _spend_rows_once(run_query, wanted)).keys() and time.monotonic() < deadline:
+    while (
+        not wanted <= (found := await _spend_rows_once(run_query, wanted, since)).keys() and time.monotonic() < deadline
+    ):
         await asyncio.sleep(poll_seconds)
     return found
 
@@ -2695,10 +2718,15 @@ def _test_model_endpoints(
 
 
 async def _logged_cost(logging_obj: LiteLLMLoggingObj, fallback: float | None) -> float | None:
-    """The cost the success handler logged for the call, which lands a moment after the call returns."""
+    """The cost the success handler logged for the call, which lands a moment after the call returns. The wait
+    ends as soon as the handler has built the call's standard logging payload, so an unpriced call does not hold
+    every turn for the full budget."""
     deadline: Final = time.monotonic() + TEST_MODEL_COST_WAIT_SECONDS
-    while (logged := _OBJECT_MAPPING.validate_python(logging_obj.model_call_details).get("response_cost")) is None and (
-        time.monotonic() < deadline
+    while (
+        (logged := (details := _OBJECT_MAPPING.validate_python(logging_obj.model_call_details)).get("response_cost"))
+        is None
+        and "standard_logging_object" not in details
+        and time.monotonic() < deadline
     ):
         await asyncio.sleep(0.05)
     return float(logged) if isinstance(logged, int | float) else fallback
@@ -2714,7 +2742,8 @@ def _turn_caller(
     """The two ways a turn reaches the provider: pinned to a saved deployment through the router, so the
     deployment's own credentials and pricing apply, or straight through litellm with the submitted params.
     Either way the call carries the proxy's own logging object, the way every proxied request does, so the
-    cost read back is the one the spend log row gets."""
+    cost read back is the one the spend log row gets. Every turn skips LiteLLM's response cache, so the reply,
+    and the cache-miss retry of the same prompt, come from the provider and not from a stored answer."""
 
     async def call(script: EndpointScript, body: Mapping[str, object], call_id: str) -> CalledTurn:
         target_params: Final = {"model": pinned_model_id} if pinned_model_id is not None else call_params
@@ -2728,6 +2757,7 @@ def _turn_caller(
                 script.metadata_key: {**proxy_metadata, "tags": [_test_model_tag(call_id)]},
                 "litellm_call_id": call_id,
                 "drop_params": True,
+                "cache": {"no-cache": True},
             },
         )
         prepared["litellm_logging_obj"] = logging_obj
@@ -2785,7 +2815,9 @@ async def test_model(
     answer from turn 2. Three checks come out of it:
 
     - `conversation`: every turn parses, turn 1 calls the tool, turn 2 uses the tool result, turn 3 carries it on
-    - `prompt_caching`: turns 2 and 3 report cached input tokens, with no cache markers in the request
+    - `prompt_caching`: turns 2 and 3 report cached input tokens and turn 3 reads at least as many as turn 2.
+      Chat Completions and Responses send no cache markers, so the provider has to cache the repeated prefix on
+      its own. Messages sends `cache_control` on the latest user message the way Claude Code does
     - `cost_tracking`: every turn carries a non-zero response cost, and the spend log row lands when a
       database is connected
 
@@ -2852,12 +2884,17 @@ async def test_model(
                 status_code=status.HTTP_403_FORBIDDEN, detail={"error": "Only proxy admins can test models."}
             )
 
-        call_params: Final = {key: value for key, value in resolved.litellm_params.items() if key != "mode"}
+        call_params: Final = {
+            **_custom_pricing_fields(resolved.resolved_model_info),
+            **{key: value for key, value in resolved.litellm_params.items() if key != "mode"},
+        }
         mode: Final = _string_mode_or_bad_request(resolved.litellm_params.get("mode")) or resolve_health_check_mode(
             resolved.probe_model_info, _OBJECT_MAPPING.validate_python(call_params)
         )
         selected: Final = _test_model_endpoints(endpoints, mode)
-        overrides_saved_deployment: Final = any(key != "model" for key in resolved.request_litellm_params)
+        overrides_saved_deployment: Final = any(
+            key != "model" or value != resolved.stored_model for key, value in resolved.request_litellm_params.items()
+        )
         pinned_model_id: Final = (
             resolved.deployment_by_id.model_info.id
             if resolved.deployment_by_id is not None and not overrides_saved_deployment
@@ -2880,9 +2917,10 @@ async def test_model(
             proxy_metadata=proxy_metadata,
         )
         run_query: Final[_RawQuery | None] = prisma_client.db.query_raw if prisma_client is not None else None
+        started_at: Final = datetime.now(timezone.utc) - timedelta(minutes=1)
 
         async def spend_lookup(request_ids: Sequence[str]) -> Mapping[str, float]:
-            return await _spend_log_rows(run_query, request_ids) if run_query is not None else {}
+            return await _spend_log_rows(run_query, request_ids, since=started_at) if run_query is not None else {}
 
         results: Final[Sequence[EndpointResult]] = await asyncio.gather(
             *(

litellm/proxy/health_endpoints/model_checks.py

@@ -16,6 +16,7 @@ tool result is used, the follow-up is answered from the earlier turn), prompt ca
 from __future__ import annotations
 
 import json
+import math
 import time
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass
@@ -316,7 +317,10 @@ def _first_choice_message(plain: Mapping[str, object]) -> object:
         choices: Final = _CHOICES.validate_python(plain.get("choices"))
     except ValidationError:
         return dict(_EMPTY_ASSISTANT)
-    return choices[0].get("message", dict(_EMPTY_ASSISTANT)) if choices else dict(_EMPTY_ASSISTANT)
+    message: Final = choices[0].get("message") if choices else None
+    if not isinstance(message, dict):
+        return dict(_EMPTY_ASSISTANT)
+    return {**_EMPTY_ASSISTANT, **_PLAIN_OBJECT.validate_python(message)}
 
 
 def _parse_chat(response: object) -> ParsedTurn | str:
@@ -595,6 +599,12 @@ def _cost_verdict(turn_number: int, cost: float | None) -> CheckResult | None:
     return None
 
 
+def _spend_disagreement(turn: TurnOutcome, logged: float) -> str:
+    if turn.cost is None or math.isclose(logged, turn.cost, rel_tol=1e-6, abs_tol=1e-12):
+        return ""
+    return f"request_id {turn.call_id} row ${logged:.6f} vs response_cost ${turn.cost:.6f}"
+
+
 def _cost_tracking_check(turns: Sequence[TurnOutcome], spend_rows: Mapping[str, float] | None) -> CheckResult:
     if not turns:
         return CheckResult("cost_tracking", "fail", "not checked, no turn returned a response")
@@ -615,6 +625,14 @@ def _cost_tracking_check(turns: Sequence[TurnOutcome], spend_rows: Mapping[str,
             "fail",
             f"{priced}, but no LiteLLM_SpendLogs row landed for request_id {', '.join(missing)}",
         )
+    disagreeing: Final = tuple(_spend_disagreement(turn, spend_rows[turn.call_id]) for turn in turns)
+    if any(disagreeing):
+        return CheckResult(
+            "cost_tracking",
+            "fail",
+            f"{priced}, but the spend log row disagrees with the response cost ("
+            f"{', '.join(note for note in disagreeing if note)})",
+        )
     logged: Final = ", ".join(f"${spend_rows[turn.call_id]:.6f}" for turn in turns)
     return CheckResult("cost_tracking", "pass", f"{priced}, spend log rows written ({logged})")
 

tests/integration/management/test_model_health_check.py

@@ -1581,4 +1581,67 @@ def test_health_test_model_names_the_check_that_failed_for_an_unsaved_model_that
         ), response.text
         assert checks["cost_tracking"][0] == "fail", response.text
         assert "no price resolved for this model" in checks["cost_tracking"][1], response.text
-        assert len(wire.drain()) >= 2, response.text
+        assert len(wire.drain()) == 5, response.text
+
+
+@pytest.mark.timeout(240)
+def test_health_test_model_bills_an_unsaved_model_at_the_custom_prices_in_its_model_info(gateway: Gateway) -> None:
+    with wire_server(lambda request: _test_model_chat_reply(request, cached_tokens=4800)) as wire:
+        response: Final = gateway.request(
+            "POST",
+            "/health/test_model",
+            {
+                "litellm_params": {
+                    "model": f"openai/priced-{uuid.uuid4().hex}",
+                    "api_base": f"{wire.url}/v1",
+                    "api_key": "synthetic-priced-key",
+                },
+                "model_info": {"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-6},
+                "endpoints": ["/v1/chat/completions"],
+            },
+        )
+
+        assert response.status_code == 200, response.text
+        payload: Final = JSON_OBJECT.validate_json(response.content)
+        assert payload["status"] == "success", response.text
+        checks: Final = _check_statuses(object_value(list_value(payload["endpoints"])[0]))
+        assert checks["cost_tracking"][0] == "pass", response.text
+        assert checks["cost_tracking"][1].startswith("turn 1 $0.005024, "), response.text
+        assert "spend log rows written" in checks["cost_tracking"][1], response.text
+        assert len(wire.drain()) == 3, response.text
+
+
+@pytest.mark.timeout(240)
+def test_health_test_model_with_a_different_model_than_the_saved_deployment_calls_that_model(gateway: Gateway) -> None:
+    with (
+        gateway.scenario() as scenario,
+        wire_server(lambda request: _test_model_chat_reply(request, cached_tokens=4800)) as wire,
+    ):
+        model_name: Final = scenario.model(api_base=f"{wire.url}/v1", api_key="synthetic-override-key")
+        model_id: Final = next(
+            string_value(object_value(info["model_info"])["id"])
+            for info in (object_value(entry) for entry in list_value(gateway.get("/v1/model/info")["data"]))
+            if info["model_name"] == model_name
+        )
+        override: Final = f"openai/override-{uuid.uuid4().hex}"
+
+        response: Final = gateway.request(
+            "POST",
+            "/health/test_model",
+            {"model_id": model_id, "litellm_params": {"model": override}, "endpoints": ["/v1/chat/completions"]},
+        )
+
+        assert response.status_code == 200, response.text
+        payload: Final = JSON_OBJECT.validate_json(response.content)
+        assert payload["model"] == override, response.text
+        assert payload["model_id"] is None, response.text
+        assert _check_statuses(object_value(list_value(payload["endpoints"])[0]))["conversation"][0] == "pass", (
+            response.text
+        )
+        requests: Final = wire.drain()
+        assert [JSON_OBJECT.validate_json(request.body)["model"] for request in requests] == [
+            override.removeprefix("openai/")
+        ] * 3, response.text
+        assert [request.headers["authorization"] for request in requests] == ["Bearer synthetic-override-key"] * 3, (
+            response.text
+        )

tests/unit/proxy/health_endpoints/test_model_checks.py

@@ -1,6 +1,6 @@
 import asyncio
 import json
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from typing import Final
 
 import pytest
@@ -11,6 +11,7 @@ from litellm.proxy.health_endpoints.model_checks import (
     RESPONSES_SCRIPT,
     WEATHER_TOOL_RESULT,
     CalledTurn,
+    EndpointResult,
     EndpointScript,
     endpoint_result_payload,
     overall_status,
@@ -150,7 +151,9 @@ def _ids() -> Sequence[str]:
     return ("call-1", "call-2", "call-3", "call-4", "call-5")
 
 
-def _run(script: EndpointScript, provider: _ScriptedProvider, spend_rows: Mapping[str, float] | None = None):
+def _run(
+    script: EndpointScript, provider: _ScriptedProvider, spend_rows: Mapping[str, float] | None = None
+) -> EndpointResult:
     ids: Final = iter(_ids())
 
     async def spend_lookup(request_ids: Sequence[str]) -> Mapping[str, float]:
@@ -166,7 +169,7 @@ def _run(script: EndpointScript, provider: _ScriptedProvider, spend_rows: Mappin
     )
 
 
-def _checks(result) -> Mapping[str, tuple[str, str]]:
+def _checks(result: EndpointResult) -> Mapping[str, tuple[str, str]]:
     return {check.name: (check.status, check.detail) for check in result.checks}
 
 
@@ -179,7 +182,9 @@ def _checks(result) -> Mapping[str, tuple[str, str]]:
     ],
     ids=["chat", "responses", "messages"],
 )
-def test_every_endpoint_passes_all_three_checks_and_feeds_each_reply_back_in_its_own_shape(script, reply, key):
+def test_every_endpoint_passes_all_three_checks_and_feeds_each_reply_back_in_its_own_shape(
+    script: EndpointScript, reply: Callable[..., object], key: str
+) -> None:
     provider: Final = _ScriptedProvider(
         CalledTurn(reply(tool_call=True), 0.004),
         [
@@ -229,6 +234,7 @@ def test_every_endpoint_passes_all_three_checks_and_feeds_each_reply_back_in_its
         tool_message: Final = turn_two["messages"][-1]
         assert tool_message == {"role": "tool", "tool_call_id": _TOOL_CALL_ID, "content": WEATHER_TOOL_RESULT}
         assert turn_two["messages"][-2]["tool_calls"][0]["id"] == _TOOL_CALL_ID
+        assert "content" in turn_two["messages"][-2] and turn_two["messages"][-2]["content"] is None, turn_two
         assert turn_three["messages"][-2]["role"] == "system", turn_three
         assert turn_three["messages"][-3]["role"] == "assistant", turn_three
         assert turn_three["messages"][-3]["content"] == _TURN_TWO_ANSWER, turn_three
@@ -465,6 +471,26 @@ def test_a_missing_spend_log_row_fails_cost_tracking_even_when_the_response_was_
     )
 
 
+def test_a_spend_log_row_that_disagrees_with_the_response_cost_fails_cost_tracking() -> None:
+    provider: Final = _ScriptedProvider(
+        CalledTurn(_chat_reply(tool_call=True), 0.004),
+        [
+            CalledTurn(_chat_reply(tool_call=False, text="72F", cached_tokens=4800), 0.0005),
+            CalledTurn(_chat_reply(tool_call=False, text="72F", cached_tokens=4800), 0.0004),
+        ],
+    )
+
+    checks: Final = _checks(
+        _run(CHAT_COMPLETIONS_SCRIPT, provider, spend_rows={"call-1": 0.004, "call-2": 0.0, "call-3": 0.0004})
+    )
+
+    assert checks["cost_tracking"] == (
+        "fail",
+        "turn 1 $0.004000, turn 2 $0.000500, turn 3 $0.000400, but the spend log row disagrees with the response "
+        "cost (request_id call-2 row $0.000000 vs response_cost $0.000500)",
+    )
+
+
 def test_a_response_that_is_not_the_endpoint_shape_fails_the_conversation_with_the_parse_error():
     provider: Final = _ScriptedProvider(CalledTurn({"id": "x", "object": "nonsense"}, 0.001), [])