From 2b13898be4959f51255d7a40978cc6515a522a8a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 20:08:21 -0800 Subject: [PATCH] fix(bedrock): check every file id field, skip MCP tool calls, accept registered S3 media types --- .../cloud_storage_security.py | 7 +-- .../prompt_templates/factory.py | 18 +++++-- litellm/proxy/common_request_processing.py | 13 ++--- ...llm_core_utils_prompt_templates_factory.py | 47 +++++++++++++++++++ .../test_cloud_storage_security.py | 21 +++++++++ .../proxy/test_common_request_processing.py | 44 +++++++++++++++++ 6 files changed, 137 insertions(+), 13 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 530f789b34a..c4b0cd3d292 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -83,7 +83,8 @@ _MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", " _ENCODED_FILE_ID_PREFIXES: Final = ("file-", "batch_") _EMBEDDED_S3_URI: Final = re.compile(r"s3://[^;]*") _MODEL_FILE_ID_MAPPING_KEY: Final = "model_file_id_mapping" -_KEYS_THAT_MAY_HIDE_AN_S3_URI: Final = frozenset({"file_id", _MODEL_FILE_ID_MAPPING_KEY}) +_KEYS_THAT_MAY_HOLD_A_FILE_ID: Final = frozenset({"file_id", "file_url", "file_uri", "fileUri"}) +_KEYS_THAT_MAY_HIDE_AN_S3_URI: Final = _KEYS_THAT_MAY_HOLD_A_FILE_ID | {_MODEL_FILE_ID_MAPPING_KEY} def is_s3_uri(value: str) -> bool: @@ -145,13 +146,13 @@ def _string_leaf(node: _BodyNode) -> str | None: def _media_file_id(node: _BodyNode) -> str | None: - return _string_leaf(node) if node.inside_media_part and node.key == "file_id" else None + return _string_leaf(node) if node.inside_media_part and node.key in _KEYS_THAT_MAY_HOLD_A_FILE_ID else None def _media_uri_candidates(trusted_file_ids: frozenset[str], node: _BodyNode) -> Iterator[str]: if not node.inside_media_part or not isinstance(node.value, str): return iter(()) - if node.key == "file_id": + if node.key in _KEYS_THAT_MAY_HOLD_A_FILE_ID: return iter((node.value,)) if node.value in trusted_file_ids else _with_embedded_s3_uris(node.value) return iter((node.value,) if node.key in _MEDIA_URI_KEYS else ()) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 4e6e72292ab..704c9852030 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2411,6 +2411,10 @@ def anthropic_process_openai_file_message( raise Exception(f"Either file_data or file_id must be present in the file message: {message}") +_S3_MIME_SUBTYPE_FORMATS: Final = MappingProxyType( + {"jpg": "jpeg", "quicktime": "mov", "x-matroska": "mkv", "x-flv": "flv", "x-ms-wmv": "wmv", "3gpp": "3gp"} +) +_CONVERSE_VIDEO_FORMAT_NAMES: Final = MappingProxyType({"3gp": "three_gp"}) _EMPTY_TEXT_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]" @@ -4238,7 +4242,11 @@ class BedrockImageProcessor: ) ) elif is_video: - return BedrockContentBlock(video=BedrockVideoBlock(source=source, format=image_format)) + return BedrockContentBlock( + video=BedrockVideoBlock( + source=source, format=_CONVERSE_VIDEO_FORMAT_NAMES.get(image_format, image_format) + ) + ) else: return BedrockContentBlock(image=BedrockImageBlock(source=source, format=image_format)) @@ -4259,14 +4267,15 @@ class BedrockImageProcessor: @classmethod def _create_s3_block(cls, uri: str, format: str | None, bucket_owner: str | None) -> BedrockContentBlock: - mime_type: Final = format or cls._s3_key_mime_type(uri) - if mime_type is None: + declared_type: Final = format or cls._s3_key_mime_type(uri) + if declared_type is None: raise litellm.BadRequestError( message=f"Cannot tell the media type of {uri} from its file extension. " "Set the part's `format` to its MIME type, e.g. `image/png`.", model="", llm_provider="bedrock", ) + mime_type: Final = declared_type.partition(";")[0].strip().lower() if mime_type.startswith("audio/"): raise litellm.BadRequestError( message=f"Bedrock Converse does not accept {uri} as {mime_type}. " @@ -4275,7 +4284,8 @@ class BedrockImageProcessor: llm_provider="bedrock", ) try: - image_format: Final = cls.validate_format(mime_type, mime_type.partition("/")[2]) + subtype: Final = mime_type.partition("/")[2] + image_format: Final = cls.validate_format(mime_type, _S3_MIME_SUBTYPE_FORMATS.get(subtype, subtype)) except ValueError as e: raise litellm.BadRequestError( message=f"Bedrock Converse does not accept {uri} as {mime_type}. {e}", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9abff0c781e..3f0df778610 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2734,12 +2734,13 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) - await refuse_s3_media_the_key_may_not_read( - raw_request_fields, - user_api_key_dict, - _REQUEST_FIELDS.validate_python(general_settings), - proxy_logging_obj.get_proxy_hook("managed_files"), - ) + if route_type != "call_mcp_tool": + await refuse_s3_media_the_key_may_not_read( + raw_request_fields, + user_api_key_dict, + _REQUEST_FIELDS.validate_python(general_settings), + proxy_logging_obj.get_proxy_hook("managed_files"), + ) self.data = await add_litellm_data_to_request( data=self.data, request=request, diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index bcd5e146ba5..1c4070d8e9d 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3406,6 +3406,53 @@ async def test_bedrock_converse_s3_media_is_sent_as_an_s3_location_without_a_dow assert await _converse_user_content(part, run_async) == [expected] +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("declared", "block", "converse_format"), + [ + pytest.param("image/jpg", "image", "jpeg", id="image-jpg"), + pytest.param("IMAGE/PNG", "image", "png", id="upper-case"), + pytest.param("image/png; charset=binary", "image", "png", id="image-with-parameter"), + pytest.param("video/quicktime", "video", "mov", id="quicktime"), + pytest.param("video/x-matroska", "video", "mkv", id="matroska"), + pytest.param("video/x-flv", "video", "flv", id="flv"), + pytest.param("video/x-ms-wmv", "video", "wmv", id="wmv"), + pytest.param("video/3gpp", "video", "three_gp", id="3gpp"), + ], +) +async def test_bedrock_converse_s3_media_format_accepts_a_registered_mime_type( + declared: str, block: str, converse_format: str, run_async: bool +): + uri: Final = "s3://media-bucket/uploads/7f3a9c" + + content: Final = await _converse_user_content( + {"type": "image_url", "image_url": {"url": uri, "format": declared}}, run_async + ) + + assert content == [{block: {"format": converse_format, "source": _s3_source(uri)}}] + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +async def test_bedrock_converse_s3_3gp_key_uses_the_converse_format_name(run_async: bool): + uri: Final = "s3://media-bucket/clips/demo.3gp" + + content: Final = await _converse_user_content({"type": "image_url", "image_url": {"url": uri}}, run_async) + + assert content == [{"video": {"format": "three_gp", "source": _s3_source(uri)}}] + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +async def test_bedrock_converse_s3_text_with_a_charset_is_a_txt_document(run_async: bool): + uri: Final = "s3://media-bucket/docs/7f3a9c" + + content: Final = await _converse_user_content( + {"type": "file", "file": {"file_id": uri, "format": "text/plain; charset=utf-8"}}, run_async + ) + + assert content[0]["document"]["format"] == "txt" + assert content[0]["document"]["source"] == _s3_source(uri) + + @pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) async def test_bedrock_converse_s3_pdf_is_a_document_named_after_its_uri(run_async: bool): report: Final = "s3://media-bucket/docs/report.pdf" diff --git a/tests/unit/litellm_core_utils/test_cloud_storage_security.py b/tests/unit/litellm_core_utils/test_cloud_storage_security.py index 3eb345a4b6b..5c2e1bc26a8 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -115,6 +115,18 @@ def _gemini_turn(*parts: object) -> dict[str, object]: _responses_turn({"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}), id="responses-unified-file-id", ), + pytest.param( + _responses_turn({"type": "input_file", "file_url": _UNIFIED_BATCH_OUTPUT_ID}), + id="responses-unified-file-id-in-file-url", + ), + pytest.param( + _gemini_turn({"file_data": {"mime_type": "application/pdf", "file_uri": _UNIFIED_BATCH_OUTPUT_ID}}), + id="gemini-unified-file-id-in-file-uri", + ), + pytest.param( + _gemini_turn({"fileData": {"mimeType": "application/pdf", "fileUri": _MODEL_EMBEDDED_UPLOAD_ID}}), + id="gemini-model-embedded-file-id-in-camel-case-file-uri", + ), pytest.param( { **_user_turn({"type": "file", "file": {"file_id": "alias"}}), @@ -393,6 +405,15 @@ def test_media_file_ids_lists_the_file_ids_of_media_parts_only(): assert media_file_ids(request_body) == ("file-a",) +def test_media_file_ids_lists_every_field_a_bridge_reads_as_a_file_id(): + request_body: Final = { + **_responses_turn({"type": "input_file", "file_url": "file-a"}), + **_gemini_turn({"file_data": {"file_uri": "file-b"}}, {"fileData": {"fileUri": "file-c"}}), + } + + assert set(media_file_ids(request_body)) == {"file-a", "file-b", "file-c"} + + def test_s3_media_uris_names_each_s3_object_once_in_the_order_it_first_appears(): other_upload: Final = "s3://media-bucket/dir/cat.png" request_body: Final = _user_turn( diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index e5c99c932d6..323a29caa3a 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -270,9 +270,21 @@ def _responses_file(file_id: str) -> dict[str, object]: return {"input": [{"role": "user", "content": [{"type": "input_file", "file_id": file_id}]}]} +def _responses_file_url(file_id: str) -> dict[str, object]: + return {"input": [{"role": "user", "content": [{"type": "input_file", "file_url": file_id}]}]} + + +def _gemini_file_uri(file_id: str) -> dict[str, object]: + return { + "contents": [{"role": "user", "parts": [{"file_data": {"mime_type": "application/pdf", "file_uri": file_id}}]}] + } + + _S3_FILE_ID_SURFACES: Final = ( pytest.param("acompletion", "/v1/chat/completions", _chat_file, id="chat-completions"), pytest.param("aresponses", "/v1/responses", _responses_file, id="responses"), + pytest.param("aresponses", "/v1/responses", _responses_file_url, id="responses-file-url"), + pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_file_uri, id="gemini-file-uri"), ) _UPLOADED_OBJECT: Final = "s3://media-bucket/litellm-bedrock-files/upload.pdf" _UPLOADED_FILE_ID: Final = ( @@ -513,6 +525,10 @@ async def _send_s3_media( ) +class _StopAfterTheS3Check(Exception): + pass + + class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response( @@ -775,6 +791,34 @@ class TestProxyBaseLLMRequestProcessing: pytest.fail("litellm_call_id is not a valid UUID") assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("route_type", "refused"), [("call_mcp_tool", False), ("acompletion", True)], ids=["mcp-tool", "chat"] + ) + async def test_s3_media_check_covers_model_requests_and_leaves_mcp_tool_arguments_alone( + self, monkeypatch, route_type: Literal["call_mcp_tool", "acompletion"], refused: bool + ): + tool_argument: Final = {"type": "image_url", "image_url": {"url": "s3://other-bucket/cat.png"}} + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"name": "describe", "arguments": tool_argument}) + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + reached_request_setup: Final = AsyncMock(side_effect=_StopAfterTheS3Check) + monkeypatch.setattr( + litellm.proxy.common_request_processing, "add_litellm_data_to_request", reached_request_setup + ) + + with pytest.raises(HTTPException if refused else _StopAfterTheS3Check): + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth( + api_key="sk-test", user_id="user-a", user_role=LitellmUserRoles.INTERNAL_USER + ), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=MagicMock(spec=ProxyConfig), + route_type=route_type, + ) + @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", [{"bad": "value"}, ["gpt-5.2"], 1]) async def test_common_processing_pre_call_logic_rejects_a_non_string_model_with_400(