From 60e16040facbe70c6433c807fc3a8eb2aa962e9f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:59:37 -0700 Subject: [PATCH] fix(vertex_ai): authenticate Vertex pass-through requests once vertex_proxy_route already resolves the caller through Depends(user_api_key_auth), but _base_vertex_proxy_route ignored it and authenticated the request again. A custom_auth function ran twice per request, and a key with max_budget reserved its estimated cost twice, so the unreleased first reservation exhausted the budget after one or two calls. The base route now uses the caller it is given and authenticates once itself only when none is passed, which is the discovery route's case. --- .../llm_passthrough_endpoints.py | 23 ++++----- .../test_llm_pass_through_endpoints.py | 48 +++++++------------ 2 files changed, 26 insertions(+), 45 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..679ae3daa78 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2956,7 +2956,7 @@ async def _base_vertex_proxy_route( fastapi_response: FastAPI response object get_vertex_pass_through_handler: Handler for the specific Vertex AI service general_settings: Proxy general settings, read per request for the anthropic-beta values to drop - user_api_key_dict: User API key authentication dict + user_api_key_dict: The caller the route already authenticated, or None to authenticate here router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore) """ from litellm.llms.vertex_ai.common_utils import ( @@ -2970,22 +2970,15 @@ async def _base_vertex_proxy_route( encoded_endpoint = httpx.URL(endpoint).path verbose_proxy_logger.debug("requested endpoint %s", endpoint) headers: Mapping[str, str] = {} - api_key_to_use = get_litellm_virtual_key(request=request) - user_api_key_dict = await user_api_key_auth( - request=request, - api_key=api_key_to_use, + caller: Final = ( + user_api_key_dict + if user_api_key_dict is not None + else await user_api_key_auth(request=request, api_key=get_litellm_virtual_key(request=request)) ) - - if user_api_key_dict is None: - api_key_to_use = get_litellm_virtual_key(request=request) - user_api_key_dict = await user_api_key_auth( - request=request, - api_key=api_key_to_use, - ) if router_credentials is not None: await assert_user_can_access_vector_store( vector_store=router_credentials, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=caller, ) vertex_project: str | None = get_vertex_project_id_from_url(endpoint) @@ -3037,7 +3030,7 @@ async def _base_vertex_proxy_route( router_credentials=router_credentials, vertex_project=vertex_project, vertex_location=vertex_location, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=caller, ) base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) @@ -3083,7 +3076,7 @@ async def _base_vertex_proxy_route( received_value: Final = await endpoint_func( request, fastapi_response, - user_api_key_dict, + caller, ) except ProxyException as e: if headers_passed_through: diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 1e9573d19a3..376b2f18d3b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -528,17 +528,12 @@ class TestVertexAIPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): # Mock credentials object with necessary attributes mock_credentials = Mock() mock_credentials.token = default_credentials mock_load_auth.return_value = (mock_credentials, default_project) - mock_auth.return_value = MagicMock() # Mock the vertex handler mock_handler = Mock() @@ -556,6 +551,7 @@ class TestVertexAIPassThroughHandler: endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), ) except Exception as e: traceback.print_exc() @@ -630,20 +626,16 @@ class TestVertexAIPassThroughHandler: mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): mock_ensure_token.return_value = ("test-auth-header", test_project) mock_get_token.return_value = (test_token, "") - mock_auth.return_value = UserAPIKeyAuth(api_key="sk-test-creds") with pytest.raises(HTTPException) as exc_info: await vertex_proxy_route( endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-creds"), ) assert exc_info.value.status_code == 401 @@ -655,8 +647,9 @@ class TestVertexAIPassThroughHandler: Critical This is how Vertex AI JS SDK will Auth to Litellm Proxy: the virtual key - arrives in x-litellm-api-key and must reach user_api_key_auth. With no Vertex - credential configured, that virtual key must not be forwarded to Google, so + arrives in x-litellm-api-key and authenticates the caller (covered through the + real auth dependency in TestVertexPassthroughAuthenticatesOncePerRequest). With no + Vertex credential configured, that virtual key must not be forwarded to Google, so the request fails with a clean 401 (LIT-5997). """ # Mock dependencies @@ -665,26 +658,21 @@ class TestVertexAIPassThroughHandler: mock_request.method = "POST" mock_response = Mock() - with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth: - mock_auth.return_value = {"api_key": "test-key-123"} + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_pass_through: + mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" - ) as mock_pass_through: - mock_pass_through.return_value = AsyncMock(return_value={"status": "success"}) + with pytest.raises(HTTPException) as exc_info: + await vertex_proxy_route( + endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key-123"), + ) - with pytest.raises(HTTPException) as exc_info: - await vertex_proxy_route( - endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", - request=mock_request, - fastapi_response=mock_response, - ) - - assert exc_info.value.status_code == 401 - mock_pass_through.assert_not_called() - mock_auth.assert_called_once() - call_args = mock_auth.call_args[1] - assert call_args["api_key"] == "Bearer test-key-123" + assert exc_info.value.status_code == 401 + mock_pass_through.assert_not_called() def test_vertex_passthrough_handler_multimodal_embedding_response(self): """