diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0182c27fac4..fde5be049eb 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -36405,6 +36405,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -36444,6 +36449,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -36483,6 +36493,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -36522,6 +36537,11 @@ "description": "Validation Error" } }, + "security": [ + { + "APIKeyHeader": [] + } + ], "summary": "Vertex Discovery Proxy Route", "tags": [ "llm_passthrough" @@ -36561,6 +36581,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 1eed00cc72c..3653fcee112 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias import httpx import openai -from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket +from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security, WebSocket from fastapi.responses import StreamingResponse from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.websockets import WebSocketState @@ -97,6 +97,8 @@ from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports _get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + api_key_header, + custom_litellm_key_header, get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, @@ -133,7 +135,6 @@ from litellm.proxy.common_utils.token_count_log import ( logged_provider_token_count, rejected_token_count, ) -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, @@ -3097,7 +3098,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, ): """ @@ -3113,7 +3114,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 authenticated router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore) """ from litellm.llms.vertex_ai.common_utils import ( @@ -3127,18 +3128,6 @@ 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, - ) - - 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, @@ -3250,6 +3239,18 @@ 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, # pyright: ignore[reportArgumentType] # its own Security default is None when unset + custom_litellm_key_header=litellm_api_key, + ) + + @router.api_route( "/vertex_ai/discovery/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -3259,6 +3260,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. @@ -3297,6 +3299,7 @@ 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, ) @@ -3316,7 +3319,7 @@ async def vertex_proxy_route( endpoint: str, request: Request, fastapi_response: Response, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)], ): """ Call LiteLLM proxy via Vertex AI SDK. 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 11b07472eaa..2fa3473d1c0 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 @@ -296,12 +296,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, @@ -312,8 +306,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() @@ -398,12 +390,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, @@ -414,8 +400,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() @@ -500,17 +484,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() @@ -528,6 +507,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() @@ -602,20 +582,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 @@ -627,8 +603,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 @@ -637,26 +614,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): """ @@ -1222,12 +1194,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, @@ -1238,8 +1204,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() @@ -1255,6 +1219,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 @@ -1269,36 +1234,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 @@ -4545,8 +4503,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( @@ -8703,6 +8659,228 @@ class TestAzureSpeechRawBodyThroughRealAuth: assert "Invalid JSON payload" in response.text +class _CountingCustomAuth: + def __init__(self, only_accepts: str | None = None) -> None: + self.only_accepts: Final = only_accepts + self.api_keys_seen: tuple[str | None, ...] = () + + async def __call__(self, request: Request, api_key: str | None) -> UserAPIKeyAuth: + self.api_keys_seen = (*self.api_keys_seen, api_key) + if self.only_accepts is not None and api_key != self.only_accepts: + raise ValueError("unknown caller") + return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}") + + +class _ServedCallerRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.served_user_ids: tuple[str | None, ...] = () + + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral + ) -> dict: + self.served_user_ids = (*self.served_user_ids, user_api_key_dict.user_id) + return data + + +class TestVertexPassthroughAuthenticatesOncePerRequest: + """Each Vertex pass-through request runs the proxy's auth once and is served as the caller that auth resolved""" + + PROJECT: Final = "vertex-proj" + VIRTUAL_KEY: Final = "sk-vertex-caller" + GENERATE_CONTENT: Final = ( + "v1/projects/vertex-proj/locations/us-central1/publishers/google/models/gemini-3.8-flash:generateContent" + ) + DISCOVERY_SEARCH: Final = ( + "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, + monkeypatch: pytest.MonkeyPatch, + path: str, + headers: Mapping[str, str], + custom_auth: _CountingCustomAuth, + credential_locations: tuple[str, ...], + ) -> tuple[httpx.Response, _ServedCallerRecorder]: + from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( + PassthroughEndpointRouter, + ) + from litellm.proxy.proxy_server import app + + router: Final = PassthroughEndpointRouter() + for location in credential_locations: + router.add_vertex_credentials( + project_id=self.PROJECT, location=location, vertex_credentials="vertex-service-account" + ) + recorder: Final = _ServedCallerRecorder() + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router", router + ) + monkeypatch.delitem(app.dependency_overrides, user_api_key_auth, raising=False) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + with ( + patch.multiple( # test-quality-ok: the real user_api_key_auth reads proxy_server module globals (custom auth, master_key, caches) that have no injection seam + "litellm.proxy.proxy_server", **{**_azure_speech_real_auth_attrs(), "user_custom_auth": custom_auth} + ), + patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth", + return_value=(Mock(token="google-access-token"), self.PROJECT), + ), + ): + response: Final = TestClient(app).post( + path, json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, headers=dict(headers) + ) + return response, recorder + + @pytest.mark.parametrize("route_prefix", ["/vertex_ai", "/vertex-ai"]) + def test_generate_content_runs_custom_auth_once_and_is_served_as_its_caller( + self, monkeypatch: pytest.MonkeyPatch, route_prefix: str + ) -> None: + 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=self.GEMINI_ANSWER)) + + response, recorder = self._post( + monkeypatch, + f"{route_prefix}/{self.GENERATE_CONTENT}", + {"Authorization": f"Bearer {self.VIRTUAL_KEY}"}, + custom_auth, + ("us-central1",), + ) + + 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 vertex.call_count == 1 + + def test_discovery_search_runs_custom_auth_once_and_is_served_as_its_caller( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + custom_auth: Final = _CountingCustomAuth() + with respx.mock(assert_all_called=True) as upstream: + discovery: Final = upstream.post(f"https://discoveryengine.googleapis.com/{self.DISCOVERY_SEARCH}").mock( + return_value=httpx.Response(200, json={"results": []}) + ) + + response, recorder = self._post( + monkeypatch, + f"/vertex_ai/discovery/{self.DISCOVERY_SEARCH}", + {"Authorization": f"Bearer {self.VIRTUAL_KEY}"}, + custom_auth, + ("global",), + ) + + assert (response.status_code, response.json()) == (200, {"results": []}) + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + 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: + custom_auth: Final = _CountingCustomAuth() + with respx.mock(assert_all_called=False) as upstream: + google: Final = upstream.route().mock(return_value=httpx.Response(200)) + + response, _ = self._post( + monkeypatch, + f"/vertex_ai/{self.GENERATE_CONTENT}", + {"x-litellm-api-key": self.VIRTUAL_KEY}, + custom_auth, + (), + ) + + assert response.status_code == 401, response.text + assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,) + assert not google.called + + @pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"]) + def test_key_outside_litellm_headers_is_rejected_and_never_reaches_google_without_a_vertex_credential( + self, monkeypatch: pytest.MonkeyPatch, header: str + ) -> None: + custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY) + with respx.mock(assert_all_called=False) as upstream: + google: Final = upstream.route().mock(return_value=httpx.Response(200)) + + response, recorder = self._post( + monkeypatch, f"/vertex_ai/{self.GENERATE_CONTENT}", {header: self.VIRTUAL_KEY}, custom_auth, () + ) + + assert response.status_code == 401, response.text + assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen + assert recorder.served_user_ids == () + assert not google.called + + + @pytest.mark.parametrize("path", [f"/vertex_ai/{GENERATE_CONTENT}", f"/vertex_ai/discovery/{DISCOVERY_SEARCH}"]) + def test_request_without_a_key_hands_custom_auth_none( + self, monkeypatch: pytest.MonkeyPatch, path: str + ) -> None: + custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY) + with respx.mock(assert_all_called=False) as upstream: + google: Final = upstream.route().mock(return_value=httpx.Response(200)) + + response, recorder = self._post(monkeypatch, path, {}, custom_auth, ("us-central1", "global")) + + assert response.status_code == 401, response.text + assert custom_auth.api_keys_seen == (None,) + assert recorder.served_user_ids == () + assert not google.called + class TestTypeSafePassthroughRoute: @staticmethod def _request(body: object, query_params: Mapping[str, str] | None = None) -> MagicMock: 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 827693168b7..910f360139c 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 @@ -67,10 +67,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 @@ -84,7 +80,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( @@ -93,6 +88,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 @@ -375,10 +371,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, @@ -397,7 +389,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}", @@ -405,6 +396,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) @@ -825,11 +813,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() @@ -839,6 +822,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 434ea6f3f97..5615db8714f 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 @@ -442,19 +442,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 @@ -542,6 +539,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}