From d0061fbea7ddf768f674954ca9fa4774c4d5bae5 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Mon, 11 May 2026 13:56:16 +0800 Subject: [PATCH] fix(vertex_ai/gemini): address mypy + gcs credential review feedback - mypy: cast image_url/file dict to Dict[str, Any] before reading mime_type/content_type, since those keys are not declared on the ChatCompletionImageUrlObject / ChatCompletionFileObjectFile TypedDicts and mypy was inferring `object` for the fallback-or chain, failing the lint job. - security (veria-ai review): in _get_gcs_object_content_type, only call VertexBase.get_access_token when the caller explicitly passes vertex_project/vertex_credentials. Without explicit credentials we now skip the authenticated call entirely and only attempt an anonymous metadata lookup, so a Gemini API-key / Google AI Studio caller can no longer trick the proxy into querying GCS with the server's default Google credentials. - Add regression test covering the no-credentials path (no ADC token requested, no Authorization header on the outbound metadata call). --- .../llms/vertex_ai/gemini/transformation.py | 31 ++++---- .../llms/vertex_ai/test_vertex.py | 72 +++++++++++++++---- 2 files changed, 77 insertions(+), 26 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 5325e06aa53..07eb03008d2 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -230,6 +230,11 @@ def _get_gcs_object_content_type( ) -> Optional[str]: """ Resolve content type from GCS object metadata. + + 仅当调用方显式传入 Vertex 凭据时才附带 Bearer token，避免在 Gemini + API key（Google AI Studio）路径上自动使用服务端的默认 Google 凭据去 + 访问 GCS，防止被用作探测私有 GCS 对象的 oracle。 + 无显式凭据时只做匿名请求，仅对公开可读对象有效。 """ try: bucket, object_name = _parse_gs_uri(image_url) @@ -242,14 +247,14 @@ def _get_gcs_object_content_type( explicit_vertex_auth_provided = ( vertex_project is not None or vertex_credentials is not None ) - try: - access_token, _ = _get_vertex_base().get_access_token( - credentials=vertex_credentials, - project_id=vertex_project, - ) - headers["Authorization"] = f"Bearer {access_token}" - except Exception as e: - if explicit_vertex_auth_provided: + if explicit_vertex_auth_provided: + try: + access_token, _ = _get_vertex_base().get_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + ) + headers["Authorization"] = f"Bearer {access_token}" + except Exception as e: raise litellm.BadRequestError( message=( "Unable to fetch GCS metadata with provided Vertex credentials/project. " @@ -258,7 +263,6 @@ def _get_gcs_object_content_type( model=None, llm_provider="vertex_ai", ) - # 未显式提供 Vertex 凭据时，metadata 对公开对象仍可能可读，继续无 token 尝试。 # 通过 httpx.URL 固定 scheme/host，并对 bucket、object 都做 URL 编码， # 避免被 CodeQL 误判为可能拼接出任意主机 URL 的 SSRF。 @@ -551,12 +555,13 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 model=model, llm_provider="vertex_ai", ) + image_url_dict = cast(Dict[str, Any], raw_image_url) format = ( - raw_image_url.get("format") - or raw_image_url.get("mime_type") - or raw_image_url.get("content_type") + image_url_dict.get("format") + or image_url_dict.get("mime_type") + or image_url_dict.get("content_type") ) - detail = raw_image_url.get("detail") + detail = image_url_dict.get("detail") media_resolution_enum = ( _convert_detail_to_media_resolution_enum(detail) ) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 8cda11c1f45..c24b03e8c71 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -1389,9 +1389,14 @@ def test_get_gcs_object_content_type_uses_shared_vertex_base_instance(): mock_http_response.json.return_value = {"contentType": "image/png"} mock_http_response.raise_for_status.return_value = None - with patch.object( - gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base - ), patch("litellm.llms.vertex_ai.gemini.transformation.httpx.get") as mock_http_get: + with ( + patch.object( + gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base + ), + patch( + "litellm.llms.vertex_ai.gemini.transformation.httpx.get" + ) as mock_http_get, + ): mock_http_get.return_value = mock_http_response content_type = gemini_transformation._get_gcs_object_content_type( image_url="gs://my-bucket/path/to/image-without-extension", @@ -1428,7 +1433,9 @@ def test_process_gemini_media_rejects_unsupported_metadata_mime_type(): "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", return_value="application/octet-stream", ): - with pytest.raises(litellm.BadRequestError, match="File type not supported by gemini"): + with pytest.raises( + litellm.BadRequestError, match="File type not supported by gemini" + ): _process_gemini_media("gs://bucket/image-without-extension") @@ -1437,7 +1444,9 @@ def test_get_gcs_object_content_type_fails_fast_with_explicit_credentials(): mock_vertex_base = MagicMock() mock_vertex_base.get_access_token.side_effect = Exception("token failure") - with patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base): + with patch.object( + gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base + ): with pytest.raises( litellm.BadRequestError, match="Unable to fetch GCS metadata with provided Vertex credentials/project", @@ -1449,6 +1458,41 @@ def test_get_gcs_object_content_type_fails_fast_with_explicit_credentials(): ) +def test_get_gcs_object_content_type_without_credentials_skips_auth(): + """未显式提供 Vertex 凭据时，不应使用服务端默认凭据去访问 GCS。 + + 这是为了避免 Gemini API key（Google AI Studio）路径被用作探测私有 + GCS 对象的 oracle（veria-ai 审阅反馈）。 + """ + from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation + + mock_vertex_base = MagicMock() + mock_http_response = MagicMock() + mock_http_response.json.return_value = {"contentType": "image/jpeg"} + mock_http_response.raise_for_status.return_value = None + + with ( + patch.object( + gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base + ), + patch( + "litellm.llms.vertex_ai.gemini.transformation.httpx.get" + ) as mock_http_get, + ): + mock_http_get.return_value = mock_http_response + content_type = gemini_transformation._get_gcs_object_content_type( + image_url="gs://public-bucket/public-object" + ) + + # 不应触发 get_access_token（避免使用服务端默认凭据） + mock_vertex_base.get_access_token.assert_not_called() + # 匿名请求仍会发出，用于公开可读对象 + mock_http_get.assert_called_once() + call_kwargs = mock_http_get.call_args.kwargs + assert "Authorization" not in call_kwargs.get("headers", {}) + assert content_type == "image/jpeg" + + def test_async_transform_request_body_does_not_block_event_loop(): """当同步 GCS metadata 查询阻塞时，async_transform_request_body 不应阻塞事件循环。""" import asyncio @@ -1513,14 +1557,16 @@ def test_async_transform_request_body_does_not_block_event_loop(): await transform_task return sleep_elapsed - with patch.object( - gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base - ), patch.object( - gemini_transformation.httpx, "get", side_effect=slow_http_get - ), patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." - "ContextCachingEndpoints.async_check_and_create_cache", - new=fake_check_and_create_cache, + with ( + patch.object( + gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base + ), + patch.object(gemini_transformation.httpx, "get", side_effect=slow_http_get), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." + "ContextCachingEndpoints.async_check_and_create_cache", + new=fake_check_and_create_cache, + ), ): sleep_elapsed = asyncio.run(run_scenario())