feat(proxy): lift every earlier billed timed-out retry attempt onto the failed request (#47490)

* feat(proxy): lift every earlier billed timed-out retry attempt onto the failed request

_billed_earlier_attempts reads the router's previous_models records, keeps the Timeout ones other than the final attempt (matched on deployment_id plus attempted_retries, since the count restarts after a fallback hop and the fallback handler re-logs the exhausting attempt) and prices each at the deployment that served it. post_call_failure_hook lifts the tuple as billed_failed_attempts; nothing writes it yet. previous_models joins the router-reserved metadata fields so a client-planted record never reaches it

* fix(proxy): price earlier timed-out attempts from the prompt estimate only

The final row may lift a mid-stream recovered usage (prompt plus delivered chunks) or a failed_attempts_usage sum; earlier attempts never produced output, so they are priced from _estimated_prompt_usage, and when failed_attempts_usage already bills them into the final row none is lifted again

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-20 09:08:11 -07:00 • committed by GitHub
parent 047cb4a8e2
commit 7a4a255d1d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 282 additions and 15 deletions

View file

@ -467,6 +467,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{
"attempted_fallbacks",
"original_model_group",
"previous_models",
"request_retry_count",
CLIENT_OUTPUT_CEILING_METADATA_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,

View file

@ -1031,9 +1031,7 @@ def _failure_usage_to_lift(
at the rate the router registered under ``deployment_id`` (``Logging.get_router_model_id``)
when the deployment declares one, else the cost map's rate for ``base_model`` (an Azure
deployment's ``model_info.base_model``) or the deployment model; any other post-dispatch
failure stays at zero cost. The raw request body backfills
the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a
chat-completions provider) without filling optional_params. Returns the
failure stays at zero cost. Returns the
(combined_usage_object, response_cost) pair to lift, or None."""
recovered_usage: Final = model_call_details.get("combined_usage_object")
failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ()))
@ -1047,6 +1045,22 @@ def _failure_usage_to_lift(
# The proxy-only error handler stamps this flag on a client disconnect too, even though the
# request had already reached the provider, so a billed failure keeps its estimate.
return None
estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body)
if estimated_usage is None:
return None
if not provider_billed:
return estimated_usage, 0.0
return estimated_usage, _estimated_input_cost(
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
)
def _estimated_prompt_usage(
model_call_details: Mapping[str, object], request_body: Mapping[str, object]
) -> Usage | None:
"""The input-only Usage estimate for the prompt this request sent, or None when the call type has no
estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint
(e.g. /v1/messages on a chat-completions provider) without filling optional_params."""
if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES:
return None
optional_params: Final = model_call_details.get("optional_params")
@ -1056,17 +1070,118 @@ def _failure_usage_to_lift(
else None
)
system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions")
estimated_usage: Final = _estimate_dispatched_failure_usage(
return _estimate_dispatched_failure_usage(
model=str(model_call_details.get("model") or ""),
request_input=model_call_details.get("messages"),
system_input=system_input,
)
if estimated_usage is None:
return None
if not provider_billed:
return estimated_usage, 0.0
return estimated_usage, _estimated_input_cost(
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
@dataclass(frozen=True)
class BilledFailedAttempt:
"""An earlier retry attempt the provider also billed, written as its own spend row."""
attempted_retries: int
deployment_id: str | None
usage: Usage
response_cost: float
_NO_RETRY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
def _retry_metadata_bucket(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]:
"""The request metadata bucket the router writes retry state into, ``litellm_metadata`` when the
route carries one (/messages, /responses) and ``metadata`` otherwise, same order as the deployment id."""
if litellm_params is None:
return _NO_RETRY_METADATA
for metadata_key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(metadata_key)
if isinstance(metadata, dict):
return cast("Mapping[str, object]", metadata) # cast-ok: untyped router dict, keys read with .get
return _NO_RETRY_METADATA
def _attempt_number(record: Mapping[str, object], index: int) -> int:
"""The attempt a retry record describes: its own ``attempted_retries`` count, else its position."""
recorded: Final = record.get("attempted_retries")
return recorded if isinstance(recorded, int) else index
def _retry_attempt_records(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]:
"""The router's ``RetryAttemptRecord`` breadcrumbs, one per failed attempt, read from the same bucket."""
if litellm_params is None:
return ()
for metadata_key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(metadata_key)
if not isinstance(metadata, dict):
continue
records = cast("Mapping[str, object]", metadata).get("previous_models") # cast-ok: untyped router dict
if isinstance(records, (list, tuple)):
return tuple(
cast("Mapping[str, object]", record) # cast-ok: RetryAttemptRecord TypedDict, read with .get
for record in cast("Sequence[object]", records) # cast-ok: untyped router tuple
if isinstance(record, dict)
)
return ()
def _record_deployment_id(record: Mapping[str, object]) -> str | None:
deployment_id: Final = record.get("deployment_id")
return deployment_id if isinstance(deployment_id, str) else None
def _attempt_identity(record: Mapping[str, object]) -> tuple[str | None, int | None]:
"""What tells one provider call from another across the request: the deployment that served it
and the retry count within its model group. Two records with the same identity describe the
same call (the fallback handler re-logs the attempt that exhausted the previous group)."""
attempted_retries: Final = record.get("attempted_retries")
return _record_deployment_id(record), attempted_retries if isinstance(attempted_retries, int) else None
def _timed_out_earlier_records(
litellm_params: Mapping[str, object] | None, final_identity: tuple[str | None, int | None]
) -> tuple[Mapping[str, object], ...]:
"""The ``Timeout`` breadcrumbs for calls other than the final attempt, one per provider call."""
by_identity: Final = {
_attempt_identity(record): record
for record in _retry_attempt_records(litellm_params)
if record.get("exception_type") == litellm.Timeout.__name__
}
return tuple(record for identity, record in by_identity.items() if identity != final_identity)
def _billed_earlier_attempts(
model_call_details: Mapping[str, object],
litellm_params: Mapping[str, object] | None,
estimated_usage: Usage,
base_model: str | None,
deployment_id: str | None,
) -> tuple[BilledFailedAttempt, ...]:
"""Every earlier retry attempt the provider also billed. The router records one
``RetryAttemptRecord`` per failed attempt in ``metadata.previous_models``, the final attempt
included, which the failure row itself bills. The final attempt is the record with the
request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback
hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a
``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final
attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the
model). Only the exception class survives on the record, so an earlier connect timeout is
billed too; any other failure class is not."""
final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries")
final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None)
return tuple(
BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
model_call_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=base_model,
),
)
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
)
@ -1265,13 +1380,26 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
return _EMPTY_LIFT
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
_litellm_params: Final = getattr(_logging_obj, "litellm_params", None)
_base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details)
_deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None
_usage_to_lift: Final = _failure_usage_to_lift(
model_call_details=_model_call_details,
request_body=request_data,
dispatched=_first_handoff is not None,
provider_billed=provider_billed,
deployment_id=_logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None,
base_model=get_base_model_from_metadata(model_call_details=_model_call_details),
deployment_id=_deployment_id,
base_model=_base_model,
)
_prompt_estimate: Final = (
_estimated_prompt_usage(_model_call_details, request_data)
if provider_billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
else None
)
_earlier_attempts: Final = (
_billed_earlier_attempts(_model_call_details, _litellm_params, _prompt_estimate, _base_model, _deployment_id)
if _prompt_estimate is not None
else ()
)
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
@ -1280,6 +1408,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]),
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
("standard_logging_object", _model_call_details.get("standard_logging_object")),
("billed_failed_attempts", _earlier_attempts or None),
)
return MappingProxyType({key: value for key, value in _entries if value is not None})

View file

@ -3330,9 +3330,7 @@ def test_x_litellm_trace_id_keeps_the_body_session_beside_it_whatever_the_log_se
assert data["litellm_trace_id"] == "req-trace-1"
@pytest.mark.parametrize(
"session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"]
)
@pytest.mark.parametrize("session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"])
@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"])
def test_session_headers_still_beat_the_body_session_with_correlation_logging(
monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool
@ -8282,6 +8280,7 @@ _PLANTED_STAMPS = {
"attempted_fallbacks": 99,
"original_model_group": "spoofed-group",
"request_retry_count": -100,
"previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}],
"_client_output_ceiling": {"api_base": "https://attacker.example"},
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9,
"client_key": "client_value",
@ -8316,6 +8315,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
assert "original_model_group" not in updated["metadata"]
assert "_client_output_ceiling" not in updated["metadata"]
assert "request_retry_count" not in updated["metadata"]
assert "previous_models" not in updated["metadata"]
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"

View file

@ -2135,6 +2135,143 @@ def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the
assert response_cost == pytest.approx(usage.prompt_tokens * 0.25)
def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(monkeypatch):
"""The router leaves one RetryAttemptRecord per failed attempt in metadata.previous_models, the
final attempt included; the earlier Timeout ones are billed (each at the deployment that served
it), a 429 attempt is not, and the final attempt is left to the failure row itself."""
from litellm.proxy.utils import _billed_earlier_attempts
monkeypatch.setitem(
litellm.model_cost,
"dep-lit145-b",
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
)
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
public_rate, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=24)
attempts = _billed_earlier_attempts(
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
litellm_params={
"metadata": {
"attempted_retries": 3,
"previous_models": (
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
{"deployment_id": "dep-lit145-b", "exception_type": "RateLimitError", "attempted_retries": 1},
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 2},
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 3},
),
}
},
estimated_usage=usage,
base_model=None,
deployment_id="dep-lit145-a",
)
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [(0, "dep-lit145-a"), (2, "dep-lit145-b")]
assert attempts[0].response_cost == pytest.approx(public_rate)
assert attempts[0].response_cost > 0
assert attempts[1].response_cost == pytest.approx(24 * 0.25)
assert all(a.usage is usage for a in attempts)
def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage():
"""A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks
delivered); the earlier timed-out attempts never produced output, so they are priced from the
input-only estimate and not from that recovered usage."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.proxy.utils import _estimated_prompt_usage, _failure_fields_to_lift
logging_obj = _dispatched_failure_logging_obj(
FAILURE_USAGE_MODEL,
get_litellm_params(
metadata={
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-final"},
"previous_models": (
{"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0},
),
}
),
)
recovered = litellm.Usage(prompt_tokens=500, completion_tokens=40, total_tokens=540)
logging_obj.model_call_details["combined_usage_object"] = recovered
logging_obj.model_call_details["response_cost"] = 0.01
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
estimate = _estimated_prompt_usage(logging_obj.model_call_details, {})
assert lifted["combined_usage_object"] is recovered
assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [
(estimate.prompt_tokens, 0)
]
assert estimate.prompt_tokens < recovered.prompt_tokens
def test_attempts_already_summed_into_the_final_row_are_not_lifted_again():
"""``Logging.carry_failed_attempt_usage`` moves a cut stream's usage into ``failed_attempts_usage``
and the final row already bills every entry of it, so no earlier attempt gets a row of its own."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.proxy.utils import _failure_fields_to_lift
logging_obj = _dispatched_failure_logging_obj(
FAILURE_USAGE_MODEL,
get_litellm_params(
metadata={
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-final"},
"previous_models": (
{"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0},
),
}
),
)
logging_obj.model_call_details["failed_attempts_usage"] = (
(litellm.Usage(prompt_tokens=24, completion_tokens=3, total_tokens=27), 0.001),
)
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
assert lifted["combined_usage_object"].prompt_tokens == 24
assert "billed_failed_attempts" not in lifted
def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed():
"""After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and
re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be
told from the previous group's attempt with the same count (dep-a, retry 1), and the re-logged
record must bill once."""
from litellm.proxy.utils import _billed_earlier_attempts
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
attempts = _billed_earlier_attempts(
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
litellm_params={
"metadata": {
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-b"},
"previous_models": (
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 0},
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 1},
),
}
},
estimated_usage=usage,
base_model=None,
deployment_id="dep-lit145-b",
)
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [
(0, "dep-lit145-a"),
(1, "dep-lit145-a"),
(0, "dep-lit145-b"),
]
assert all(a.response_cost > 0 for a in attempts)
def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object:
"""A Logging object as the failure hook sees it: litellm_params in the real get_litellm_params()
shape (every key present, ``litellm_metadata`` None on chat completions) and a first-handoff