fix(router): keep the request's metadata on rust core fallback hop kwargs (#50263)

The fallback loggers got the walk's own metadata in place of the request's, so prometheus fallback counters lost the caller's key and team labels. Lay the walk's stamps over a copy of the request's bucket, as run_async_fallback did
This commit is contained in:
Nate Armstrong 2026-11-05 15:46:44 -08:00 • committed by GitHub
parent 3b20b48473
commit 41cbbfeb8d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 59 additions and 2 deletions

View file

@ -264,10 +264,18 @@ class WalkMetadata:
def hop_kwargs(self, event: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a fresh dict for the loggers
"""The kwargs the level held after a hop (``kwargs["model"] = target``), as the
fallback loggers see them."""
fallback loggers see them: the request's own metadata (the proxy's ``user_api_key_*``
labels) with the level's walk stamps laid over it, as ``run_async_fallback``'s
``_hop_metadata`` copies it."""
scope: Final = self._scope(event.get("scope"))
model: Final = scope.metadata.get("model_group", event.get("to"))
return {**self._request, **scope.kwargs, "model": model, self._field: dict(scope.metadata)}
request_bucket: Final = as_str_mapping(self._request.get(self._field)) or {}
return {
**self._request,
**scope.kwargs,
"model": model,
self._field: {**request_bucket, **scope.metadata},
}
def log_retry(self, event: Mapping[str, object], error: BaseException | None) -> None:
"""A ``retry_logged`` event: the failure joins the scope's ``previous_models``."""

View file

@ -0,0 +1,49 @@
import pytest
from litellm.router_rust.invoke.walk import WalkMetadata
def _no_base_model(_model_id: str) -> str | None:
return None
@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"])
def test_hop_kwargs_keep_the_requests_key_labels_under_the_walk_stamps(field: str) -> None:
"""The fallback loggers (prometheus' fallback counters) read the caller's key and team off
the hop's metadata, so a hop keeps the request's bucket and lays its own stamps over it."""
request = {
"model": "primary",
field: {
"user_api_key_hash": "hashed-key",
"user_api_key_alias": "key-alias",
"user_api_key_team_id": "team-id",
"user_api_key_team_alias": "team-alias",
"attempted_fallbacks": "stale",
},
}
walk = WalkMetadata(request, field, _no_base_model)
walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5})
kwargs = walk.hop_kwargs({"scope": 1, "to": "fallback"})
assert kwargs["model"] == "fallback"
metadata = kwargs[field]
assert isinstance(metadata, dict)
assert metadata["user_api_key_hash"] == "hashed-key"
assert metadata["user_api_key_alias"] == "key-alias"
assert metadata["user_api_key_team_id"] == "team-id"
assert metadata["user_api_key_team_alias"] == "team-alias"
assert metadata["model_group"] == "fallback"
assert metadata["original_model_group"] == "primary"
assert metadata["attempted_fallbacks"] == 1
assert request[field]["attempted_fallbacks"] == "stale"
assert "model_group" not in request[field]
def test_hop_kwargs_without_request_metadata_hold_the_walk_stamps() -> None:
walk = WalkMetadata({"model": "primary"}, "metadata", _no_base_model)
walk.hop({"parent": 0, "scope": 1, "from": "primary", "to": "fallback", "depth": 1, "max_fallbacks": 5})
metadata = walk.hop_kwargs({"scope": 1, "to": "fallback"})["metadata"]
assert metadata == {"original_model_group": "primary", "model_group": "fallback", "attempted_fallbacks": 1}