Merge pull request #38111 from BerriAI/litellm_fix_passthrough_router_budget_leak

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

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:

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)