fix(passthrough): attribute spend and release budget reservation on router-model /vllm and /azure routes

The /vllm and /azure router-model passthrough branches called
llm_router.allm_passthrough_route directly with no request metadata,
so the cost callback saw no user_api_key and no
user_api_key_budget_reservation. Spend for a budgeted virtual key hit
neither the key's spend nor the spend logs, and the reservation minted
at auth into the shared Redis counter was never released, drifting the
counter up until the key falsely tripped BudgetExceededError.

Thread the authenticated key's attribution metadata into both calls via
the same builder add_litellm_data_to_request uses, so the cost callback
attributes spend and reconciles the reservation. Regression tests cover
both branches.
This commit is contained in:
mateo-berri 2026-08-24 11:16:57 -07:00
parent 3122600e21
commit d23069e907
2 changed files with 109 additions and 1 deletions

View file

@ -9,7 +9,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
import json
import os
import re
from collections.abc import Callable
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Any, Final, cast
@ -106,6 +106,28 @@ def is_passthrough_request_streaming(request_body: object) -> bool:
return bool(request_body.get("stream", False))
def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, Any]:
"""
Build the request metadata carrying key-level spend attribution and the
pre-call budget reservation for a router-model passthrough request.
Router-model passthrough branches call ``allm_passthrough_route`` directly,
bypassing ``add_litellm_data_to_request``. Without this metadata the cost
callback cannot attribute spend to the calling key and never releases the
budget reservation minted at auth time, so the shared spend counter drifts
up until the key falsely trips ``BudgetExceededError``.
"""
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
request_data: Final = {"metadata": {}} # mutable-ok: attribution builder + litellm mutate this dict in place
LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data=request_data,
user_api_key_dict=user_api_key_dict,
_metadata_variable_name="metadata",
)
return request_data["metadata"]
async def llm_passthrough_factory_proxy_route(
custom_llm_provider: str,
endpoint: str,
@ -346,6 +368,7 @@ async def vllm_proxy_route(
params=None,
headers=None,
cookies=None,
metadata=get_passthrough_router_request_metadata(user_api_key_dict),
),
)
@ -1475,6 +1498,7 @@ async def azure_proxy_route(
params=None,
headers=None,
cookies=None,
metadata=get_passthrough_router_request_metadata(user_api_key_dict),
)
if is_streaming_request:

View file

@ -4055,3 +4055,87 @@ class TestVertexAILiveWebsocketPassthrough:
assert "use_in_pass_through" in close_kwargs["reason"]
assert "default_vertex_config" in close_kwargs["reason"]
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
class TestPassthroughRouterModelBudgetReservation:
"""
Router-model passthrough on /vllm and /azure must thread the calling key's
metadata into ``allm_passthrough_route``. Without ``user_api_key`` the spend
is attributed to nobody, and without ``user_api_key_budget_reservation`` the
pre-call reservation is never released, so the shared spend counter drifts up
until the key falsely trips a 429 BudgetExceededError (LIT-5470).
"""
def _key_with_reservation(self) -> UserAPIKeyAuth:
reservation = {
"reserved_cost": 0.5,
"entries": [{"counter_key": "spend:key:hashed-token", "reserved_cost": 0.5}],
}
return UserAPIKeyAuth(
api_key="hashed-token",
user_id="u1",
team_id="t1",
budget_reservation=reservation,
)
def _request(self) -> MagicMock:
request = MagicMock(spec=Request)
request.method = "POST"
request.headers = {"content-type": "application/json"}
request.query_params = {}
return request
def _install_recording_router(self, monkeypatch, body: dict) -> list[dict]:
import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep
import litellm.proxy.proxy_server as proxy_server
captured: list[dict] = []
class RecordingRouter:
async def allm_passthrough_route(self, **kwargs):
captured.append(kwargs)
return httpx.Response(200, json={"ok": True})
async def fake_get_request_body(_request):
return body
monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter())
monkeypatch.setattr(ep, "get_request_body", fake_get_request_body)
monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True)
return captured
def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None:
assert len(captured) == 1, "the router-model branch must dispatch exactly once"
metadata = captured[0]["metadata"]
assert metadata["user_api_key"] == user_api_key_dict.api_key
assert metadata["user_api_key_budget_reservation"] is user_api_key_dict.budget_reservation
assert metadata["user_api_key_user_id"] == user_api_key_dict.user_id
assert metadata["user_api_key_team_id"] == user_api_key_dict.team_id
@pytest.mark.asyncio
async def test_vllm_router_model_threads_key_metadata(self, monkeypatch):
user_api_key_dict = self._key_with_reservation()
captured = self._install_recording_router(monkeypatch, {"model": "router-model", "stream": False})
await vllm_proxy_route(
endpoint="/chat/completions",
request=self._request(),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=user_api_key_dict,
)
self._assert_metadata_carries_attribution(captured, user_api_key_dict)
@pytest.mark.asyncio
async def test_azure_router_model_threads_key_metadata(self, monkeypatch):
user_api_key_dict = self._key_with_reservation()
captured = self._install_recording_router(monkeypatch, {"model": "gpt-5", "stream": False})
await azure_proxy_route(
endpoint="openai/deployments/gpt-5/chat/completions",
request=self._request(),
fastapi_response=MagicMock(spec=Response),
user_api_key_dict=user_api_key_dict,
)
self._assert_metadata_carries_attribution(captured, user_api_key_dict)