diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 545b317055a..85dba9c1469 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -34564,6 +34564,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34603,6 +34608,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34642,6 +34652,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34681,6 +34696,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -34720,6 +34740,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c64f3a9db2a..3dccd2b68be 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -105,7 +105,6 @@ from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) -from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, without_injected_usage_frames, @@ -2942,7 +2941,7 @@ async def _base_vertex_proxy_route( fastapi_response: Response, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, general_settings: Mapping[str, object], - user_api_key_dict: UserAPIKeyAuth | None = None, + user_api_key_dict: UserAPIKeyAuth, router_credentials: LiteLLM_ManagedVectorStore | None = None, ): """ @@ -2958,7 +2957,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: The caller the route already authenticated, or None to authenticate here + user_api_key_dict: The caller the route authenticated router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore) """ from litellm.llms.vertex_ai.common_utils import ( @@ -2972,15 +2971,10 @@ async def _base_vertex_proxy_route( encoded_endpoint = httpx.URL(endpoint).path verbose_proxy_logger.debug("requested endpoint %s", endpoint) headers: Mapping[str, str] = {} - 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 router_credentials is not None: await assert_user_can_access_vector_store( vector_store=router_credentials, - user_api_key_dict=caller, + user_api_key_dict=user_api_key_dict, ) vertex_project: str | None = get_vertex_project_id_from_url(endpoint) @@ -3032,7 +3026,7 @@ async def _base_vertex_proxy_route( router_credentials=router_credentials, vertex_project=vertex_project, vertex_location=vertex_location, - user_api_key_dict=caller, + user_api_key_dict=user_api_key_dict, ) base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location) @@ -3078,7 +3072,7 @@ async def _base_vertex_proxy_route( received_value: Final = await endpoint_func( request, fastapi_response, - caller, + user_api_key_dict, ) except ProxyException as e: if headers_passed_through: @@ -3088,6 +3082,16 @@ async def _base_vertex_proxy_route( return received_value +async def _authenticate_vertex_caller( + request: Request, + authorization: Annotated[str | None, Security(api_key_header)], + litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)], +) -> UserAPIKeyAuth: + return await user_api_key_auth( + request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key + ) + + @router.api_route( "/vertex_ai/discovery/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -3097,6 +3101,7 @@ async def vertex_discovery_proxy_route( endpoint: str, request: Request, fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)], ): """ Call any vertex discovery endpoint using the proxy. @@ -3135,20 +3140,11 @@ async def vertex_discovery_proxy_route( fastapi_response=fastapi_response, get_vertex_pass_through_handler=discovery_handler, general_settings=_proxy_general_settings(), + user_api_key_dict=user_api_key_dict, router_credentials=vector_store_credentials, ) -async def _authenticate_vertex_caller( - request: Request, - authorization: Annotated[str | None, Security(api_key_header)], - litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)], -) -> UserAPIKeyAuth: - return await user_api_key_auth( - request=request, api_key=authorization or "", custom_litellm_key_header=litellm_api_key - ) - - @router.api_route( "/vertex-ai/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], 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 54874b51478..aa26145c406 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 @@ -324,12 +324,6 @@ 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.get_litellm_virtual_key" - ) as mock_get_virtual_key, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_user_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, @@ -340,8 +334,6 @@ class TestVertexAIPassThroughHandler: # Setup mocks mock_load_auth.return_value = (mock_credentials, test_project) - mock_get_virtual_key.return_value = "Bearer test-key" - mock_user_auth.return_value = {"api_key": "test-key"} # Mock the vertex handler mock_handler = Mock() @@ -426,12 +418,6 @@ 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.get_litellm_virtual_key" - ) as mock_get_virtual_key, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_user_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, @@ -442,8 +428,6 @@ class TestVertexAIPassThroughHandler: # Setup mocks mock_load_auth.return_value = (mock_credentials, test_project) - mock_get_virtual_key.return_value = "Bearer test-key" - mock_user_auth.return_value = {"api_key": "test-key"} # Mock the vertex handler for global location mock_handler = Mock() @@ -1238,12 +1222,6 @@ class TestVertexAIDiscoveryPassThroughHandler: 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.get_litellm_virtual_key" - ) as mock_get_virtual_key, - mock.patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth" - ) as mock_user_auth, mock.patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler" ) as mock_get_handler, @@ -1254,8 +1232,6 @@ class TestVertexAIDiscoveryPassThroughHandler: # Setup mocks mock_load_auth.return_value = (mock_credentials, test_project) - mock_get_virtual_key.return_value = "Bearer test-key" - mock_user_auth.return_value = {"api_key": "test-key"} # Mock the discovery handler mock_handler = Mock() @@ -1271,6 +1247,7 @@ class TestVertexAIDiscoveryPassThroughHandler: endpoint=endpoint, request=mock_request, fastapi_response=mock_response, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), ) # Verify create_pass_through_route was called with correct arguments @@ -1285,36 +1262,29 @@ class TestVertexAIDiscoveryPassThroughHandler: @pytest.mark.asyncio async def test_vertex_discovery_proxy_route_api_key_auth(self): """ - 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 + With no Vertex credential configured, the caller's virtual key must not be forwarded to Google, so the request fails with a clean 401 (LIT-5997). """ - # Mock dependencies mock_request = Mock() mock_request.headers = {"x-litellm-api-key": "test-key-123"} 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_discovery_proxy_route( + endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search", + 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_discovery_proxy_route( - endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search", - 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() @pytest.mark.asyncio @@ -4375,8 +4345,6 @@ class TestVertexRawPredictStreamingClassification: return_value=(mock_credentials, "test-project"), ), mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), - mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"), - mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(api_key="test-key"))), mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler), ): await vertex_proxy_route( @@ -8163,6 +8131,10 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: "v1/projects/vertex-proj/locations/global/collections/default_collection/engines/support-engine" "/servingConfigs/default_search:search" ) + GEMINI_ANSWER: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + } def _post( self, @@ -8211,17 +8183,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: custom_auth: Final = _CountingCustomAuth() vertex_url: Final = f"https://us-central1-aiplatform.googleapis.com/{self.GENERATE_CONTENT}" with respx.mock(assert_all_called=True) as upstream: - vertex: Final = upstream.post(vertex_url).mock( - return_value=httpx.Response( - 200, - json={ - "candidates": [ - {"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"} - ], - "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, - }, - ) - ) + vertex: Final = upstream.post(vertex_url).mock(return_value=httpx.Response(200, json=self.GEMINI_ANSWER)) response, recorder = self._post( monkeypatch, @@ -8258,6 +8220,51 @@ class TestVertexPassthroughAuthenticatesOncePerRequest: assert recorder.served_user_ids == ["custom-auth-caller-1"] assert discovery.call_count == 1 + @pytest.mark.parametrize("scheme", ["bearer", "Basic"]) + @pytest.mark.parametrize( + "route, upstream_url, credential_location, answer", + [ + ( + f"/vertex_ai/{GENERATE_CONTENT}", + f"https://us-central1-aiplatform.googleapis.com/{GENERATE_CONTENT}", + "us-central1", + GEMINI_ANSWER, + ), + ( + f"/vertex_ai/discovery/{DISCOVERY_SEARCH}", + f"https://discoveryengine.googleapis.com/{DISCOVERY_SEARCH}", + "global", + {"results": []}, + ), + ], + ids=["generate_content", "discovery_search"], + ) + def test_x_litellm_api_key_with_any_auth_scheme_authenticates_once_on_every_vertex_route( + self, + monkeypatch: pytest.MonkeyPatch, + scheme: str, + route: str, + upstream_url: str, + credential_location: str, + answer: Mapping[str, object], + ) -> None: + custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY) + with respx.mock(assert_all_called=True) as upstream: + google: Final = upstream.post(upstream_url).mock(return_value=httpx.Response(200, json=answer)) + + response, recorder = self._post( + monkeypatch, + route, + {"x-litellm-api-key": f"{scheme} {self.VIRTUAL_KEY}"}, + custom_auth, + (credential_location,), + ) + + assert response.status_code == 200, response.text + assert custom_auth.api_keys_seen == [self.VIRTUAL_KEY] + assert recorder.served_user_ids == ["custom-auth-caller-1"] + assert google.call_count == 1 + def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential( self, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index ed14eeda6a2..c5b850ea2a9 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/unit/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -66,10 +66,6 @@ async def test_vertex_passthrough_load_balancing(): patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): # Setup additional mocks to avoid side effects @@ -83,7 +79,6 @@ async def test_vertex_passthrough_load_balancing(): mock_endpoint_func = AsyncMock() mock_create_route.return_value = mock_endpoint_func - mock_auth.return_value = {} # Execute await _base_vertex_proxy_route( @@ -92,6 +87,7 @@ async def test_vertex_passthrough_load_balancing(): fastapi_response=mock_response, get_vertex_pass_through_handler=mock_handler, general_settings={}, + user_api_key_dict=UserAPIKeyAuth(), ) # Verify @@ -376,10 +372,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - patch( # test-quality-ok: the route calls auth directly rather than through Depends - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", new_callable=AsyncMock, @@ -398,7 +390,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( "global", ) mock_create_route.return_value = AsyncMock() - mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key") await _base_vertex_proxy_route( endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", @@ -406,6 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( fastapi_response=MagicMock(), get_vertex_pass_through_handler=VertexAIPassThroughHandler(), general_settings={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"), ) upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] @@ -723,10 +715,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" ) as mock_create_route, - patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - ) as mock_auth, ): mock_pt_router.get_vertex_credentials.return_value = MagicMock() @@ -738,7 +726,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): ) mock_endpoint_func = AsyncMock() mock_create_route.return_value = mock_endpoint_func - mock_auth.return_value = {} mock_handler.get_default_base_target_url.return_value = ( "https://global-aiplatform.googleapis.com" @@ -750,6 +737,7 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url(): fastapi_response=mock_response, get_vertex_pass_through_handler=mock_handler, general_settings={}, + user_api_key_dict=UserAPIKeyAuth(), ) # Verify the router was called with the custom model name (extracted from URL) @@ -826,11 +814,6 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", return_value=relay_returning_logging_kwargs, ), - patch( # test-quality-ok: the route calls auth directly rather than through Depends - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new_callable=AsyncMock, - return_value=UserAPIKeyAuth(api_key="hashed-key"), - ), ): mock_pt_router.get_vertex_credentials.return_value = MagicMock() @@ -840,6 +823,7 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment fastapi_response=MagicMock(), get_vertex_pass_through_handler=mock_handler, general_settings={}, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), ) assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep" diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py index 29efc41b484..5b356c43245 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py @@ -441,19 +441,16 @@ async def test_vertex_discovery_denies_other_team_vector_store_credentials(): "team_id": "team-b", } - with patch( - "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", - new=AsyncMock(return_value=UserAPIKeyAuth(team_id="team-a")), - ): - with pytest.raises(HTTPException) as exc_info: - await _base_vertex_proxy_route( - endpoint="projects/p/locations/us-central1/dataStores/vs_other_team", - request=request, - fastapi_response=Response(), - get_vertex_pass_through_handler=MagicMock(), - general_settings={}, - router_credentials=vector_store_credentials, - ) + with pytest.raises(HTTPException) as exc_info: + await _base_vertex_proxy_route( + endpoint="projects/p/locations/us-central1/dataStores/vs_other_team", + request=request, + fastapi_response=Response(), + get_vertex_pass_through_handler=MagicMock(), + general_settings={}, + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + router_credentials=vector_store_credentials, + ) assert exc_info.value.status_code == 403 @@ -541,6 +538,7 @@ async def test_vertex_discovery_allows_unregistered_provider_native_datastore_id endpoint="projects/p/locations/us-central1/dataStores/vs_unknown", request=request, fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(), ) assert response == {"ok": True}