diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7ce41c1d5b6..55d6b88363f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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,36 @@ 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``. + + The payload rides the ``litellm_metadata`` bucket, not ``metadata``: the + router hop ``_ageneric_api_call_with_fallbacks`` canonicalises this call + type into ``litellm_metadata``, and the cost callback reads spend + attribution from that bucket while only backfilling ``user_api_key*`` keys + from ``metadata``. Passing ``metadata=`` would silently drop the secondary + attribution fields the helper sets (``agent_id``, + ``user_api_end_user_max_budget``) before the callback ever sees them. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + request_data: Final = {"litellm_metadata": {}} # mutable-ok: builder + litellm mutate this 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="litellm_metadata", + ) + return request_data["litellm_metadata"] + + async def llm_passthrough_factory_proxy_route( custom_llm_provider: str, endpoint: str, @@ -346,6 +376,7 @@ async def vllm_proxy_route( params=None, headers=None, cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), ), ) @@ -1475,6 +1506,7 @@ async def azure_proxy_route( params=None, headers=None, cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), ) if is_streaming_request: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 97b77ec95e0..22843ceae40 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -4056,3 +4056,95 @@ 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, + agent_id="agent-xyz", + end_user_max_budget=42.0, + ) + + 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" + assert captured[0].get("metadata") is None, ( + "attribution must ride the litellm_metadata bucket the router canonicalizes on; " + "the plain metadata bucket is dropped for every non-user_api_key field" + ) + litellm_metadata = captured[0]["litellm_metadata"] + assert litellm_metadata["user_api_key"] == user_api_key_dict.api_key + assert litellm_metadata["user_api_key_budget_reservation"] is user_api_key_dict.budget_reservation + assert litellm_metadata["user_api_key_user_id"] == user_api_key_dict.user_id + assert litellm_metadata["user_api_key_team_id"] == user_api_key_dict.team_id + assert litellm_metadata["agent_id"] == user_api_key_dict.agent_id + assert litellm_metadata["user_api_end_user_max_budget"] == user_api_key_dict.end_user_max_budget + + @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)