diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7ce41c1d5b6..fc891978ad5 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,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: 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 ac140abe31f..032cb360d4c 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 @@ -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)