From a0c3855cc50c959b52c2299d1130013caaced58e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 18:50:13 -0800 Subject: [PATCH 01/24] fix(bedrock): send s3:// image, document and video URLs to Converse as S3 sources An s3:// URL in an image_url or file part now becomes a Converse s3Location source instead of failing with "Unsupported image type". The format comes from the part's format or the key's extension, an optional bucket_owner is sent as bucketOwner, and an unknown or unsupported format answers 400. An image in an assistant turn no longer crashes the async path. Backport to lts/1.108.x, adapted to the line. Cherry-picked from fe631f3fb23d2e87916bf927478a35b3b7971280 (litellm_bedrock_converse_s3_image_sources). --- .../prompt_templates/factory.py | 124 +++++++++-- litellm/proxy/_lazy_openapi_snapshot.json | 8 + .../transformation.py | 11 +- litellm/types/llms/bedrock.py | 10 +- litellm/types/llms/openai.py | 2 + ...llm_core_utils_prompt_templates_factory.py | 207 ++++++++++++++++-- .../test_litellm_completion_responses.py | 30 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 8 files changed, 355 insertions(+), 41 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 2613e0bbd28..e07a4394162 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -8,11 +8,13 @@ import re import xml.etree.ElementTree as ET from collections.abc import Container, Iterable, Iterator, Mapping, Sequence from enum import Enum +from pathlib import PurePosixPath from types import MappingProxyType -from typing import Any, Final, TypeAlias, TypedDict, cast, overload +from typing import Any, Final, TypeAlias, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, JsonValue, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm import litellm.types @@ -31,7 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ) from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client, header_value -from litellm.types.files import get_file_extension_from_mime_type +from litellm.types.files import get_file_extension_from_mime_type, get_file_mime_type_from_extension from litellm.types.llms.anthropic import * from litellm.types.llms.bedrock import CachePointBlock from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock @@ -3647,6 +3649,7 @@ from litellm.types.llms.bedrock import ( from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock from litellm.types.llms.bedrock import DocumentBlock as BedrockDocumentBlock from litellm.types.llms.bedrock import ImageBlock as BedrockImageBlock +from litellm.types.llms.bedrock import S3Location as BedrockS3Location from litellm.types.llms.bedrock import SourceBlock as BedrockSourceBlock from litellm.types.llms.bedrock import ToolBlock as BedrockToolBlock from litellm.types.llms.bedrock import ToolResultBlock as BedrockToolResultBlock @@ -3819,10 +3822,8 @@ class BedrockImageProcessor: return valid_extensions[0] @staticmethod - def _create_bedrock_block(image_bytes: str, mime_type: str, image_format: str) -> BedrockContentBlock: + def _create_bedrock_block(source: BedrockSourceBlock, mime_type: str, image_format: str) -> BedrockContentBlock: """Create appropriate Bedrock content block based on mime type.""" - _blob: Final = BedrockSourceBlock(bytes=image_bytes) - document_types: Final = ["application", "text"] is_document: Final = any(mime_type.startswith(doc_type) for doc_type in document_types) @@ -3832,12 +3833,10 @@ class BedrockImageProcessor: HASH_SAMPLE_BYTES: Final = 64 * 1024 # hash up to 64 KB of data if is_document: - # --- Prepare normalized bytes for hashing (without modifying original) --- - if isinstance(image_bytes, str): - # Remove whitespace/newlines so base64 variations hash identically - normalized = "".join(image_bytes.split()).encode("utf-8") - else: - normalized = image_bytes + s3_location: Final = source.get("s3Location") + hashed_content: Final = s3_location["uri"] if s3_location is not None else source.get("bytes") or "" + # Remove whitespace/newlines so base64 variations hash identically + normalized: Final = "".join(hashed_content.split()).encode("utf-8") # --- Use only the first 64 KB for speed --- if len(normalized) <= HASH_SAMPLE_BYTES: @@ -3857,20 +3856,63 @@ class BedrockImageProcessor: # --- Return content block --- return BedrockContentBlock( document=BedrockDocumentBlock( - source=_blob, + source=source, format=image_format, name=document_name, ) ) elif is_video: - return BedrockContentBlock(video=BedrockVideoBlock(source=_blob, format=image_format)) + return BedrockContentBlock(video=BedrockVideoBlock(source=source, format=image_format)) else: - return BedrockContentBlock(image=BedrockImageBlock(source=_blob, format=image_format)) + return BedrockContentBlock(image=BedrockImageBlock(source=source, format=image_format)) + + @staticmethod + def _s3_key_mime_type(uri: str) -> str | None: + extension: Final = PurePosixPath(uri).suffix.removeprefix(".").lower() + if not extension: + return None + bedrock_format: Final = "jpeg" if extension == "jpg" else extension + config: Final = litellm.AmazonConverseConfig() + if bedrock_format in config.get_supported_image_types(): + return f"image/{bedrock_format}" + if bedrock_format in config.get_supported_video_types(): + return f"video/{bedrock_format}" + try: + return get_file_mime_type_from_extension(extension) + except ValueError: + return None @classmethod - def process_image_sync(cls, image_url: str, format: str | None = None) -> BedrockContentBlock: + 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: + 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", + ) + try: + image_format: Final = cls.validate_format(mime_type, mime_type.partition("/")[2]) + except ValueError as e: + raise litellm.BadRequestError( + message=f"Bedrock Converse does not accept {uri} as {mime_type}. {e}", + model="", + llm_provider="bedrock", + ) from e + s3_location: Final = ( + BedrockS3Location(uri=uri) if bucket_owner is None else BedrockS3Location(uri=uri, bucketOwner=bucket_owner) + ) + return cls._create_bedrock_block(BedrockSourceBlock(s3Location=s3_location), mime_type, image_format) + + @classmethod + def process_image_sync( + cls, image_url: str, format: str | None = None, bucket_owner: str | None = None + ) -> BedrockContentBlock: """Synchronous image processing.""" + if image_url.startswith("s3://"): + return cls._create_s3_block(image_url, format, bucket_owner) if "base64" in image_url: img_bytes, mime_type, image_format = cls._parse_base64_image(image_url) elif "http://" in image_url or "https://" in image_url: @@ -3884,12 +3926,16 @@ class BedrockImageProcessor: image_format = mime_type.split("/")[1] image_format = cls.validate_format(mime_type, image_format) - return cls._create_bedrock_block(img_bytes, mime_type, image_format) + return cls._create_bedrock_block(BedrockSourceBlock(bytes=img_bytes), mime_type, image_format) @classmethod - async def process_image_async(cls, image_url: str, format: str | None) -> BedrockContentBlock: + async def process_image_async( + cls, image_url: str, format: str | None = None, bucket_owner: str | None = None + ) -> BedrockContentBlock: """Asynchronous image processing.""" + if image_url.startswith("s3://"): + return cls._create_s3_block(image_url, format, bucket_owner) if "base64" in image_url: img_bytes, mime_type, image_format = cls._parse_base64_image(image_url) elif "http://" in image_url or "https://" in image_url: @@ -3903,7 +3949,7 @@ class BedrockImageProcessor: image_format = mime_type.split("/")[1] image_format = cls.validate_format(mime_type, image_format) - return cls._create_bedrock_block(img_bytes, mime_type, image_format) + return cls._create_bedrock_block(BedrockSourceBlock(bytes=img_bytes), mime_type, image_format) def _convert_to_bedrock_tool_call_invoke( @@ -4045,6 +4091,18 @@ def _append_bedrock_tool_result_media_block( ) +class _S3MediaOwner(TypedDict, total=False): + bucket_owner: ReadOnly[object] + + +_S3_MEDIA_OWNER: Final = TypeAdapter(_S3MediaOwner) + + +def _s3_bucket_owner(media: _S3MediaOwner) -> str | None: + owner: Final = media.get("bucket_owner") + return owner if isinstance(owner, str) else None + + def _append_bedrock_tool_result_image_url_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], content: dict, @@ -4055,9 +4113,15 @@ def _append_bedrock_tool_result_image_url_block( format = content["image_url"].get("format") else: image_url = content["image_url"] + bucket_owner: Final = ( + _s3_bucket_owner(_S3_MEDIA_OWNER.validate_python(content["image_url"])) + if isinstance(content["image_url"], dict) + else None + ) processed_block: Final = BedrockImageProcessor.process_image_sync( image_url=image_url, format=format, + bucket_owner=bucket_owner, ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "image_url") @@ -4080,6 +4144,7 @@ def _append_bedrock_tool_result_file_block( processed_block: Final = BedrockImageProcessor.process_image_sync( image_url=cast(str, file_id or file_data), format=file_obj.get("format"), + bucket_owner=_s3_bucket_owner(_S3_MEDIA_OWNER.validate_python(file_obj)), ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "file") @@ -4699,8 +4764,13 @@ class BedrockConverseMessagesProcessor: format = element["image_url"].get("format") else: image_url = element["image_url"] + bucket_owner: str | None = ( + element["image_url"].get("bucket_owner") + if isinstance(element["image_url"], dict) + else None + ) _part = await BedrockImageProcessor.process_image_async( - image_url=image_url, format=format + image_url=image_url, format=format, bucket_owner=bucket_owner ) _parts.append(_part) elif element["type"] == "file": @@ -4960,7 +5030,9 @@ class BedrockConverseMessagesProcessor: llm_provider="bedrock", ) format: Final = file_message.get("format") - return BedrockImageProcessor.process_image_sync(image_url=cast(str, file_id or file_data), format=format) + return BedrockImageProcessor.process_image_sync( + image_url=cast(str, file_id or file_data), format=format, bucket_owner=file_message.get("bucket_owner") + ) _process_file_message = process_file_message @@ -4984,7 +5056,9 @@ class BedrockConverseMessagesProcessor: model="", llm_provider="bedrock", ) - return await BedrockImageProcessor.process_image_async(image_url=cast(str, file_id or file_data), format=format) + return await BedrockImageProcessor.process_image_async( + image_url=cast(str, file_id or file_data), format=format, bucket_owner=file_message.get("bucket_owner") + ) @staticmethod def _process_document_message(element: dict) -> BedrockContentBlock: @@ -5111,9 +5185,13 @@ def _bedrock_converse_messages_pt( format = element["image_url"].get("format") else: image_url = element["image_url"] + bucket_owner: str | None = ( + element["image_url"].get("bucket_owner") + if isinstance(element["image_url"], dict) + else None + ) _part = BedrockImageProcessor.process_image_sync( - image_url=image_url, - format=format, + image_url=image_url, format=format, bucket_owner=bucket_owner ) _parts.append(_part) elif element["type"] == "file": diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 467cb2d624c..dbd9eae0fbc 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -46591,6 +46591,10 @@ }, "ChatCompletionFileObjectFile": { "properties": { + "bucket_owner": { + "title": "Bucket Owner", + "type": "string" + }, "detail": { "title": "Detail", "type": "string" @@ -46697,6 +46701,10 @@ }, "ChatCompletionImageUrlObject": { "properties": { + "bucket_owner": { + "title": "Bucket Owner", + "type": "string" + }, "detail": { "title": "Detail", "type": "string" diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2e7fb6da194..13c9388cc21 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2257,6 +2257,8 @@ class LiteLLMCompletionResponsesConfig: file_dict["file_data"] = item["file_data"] if item.get("filename"): file_dict["filename"] = item["filename"] + if item.get("bucket_owner"): + file_dict["bucket_owner"] = item["bucket_owner"] new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: @@ -2270,8 +2272,13 @@ class LiteLLMCompletionResponsesConfig: """ Transform a Responses API input_image item to a Chat Completion image item """ - image_url_obj: Final = ChatCompletionImageUrlObject( - url=item.get("image_url") or "", detail=item.get("detail") or "auto" + url: Final = item.get("image_url") or "" + detail: Final = item.get("detail") or "auto" + bucket_owner: Final = item.get("bucket_owner") + image_url_obj: Final = ( + ChatCompletionImageUrlObject(url=url, detail=detail, bucket_owner=bucket_owner) + if bucket_owner + else ChatCompletionImageUrlObject(url=url, detail=detail) ) return ChatCompletionImageObject(type="image_url", image_url=image_url_obj) diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index ff4f2c4fd1e..5425391ee64 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -22,8 +22,14 @@ class SystemContentBlock(TypedDict, total=False): cachePoint: CachePointBlock -class SourceBlock(TypedDict): - bytes: str | None # base 64 encoded string +class S3Location(TypedDict): + uri: ReadOnly[str] + bucketOwner: NotRequired[ReadOnly[str]] + + +class SourceBlock(TypedDict, total=False): + bytes: ReadOnly[str | None] # base 64 encoded string + s3Location: ReadOnly[S3Location] BedrockImageTypes = Literal["png", "jpeg", "gif", "webp"] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8cc9e912e80..a1e425634cd 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -796,6 +796,7 @@ class ChatCompletionImageUrlObject(TypedDict, total=False): url: Required[str] detail: str format: str + bucket_owner: ReadOnly[str] class ChatCompletionImageObject(TypedDict): @@ -842,6 +843,7 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): format: str detail: str # For video/image resolution control (low, medium, high, ultra_high) video_metadata: dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) + bucket_owner: ReadOnly[str] class ChatCompletionFileObject(TypedDict): 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 ad216548d08..8ab8dd70910 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 @@ -68,6 +68,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock +from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock from litellm.types.llms.openai import AllMessageValues @@ -2162,10 +2163,10 @@ def test_bedrock_create_bedrock_block_deterministic_document_hash(): # Create two blocks with same content block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) # Both should have the same document name @@ -2190,10 +2191,10 @@ def test_bedrock_create_bedrock_block_different_content_different_hash(): # Create blocks block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content1, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content1}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content2, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content2}, mime_type="application/pdf", image_format="pdf" ) # Should have different document names @@ -2216,15 +2217,15 @@ def test_bedrock_create_bedrock_block_normalized_base64(): # Create blocks block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_with_newlines, + source={"bytes": base64_with_newlines}, mime_type="application/pdf", image_format="pdf", ) block3 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_with_spaces, + source={"bytes": base64_with_spaces}, mime_type="application/pdf", image_format="pdf", ) @@ -2251,10 +2252,10 @@ def test_bedrock_create_bedrock_block_large_file_sampling(): # Create blocks block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content1, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content1}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content2, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content2}, mime_type="application/pdf", image_format="pdf" ) # Should have different names because total length is different @@ -2273,7 +2274,7 @@ def test_bedrock_create_bedrock_block_very_large_file(): # Create block block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) # Should have a valid document name @@ -2295,7 +2296,7 @@ def test_bedrock_create_bedrock_block_image_type(): # Create block block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="image/png", image_format="png" + source={"bytes": base64_content}, mime_type="image/png", image_format="png" ) # Should be an image block, not a document block @@ -2316,7 +2317,7 @@ def test_bedrock_create_bedrock_block_video_type(): # Create block block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="video/mp4", image_format="mp4" + source={"bytes": base64_content}, mime_type="video/mp4", image_format="mp4" ) # Should be a video block, not a document block @@ -2337,7 +2338,7 @@ def test_bedrock_create_bedrock_block_document_name_format(): base64_content = base64.b64encode(pdf_content).decode("utf-8") block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) document_name = block["document"]["name"] @@ -2366,7 +2367,7 @@ def test_bedrock_create_bedrock_block_different_document_formats(): for content, mime_type, format_type in test_cases: base64_content = base64.b64encode(content + b"\x00" * 100).decode("utf-8") block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type=mime_type, image_format=format_type + source={"bytes": base64_content}, mime_type=mime_type, image_format=format_type ) assert block.get("document") is not None @@ -2375,6 +2376,184 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["format"] == format_type +_S3_MEDIA_MODEL: Final = "us.amazon.nova-2-lite-v1:0" +_S3_BUCKET_OWNER: Final = "111122223333" + + +async def _converse_user_content(part: Mapping[str, object], run_async: bool) -> Sequence[BedrockContentBlock]: + messages: Final = [{"role": "user", "content": [{"type": "text", "text": "Describe this."}, part]}] + with respx.mock: + converted: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + if run_async + else _bedrock_converse_messages_pt(messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock") + ) + return converted[0]["content"][1:] + + +def _s3_source(uri: str, bucket_owner: str | None = None) -> Mapping[str, object]: + return {"s3Location": {"uri": uri} if bucket_owner is None else {"uri": uri, "bucketOwner": bucket_owner}} + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("part", "expected"), + [ + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.png"}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png")}}, + id="image-format-from-extension", + ), + pytest.param( + {"type": "image_url", "image_url": "s3://media-bucket/dir/CAT.JPG"}, + {"image": {"format": "jpeg", "source": _s3_source("s3://media-bucket/dir/CAT.JPG")}}, + id="string-url-uppercase-extension", + ), + pytest.param( + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}, + }, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}}, + id="bucket-owner", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.png", "format": "image/webp"}}, + {"image": {"format": "webp", "source": _s3_source("s3://media-bucket/dir/cat.png")}}, + id="explicit-format-wins", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/base64/cat.png"}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/base64/cat.png")}}, + id="base64-in-key", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/clips/demo.mp4"}}, + {"video": {"format": "mp4", "source": _s3_source("s3://media-bucket/clips/demo.mp4")}}, + id="video", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/clips/demo.mov"}}, + {"video": {"format": "mov", "source": _s3_source("s3://media-bucket/clips/demo.mov")}}, + id="video-extension-is-the-bedrock-format", + ), + pytest.param( + {"type": "file", "file": {"file_id": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}}, + id="file-id", + ), + ], +) +async def test_bedrock_converse_s3_media_is_sent_as_an_s3_location_without_a_download( + part: Mapping[str, object], expected: Mapping[str, object], run_async: bool +): + assert await _converse_user_content(part, run_async) == [expected] + + +@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" + invoice: Final = "s3://media-bucket/docs/invoice.pdf" + + report_content: Final = await _converse_user_content({"type": "file", "file": {"file_id": report}}, run_async) + report_again: Final = await _converse_user_content({"type": "file", "file": {"file_id": report}}, run_async) + invoice_content: Final = await _converse_user_content({"type": "file", "file": {"file_id": invoice}}, run_async) + + report_name: Final = report_content[0]["document"]["name"] + assert report_content == [{"document": {"format": "pdf", "name": report_name, "source": _s3_source(report)}}] + assert re.fullmatch(r"DocumentPDFmessages_[0-9a-f]{16}_pdf", report_name) + assert report_again == report_content + assert invoice_content[0]["document"]["name"] != report_name + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("uri", ["s3://media-bucket/uploads/7f3a9c", "s3://media-bucket/uploads/blob.unknownext"]) +async def test_bedrock_converse_s3_media_without_a_known_extension_or_format_is_a_bad_request( + uri: str, run_async: bool +): + with pytest.raises(litellm.BadRequestError) as excinfo: + await _converse_user_content({"type": "image_url", "image_url": {"url": uri}}, run_async) + + assert uri in excinfo.value.message + assert "`format`" in excinfo.value.message + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("image_url", "named_format"), + [ + pytest.param({"url": "s3://media-bucket/scans/page.tiff"}, "page.tiff", id="tiff-key"), + pytest.param( + {"url": "s3://media-bucket/scans/page.png", "format": "image/tiff"}, "image/tiff", id="tiff-format" + ), + ], +) +async def test_bedrock_converse_s3_media_in_a_format_bedrock_does_not_take_is_a_bad_request( + image_url: Mapping[str, str], named_format: str, run_async: bool +): + with pytest.raises(litellm.BadRequestError) as excinfo: + await _converse_user_content({"type": "image_url", "image_url": image_url}, run_async) + + assert image_url["url"] in excinfo.value.message + assert named_format in excinfo.value.message + + +@pytest.mark.parametrize( + "part", + [ + pytest.param( + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}, + }, + id="image-url", + ), + pytest.param( + {"type": "file", "file": {"file_id": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}}, + id="file", + ), + ], +) +def test_bedrock_converse_tool_result_s3_media_keeps_its_bucket_owner(part: Mapping[str, object]): + messages: Final = [ + {"role": "user", "content": "Describe the screenshot."}, + {"tool_call_id": "tooluse_s3_1", "role": "tool", "name": "fetch_screenshot", "content": [part]}, + ] + + with respx.mock: + converted: Final = _bedrock_converse_messages_pt( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + + assert converted[-1]["content"][-1]["toolResult"]["content"] == [ + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}} + ] + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +async def test_bedrock_converse_assistant_s3_image_is_sent_as_an_s3_location(run_async: bool): + messages: Final = [ + {"role": "user", "content": "Draw a cat."}, + {"role": "assistant", "content": [{"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.png"}}]}, + {"role": "user", "content": "Make it bigger."}, + ] + + with respx.mock: + converted: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + if run_async + else _bedrock_converse_messages_pt(messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock") + ) + + assert converted[1]["content"] == [ + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png")}} + ] + + def test_bedrock_nova_web_search_options_mapping(): """ Test that web_search_options is correctly mapped to Nova grounding. diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 79d5ec5d877..5fb908ac42f 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -263,6 +263,36 @@ class TestLiteLLMCompletionResponsesConfig: assert "extra_field" not in result assert "another_field" not in result + @pytest.mark.parametrize( + ("item", "expected"), + [ + pytest.param( + {"type": "input_image", "image_url": "s3://media-bucket/dir/cat.png", "bucket_owner": "111122223333"}, + { + "type": "image_url", + "image_url": { + "url": "s3://media-bucket/dir/cat.png", + "detail": "auto", + "bucket_owner": "111122223333", + }, + }, + id="input-image", + ), + pytest.param( + {"type": "input_file", "file_id": "s3://media-bucket/docs/report.pdf", "bucket_owner": "111122223333"}, + { + "type": "file", + "file": {"file_id": "s3://media-bucket/docs/report.pdf", "bucket_owner": "111122223333"}, + }, + id="input-file", + ), + ], + ) + def test_transform_input_media_item_keeps_bucket_owner(self, item: Mapping[str, object], expected: object): + assert LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([item]) == [ + expected + ] + def test_transform_chat_completion_response_with_reasoning_content(self): """Test that reasoning content is preserved in the full transformation pipeline""" # Setup diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 89cc9157359..6d7f0f36014 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28710,6 +28710,8 @@ export interface components { }; /** ChatCompletionFileObjectFile */ ChatCompletionFileObjectFile: { + /** Bucket Owner */ + bucket_owner?: string; /** Detail */ detail?: string; /** File Data */ @@ -28751,6 +28753,8 @@ export interface components { }; /** ChatCompletionImageUrlObject */ ChatCompletionImageUrlObject: { + /** Bucket Owner */ + bucket_owner?: string; /** Detail */ detail?: string; /** Format */ From b5c98922eba75fc357f1add1f0e48074d53e9407 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 01:20:46 -0800 Subject: [PATCH 02/24] fix(proxy): limit non-admin s3:// media to allowed buckets and refuse managed prefixes A non-admin key may only send an s3:// media URL whose bucket is listed in general_settings.allowed_client_side_s3_buckets (or when client-side credentials are allowed), and never one under litellm's managed prefixes. Proxy admins are not limited. A bucket_owner on an s3:// URL that is not a non-empty string is a 400, the Responses bridge forwards the owner only for s3:// URLs, audio on S3 is a 400, and video in a tool result is sent to Bedrock. Backport to lts/1.108.x, adapted to the line. Cherry-picked from 2c7b965ff4db4dd7e64b67107c704ea019e44009 (litellm_bedrock_converse_s3_image_sources). --- .../cloud_storage_security.py | 72 +++++++- litellm/litellm_core_utils/core_helpers.py | 6 +- .../prompt_templates/factory.py | 108 +++++++----- litellm/proxy/_lazy_openapi_snapshot.json | 8 - litellm/proxy/_types.py | 4 + litellm/proxy/auth/auth_utils.py | 11 +- litellm/proxy/common_request_processing.py | 37 +++- .../transformation.py | 14 +- litellm/types/llms/bedrock.py | 7 +- litellm/types/llms/openai.py | 2 - ...llm_core_utils_prompt_templates_factory.py | 92 +++++++++- .../test_cloud_storage_security.py | 165 ++++++++++++++++++ .../chat/test_converse_transformation.py | 37 ++++ .../test_generic_guardrail_api.py | 46 +++++ .../proxy/test_common_request_processing.py | 127 +++++++++++++- ...test_proxy_server_endpoints_and_startup.py | 24 +++ .../test_litellm_completion_responses.py | 45 +++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 +- 18 files changed, 742 insertions(+), 72 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index c913e4de02e..88ddb14f498 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,11 +1,13 @@ import posixpath import re -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import accumulate, chain, repeat from types import MappingProxyType from typing import Any, Final, cast from urllib.parse import quote, unquote from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping VERTEX_AI_MANAGED_GCS_PREFIX: Final = "litellm-vertex-files/" BEDROCK_MANAGED_S3_BATCH_PREFIX: Final = "litellm-bedrock-files-" @@ -18,6 +20,15 @@ BEDROCK_MANAGED_S3_PREFIXES: Final = ( ) MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://") _MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({})) +MAX_FILE_ID_DECODE_PASSES: Final = 8 + + +def _unquote_once(current: str, _: str) -> str: + return unquote(current) + + +def _decodings(value: str) -> tuple[str, ...]: + return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), _unquote_once)) def is_managed_cloud_storage_uri(file_id: str) -> bool: @@ -31,6 +42,65 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: return isinstance(file_id, str) and file_id.startswith(MANAGED_CLOUD_STORAGE_SCHEMES) +_S3_SCHEME: Final = "s3://" +_FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3: Final = ("s", "%") +_CLIENT_MEDIA_FIELDS: Final = ("messages", "input", "system") +_MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data"}) + + +def _media_uris(value: object, key: str | None = None) -> Iterator[str]: + if isinstance(value, str): + return iter((value,) if key in _MEDIA_URI_KEYS else ()) + fields: Final = as_str_mapping(value) + if fields is not None: + return chain.from_iterable(_media_uris(child, child_key) for child_key, child in fields.items()) + items: Final = as_object_list(value) + return chain.from_iterable(map(_media_uris, items)) if items is not None else iter(()) + + +def _s3_bucket(uri: str) -> str | None: + return uri.removeprefix(_S3_SCHEME).partition("/")[0] if uri.startswith(_S3_SCHEME) else None + + +def _s3_object_key(uri: str) -> str | None: + if not uri.startswith(_S3_SCHEME): + return None + _, separator, object_key = uri.removeprefix(_S3_SCHEME).partition("/") + return object_key if separator else None + + +def _outside_managed_prefixes(object_key: str) -> bool: + return not any(prefix in object_key for prefix in BEDROCK_MANAGED_S3_PREFIXES) + + +def _refused_in_some_decoding( + uri: str, s3_part: Callable[[str], str | None], is_allowed: Callable[[str], bool] +) -> bool: + decodings: Final = _decodings(uri) + parts: Final = tuple(part for part in map(s3_part, decodings) if part is not None) + settled: Final = decodings[-1] == decodings[-2] + return bool(parts) and (not settled or not all(map(is_allowed, parts))) + + +def _is_s3_uri(uri: str) -> bool: + return uri.startswith(_FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3) and any( + decoding.startswith(_S3_SCHEME) for decoding in _decodings(uri) + ) + + +def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: + media_uris: Final = chain.from_iterable(_media_uris(request_body.get(field)) for field in _CLIENT_MEDIA_FIELDS) + return tuple(filter(_is_s3_uri, media_uris)) + + +def is_managed_s3_uri(uri: str) -> bool: + return _refused_in_some_decoding(uri, _s3_object_key, _outside_managed_prefixes) + + +def names_an_unlisted_s3_bucket(uri: str, is_listed: Callable[[str], bool]) -> bool: + return _refused_in_some_decoding(uri, _s3_bucket, is_listed) + + _SAFE_OBJECT_COMPONENT_PATTERN: Final = re.compile(r"[^A-Za-z0-9._-]+") diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index e706bcadef8..9e96da43b00 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -3,7 +3,7 @@ import copy import logging import re -from collections.abc import Collection, Iterable, Iterator, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable @@ -56,6 +56,10 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def as_object_list(value: object) -> Sequence[object] | None: + return value if isinstance(value, list) else None + + @runtime_checkable class _SupportsIter(Protocol): def __iter__(self) -> Iterator[object]: ... diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index e07a4394162..0e5393d1e37 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -10,11 +10,10 @@ from collections.abc import Container, Iterable, Iterator, Mapping, Sequence from enum import Enum from pathlib import PurePosixPath from types import MappingProxyType -from typing import Any, Final, TypeAlias, cast, overload +from typing import Any, Final, TypeAlias, TypedDict, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, JsonValue, TypeAdapter -from typing_extensions import ReadOnly, TypedDict import litellm import litellm.types @@ -22,6 +21,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -3871,12 +3871,11 @@ class BedrockImageProcessor: extension: Final = PurePosixPath(uri).suffix.removeprefix(".").lower() if not extension: return None - bedrock_format: Final = "jpeg" if extension == "jpg" else extension config: Final = litellm.AmazonConverseConfig() - if bedrock_format in config.get_supported_image_types(): - return f"image/{bedrock_format}" - if bedrock_format in config.get_supported_video_types(): - return f"video/{bedrock_format}" + if extension in config.get_supported_image_types(): + return f"image/{extension}" + if extension in config.get_supported_video_types(): + return f"video/{extension}" try: return get_file_mime_type_from_extension(extension) except ValueError: @@ -3892,6 +3891,13 @@ class BedrockImageProcessor: model="", llm_provider="bedrock", ) + if mime_type.startswith("audio/"): + raise litellm.BadRequestError( + message=f"Bedrock Converse does not accept {uri} as {mime_type}. " + "Audio is not one of its image, document, or video formats.", + model="", + llm_provider="bedrock", + ) try: image_format: Final = cls.validate_format(mime_type, mime_type.partition("/")[2]) except ValueError as e: @@ -4075,13 +4081,15 @@ def _convert_to_bedrock_tool_call_invoke( def _append_bedrock_tool_result_media_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], processed_block: BedrockContentBlock, - content: dict, + content: Mapping[str, object], content_type: str, ) -> None: if "image" in processed_block: tool_result_content_blocks.append(BedrockToolResultContentBlock(image=processed_block["image"])) elif "document" in processed_block: tool_result_content_blocks.append(BedrockToolResultContentBlock(document=processed_block["document"])) + elif "video" in processed_block: + tool_result_content_blocks.append(BedrockToolResultContentBlock(video=processed_block["video"])) else: verbose_logger.warning( "Bedrock Converse: unrecognized BedrockContentBlock keys %s for %s tool-result block %s; dropping.", @@ -4091,21 +4099,23 @@ def _append_bedrock_tool_result_media_block( ) -class _S3MediaOwner(TypedDict, total=False): - bucket_owner: ReadOnly[object] - - -_S3_MEDIA_OWNER: Final = TypeAdapter(_S3MediaOwner) - - -def _s3_bucket_owner(media: _S3MediaOwner) -> str | None: - owner: Final = media.get("bucket_owner") - return owner if isinstance(owner, str) else None +def _s3_bucket_owner(media: object, url: str) -> str | None: + fields: Final = as_str_mapping(media) + owner: Final = fields.get("bucket_owner") if fields is not None else None + if owner is None or (isinstance(owner, str) and owner): + return owner + if not url.startswith("s3://"): + return None + raise litellm.BadRequestError( + message=f"`bucket_owner` for {url} must be the AWS account id as a non-empty string, not {owner!r}.", + model="", + llm_provider="bedrock", + ) def _append_bedrock_tool_result_image_url_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], - content: dict, + content: ChatCompletionImageObject, ) -> None: format: str | None = None if isinstance(content["image_url"], dict): @@ -4113,26 +4123,21 @@ def _append_bedrock_tool_result_image_url_block( format = content["image_url"].get("format") else: image_url = content["image_url"] - bucket_owner: Final = ( - _s3_bucket_owner(_S3_MEDIA_OWNER.validate_python(content["image_url"])) - if isinstance(content["image_url"], dict) - else None - ) processed_block: Final = BedrockImageProcessor.process_image_sync( image_url=image_url, format=format, - bucket_owner=bucket_owner, + bucket_owner=_s3_bucket_owner(content["image_url"], image_url), ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "image_url") def _append_bedrock_tool_result_file_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], - content: dict, + content: ChatCompletionFileObject, ) -> None: # Match the user-message path (_process_file_message): accept either # file_data (base64 data URI) or file_id (server-side reference / URL). - file_obj: Final = content.get("file") or {} + file_obj: Final = content.get("file") or ChatCompletionFileObjectFile() file_data: Final = file_obj.get("file_data") file_id: Final = file_obj.get("file_id") if file_data is None and file_id is None: @@ -4141,10 +4146,11 @@ def _append_bedrock_tool_result_file_block( model="", llm_provider="bedrock", ) + url: Final = cast(str, file_id or file_data) processed_block: Final = BedrockImageProcessor.process_image_sync( - image_url=cast(str, file_id or file_data), + image_url=url, format=file_obj.get("format"), - bucket_owner=_s3_bucket_owner(_S3_MEDIA_OWNER.validate_python(file_obj)), + bucket_owner=_s3_bucket_owner(file_obj, url), ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "file") @@ -4764,13 +4770,10 @@ class BedrockConverseMessagesProcessor: format = element["image_url"].get("format") else: image_url = element["image_url"] - bucket_owner: str | None = ( - element["image_url"].get("bucket_owner") - if isinstance(element["image_url"], dict) - else None - ) _part = await BedrockImageProcessor.process_image_async( - image_url=image_url, format=format, bucket_owner=bucket_owner + image_url=image_url, + format=format, + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), ) _parts.append(_part) elif element["type"] == "file": @@ -4913,7 +4916,15 @@ class BedrockConverseMessagesProcessor: image_url = element["image_url"]["url"] else: image_url = element["image_url"] - assistants_part = await BedrockImageProcessor.process_image_async(image_url=image_url) + assistants_part = await BedrockImageProcessor.process_image_async( + image_url=image_url, + format=( + element["image_url"].get("format") + if isinstance(element["image_url"], dict) + else None + ), + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), + ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( @@ -5030,8 +5041,9 @@ class BedrockConverseMessagesProcessor: llm_provider="bedrock", ) format: Final = file_message.get("format") + url: Final = cast(str, file_id or file_data) return BedrockImageProcessor.process_image_sync( - image_url=cast(str, file_id or file_data), format=format, bucket_owner=file_message.get("bucket_owner") + image_url=url, format=format, bucket_owner=_s3_bucket_owner(file_message, url) ) _process_file_message = process_file_message @@ -5056,8 +5068,9 @@ class BedrockConverseMessagesProcessor: model="", llm_provider="bedrock", ) + url: Final = cast(str, file_id or file_data) return await BedrockImageProcessor.process_image_async( - image_url=cast(str, file_id or file_data), format=format, bucket_owner=file_message.get("bucket_owner") + image_url=url, format=format, bucket_owner=_s3_bucket_owner(file_message, url) ) @staticmethod @@ -5185,13 +5198,10 @@ def _bedrock_converse_messages_pt( format = element["image_url"].get("format") else: image_url = element["image_url"] - bucket_owner: str | None = ( - element["image_url"].get("bucket_owner") - if isinstance(element["image_url"], dict) - else None - ) _part = BedrockImageProcessor.process_image_sync( - image_url=image_url, format=format, bucket_owner=bucket_owner + image_url=image_url, + format=format, + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), ) _parts.append(_part) elif element["type"] == "file": @@ -5337,7 +5347,15 @@ def _bedrock_converse_messages_pt( image_url = element["image_url"]["url"] else: image_url = element["image_url"] - assistants_part = BedrockImageProcessor.process_image_sync(image_url=image_url) + assistants_part = BedrockImageProcessor.process_image_sync( + image_url=image_url, + format=( + element["image_url"].get("format") + if isinstance(element["image_url"], dict) + else None + ), + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), + ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index dbd9eae0fbc..467cb2d624c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -46591,10 +46591,6 @@ }, "ChatCompletionFileObjectFile": { "properties": { - "bucket_owner": { - "title": "Bucket Owner", - "type": "string" - }, "detail": { "title": "Detail", "type": "string" @@ -46701,10 +46697,6 @@ }, "ChatCompletionImageUrlObject": { "properties": { - "bucket_owner": { - "title": "Bucket Owner", - "type": "string" - }, "detail": { "title": "Detail", "type": "string" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5e812eec477..f29dab18e33 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2736,6 +2736,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): description="How long a blocked source address, or source address and username, stays blocked. Every attempt from a blocked key, right or wrong, is refused with 429 before the password is checked; the block is not extended by refused attempts. Set under `general_settings` in config.yaml. Defaults to 300", ) allowed_routes: list | None = Field(None, description="Proxy API Endpoints you want users to be able to access") + allowed_client_side_s3_buckets: tuple[str, ...] | None = Field( + None, + description="S3 bucket names that keys other than proxy admins may name in `s3://` image and file URLs. Any other bucket is rejected unless client-side credentials are allowed", + ) reject_clientside_metadata_tags: bool | None = Field( None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 98bb0985f4a..e058091b69b 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -9,7 +9,7 @@ from types import MappingProxyType from typing import Any, Final, Protocol from fastapi import HTTPException, Request, status -from pydantic import PositiveInt, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, PositiveInt, TypeAdapter, ValidationError import litellm from litellm import Router, constants, provider_list @@ -462,6 +462,15 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( ) +class ClientSideStorageAllowlist(BaseModel): + model_config = ConfigDict(frozen=True) + + allowed_client_side_s3_buckets: tuple[str, ...] | None = None + + def admits_s3_bucket(self, bucket: str) -> bool: + return bucket in (self.allowed_client_side_s3_buckets or ()) + + def _check_banned_params( body: Mapping[str, object], general_settings: dict, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c6562dc2273..b148ff1bc2c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -58,6 +58,11 @@ from litellm.litellm_core_utils.bug_report import ( should_report_bug, strip_bug_report_notice, ) +from litellm.litellm_core_utils.cloud_storage_security import ( + is_managed_s3_uri, + names_an_unlisted_s3_bucket, + s3_media_uris, +) from litellm.litellm_core_utils.core_helpers import ( get_or_create_metadata_bucket, independent_snapshot, @@ -87,14 +92,14 @@ from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields -from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, is_model_cost_zero, request_skips_budget_checks, tag_max_budget_check_for_tags, ) -from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route +from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -855,6 +860,33 @@ def _request_model(data: Mapping[str, object]) -> str | list[str] | None: return None +def _refuse_s3_media_the_key_may_not_read( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], +) -> None: + s3_uris: Final = s3_media_uris(data) + if not s3_uris or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + managed_uris: Final = tuple(filter(is_managed_s3_uri, s3_uris)) + if managed_uris: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " + f"({', '.join(managed_uris)}). Use the LiteLLM managed file id returned when the file was created.", + ) + if general_settings.get("allow_client_side_credentials") is True: + return + is_listed: Final = ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket + unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed)) + if unlisted_uris: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Raw s3:// URIs can only name a bucket listed in `general_settings.allowed_client_side_s3_buckets` " + f"({', '.join(unlisted_uris)}). Ask the proxy admin to list the bucket there.", + ) + + async def _enforce_guardrail_added_tag_budgets( data: Mapping[str, object], tags_before_guardrails: frozenset[str], @@ -2277,6 +2309,7 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) + _refuse_s3_media_the_key_may_not_read(self.data, user_api_key_dict, general_settings) self.data = await add_litellm_data_to_request( data=self.data, request=request, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 13c9388cc21..965bdd29b21 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -62,6 +62,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ) from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex from litellm.types.llms.base import CachedTokensDetails +from litellm.types.llms.bedrock import S3MediaImageUrlObject from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -226,6 +227,10 @@ _TEXT_ADAPTER: Final = TypeAdapter(str) _RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) +def _is_s3_url(url: object) -> bool: + return isinstance(url, str) and url.startswith("s3://") + + @runtime_checkable class _SupportsIter(Protocol): def __iter__(self) -> Iterator[object]: ... @@ -2257,8 +2262,9 @@ class LiteLLMCompletionResponsesConfig: file_dict["file_data"] = item["file_data"] if item.get("filename"): file_dict["filename"] = item["filename"] - if item.get("bucket_owner"): - file_dict["bucket_owner"] = item["bucket_owner"] + bucket_owner: Final = item.get("bucket_owner") + if bucket_owner is not None and _is_s3_url(file_id or item.get("file_data")): + file_dict["bucket_owner"] = bucket_owner new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: @@ -2276,8 +2282,8 @@ class LiteLLMCompletionResponsesConfig: detail: Final = item.get("detail") or "auto" bucket_owner: Final = item.get("bucket_owner") image_url_obj: Final = ( - ChatCompletionImageUrlObject(url=url, detail=detail, bucket_owner=bucket_owner) - if bucket_owner + S3MediaImageUrlObject(url=url, detail=detail, bucket_owner=bucket_owner) + if bucket_owner is not None and _is_s3_url(url) else ChatCompletionImageUrlObject(url=url, detail=detail) ) diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5425391ee64..4a876dabc17 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -9,7 +9,7 @@ from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, overri from litellm.types.llms.base import LiteLLMBaseModel -from .openai import ChatCompletionToolCallChunk +from .openai import ChatCompletionImageUrlObject, ChatCompletionToolCallChunk class CachePointBlock(TypedDict, total=False): @@ -32,6 +32,10 @@ class SourceBlock(TypedDict, total=False): s3Location: ReadOnly[S3Location] +class S3MediaImageUrlObject(ChatCompletionImageUrlObject, total=False): + bucket_owner: ReadOnly[object] + + BedrockImageTypes = Literal["png", "jpeg", "gif", "webp"] @@ -74,6 +78,7 @@ class SearchResultBlock(TypedDict, total=False): class ToolResultContentBlock(TypedDict, total=False): image: ImageBlock document: DocumentBlock + video: ReadOnly[VideoBlock] searchResult: SearchResultBlock json: dict text: str diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index a1e425634cd..8cc9e912e80 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -796,7 +796,6 @@ class ChatCompletionImageUrlObject(TypedDict, total=False): url: Required[str] detail: str format: str - bucket_owner: ReadOnly[str] class ChatCompletionImageObject(TypedDict): @@ -843,7 +842,6 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): format: str detail: str # For video/image resolution control (low, medium, high, ultra_high) video_metadata: dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) - bucket_owner: ReadOnly[str] class ChatCompletionFileObject(TypedDict): 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 8ab8dd70910..5035900e2bb 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 @@ -2411,6 +2411,21 @@ def _s3_source(uri: str, bucket_owner: str | None = None) -> Mapping[str, object {"image": {"format": "jpeg", "source": _s3_source("s3://media-bucket/dir/CAT.JPG")}}, id="string-url-uppercase-extension", ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.jpg"}}, + {"image": {"format": "jpeg", "source": _s3_source("s3://media-bucket/dir/cat.jpg")}}, + id="jpg-is-jpeg", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/CAT.PNG"}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/CAT.PNG")}}, + id="uppercase-image-extension", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/clips/DEMO.MOV"}}, + {"video": {"format": "mov", "source": _s3_source("s3://media-bucket/clips/DEMO.MOV")}}, + id="uppercase-video-extension", + ), pytest.param( { "type": "image_url", @@ -2488,6 +2503,10 @@ async def test_bedrock_converse_s3_media_without_a_known_extension_or_format_is_ pytest.param( {"url": "s3://media-bucket/scans/page.png", "format": "image/tiff"}, "image/tiff", id="tiff-format" ), + pytest.param({"url": "s3://media-bucket/voice/note.mp3"}, "audio/mpeg", id="mp3-key"), + pytest.param( + {"url": "s3://media-bucket/voice/note.bin", "format": "audio/mp4"}, "audio/mp4", id="audio-format" + ), ], ) async def test_bedrock_converse_s3_media_in_a_format_bedrock_does_not_take_is_a_bad_request( @@ -2500,6 +2519,42 @@ async def test_bedrock_converse_s3_media_in_a_format_bedrock_does_not_take_is_a_ assert named_format in excinfo.value.message +_PNG_BASE64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" + + +def _owned_image_url_part(url: str, bucket_owner: object) -> Mapping[str, object]: + return {"type": "image_url", "image_url": {"url": url, "bucket_owner": bucket_owner}} + + +def _owned_file_part(url: str, bucket_owner: object) -> Mapping[str, object]: + return {"type": "file", "file": {"file_id": url, "bucket_owner": bucket_owner}} + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("bucket_owner", [111122223333, ["111122223333"], ""], ids=["int", "list", "empty-string"]) +@pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) +async def test_bedrock_converse_s3_media_with_a_bucket_owner_that_is_not_an_account_id_is_a_bad_request( + part_for: Callable[[str, object], Mapping[str, object]], bucket_owner: object, run_async: bool +): + with pytest.raises(litellm.BadRequestError) as excinfo: + await _converse_user_content(part_for("s3://media-bucket/dir/cat.png", bucket_owner), run_async) + + assert "`bucket_owner`" in excinfo.value.message + assert "s3://media-bucket/dir/cat.png" in excinfo.value.message + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) +async def test_bedrock_converse_ignores_a_bucket_owner_on_media_that_is_not_on_s3( + part_for: Callable[[str, object], Mapping[str, object]], run_async: bool +): + content: Final = await _converse_user_content( + part_for(f"data:image/png;base64,{_PNG_BASE64}", 111122223333), run_async + ) + + assert content == [{"image": {"format": "png", "source": {"bytes": _PNG_BASE64}}}] + + @pytest.mark.parametrize( "part", [ @@ -2532,11 +2587,44 @@ def test_bedrock_converse_tool_result_s3_media_keeps_its_bucket_owner(part: Mapp ] +@pytest.mark.parametrize("bucket_owner", [111122223333, ""], ids=["int", "empty-string"]) +@pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) +def test_bedrock_converse_tool_result_s3_media_with_a_bucket_owner_that_is_not_an_account_id_is_a_bad_request( + part_for: Callable[[str, object], Mapping[str, object]], bucket_owner: object +): + messages: Final = [ + {"role": "user", "content": "Describe the screenshot."}, + { + "tool_call_id": "tooluse_s3_1", + "role": "tool", + "name": "fetch_screenshot", + "content": [part_for("s3://media-bucket/dir/cat.png", bucket_owner)], + }, + ] + + with respx.mock, pytest.raises(litellm.BadRequestError) as excinfo: + _bedrock_converse_messages_pt(messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock") + + assert "`bucket_owner`" in excinfo.value.message + + @pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) async def test_bedrock_converse_assistant_s3_image_is_sent_as_an_s3_location(run_async: bool): messages: Final = [ {"role": "user", "content": "Draw a cat."}, - {"role": "assistant", "content": [{"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.png"}}]}, + { + "role": "assistant", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "s3://media-bucket/dir/cat.png", + "format": "image/webp", + "bucket_owner": _S3_BUCKET_OWNER, + }, + } + ], + }, {"role": "user", "content": "Make it bigger."}, ] @@ -2550,7 +2638,7 @@ async def test_bedrock_converse_assistant_s3_image_is_sent_as_an_s3_location(run ) assert converted[1]["content"] == [ - {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png")}} + {"image": {"format": "webp", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}} ] 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 c3a2511a263..ea4953cb7c8 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -1,5 +1,15 @@ +from functools import reduce +from typing import Final +from urllib.parse import quote + +import pytest + from litellm.litellm_core_utils.cloud_storage_security import ( + MAX_FILE_ID_DECODE_PASSES, is_managed_cloud_storage_uri, + is_managed_s3_uri, + names_an_unlisted_s3_bucket, + s3_media_uris, ) @@ -13,3 +23,158 @@ def test_is_managed_cloud_storage_uri_ignores_provider_and_unified_ids(): assert not is_managed_cloud_storage_uri("file-abc123") assert not is_managed_cloud_storage_uri("bGl0ZWxsbV9wcm94eQ==") assert not is_managed_cloud_storage_uri("") + + +def _quote_once(current: str, _: int) -> str: + return quote(current, safe="") + + +def _quoted_times(value: str, times: int) -> str: + return reduce(_quote_once, range(times), value) + + +_MANAGED_UPLOAD: Final = "s3://media-bucket/litellm-bedrock-files/other-team.pdf" + + +def _user_turn(*parts: object) -> dict[str, object]: + return {"messages": [{"role": "user", "content": list(parts)}]} + + +def _responses_turn(*parts: object) -> dict[str, object]: + return {"input": [{"role": "user", "content": list(parts)}]} + + +@pytest.mark.parametrize( + "request_body", + ( + pytest.param(_user_turn({"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}), id="chat-image-url"), + pytest.param(_user_turn({"type": "image_url", "image_url": _MANAGED_UPLOAD}), id="chat-image-url-string"), + pytest.param(_user_turn({"type": "file", "file": {"file_id": _MANAGED_UPLOAD}}), id="chat-file-id"), + pytest.param(_user_turn({"type": "file", "file": {"file_data": _MANAGED_UPLOAD}}), id="chat-file-data"), + pytest.param( + { + "messages": [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}], + } + ] + }, + id="chat-tool-result", + ), + pytest.param(_responses_turn({"type": "input_image", "image_url": _MANAGED_UPLOAD}), id="responses-image"), + pytest.param(_responses_turn({"type": "input_file", "file_url": _MANAGED_UPLOAD}), id="responses-file-url"), + pytest.param(_responses_turn({"type": "input_file", "file_id": _MANAGED_UPLOAD}), id="responses-file-id"), + pytest.param( + { + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [{"type": "input_image", "image_url": _MANAGED_UPLOAD}], + } + ] + }, + id="responses-function-call-output", + ), + pytest.param( + _user_turn({"type": "image", "source": {"type": "url", "url": _MANAGED_UPLOAD}}), id="anthropic-image" + ), + pytest.param( + _user_turn({"type": "document", "source": {"type": "url", "url": _MANAGED_UPLOAD}}), + id="anthropic-document", + ), + pytest.param( + {"system": [{"type": "image", "source": {"type": "url", "url": _MANAGED_UPLOAD}}], "messages": []}, + id="anthropic-system", + ), + ), +) +def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]): + assert s3_media_uris(request_body) == (_MANAGED_UPLOAD,) + + +@pytest.mark.parametrize( + "uri", + ( + "s3://media-bucket/litellm-bedrock-files/x.pdf", + "s3://media-bucket/litellm-bedrock-files-batch-1/input.jsonl", + "s3://media-bucket/litellm-batch-outputs/x.jsonl.out", + "s3://media-bucket/tenant-a/litellm-bedrock-files/x.pdf", + "s3://media-bucket//litellm-bedrock-files/x.pdf", + "s3://media-bucket/./litellm-bedrock-files/x.pdf", + "s3://media-bucket/litellm%2Dbedrock-files/x.pdf", + "s3://media-bucket/litellm%252Dbedrock-files/x.pdf", + "s3://media-bucket%2Flitellm-bedrock-files/x.pdf", + "s3%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf", + "%73%33%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf", + "s3:///litellm-bedrock-files/x.pdf", + ), +) +def test_is_managed_s3_uri_sees_a_managed_prefix_anywhere_in_the_decoded_key(uri: str): + assert s3_media_uris(_user_turn({"type": "file", "file": {"file_id": uri}})) == (uri,) + assert is_managed_s3_uri(uri) + + +def test_is_managed_s3_uri_fails_closed_on_a_key_encoded_past_the_pass_cap(): + uri: Final = "s3://media-bucket/" + _quoted_times("litellm-bedrock-files/x.pdf", MAX_FILE_ID_DECODE_PASSES + 2) + + assert s3_media_uris(_user_turn({"type": "image_url", "image_url": {"url": uri}})) == (uri,) + assert is_managed_s3_uri(uri) + + +@pytest.mark.parametrize( + "uri", + ( + pytest.param("s3://media-bucket/dir/cat.png", id="unmanaged-key"), + pytest.param("s3://litellm-bedrock-files-bucket/cat.png", id="managed-prefix-in-bucket-name-only"), + pytest.param("s3://media-bucket", id="no-key"), + ), +) +def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): + assert not is_managed_s3_uri(uri) + + +@pytest.mark.parametrize( + "request_body", + ( + pytest.param( + _user_turn({"type": "image_url", "image_url": {"url": "https://example.test/litellm-bedrock-files/x.png"}}), + id="https-url", + ), + pytest.param( + _user_turn({"type": "image_url", "image_url": {"url": "data:image/png;base64,czM6Ly9idWNrZXQva2V5"}}), + id="data-url", + ), + pytest.param(_user_turn({"type": "text", "text": _MANAGED_UPLOAD}), id="text-part"), + pytest.param({"messages": [{"role": "user", "content": _MANAGED_UPLOAD}]}, id="string-content"), + pytest.param({"input": _MANAGED_UPLOAD}, id="responses-string-input"), + pytest.param({"system": _MANAGED_UPLOAD, "messages": []}, id="string-system"), + pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="non-media-field"), + ), +) +def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_body: dict[str, object]): + assert s3_media_uris(request_body) == () + + +@pytest.mark.parametrize( + ("uri", "names_an_unlisted_bucket"), + ( + pytest.param("s3://listed-bucket/dir/cat.png", False, id="listed"), + pytest.param("s3://listed-bucket/dir/a%20cat.png", False, id="listed-with-an-encoded-key"), + pytest.param("s3://other-bucket/dir/cat.png", True, id="unlisted"), + pytest.param("s3://listed-bucket.example/cat.png", True, id="listed-name-as-a-prefix"), + pytest.param("s3://listed%2Dbucket/cat.png", True, id="encoded-bucket"), + pytest.param("s3://other-bucket%2Flisted-bucket/cat.png", True, id="encoded-slash-in-the-bucket"), + pytest.param("s3:///listed-bucket/cat.png", True, id="empty-bucket"), + pytest.param("s3%3A%2F%2Fother-bucket%2Fcat.png", True, id="encoded-scheme"), + pytest.param( + "s3://listed-bucket/" + _quoted_times("a cat.png", MAX_FILE_ID_DECODE_PASSES + 2), + True, + id="encoded-past-the-pass-cap", + ), + ), +) +def test_names_an_unlisted_s3_bucket_checks_the_bucket_of_every_decoding(uri: str, names_an_unlisted_bucket: bool): + assert names_an_unlisted_s3_bucket(uri, frozenset({"listed-bucket"}).__contains__) is names_an_unlisted_bucket diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 6c996a7ce75..9014a88f6dd 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -356,6 +356,43 @@ def test_bedrock_system_prompt_variants_keep_continue_message( ] +def test_bedrock_tool_result_s3_video_reaches_the_converse_request_as_a_video_block() -> None: + clip: Final = "s3://media-bucket/clips/demo.mp4" + request: Final = AmazonConverseConfig()._transform_request( + model="us.amazon.nova-2-lite-v1:0", + messages=[ + {"role": "user", "content": "What happens in the clip?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "fetch_clip", "arguments": "{}"}} + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "text", "text": "video attached"}, {"type": "file", "file": {"file_id": clip}}], + }, + ], + optional_params={ + "tools": [ + { + "type": "function", + "function": {"name": "fetch_clip", "parameters": {"type": "object", "properties": {}}}, + } + ] + }, + litellm_params={}, + headers={}, + ) + + assert request["messages"][-1]["content"][0]["toolResult"]["content"] == [ + {"text": "video attached"}, + {"video": {"format": "mp4", "source": {"s3Location": {"uri": clip}}}}, + ] + + @pytest.mark.respx(assert_all_called=True) def test_bedrock_json_schema_description_is_sent_in_request(respx_mock: respx.Router) -> None: description: Final = "Find the meaning inside a poem" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 19eb3c71a5b..45d217dd5da 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -5,6 +5,7 @@ This test file tests the Generic Guardrail API implementation, specifically focusing on metadata extraction and passing. """ +import json import os from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -16,6 +17,7 @@ import litellm from litellm import ModelResponse from litellm._version import version as litellm_version from litellm.exceptions import GuardrailRaisedException, Timeout +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPI, @@ -26,6 +28,26 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guar from litellm.types.utils import Choices, Message +async def _structured_messages_sent(guardrail: GenericGuardrailAPI, rows: list[dict]) -> list: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: records every body the guardrail endpoint received + + async def answer(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.content) + return httpx.Response(200, json={"action": "NONE"}, request=request) + + endpoint: Final = AsyncHTTPHandler(transport=httpx.MockTransport(answer)) + try: + with patch.object(guardrail, "async_handler", endpoint): + await guardrail.apply_guardrail( + inputs={"texts": ["What is written here?"], "structured_messages": rows}, + request_data={}, + input_type="request", + ) + finally: + await endpoint.close() + return [json.loads(body)["structured_messages"] for body in sent_bodies] + + @pytest.fixture def generic_guardrail(): """Create a GenericGuardrailAPI instance for testing""" @@ -672,6 +694,30 @@ class TestStructuredMessagesInResponse: assert "structured_messages" not in guardrailed_inputs assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."] + @pytest.mark.asyncio + async def test_an_image_with_a_bucket_owner_that_is_not_a_string_keeps_its_turn_for_the_guardrail( + self, generic_guardrail + ): + row: Final = { + "role": "user", + "content": [ + {"type": "text", "text": "What is written here?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA", "bucket_owner": 111122223333}}, + ], + } + + assert await _structured_messages_sent(generic_guardrail, [row]) == [ + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is written here?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}, + ], + } + ] + ] + @pytest.mark.asyncio @pytest.mark.parametrize( "structured_messages", diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8da89598eff..fd0aec87b2d 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -63,9 +63,10 @@ from litellm.proxy.common_request_processing import ( from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException +from litellm.proxy._types import Litellm_EntityType, LitellmUserRoles, ProxyErrorTypes, ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck from litellm.proxy.hooks.model_max_budget_limiter import ( @@ -204,6 +205,65 @@ def test_include_guardrail_response_requested_is_false_without_exact_true(): assert include_guardrail_response_requested({}) is False +def _chat_image(uri: str) -> dict[str, object]: + return {"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": uri}}]}]} + + +def _responses_image(uri: str) -> dict[str, object]: + return {"input": [{"role": "user", "content": [{"type": "input_image", "image_url": uri}]}]} + + +def _anthropic_image(uri: str) -> dict[str, object]: + return { + "max_tokens": 16, + "messages": [{"role": "user", "content": [{"type": "image", "source": {"type": "url", "url": uri}}]}], + } + + +_S3_MEDIA_SURFACES: Final = ( + pytest.param("acompletion", "/v1/chat/completions", _chat_image, id="chat-completions"), + pytest.param("aresponses", "/v1/responses", _responses_image, id="responses"), + pytest.param("anthropic_messages", "/v1/messages", _anthropic_image, id="anthropic-messages"), +) + + +async def _send_s3_media( + route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + path: str, + body: dict[str, object], + user_role: LitellmUserRoles, + general_settings: dict[str, object], +) -> object: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + router: Final = Router( + model_list=[ + { + "model_name": "nova", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "mock_response": "a cat"}, + } + ] + ) + return await ProxyBaseLLMRequestProcessing(data={"model": "nova", **body}).base_process_llm_request( + request=Request( + { + "type": "http", + "method": "POST", + "path": path, + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + } + ), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_role=user_role), + route_type=route_type, + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + general_settings=general_settings, + proxy_config=RealProxyConfig(), + llm_router=router, + ) + + class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response( @@ -509,6 +569,71 @@ class TestProxyBaseLLMRequestProcessing: assert exc_info.value.param == "model" add_litellm_data_to_request.assert_not_awaited() + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_MEDIA_SURFACES) + @pytest.mark.parametrize( + ("uri", "general_settings", "named_in_refusal"), + [ + pytest.param( + "s3://media-bucket/cat.png", + {}, + "general_settings.allowed_client_side_s3_buckets", + id="no-bucket-listed", + ), + pytest.param( + "s3://other-bucket/cat.png", + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + "general_settings.allowed_client_side_s3_buckets", + id="unlisted-bucket", + ), + pytest.param( + "s3://media-bucket/litellm-bedrock-files/other-team.png", + {"allowed_client_side_s3_buckets": ["media-bucket"], "allow_client_side_credentials": True}, + "LiteLLM-managed storage", + id="managed-key-in-a-listed-bucket", + ), + ], + ) + async def test_non_admin_s3_media_is_refused_while_a_proxy_admin_key_passes( + self, + route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + path: str, + body_for: Callable[[str], dict[str, object]], + uri: str, + general_settings: dict[str, object], + named_in_refusal: str, + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media(route_type, path, body_for(uri), LitellmUserRoles.INTERNAL_USER, general_settings) + admin_response: Final = await _send_s3_media(route_type, path, body_for(uri), LitellmUserRoles.PROXY_ADMIN, {}) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert uri in refused.value.detail + assert named_in_refusal in refused.value.detail + assert "a cat" in str(admin_response) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_MEDIA_SURFACES) + @pytest.mark.parametrize( + "general_settings", + [ + pytest.param({"allowed_client_side_s3_buckets": ["media-bucket"]}, id="listed-bucket"), + pytest.param({"allow_client_side_credentials": True}, id="client-side-credentials-allowed"), + ], + ) + async def test_non_admin_s3_media_passes_in_a_bucket_the_admin_allowed( + self, + route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + path: str, + body_for: Callable[[str], dict[str, object]], + general_settings: dict[str, object], + ) -> None: + response: Final = await _send_s3_media( + route_type, path, body_for("s3://media-bucket/dir/cat.png"), LitellmUserRoles.INTERNAL_USER, general_settings + ) + + assert "a cat" in str(response) + @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True]) @pytest.mark.parametrize( diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 114f5123dd6..f929a8c730d 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -12519,6 +12519,30 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch): assert before["some_api_key"] != "sk-stored-secret" +@pytest.mark.asyncio +async def test_update_config_general_settings_stores_a_client_side_s3_bucket_the_gate_admits(monkeypatch): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist + from litellm.proxy.proxy_server import update_config_general_settings + + fake: Final = _fake_prisma_with_config({}) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake) + monkeypatch.setattr(litellm, "store_audit_logs", False) + + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="allowed_client_side_s3_buckets", + field_value=["media-bucket"], + config_type="general_settings", + ), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + stored: Final = json.loads(fake.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"]) + assert ClientSideStorageAllowlist.model_validate(stored).admits_s3_bucket("media-bucket") + + @pytest.mark.asyncio async def test_delete_config_general_settings_is_visible_to_the_next_read(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 5fb908ac42f..3e30f38b82f 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -286,6 +286,22 @@ class TestLiteLLMCompletionResponsesConfig: }, id="input-file", ), + pytest.param( + {"type": "input_image", "image_url": "s3://media-bucket/dir/cat.png", "bucket_owner": 111122223333}, + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "detail": "auto", "bucket_owner": 111122223333}, + }, + id="input-image-owner-left-for-bedrock-to-check", + ), + pytest.param( + {"type": "input_image", "image_url": "s3://media-bucket/dir/cat.png", "bucket_owner": ""}, + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "detail": "auto", "bucket_owner": ""}, + }, + id="input-image-empty-owner-left-for-bedrock-to-check", + ), ], ) def test_transform_input_media_item_keeps_bucket_owner(self, item: Mapping[str, object], expected: object): @@ -293,6 +309,35 @@ class TestLiteLLMCompletionResponsesConfig: expected ] + @pytest.mark.parametrize( + ("item", "expected"), + [ + pytest.param( + {"type": "input_image", "image_url": "data:image/png;base64,AA", "bucket_owner": "111122223333"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA", "detail": "auto"}}, + id="data-url-image", + ), + pytest.param( + {"type": "input_image", "image_url": "https://example.test/cat.png", "bucket_owner": 111122223333}, + {"type": "image_url", "image_url": {"url": "https://example.test/cat.png", "detail": "auto"}}, + id="https-image-with-an-int-owner", + ), + pytest.param( + { + "type": "input_file", + "file_data": "data:application/pdf;base64,AA", + "bucket_owner": ["111122223333"], + }, + {"type": "file", "file": {"file_data": "data:application/pdf;base64,AA"}}, + id="data-url-file-with-a-list-owner", + ), + ], + ) + def test_transform_input_media_item_drops_bucket_owner_off_s3(self, item: Mapping[str, object], expected: object): + assert LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([item]) == [ + expected + ] + def test_transform_chat_completion_response_with_reasoning_content(self): """Test that reasoning content is preserved in the full transformation pipeline""" # Setup diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6d7f0f36014..a532d9e8d40 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28710,8 +28710,6 @@ export interface components { }; /** ChatCompletionFileObjectFile */ ChatCompletionFileObjectFile: { - /** Bucket Owner */ - bucket_owner?: string; /** Detail */ detail?: string; /** File Data */ @@ -28753,8 +28751,6 @@ export interface components { }; /** ChatCompletionImageUrlObject */ ChatCompletionImageUrlObject: { - /** Bucket Owner */ - bucket_owner?: string; /** Detail */ detail?: string; /** Format */ @@ -29719,6 +29715,11 @@ export interface components { * @description If True, lets keys address Responses API ids that this proxy did not issue (raw provider ids, or ids issued before response-id encryption was configured). Such an id carries no owner, so no ownership check can run on it; ids this proxy did issue keep full ownership enforcement. Off by default, in which case an unrecognized response id is rejected with 403 */ allow_unmanaged_response_ids?: boolean | null; + /** + * Allowed Client Side S3 Buckets + * @description S3 bucket names that keys other than proxy admins may name in `s3://` image and file URLs. Any other bucket is rejected unless client-side credentials are allowed + */ + allowed_client_side_s3_buckets?: string[] | null; /** * Allowed File Extensions * @description the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied From 544e20570dbd02efefcce2124baee6670ea2b750 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 23:02:55 -0800 Subject: [PATCH 03/24] style: clear the ruff reds lts/1.108.x already has Sorts two imports ruff flags (I001) in redis_update_buffer.py and utils.py, the same hunks #50348 carries, and drops a duplicate GenericStreamingChunk import (F811) in a streaming test. No behavior change. --- litellm/proxy/db/db_transaction_queue/redis_update_buffer.py | 1 - litellm/utils.py | 2 +- .../litellm_core_utils/test_streaming_chunk_builder_utils.py | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 22f55ea9153..8fbaa7e3a9a 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -13,7 +13,6 @@ from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import TypeAdapter - from redis.exceptions import RedisError from litellm._internal_context import with_service_target diff --git a/litellm/utils.py b/litellm/utils.py index 51dc6db456a..71c9ca0a16e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -434,11 +434,11 @@ if TYPE_CHECKING: ) from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_llm import CustomLLM from litellm.router_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy, reset_retry_policy, ) - from litellm.llms.custom_llm import CustomLLM from litellm.types.llms.anthropic import ( ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index a46c8bed4ad..716c0af236c 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -18,7 +18,6 @@ from litellm.types.utils import ( ChatCompletionMessageToolCall, Delta, Function, - GenericStreamingChunk, ModelResponseStream, PromptTokensDetails, ServerToolUse, From ef3a93b0cdfe43444f6da3fcd643dc7aaa0c3862 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 02:53:33 -0800 Subject: [PATCH 04/24] fix(proxy): check s3:// media in every request field and refuse it in Bedrock batch files The s3:// media check now reads the whole request body instead of a fixed list of fields, and it also reads the Gemini file_uri and fileUri keys. A Bedrock batch input line with an s3:// media URL is a 400 before upload that says "LiteLLM does not accept s3:// media URLs in Bedrock batch files", since Bedrock batch files do not take s3:// media URLs yet. Backport to lts/1.108.x, adapted to the line. The line has no per-line batch parser, so the batch check runs at the top of the per-record loop that every Bedrock batch entry point shares. Cherry-picked from d437577e3dd4a7641d926a59776c8f8bd4bc7405 and b089fb9f903f3b568646bd442f63603a4887cd08 (litellm_bedrock_converse_s3_image_sources). --- .../cloud_storage_security.py | 6 +- litellm/llms/bedrock/files/transformation.py | 29 +++++++- .../test_cloud_storage_security.py | 19 ++++- .../test_bedrock_files_transformation.py | 72 +++++++++++++++++++ .../proxy/test_common_request_processing.py | 26 +++++-- 5 files changed, 140 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 88ddb14f498..e06cadcaa93 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -44,8 +44,7 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: _S3_SCHEME: Final = "s3://" _FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3: Final = ("s", "%") -_CLIENT_MEDIA_FIELDS: Final = ("messages", "input", "system") -_MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data"}) +_MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data", "file_uri", "fileUri"}) def _media_uris(value: object, key: str | None = None) -> Iterator[str]: @@ -89,8 +88,7 @@ def _is_s3_uri(uri: str) -> bool: def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: - media_uris: Final = chain.from_iterable(_media_uris(request_body.get(field)) for field in _CLIENT_MEDIA_FIELDS) - return tuple(filter(_is_s3_uri, media_uris)) + return tuple(filter(_is_s3_uri, _media_uris(request_body))) def is_managed_s3_uri(uri: str) -> bool: diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4d40e0dec85..3ab84c6d58c 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_UPLOAD_PREFIX, build_managed_cloud_object_name, encode_s3_object_key_for_url, + s3_media_uris, sanitize_cloud_object_component, should_allow_legacy_cloud_file_ids, split_configured_cloud_bucket_name, @@ -321,8 +322,7 @@ def _validate_file_id_against_configured_buckets( _REJECTED_FILE_ID_REQUEST_URL: Final = "https://litellm.ai" -def _rejected_file_id(reason: ValueError) -> BedrockError: - message: Final = str(reason) +def _rejected_request(message: str) -> BedrockError: return BedrockError( status_code=400, message=message, @@ -334,6 +334,29 @@ def _rejected_file_id(reason: ValueError) -> BedrockError: ) +def _rejected_file_id(reason: ValueError) -> BedrockError: + return _rejected_request(str(reason)) + + +def _batch_line_error(line_number: int, problem: str) -> BedrockError: + return _rejected_request(f"Batch input file line {line_number} {problem}") + + +def _json_object(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None + + +def _batch_body_problem(body: object) -> str | None: + json_body: Final = _json_object(body) + s3_uris: Final = s3_media_uris(json_body) if json_body is not None else () + if not s3_uris: + return None + return ( + f"has s3:// media URLs ({', '.join(s3_uris)}). " + "LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data" + ) + + def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: configured_bucket_names: Final = get_configured_s3_bucket_names(litellm_params) allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params) @@ -1083,6 +1106,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): for idx, _openai_jsonl_content in enumerate(openai_jsonl_content): # Extract the request body from OpenAI format openai_body = _openai_jsonl_content.get("body", {}) + if (body_problem := _batch_body_problem(openai_body)) is not None: + raise _batch_line_error(idx + 1, body_problem) record_model = openai_body.get("model", "") resolved_model = litellm.model_alias_map.get(record_model, record_model) model_for_transform, provider = self._resolve_batch_record_model_and_provider( 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 ea4953cb7c8..e476650c0c8 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -44,6 +44,10 @@ def _responses_turn(*parts: object) -> dict[str, object]: return {"input": [{"role": "user", "content": list(parts)}]} +def _gemini_turn(*parts: object) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": list(parts)}]} + + @pytest.mark.parametrize( "request_body", ( @@ -89,6 +93,19 @@ def _responses_turn(*parts: object) -> dict[str, object]: {"system": [{"type": "image", "source": {"type": "url", "url": _MANAGED_UPLOAD}}], "messages": []}, id="anthropic-system", ), + pytest.param( + _gemini_turn({"file_data": {"mime_type": "image/png", "file_uri": _MANAGED_UPLOAD}}), + id="gemini-file-data-snake-case", + ), + pytest.param( + _gemini_turn({"fileData": {"mimeType": "image/png", "fileUri": _MANAGED_UPLOAD}}), + id="gemini-file-data-camel-case", + ), + pytest.param( + {"system_instruction": {"parts": [{"file_data": {"file_uri": _MANAGED_UPLOAD}}]}, "contents": []}, + id="gemini-system-instruction", + ), + pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="media-key-in-any-field"), ), ) def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]): @@ -151,7 +168,7 @@ def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): pytest.param({"messages": [{"role": "user", "content": _MANAGED_UPLOAD}]}, id="string-content"), pytest.param({"input": _MANAGED_UPLOAD}, id="responses-string-input"), pytest.param({"system": _MANAGED_UPLOAD, "messages": []}, id="string-system"), - pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="non-media-field"), + pytest.param({"metadata": {"note": _MANAGED_UPLOAD}}, id="non-media-key"), ), ) def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_body: dict[str, object]): diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 425beefac8a..0ec966dfb05 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -22,6 +22,7 @@ from botocore.credentials import Credentials import litellm from litellm.constants import DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation from tests.unit.proxy.conftest import httpx_transport @@ -4334,3 +4335,74 @@ async def test_async_create_file( "s3://litellm-proxy-941277531214/litellm-bedrock-files-" ) assert file_obj.filename.endswith(".jsonl") + + +_VALID_BATCH_LINE: Final = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } +) +_BATCH_S3_IMAGE: Final = "s3://media-bucket/dir/cat.png" +_BATCH_S3_MEDIA_REFUSAL: Final = ( + f"Batch input file line 2 has s3:// media URLs ({_BATCH_S3_IMAGE}). " + "LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data" +) +_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( + { + "custom_id": "req-2", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": _BATCH_S3_IMAGE}}]}], + }, + } +) +_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( + { + "custom_id": "req-2", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "input": [{"role": "user", "content": [{"type": "input_image", "image_url": _BATCH_S3_IMAGE}]}], + }, + } +) + + +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "s3_media_line", + ( + pytest.param(_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE, id="chat-s3-image"), + pytest.param(_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE, id="responses-s3-image"), + ), +) +async def test_create_file_refuses_a_batch_line_with_s3_media_with_a_400_before_uploading( + respx_mock: respx.MockRouter, s3_media_line: str +) -> None: + put_route: Final = respx_mock.put(url__regex=r"https://s3\..+").respond(status_code=200) + + with pytest.raises(BedrockError) as exc_info: + await litellm.acreate_file( + file=("batch.jsonl", f"{_VALID_BATCH_LINE}\n{s3_media_line}\n".encode(), "application/jsonl"), + purpose="batch", + custom_llm_provider="bedrock", + s3_bucket_name="litellm-batch-bucket", + aws_region_name="us-west-2", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.message == _BATCH_S3_MEDIA_REFUSAL + assert not put_route.called diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fd0aec87b2d..5c07835b3f9 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,7 +4,7 @@ import datetime import json from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence +from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeAlias from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -220,15 +220,27 @@ def _anthropic_image(uri: str) -> dict[str, object]: } +def _gemini_snake_case_image(uri: str) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": [{"file_data": {"mime_type": "image/png", "file_uri": uri}}]}]} + + +def _gemini_camel_case_image(uri: str) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": [{"fileData": {"mimeType": "image/png", "fileUri": uri}}]}]} + + +_S3MediaRoute: TypeAlias = Literal["acompletion", "aresponses", "anthropic_messages", "agenerate_content"] +_GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/nova:generateContent" _S3_MEDIA_SURFACES: Final = ( pytest.param("acompletion", "/v1/chat/completions", _chat_image, id="chat-completions"), pytest.param("aresponses", "/v1/responses", _responses_image, id="responses"), pytest.param("anthropic_messages", "/v1/messages", _anthropic_image, id="anthropic-messages"), + pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_snake_case_image, id="gemini-snake-case"), + pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_camel_case_image, id="gemini-camel-case"), ) async def _send_s3_media( - route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + route_type: _S3MediaRoute, path: str, body: dict[str, object], user_role: LitellmUserRoles, @@ -596,7 +608,7 @@ class TestProxyBaseLLMRequestProcessing: ) async def test_non_admin_s3_media_is_refused_while_a_proxy_admin_key_passes( self, - route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]], uri: str, @@ -623,13 +635,17 @@ class TestProxyBaseLLMRequestProcessing: ) async def test_non_admin_s3_media_passes_in_a_bucket_the_admin_allowed( self, - route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]], general_settings: dict[str, object], ) -> None: response: Final = await _send_s3_media( - route_type, path, body_for("s3://media-bucket/dir/cat.png"), LitellmUserRoles.INTERNAL_USER, general_settings + route_type, + path, + body_for("s3://media-bucket/dir/cat.png"), + LitellmUserRoles.INTERNAL_USER, + general_settings, ) assert "a cat" in str(response) From 2c96d603df1514d10a3b99067735818c9a34be62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 05:05:34 -0800 Subject: [PATCH 05/24] fix(proxy): check s3:// URLs only inside media parts and refuse them only on Converse batch lines The s3:// bucket check now reads URL fields only inside the content parts that a Bedrock transform turns into an S3 source, and it uses the same s3:// test as the prompt factory. Values in metadata, annotations, tool arguments and other fields pass as before. Each URL is named once in the 403 and 400 messages Bedrock batch files refuse s3:// media only on the lines that go through the Converse transform, so passthrough-provider lines upload unchanged Backport to lts/1.108.x, adapted to the line. The line has no per-line batch parser, so the per-record loop turns the Converse refusal into the line error, and the old body check is gone Cherry-picked from 2c23d2253a52798ebe23eccfbbd9528fe0d91b90 (litellm_bedrock_converse_s3_image_sources) --- .../cloud_storage_security.py | 34 +++- .../prompt_templates/factory.py | 7 +- litellm/llms/bedrock/files/transformation.py | 37 ++-- .../transformation.py | 3 +- .../test_cloud_storage_security.py | 101 +++++++++- .../test_bedrock_files_transformation.py | 75 ++++++- .../proxy/test_common_request_processing.py | 187 ++++++++++++++++++ 7 files changed, 403 insertions(+), 41 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index e06cadcaa93..7ca44b0fa6a 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -43,14 +43,36 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: _S3_SCHEME: Final = "s3://" -_FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3: Final = ("s", "%") +_S3_MEDIA_PART_TYPES: Final = frozenset( + {"image_url", "file", "input_image", "input_file", "computer_screenshot", "image", "document"} +) +_GEMINI_MEDIA_PART_KEYS: Final = frozenset({"file_data", "fileData"}) _MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data", "file_uri", "fileUri"}) -def _media_uris(value: object, key: str | None = None) -> Iterator[str]: +def is_s3_uri(value: str) -> bool: + return value.startswith(_S3_SCHEME) + + +def _media_part_uris(value: object, key: str | None = None) -> Iterator[str]: if isinstance(value, str): return iter((value,) if key in _MEDIA_URI_KEYS else ()) fields: Final = as_str_mapping(value) + if fields is not None: + return chain.from_iterable(_media_part_uris(child, child_key) for child_key, child in fields.items()) + items: Final = as_object_list(value) + return chain.from_iterable(map(_media_part_uris, items)) if items is not None else iter(()) + + +def _is_media_part(fields: Mapping[str, object], key: str | None) -> bool: + part_type: Final = fields.get("type") + return (isinstance(part_type, str) and part_type in _S3_MEDIA_PART_TYPES) or key in _GEMINI_MEDIA_PART_KEYS + + +def _media_uris(value: object, key: str | None = None) -> Iterator[str]: + fields: Final = as_str_mapping(value) + if fields is not None and _is_media_part(fields, key): + return _media_part_uris(fields) if fields is not None: return chain.from_iterable(_media_uris(child, child_key) for child_key, child in fields.items()) items: Final = as_object_list(value) @@ -81,14 +103,8 @@ def _refused_in_some_decoding( return bool(parts) and (not settled or not all(map(is_allowed, parts))) -def _is_s3_uri(uri: str) -> bool: - return uri.startswith(_FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3) and any( - decoding.startswith(_S3_SCHEME) for decoding in _decodings(uri) - ) - - def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: - return tuple(filter(_is_s3_uri, _media_uris(request_body))) + return tuple(dict.fromkeys(filter(is_s3_uri, _media_uris(request_body)))) def is_managed_s3_uri(uri: str) -> bool: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0e5393d1e37..5794e69d1a9 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -21,6 +21,7 @@ import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid from litellm.constants import REDACTED_BY_LITELLM +from litellm.litellm_core_utils.cloud_storage_security import is_s3_uri from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, @@ -3917,7 +3918,7 @@ class BedrockImageProcessor: ) -> BedrockContentBlock: """Synchronous image processing.""" - if image_url.startswith("s3://"): + if is_s3_uri(image_url): return cls._create_s3_block(image_url, format, bucket_owner) if "base64" in image_url: img_bytes, mime_type, image_format = cls._parse_base64_image(image_url) @@ -3940,7 +3941,7 @@ class BedrockImageProcessor: ) -> BedrockContentBlock: """Asynchronous image processing.""" - if image_url.startswith("s3://"): + if is_s3_uri(image_url): return cls._create_s3_block(image_url, format, bucket_owner) if "base64" in image_url: img_bytes, mime_type, image_format = cls._parse_base64_image(image_url) @@ -4104,7 +4105,7 @@ def _s3_bucket_owner(media: object, url: str) -> str | None: owner: Final = fields.get("bucket_owner") if fields is not None else None if owner is None or (isinstance(owner, str) and owner): return owner - if not url.startswith("s3://"): + if not is_s3_uri(url): return None raise litellm.BadRequestError( message=f"`bucket_owner` for {url} must be the AWS account id as a non-empty string, not {owner!r}.", diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 3ab84c6d58c..f5c4732a4fe 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -342,19 +342,12 @@ def _batch_line_error(line_number: int, problem: str) -> BedrockError: return _rejected_request(f"Batch input file line {line_number} {problem}") -def _json_object(value: object) -> Mapping[str, object] | None: - return value if isinstance(value, dict) else None - - -def _batch_body_problem(body: object) -> str | None: - json_body: Final = _json_object(body) - s3_uris: Final = s3_media_uris(json_body) if json_body is not None else () - if not s3_uris: - return None - return ( - f"has s3:// media URLs ({', '.join(s3_uris)}). " - "LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data" - ) +class _S3MediaInBatchRecord(ValueError): + def __init__(self, s3_uris: tuple[str, ...]) -> None: + super().__init__( + f"has s3:// media URLs ({', '.join(s3_uris)}). " + "LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data" + ) def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: @@ -1042,6 +1035,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): AmazonConverseConfig, ) + s3_uris: Final = s3_media_uris(openai_request_body) + if s3_uris: + raise _S3MediaInBatchRecord(s3_uris) converse_config: Final = AmazonConverseConfig() mapped_params = converse_config.map_openai_params( non_default_params=optional_params, @@ -1106,8 +1102,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): for idx, _openai_jsonl_content in enumerate(openai_jsonl_content): # Extract the request body from OpenAI format openai_body = _openai_jsonl_content.get("body", {}) - if (body_problem := _batch_body_problem(openai_body)) is not None: - raise _batch_line_error(idx + 1, body_problem) record_model = openai_body.get("model", "") resolved_model = litellm.model_alias_map.get(record_model, record_model) model_for_transform, provider = self._resolve_batch_record_model_and_provider( @@ -1125,11 +1119,14 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): openai_request_body=openai_body, model=model_for_transform ) else: - model_input = self._map_openai_to_bedrock_params( - openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), - model=model_for_transform, - provider=provider, - ) + try: + model_input = self._map_openai_to_bedrock_params( + openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), + model=model_for_transform, + provider=provider, + ) + except _S3MediaInBatchRecord as e: + raise _batch_line_error(idx + 1, str(e)) from e # Create Bedrock batch record record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 965bdd29b21..c8535aec969 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -41,6 +41,7 @@ from litellm._logging import verbose_logger from litellm.caching import InMemoryCache from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook +from litellm.litellm_core_utils.cloud_storage_security import is_s3_uri from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -228,7 +229,7 @@ _RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) def _is_s3_url(url: object) -> bool: - return isinstance(url, str) and url.startswith("s3://") + return isinstance(url, str) and is_s3_uri(url) @runtime_checkable 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 e476650c0c8..14b931714f2 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -105,7 +105,18 @@ def _gemini_turn(*parts: object) -> dict[str, object]: {"system_instruction": {"parts": [{"file_data": {"file_uri": _MANAGED_UPLOAD}}]}, "contents": []}, id="gemini-system-instruction", ), - pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="media-key-in-any-field"), + pytest.param( + { + "input": [ + { + "type": "computer_call_output", + "call_id": "call_1", + "output": {"type": "computer_screenshot", "image_url": _MANAGED_UPLOAD}, + } + ] + }, + id="responses-computer-screenshot", + ), ), ) def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]): @@ -124,8 +135,6 @@ def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: d "s3://media-bucket/litellm%2Dbedrock-files/x.pdf", "s3://media-bucket/litellm%252Dbedrock-files/x.pdf", "s3://media-bucket%2Flitellm-bedrock-files/x.pdf", - "s3%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf", - "%73%33%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf", "s3:///litellm-bedrock-files/x.pdf", ), ) @@ -169,12 +178,98 @@ def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): pytest.param({"input": _MANAGED_UPLOAD}, id="responses-string-input"), pytest.param({"system": _MANAGED_UPLOAD, "messages": []}, id="string-system"), pytest.param({"metadata": {"note": _MANAGED_UPLOAD}}, id="non-media-key"), + pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="metadata-url"), + pytest.param({"extra_body": {"url": _MANAGED_UPLOAD}}, id="extra-body-url"), + pytest.param({"name": "fetch", "arguments": {"url": _MANAGED_UPLOAD}}, id="mcp-tool-arguments"), + pytest.param( + { + "messages": [ + { + "role": "assistant", + "content": "See the source.", + "annotations": [{"type": "url_citation", "url_citation": {"url": _MANAGED_UPLOAD}}], + } + ] + }, + id="chat-url-citation", + ), + pytest.param( + { + "input": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "See the file.", + "annotations": [{"type": "file_citation", "file_id": _MANAGED_UPLOAD, "index": 0}], + } + ], + } + ] + }, + id="responses-file-citation", + ), + pytest.param( + _gemini_turn({"functionCall": {"name": "fetch", "args": {"url": _MANAGED_UPLOAD}}}), + id="gemini-function-call-camel-case", + ), + pytest.param( + _gemini_turn({"function_call": {"name": "fetch", "args": {"url": _MANAGED_UPLOAD}}}), + id="gemini-function-call-snake-case", + ), + pytest.param( + { + "tools": [ + { + "type": "function", + "function": { + "name": "fetch", + "parameters": { + "type": "object", + "properties": {"url": {"type": "string", "enum": [_MANAGED_UPLOAD]}}, + "default": {"url": _MANAGED_UPLOAD}, + }, + }, + } + ] + }, + id="tool-schema", + ), + pytest.param( + _user_turn({"type": "image_url", "image_url": {"url": "S3://media-bucket/litellm-bedrock-files/x.png"}}), + id="uppercase-scheme", + ), + pytest.param( + _user_turn( + {"type": "file", "file": {"file_id": "s3%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf"}} + ), + id="encoded-scheme", + ), + pytest.param( + _user_turn( + {"type": "file", "file": {"file_id": "%73%33%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf"}} + ), + id="fully-encoded-scheme", + ), ), ) def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_body: dict[str, object]): assert s3_media_uris(request_body) == () +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( + {"type": "image_url", "image_url": {"url": other_upload}}, + {"type": "file", "file": {"file_id": _MANAGED_UPLOAD}}, + {"type": "image_url", "image_url": {"url": other_upload}}, + ) + + assert s3_media_uris(request_body) == (other_upload, _MANAGED_UPLOAD) + + @pytest.mark.parametrize( ("uri", "names_an_unlisted_bucket"), ( diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 0ec966dfb05..df72d817edb 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -4365,6 +4365,36 @@ _CHAT_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( }, } ) +_S3_IMAGE_PART: Final = {"type": "image_url", "image_url": {"url": _BATCH_S3_IMAGE}} +_OTHER_BATCH_S3_IMAGE: Final = "s3://media-bucket/dir/dog.png" +_BATCH_TWO_S3_MEDIA_REFUSAL: Final = _BATCH_S3_MEDIA_REFUSAL.replace( + f"({_BATCH_S3_IMAGE})", f"({_BATCH_S3_IMAGE}, {_OTHER_BATCH_S3_IMAGE})" +) + + +def _nova_chat_batch_line(custom_id: str, **body: object) -> str: + return json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "bedrock/amazon.nova-pro-v1:0", **body}, + } + ) + + +_CHAT_BATCH_LINE_WITH_THE_SAME_S3_IMAGE_TWICE: Final = _nova_chat_batch_line( + "req-2", messages=[{"role": "user", "content": [_S3_IMAGE_PART, _S3_IMAGE_PART]}] +) +_CHAT_BATCH_LINE_WITH_TWO_S3_IMAGES: Final = _nova_chat_batch_line( + "req-2", + messages=[ + { + "role": "user", + "content": [_S3_IMAGE_PART, {"type": "image_url", "image_url": {"url": _OTHER_BATCH_S3_IMAGE}}], + } + ], +) _RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( { "custom_id": "req-2", @@ -4381,14 +4411,18 @@ _RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( @pytest.mark.usefixtures(httpx_transport.__name__) @pytest.mark.asyncio @pytest.mark.parametrize( - "s3_media_line", + ("s3_media_line", "expected_message"), ( - pytest.param(_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE, id="chat-s3-image"), - pytest.param(_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE, id="responses-s3-image"), + pytest.param(_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE, _BATCH_S3_MEDIA_REFUSAL, id="chat-s3-image"), + pytest.param(_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE, _BATCH_S3_MEDIA_REFUSAL, id="responses-s3-image"), + pytest.param( + _CHAT_BATCH_LINE_WITH_THE_SAME_S3_IMAGE_TWICE, _BATCH_S3_MEDIA_REFUSAL, id="chat-repeated-s3-image" + ), + pytest.param(_CHAT_BATCH_LINE_WITH_TWO_S3_IMAGES, _BATCH_TWO_S3_MEDIA_REFUSAL, id="chat-two-s3-images"), ), ) async def test_create_file_refuses_a_batch_line_with_s3_media_with_a_400_before_uploading( - respx_mock: respx.MockRouter, s3_media_line: str + respx_mock: respx.MockRouter, s3_media_line: str, expected_message: str ) -> None: put_route: Final = respx_mock.put(url__regex=r"https://s3\..+").respond(status_code=200) @@ -4404,5 +4438,36 @@ async def test_create_file_refuses_a_batch_line_with_s3_media_with_a_400_before_ ) assert exc_info.value.status_code == 400 - assert exc_info.value.message == _BATCH_S3_MEDIA_REFUSAL + assert exc_info.value.message == expected_message assert not put_route.called + + +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.asyncio +async def test_create_file_uploads_s3_values_that_no_converse_record_sends_as_an_s3_source( + respx_mock: respx.MockRouter, +) -> None: + put_route: Final = respx_mock.put(url__regex=r"https://s3\..+").respond(status_code=200) + passthrough_messages: Final = [{"role": "user", "content": [_S3_IMAGE_PART]}] + passthrough_line: Final = _nova_chat_batch_line( + "req-1", model="bedrock/openai.gpt-oss-120b-1:0", max_tokens=10, messages=passthrough_messages + ) + metadata_only_line: Final = _nova_chat_batch_line( + "req-2", messages=[{"role": "user", "content": "Hello"}], max_tokens=10, metadata={"url": _BATCH_S3_IMAGE} + ) + + await litellm.acreate_file( + file=("batch.jsonl", f"{passthrough_line}\n{metadata_only_line}\n".encode(), "application/jsonl"), + purpose="batch", + custom_llm_provider="bedrock", + s3_bucket_name="litellm-batch-bucket", + aws_region_name="us-west-2", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + ) + + passthrough_record, metadata_only_record = ( + json.loads(line) for line in put_route.calls.last.request.content.decode().splitlines() + ) + assert passthrough_record["modelInput"]["messages"] == passthrough_messages + assert metadata_only_record["recordId"] == "req-2" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5c07835b3f9..eaab813db9f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -239,6 +239,142 @@ _S3_MEDIA_SURFACES: Final = ( ) +_S3_OBJECT: Final = "s3://media-bucket/dir/cat.png" +_HELLO: Final = {"role": "user", "content": "hello"} +_S3_VALUES_OUTSIDE_MEDIA_PARTS: Final = ( + pytest.param( + "acompletion", + "/v1/chat/completions", + {"messages": [_HELLO], "metadata": {"url": _S3_OBJECT}}, + id="chat-metadata-url", + ), + pytest.param( + "acompletion", + "/v1/chat/completions", + {"messages": [_HELLO], "extra_body": {"url": _S3_OBJECT}}, + id="chat-extra-body-url", + ), + pytest.param( + "acompletion", + "/v1/chat/completions", + { + "messages": [ + _HELLO, + { + "role": "assistant", + "content": "See the source.", + "annotations": [{"type": "url_citation", "url_citation": {"url": _S3_OBJECT}}], + }, + _HELLO, + ] + }, + id="chat-url-citation", + ), + pytest.param( + "acompletion", + "/v1/chat/completions", + { + "messages": [_HELLO], + "tools": [ + { + "type": "function", + "function": { + "name": "fetch", + "parameters": { + "type": "object", + "properties": {"url": {"type": "string", "enum": [_S3_OBJECT]}}, + "default": {"url": _S3_OBJECT}, + }, + }, + } + ], + }, + id="chat-tool-schema", + ), + pytest.param( + "aresponses", + "/v1/responses", + { + "input": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "See the file.", + "annotations": [{"type": "file_citation", "file_id": _S3_OBJECT, "index": 0}], + } + ], + }, + _HELLO, + ] + }, + id="responses-file-citation", + ), + pytest.param( + "agenerate_content", + _GEMINI_GENERATE_CONTENT_PATH, + { + "contents": [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"functionCall": {"name": "fetch", "args": {"url": _S3_OBJECT}}}]}, + {"role": "user", "parts": [{"functionResponse": {"name": "fetch", "response": {"ok": True}}}]}, + ] + }, + id="gemini-function-call-camel-case", + ), + pytest.param( + "agenerate_content", + _GEMINI_GENERATE_CONTENT_PATH, + { + "contents": [ + {"role": "user", "parts": [{"text": "hello"}]}, + {"role": "model", "parts": [{"function_call": {"name": "fetch", "args": {"url": _S3_OBJECT}}}]}, + {"role": "user", "parts": [{"function_response": {"name": "fetch", "response": {"ok": True}}}]}, + ] + }, + id="gemini-function-call-snake-case", + ), +) +_S3_VALUES_ON_NON_MEDIA_ROUTES: Final = ( + pytest.param( + "call_mcp_tool", + "/mcp-rest/tools/call", + {"name": "fetch", "arguments": {"url": _S3_OBJECT}}, + id="mcp-tool-arguments", + ), + pytest.param("afile_retrieve", "/v1/files/s3-object", {"file_id": _S3_OBJECT}, id="file-retrieve"), + pytest.param("afile_content", "/v1/files/s3-object/content", {"file_id": _S3_OBJECT}, id="file-content"), + pytest.param( + "avector_store_file_create", + "/v1/vector_stores/vs_1/files", + {"vector_store_id": "vs_1", "file_id": _S3_OBJECT}, + id="vector-store-file-create", + ), + pytest.param( + "asend_message", + "/a2a/agent-1/message/send", + { + "jsonrpc": "2.0", + "id": "1", + "method": "message/send", + "params": { + "message": { + "role": "user", + "messageId": "m1", + "parts": [ + {"kind": "data", "data": {"file_url": _S3_OBJECT}}, + {"kind": "file", "file": {"uri": _S3_OBJECT, "mimeType": "image/png"}}, + ], + } + }, + }, + id="a2a-data-and-file-parts", + ), +) + + async def _send_s3_media( route_type: _S3MediaRoute, path: str, @@ -650,6 +786,57 @@ class TestProxyBaseLLMRequestProcessing: assert "a cat" in str(response) + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_OUTSIDE_MEDIA_PARTS) + async def test_non_admin_s3_values_outside_media_parts_pass_with_no_bucket_listed( + self, route_type: _S3MediaRoute, path: str, body: dict[str, object] + ) -> None: + response: Final = await _send_s3_media(route_type, path, body, LitellmUserRoles.INTERNAL_USER, {}) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_ON_NON_MEDIA_ROUTES) + async def test_non_admin_s3_values_on_routes_without_media_parts_reach_the_route( + self, route_type: str, path: str, body: dict[str, object] + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + data, _ = await ProxyBaseLLMRequestProcessing(data=copy.deepcopy(body)).common_processing_pre_call_logic( + request=Request({"type": "http", "method": "POST", "path": path, "query_string": b"", "headers": []}), + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.INTERNAL_USER), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=RealProxyConfig(), + route_type=route_type, + ) + + assert {key: data[key] for key in body} == body + + @pytest.mark.asyncio + async def test_non_admin_s3_media_refusal_names_a_repeated_object_once(self) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": _S3_OBJECT}}, + {"type": "image_url", "image_url": {"url": _S3_OBJECT}}, + ], + } + ] + }, + LitellmUserRoles.INTERNAL_USER, + {}, + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN + assert refused.value.detail.count(_S3_OBJECT) == 1 + @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True]) @pytest.mark.parametrize( From c70ad06604f73707b46cbffaded1b682b4695a64 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 06:06:33 -0800 Subject: [PATCH 06/24] test(proxy): check that a file part's filename is not read as an s3:// URL s3_media_uris reads only the URL keys inside a media part. A change that returned every string in the part passed every test, so this case pins that an s3:// filename is ignored --- .../litellm_core_utils/test_cloud_storage_security.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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 14b931714f2..093d0076234 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -237,6 +237,15 @@ def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): }, id="tool-schema", ), + pytest.param( + _user_turn( + { + "type": "file", + "file": {"file_data": "data:application/pdf;base64,JVBERi0=", "filename": _MANAGED_UPLOAD}, + } + ), + id="media-part-filename", + ), pytest.param( _user_turn({"type": "image_url", "image_url": {"url": "S3://media-bucket/litellm-bedrock-files/x.png"}}), id="uppercase-scheme", From d6b5036c1977a9f923634f8f0431edde23192068 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 17:11:17 -0800 Subject: [PATCH 07/24] fix(proxy): check s3 URIs inside encoded file ids and on queued chat completions The s3 media check now decodes model-embedded and unified file ids, since completion() turns them back into the raw s3 URI they carry. The /queue/chat/completions endpoint skips common request processing, so it now runs the same check. The media walk goes level by level instead of recursing, which the recursive_detector code-quality step requires. (cherry picked from commit 5b45c632c57fd46b9dda537b694edcda2e0f41d3) --- .../cloud_storage_security.py | 68 ++++++++++++++----- litellm/proxy/common_request_processing.py | 4 +- litellm/proxy/proxy_server.py | 2 + .../test_cloud_storage_security.py | 45 ++++++++++++ tests/unit/proxy/test_proxy_server.py | 61 +++++++++++++++++ 5 files changed, 161 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 7ca44b0fa6a..7c2da8d984d 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,7 +1,9 @@ +import base64 import posixpath import re from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate, chain, repeat +from dataclasses import dataclass +from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType from typing import Any, Final, cast from urllib.parse import quote, unquote @@ -48,20 +50,35 @@ _S3_MEDIA_PART_TYPES: Final = frozenset( ) _GEMINI_MEDIA_PART_KEYS: Final = frozenset({"file_data", "fileData"}) _MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data", "file_uri", "fileUri"}) +_ENCODED_FILE_ID_PREFIXES: Final = ("file-", "batch_") +_EMBEDDED_S3_URI: Final = re.compile(r"s3://[^;]*") def is_s3_uri(value: str) -> bool: return value.startswith(_S3_SCHEME) -def _media_part_uris(value: object, key: str | None = None) -> Iterator[str]: - if isinstance(value, str): - return iter((value,) if key in _MEDIA_URI_KEYS else ()) - fields: Final = as_str_mapping(value) - if fields is not None: - return chain.from_iterable(_media_part_uris(child, child_key) for child_key, child in fields.items()) - items: Final = as_object_list(value) - return chain.from_iterable(map(_media_part_uris, items)) if items is not None else iter(()) +def _base64url_text(value: str) -> str | None: + try: + return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4)).decode() + except ValueError: + return None + + +def _s3_uris_in(text: str) -> Iterator[str]: + return (match.group(0) for match in _EMBEDDED_S3_URI.finditer(text)) + + +def _s3_uris_embedded_in_file_id(file_id: str) -> Iterator[str]: + encodings: Final = dict.fromkeys((file_id, *(file_id.removeprefix(prefix) for prefix in _ENCODED_FILE_ID_PREFIXES))) + return chain.from_iterable(map(_s3_uris_in, filter(None, map(_base64url_text, encodings)))) + + +@dataclass(frozen=True, slots=True) +class _BodyNode: + value: object + key: str | None = None + inside_media_part: bool = False def _is_media_part(fields: Mapping[str, object], key: str | None) -> bool: @@ -69,14 +86,30 @@ def _is_media_part(fields: Mapping[str, object], key: str | None) -> bool: return (isinstance(part_type, str) and part_type in _S3_MEDIA_PART_TYPES) or key in _GEMINI_MEDIA_PART_KEYS -def _media_uris(value: object, key: str | None = None) -> Iterator[str]: - fields: Final = as_str_mapping(value) - if fields is not None and _is_media_part(fields, key): - return _media_part_uris(fields) +def _child_nodes(node: _BodyNode) -> Iterator[_BodyNode]: + fields: Final = as_str_mapping(node.value) if fields is not None: - return chain.from_iterable(_media_uris(child, child_key) for child_key, child in fields.items()) - items: Final = as_object_list(value) - return chain.from_iterable(map(_media_uris, items)) if items is not None else iter(()) + inside_media_part: Final = node.inside_media_part or _is_media_part(fields, node.key) + return (_BodyNode(child, child_key, inside_media_part) for child_key, child in fields.items()) + items: Final = as_object_list(node.value) + return (_BodyNode(item, None, node.inside_media_part) for item in items) if items is not None else iter(()) + + +def _next_level(level: tuple[_BodyNode, ...], _: None) -> tuple[_BodyNode, ...]: + return tuple(chain.from_iterable(map(_child_nodes, level))) + + +def _body_nodes(request_body: Mapping[str, object]) -> Iterator[_BodyNode]: + levels: Final = accumulate(repeat(None), _next_level, initial=(_BodyNode(request_body),)) + return chain.from_iterable(takewhile(bool, levels)) + + +def _media_uri_candidates(node: _BodyNode) -> Iterator[str]: + if not node.inside_media_part or not isinstance(node.value, str): + return iter(()) + if node.key == "file_id": + return chain((node.value,), _s3_uris_embedded_in_file_id(node.value)) + return iter((node.value,) if node.key in _MEDIA_URI_KEYS else ()) def _s3_bucket(uri: str) -> str | None: @@ -104,7 +137,8 @@ def _refused_in_some_decoding( def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: - return tuple(dict.fromkeys(filter(is_s3_uri, _media_uris(request_body)))) + media_uris: Final = chain.from_iterable(map(_media_uri_candidates, _body_nodes(request_body))) + return tuple(dict.fromkeys(filter(is_s3_uri, media_uris))) def is_managed_s3_uri(uri: str) -> bool: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b148ff1bc2c..9957eb7792e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -860,7 +860,7 @@ def _request_model(data: Mapping[str, object]) -> str | list[str] | None: return None -def _refuse_s3_media_the_key_may_not_read( +def refuse_s3_media_the_key_may_not_read( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], @@ -2309,7 +2309,7 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) - _refuse_s3_media_the_key_may_not_read(self.data, user_api_key_dict, general_settings) + refuse_s3_media_the_key_may_not_read(self.data, user_api_key_dict, general_settings) self.data = await add_litellm_data_to_request( data=self.data, request=request, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 80970cba90b..bb74b2a08af 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -445,6 +445,7 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l log_llm_api_exception, open_sse_before_first_byte, proxy_exception_from_route_error, + refuse_s3_media_the_key_may_not_read, request_litellm_call_id, resolve_litellm_call_id, should_return_raw_model_name, @@ -17679,6 +17680,7 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) + refuse_s3_media_the_key_may_not_read(data, user_api_key_dict, general_settings) # Include original request and headers in the data data["proxy_server_request"] = { 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 093d0076234..9aa44e1fccf 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -1,3 +1,4 @@ +from base64 import urlsafe_b64encode from functools import reduce from typing import Final from urllib.parse import quote @@ -36,6 +37,16 @@ def _quoted_times(value: str, times: int) -> str: _MANAGED_UPLOAD: Final = "s3://media-bucket/litellm-bedrock-files/other-team.pdf" +def _unpadded_base64url(text: str) -> str: + return urlsafe_b64encode(text.encode()).decode().rstrip("=") + + +_MODEL_EMBEDDED_UPLOAD_ID: Final = "file-" + _unpadded_base64url(f"litellm:{_MANAGED_UPLOAD};model,nova") +_UNIFIED_BATCH_OUTPUT_ID: Final = _unpadded_base64url( + f"litellm_proxy:application/json;unified_id,u1;target_model_names,nova;llm_output_file_id,{_MANAGED_UPLOAD};" +) + + def _user_turn(*parts: object) -> dict[str, object]: return {"messages": [{"role": "user", "content": list(parts)}]} @@ -67,6 +78,14 @@ def _gemini_turn(*parts: object) -> dict[str, object]: }, id="chat-tool-result", ), + pytest.param( + _user_turn({"type": "file", "file": {"file_id": _MODEL_EMBEDDED_UPLOAD_ID}}), + id="chat-model-embedded-file-id", + ), + pytest.param( + _responses_turn({"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}), + id="responses-unified-file-id", + ), pytest.param(_responses_turn({"type": "input_image", "image_url": _MANAGED_UPLOAD}), id="responses-image"), pytest.param(_responses_turn({"type": "input_file", "file_url": _MANAGED_UPLOAD}), id="responses-file-url"), pytest.param(_responses_turn({"type": "input_file", "file_id": _MANAGED_UPLOAD}), id="responses-file-id"), @@ -173,6 +192,32 @@ def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): _user_turn({"type": "image_url", "image_url": {"url": "data:image/png;base64,czM6Ly9idWNrZXQva2V5"}}), id="data-url", ), + pytest.param( + _user_turn( + { + "type": "file", + "file": { + "file_id": _unpadded_base64url( + "litellm_proxy:application/pdf;unified_id,u2;target_model_names,nova;" + "llm_output_file_id,file-abc123;" + ) + }, + } + ), + id="unified-file-id-of-a-provider-file", + ), + pytest.param( + _user_turn( + { + "type": "file", + "file": {"file_data": "data:application/pdf;base64," + _unpadded_base64url(_MANAGED_UPLOAD)}, + } + ), + id="file-data-payload-naming-an-s3-object", + ), + pytest.param( + _user_turn({"type": "text", "text": _MODEL_EMBEDDED_UPLOAD_ID}), id="encoded-file-id-outside-a-media-part" + ), pytest.param(_user_turn({"type": "text", "text": _MANAGED_UPLOAD}), id="text-part"), pytest.param({"messages": [{"role": "user", "content": _MANAGED_UPLOAD}]}, id="string-content"), pytest.param({"input": _MANAGED_UPLOAD}, id="responses-string-input"), diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 80bc3ad1a34..c3a7a9c0925 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -3789,3 +3789,64 @@ def test_live_session_preserves_pretty_printed_start_frame( assert captured.route_type == "_arealtime" assert routed_data["model"] == "gpt-live-1" assert routed_data["live_session_start"] == expected_start + + +@pytest.mark.asyncio +async def test_queue_request_refuses_a_non_admin_s3_image_outside_the_allowed_buckets( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.requests import Request + from starlette.responses import Response + + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.proxy_server import async_queue_request + + body: Final = json.dumps( + { + "model": "nova-lite", + "messages": [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "s3://another-teams-bucket/cat.png"}}], + } + ], + } + ).encode() + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/queue/chat/completions", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "root_path": "", + }, + receive, + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "nova-lite", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test", "mock_response": "ok"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "general_settings", {"allowed_client_side_s3_buckets": ["media-bucket"]}) + + with pytest.raises(ProxyException) as exc_info: + await async_queue_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.code == "403" From e2931d410919b93847e7fbf1ac065f013f2e0694 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 12:53:55 -0800 Subject: [PATCH 08/24] fix(proxy): let a key send managed file ids it may read and refuse s3 sources hidden in model_file_id_mapping The s3 media gate refused every LiteLLM managed file id whose upload sits under the managed prefix, even when the key's own user or team created it. It now asks the managed files hook which unified ids the key may read and skips only the objects those ids embed. A forged or unowned id, or a proxy without the hook, still gets a 403. The gate also reads the provider file ids in a client-supplied model_file_id_mapping, so an alias can no longer point a non-admin request at an unlisted or managed object. Markdown documents (.md, .markdown) now map to Bedrock's md document format, so an s3 Markdown source without an explicit format no longer fails with an unknown media type. On lts/1.108.x the gate still walks the body before it reads the key role, so callers that pass a specced UserAPIKeyAuth mock keep working, and it reads self.data and general_settings directly since the line has no raw_request_fields. (cherry picked from commit 37f085c966dd83e84b646eb3acbdfbaef7b6b0db) --- .../cloud_storage_security.py | 39 +++++- litellm/proxy/common_request_processing.py | 42 +++++- litellm/proxy/proxy_server.py | 4 +- litellm/types/files.py | 5 +- ...llm_core_utils_prompt_templates_factory.py | 10 ++ .../test_cloud_storage_security.py | 42 ++++++ .../proxy/test_common_request_processing.py | 131 +++++++++++++++++- 7 files changed, 257 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 7c2da8d984d..6971eecaf74 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -3,6 +3,7 @@ import posixpath import re from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from functools import partial from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType from typing import Any, Final, cast @@ -52,6 +53,7 @@ _GEMINI_MEDIA_PART_KEYS: Final = frozenset({"file_data", "fileData"}) _MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data", "file_uri", "fileUri"}) _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" def is_s3_uri(value: str) -> bool: @@ -74,6 +76,10 @@ def _s3_uris_embedded_in_file_id(file_id: str) -> Iterator[str]: return chain.from_iterable(map(_s3_uris_in, filter(None, map(_base64url_text, encodings)))) +def _with_embedded_s3_uris(file_id: str) -> Iterator[str]: + return chain((file_id,), _s3_uris_embedded_in_file_id(file_id)) + + @dataclass(frozen=True, slots=True) class _BodyNode: value: object @@ -99,19 +105,32 @@ def _next_level(level: tuple[_BodyNode, ...], _: None) -> tuple[_BodyNode, ...]: return tuple(chain.from_iterable(map(_child_nodes, level))) -def _body_nodes(request_body: Mapping[str, object]) -> Iterator[_BodyNode]: +def _body_nodes(request_body: object) -> Iterator[_BodyNode]: levels: Final = accumulate(repeat(None), _next_level, initial=(_BodyNode(request_body),)) return chain.from_iterable(takewhile(bool, levels)) -def _media_uri_candidates(node: _BodyNode) -> Iterator[str]: +def _string_leaf(node: _BodyNode) -> str | None: + return node.value if isinstance(node.value, str) else 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 + + +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": - return chain((node.value,), _s3_uris_embedded_in_file_id(node.value)) + 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 ()) +def _mapped_provider_file_ids(request_body: Mapping[str, object]) -> Iterator[str]: + mapping_leaves: Final = filter(None, map(_string_leaf, _body_nodes(request_body.get(_MODEL_FILE_ID_MAPPING_KEY)))) + return chain.from_iterable(map(_with_embedded_s3_uris, mapping_leaves)) + + def _s3_bucket(uri: str) -> str | None: return uri.removeprefix(_S3_SCHEME).partition("/")[0] if uri.startswith(_S3_SCHEME) else None @@ -136,9 +155,17 @@ def _refused_in_some_decoding( return bool(parts) and (not settled or not all(map(is_allowed, parts))) -def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: - media_uris: Final = chain.from_iterable(map(_media_uri_candidates, _body_nodes(request_body))) - return tuple(dict.fromkeys(filter(is_s3_uri, media_uris))) +def media_file_ids(request_body: Mapping[str, object]) -> tuple[str, ...]: + return tuple(dict.fromkeys(filter(None, map(_media_file_id, _body_nodes(request_body))))) + + +def s3_media_uris( + request_body: Mapping[str, object], trusted_file_ids: frozenset[str] = frozenset() +) -> tuple[str, ...]: + media_uris: Final = chain.from_iterable( + map(partial(_media_uri_candidates, trusted_file_ids), _body_nodes(request_body)) + ) + return tuple(dict.fromkeys(filter(is_s3_uri, chain(media_uris, _mapped_provider_file_ids(request_body))))) def is_managed_s3_uri(uri: str) -> bool: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9957eb7792e..1406b2e784c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -6,7 +6,7 @@ import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache -from itertools import chain +from itertools import chain, compress from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -60,6 +60,7 @@ from litellm.litellm_core_utils.bug_report import ( ) from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_s3_uri, + media_file_ids, names_an_unlisted_s3_bucket, s3_media_uris, ) @@ -130,6 +131,10 @@ from litellm.proxy.common_utils.sse_keepalive import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.native_compaction import with_proxy_compaction_executor +from litellm.proxy.openai_files_endpoints.common_utils import ( + ManagedResourceAccessChecker, + is_base64_encoded_unified_file_id, +) from litellm.proxy.route_llm_request import ( route_request, ) @@ -860,20 +865,43 @@ def _request_model(data: Mapping[str, object]) -> str | list[str] | None: return None -def refuse_s3_media_the_key_may_not_read( +async def _key_may_read_unified_file_id( + managed_files: ManagedResourceAccessChecker, user_api_key_dict: UserAPIKeyAuth, unified_file_id: str +) -> bool: + try: + return await managed_files.can_user_call_unified_file_id(unified_file_id, user_api_key_dict) + except HTTPException: + return False + + +async def _unified_file_ids_the_key_may_read( + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, managed_files: object +) -> frozenset[str]: + if not isinstance(managed_files, ManagedResourceAccessChecker): + return frozenset() + unified_file_ids: Final = tuple(filter(is_base64_encoded_unified_file_id, media_file_ids(data))) + readable: Final = await asyncio.gather( + *(_key_may_read_unified_file_id(managed_files, user_api_key_dict, file_id) for file_id in unified_file_ids) + ) + return frozenset(compress(unified_file_ids, readable)) + + +async def refuse_s3_media_the_key_may_not_read( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], + managed_files: object, ) -> None: - s3_uris: Final = s3_media_uris(data) - if not s3_uris or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return + readable_file_ids: Final = await _unified_file_ids_the_key_may_read(data, user_api_key_dict, managed_files) + s3_uris: Final = s3_media_uris(data, readable_file_ids) managed_uris: Final = tuple(filter(is_managed_s3_uri, s3_uris)) if managed_uris: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " - f"({', '.join(managed_uris)}). Use the LiteLLM managed file id returned when the file was created.", + f"({', '.join(managed_uris)}). Use a LiteLLM managed file id that this key's user or team created.", ) if general_settings.get("allow_client_side_credentials") is True: return @@ -2309,7 +2337,9 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) - refuse_s3_media_the_key_may_not_read(self.data, user_api_key_dict, general_settings) + await refuse_s3_media_the_key_may_not_read( + self.data, user_api_key_dict, 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/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bb74b2a08af..e44af2aaf6d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17680,7 +17680,9 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) - refuse_s3_media_the_key_may_not_read(data, user_api_key_dict, general_settings) + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) # Include original request and headers in the data data["proxy_server_request"] = { diff --git a/litellm/types/files.py b/litellm/types/files.py index 259a836d9ad..083f5bf77d1 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -29,6 +29,7 @@ class FileType(Enum): JSON = "JSON" M4A = "M4A" M4V = "M4V" + MD = "MD" MOV = "MOV" MP3 = "MP3" MP4 = "MP4" @@ -75,6 +76,7 @@ FILE_EXTENSIONS: Final[Mapping[FileType, list[str]]] = MappingProxyType( FileType.JSON: ["json"], FileType.M4A: ["m4a"], FileType.M4V: ["m4v"], + FileType.MD: ["md", "markdown"], FileType.MOV: ["mov"], FileType.MP3: ["mp3"], FileType.MP4: ["mp4"], @@ -122,6 +124,7 @@ FILE_MIME_TYPES: Final[Mapping[FileType, str]] = MappingProxyType( FileType.JSON: "application/json", FileType.M4A: "audio/x-m4a", FileType.M4V: "video/x-m4v", + FileType.MD: "text/markdown", FileType.MOV: "video/quicktime", FileType.MP3: "audio/mpeg", FileType.MP4: "video/mp4", @@ -238,7 +241,7 @@ def is_audio_file_type(file_type): # Text -TEXT_FILE_TYPES: Final = {FileType.CSV, FileType.HTML, FileType.RTF, FileType.TXT} +TEXT_FILE_TYPES: Final = {FileType.CSV, FileType.HTML, FileType.MD, FileType.RTF, FileType.TXT} def is_text_file_type(file_type): 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 71b97dbbd3e..13a9d129a7b 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 @@ -2802,6 +2802,16 @@ async def test_bedrock_converse_s3_pdf_is_a_document_named_after_its_uri(run_asy assert invoice_content[0]["document"]["name"] != report_name +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("uri", ["s3://media-bucket/docs/notes.md", "s3://media-bucket/docs/NOTES.markdown"]) +async def test_bedrock_converse_s3_markdown_without_a_format_is_an_md_document(uri: str, run_async: bool): + content: Final = await _converse_user_content({"type": "file", "file": {"file_id": uri}}, run_async) + + assert [{**block, "document": {**block["document"], "name": "-"}} for block in content] == [ + {"document": {"format": "md", "name": "-", "source": _s3_source(uri)}} + ] + + @pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) @pytest.mark.parametrize("uri", ["s3://media-bucket/uploads/7f3a9c", "s3://media-bucket/uploads/blob.unknownext"]) async def test_bedrock_converse_s3_media_without_a_known_extension_or_format_is_a_bad_request( 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 9aa44e1fccf..882f4983d5a 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, is_managed_s3_uri, names_an_unlisted_s3_bucket, + media_file_ids, s3_media_uris, ) @@ -86,6 +87,20 @@ 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( + { + **_user_turn({"type": "file", "file": {"file_id": "alias"}}), + "model_file_id_mapping": {"alias": {"deployment-1": _MANAGED_UPLOAD}}, + }, + id="model-file-id-mapping", + ), + pytest.param( + { + **_responses_turn({"type": "input_file", "file_id": "alias"}), + "model_file_id_mapping": {"alias": {"deployment-1": _MODEL_EMBEDDED_UPLOAD_ID}}, + }, + id="model-file-id-mapping-to-a-model-embedded-file-id", + ), pytest.param(_responses_turn({"type": "input_image", "image_url": _MANAGED_UPLOAD}), id="responses-image"), pytest.param(_responses_turn({"type": "input_file", "file_url": _MANAGED_UPLOAD}), id="responses-file-url"), pytest.param(_responses_turn({"type": "input_file", "file_id": _MANAGED_UPLOAD}), id="responses-file-id"), @@ -313,6 +328,33 @@ def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_bod assert s3_media_uris(request_body) == () +def test_s3_media_uris_skips_only_the_objects_a_trusted_file_id_embeds(): + listed_image: Final = "s3://media-bucket/dir/cat.png" + request_body: Final = _responses_turn( + {"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}, + {"type": "input_file", "file_id": _MODEL_EMBEDDED_UPLOAD_ID}, + {"type": "input_image", "image_url": listed_image}, + ) + + assert s3_media_uris(request_body, frozenset({_UNIFIED_BATCH_OUTPUT_ID})) == (_MANAGED_UPLOAD, listed_image) + assert s3_media_uris(request_body, frozenset({_UNIFIED_BATCH_OUTPUT_ID, _MODEL_EMBEDDED_UPLOAD_ID})) == ( + listed_image, + ) + + +def test_media_file_ids_lists_the_file_ids_of_media_parts_only(): + request_body: Final = { + **_user_turn( + {"type": "file", "file": {"file_id": "file-a"}}, + {"type": "text", "text": "file-b", "file_id": "file-b"}, + {"type": "file", "file": {"file_id": "file-a"}}, + ), + "metadata": {"file_id": "file-c"}, + } + + assert media_file_ids(request_body) == ("file-a",) + + 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 eaab813db9f..78cafd9b709 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,6 +2,7 @@ import asyncio import copy import datetime import json +from base64 import urlsafe_b64encode from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeAlias @@ -25,6 +26,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic +from litellm.types.utils import SpecialEnums from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -239,6 +241,48 @@ _S3_MEDIA_SURFACES: Final = ( ) +def _chat_file(file_id: str) -> dict[str, object]: + return {"messages": [{"role": "user", "content": [{"type": "file", "file": {"file_id": file_id}}]}]} + + +def _responses_file(file_id: str) -> dict[str, object]: + return {"input": [{"role": "user", "content": [{"type": "input_file", "file_id": 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"), +) +_UPLOADED_OBJECT: Final = "s3://media-bucket/litellm-bedrock-files/upload.pdf" +_UPLOADED_FILE_ID: Final = ( + urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/pdf", "u1", "nova", _UPLOADED_OBJECT, "deployment-1" + ).encode() + ) + .decode() + .rstrip("=") +) + + +class _ManagedFilesCreatedBy(CustomLogger): + def __init__(self, creators: Mapping[str, str]) -> None: + super().__init__() + self.creators: Final = creators + + async def can_user_call_unified_file_id( + self, unified_file_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + if unified_file_id not in self.creators: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"File not found: {unified_file_id}") + return self.creators[unified_file_id] == user_api_key_dict.user_id + + async def can_user_call_unified_object_id( + self, unified_object_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + return False + + _S3_OBJECT: Final = "s3://media-bucket/dir/cat.png" _HELLO: Final = {"role": "user", "content": "hello"} _S3_VALUES_OUTSIDE_MEDIA_PARTS: Final = ( @@ -381,9 +425,14 @@ async def _send_s3_media( body: dict[str, object], user_role: LitellmUserRoles, general_settings: dict[str, object], + managed_files: CustomLogger | None = None, ) -> object: from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + if managed_files is not None: + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + router: Final = Router( model_list=[ { @@ -403,9 +452,9 @@ async def _send_s3_media( } ), fastapi_response=Response(), - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_role=user_role), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_id="user-a", user_role=user_role), route_type=route_type, - proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, proxy_config=RealProxyConfig(), llm_router=router, @@ -813,6 +862,84 @@ class TestProxyBaseLLMRequestProcessing: assert {key: data[key] for key in body} == body + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_FILE_ID_SURFACES) + async def test_non_admin_managed_file_id_its_user_created_passes( + self, route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]] + ) -> None: + response: Final = await _send_s3_media( + route_type, + path, + body_for(_UPLOADED_FILE_ID), + LitellmUserRoles.INTERNAL_USER, + {}, + _ManagedFilesCreatedBy({_UPLOADED_FILE_ID: "user-a"}), + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_FILE_ID_SURFACES) + @pytest.mark.parametrize( + "managed_files", + [ + pytest.param(_ManagedFilesCreatedBy({_UPLOADED_FILE_ID: "user-b"}), id="another-user-created-it"), + pytest.param(_ManagedFilesCreatedBy({}), id="forged-id"), + pytest.param(None, id="no-managed-files-hook"), + ], + ) + async def test_non_admin_managed_file_id_the_key_may_not_read_is_refused( + self, + route_type: _S3MediaRoute, + path: str, + body_for: Callable[[str], dict[str, object]], + managed_files: CustomLogger | None, + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + route_type, path, body_for(_UPLOADED_FILE_ID), LitellmUserRoles.INTERNAL_USER, {}, managed_files + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert _UPLOADED_OBJECT in refused.value.detail + assert "LiteLLM-managed storage" in refused.value.detail + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("mapped_source", "named_in_refusal"), + [ + pytest.param( + "s3://other-bucket/private.pdf", + "general_settings.allowed_client_side_s3_buckets", + id="unlisted-bucket", + ), + pytest.param( + "s3://media-bucket/litellm-bedrock-files/other-team.pdf", "LiteLLM-managed storage", id="managed-key" + ), + ], + ) + async def test_non_admin_s3_source_behind_a_model_file_id_mapping_is_refused_while_a_proxy_admin_key_passes( + self, mapped_source: str, named_in_refusal: str + ) -> None: + body: Final = {**_chat_file("alias"), "model_file_id_mapping": {"alias": {"deployment-1": mapped_source}}} + + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + body, + LitellmUserRoles.INTERNAL_USER, + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + ) + admin_response: Final = await _send_s3_media( + "acompletion", "/v1/chat/completions", body, LitellmUserRoles.PROXY_ADMIN, {} + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert mapped_source in refused.value.detail + assert named_in_refusal in refused.value.detail + assert "a cat" in str(admin_response) + @pytest.mark.asyncio async def test_non_admin_s3_media_refusal_names_a_repeated_object_once(self) -> None: with pytest.raises(HTTPException) as refused: From 34e272a4aa4b38baea890ba772f8f481c12a5d19 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:39:32 -0800 Subject: [PATCH 09/24] fix(proxy): keep inline tool-result video dropped and let an admitted client AWS identity skip the s3 bucket list A tool result forwards a video to Bedrock only when it is an s3Location. An inline base64 video is dropped with a warning, as it was before this change, since Converse rejects video in a toolResult block A non-admin key that sends an AWS identity its deployment admits through configurable_clientside_auth_params (aws_role_name, aws_web_identity_token or aws_profile_name) skips the bucket list, because Bedrock then reads the object with that identity. The managed-prefix 403 still applies The check now skips the body walk when no string starts with s3:// and no file_id or model_file_id_mapping key is present, which keeps large image bodies cheap Adapted for lts/1.108.x. The per-frame check on WS /v1/responses is not ported, because this line has no Responses WebSocket frame guard to hook it into. The gate keeps this line's admits_s3_bucket allowlist and its order of walking the body before the admin role check (cherry picked from commit 6662cb34d768bd68a7dc34ea27c363a1bd7f6578) --- .../cloud_storage_security.py | 20 +++++ .../prompt_templates/factory.py | 2 +- litellm/proxy/auth/auth_utils.py | 21 ++++++ litellm/proxy/common_request_processing.py | 18 ++++- litellm/proxy/proxy_server.py | 2 +- ...llm_core_utils_prompt_templates_factory.py | 36 +++++++++ .../test_cloud_storage_security.py | 13 +++- .../proxy/test_common_request_processing.py | 73 ++++++++++++++++++- 8 files changed, 178 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 6971eecaf74..edf54fd375f 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -54,6 +54,7 @@ _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}) def is_s3_uri(value: str) -> bool: @@ -155,6 +156,23 @@ def _refused_in_some_decoding( return bool(parts) and (not settled or not all(map(is_allowed, parts))) +def _may_name_an_s3_uri(value: object) -> bool: + if isinstance(value, str): + return is_s3_uri(value) + fields: Final = as_str_mapping(value) + if fields is not None: + return not _KEYS_THAT_MAY_HIDE_AN_S3_URI.isdisjoint(fields) or any(map(_may_name_an_s3_uri, fields.values())) + items: Final = as_object_list(value) + return items is not None and any(map(_may_name_an_s3_uri, items)) + + +def _may_carry_s3_media(request_body: Mapping[str, object]) -> bool: + try: + return _may_name_an_s3_uri(request_body) + except RecursionError: + return True + + def media_file_ids(request_body: Mapping[str, object]) -> tuple[str, ...]: return tuple(dict.fromkeys(filter(None, map(_media_file_id, _body_nodes(request_body))))) @@ -162,6 +180,8 @@ def media_file_ids(request_body: Mapping[str, object]) -> tuple[str, ...]: def s3_media_uris( request_body: Mapping[str, object], trusted_file_ids: frozenset[str] = frozenset() ) -> tuple[str, ...]: + if not _may_carry_s3_media(request_body): + return () media_uris: Final = chain.from_iterable( map(partial(_media_uri_candidates, trusted_file_ids), _body_nodes(request_body)) ) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7da1753c04f..6da1738dcc8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4081,7 +4081,7 @@ def _append_bedrock_tool_result_media_block( tool_result_content_blocks.append(BedrockToolResultContentBlock(image=processed_block["image"])) elif "document" in processed_block: tool_result_content_blocks.append(BedrockToolResultContentBlock(document=processed_block["document"])) - elif "video" in processed_block: + elif "video" in processed_block and "s3Location" in processed_block["video"]["source"]: tool_result_content_blocks.append(BedrockToolResultContentBlock(video=processed_block["video"])) else: verbose_logger.warning( diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index e9e6157252f..a6f777d9f45 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -234,6 +234,27 @@ def _allow_model_level_clientside_configurable_parameters( ) +_CLIENT_AWS_IDENTITY_PARAMS: Final = ("aws_role_name", "aws_web_identity_token", "aws_profile_name") + + +def _sends_a_value(body: Mapping[str, object], param: str) -> bool: + value: Final = body.get(param) + return isinstance(value, str) and bool(value) + + +def deployment_admits_a_client_aws_identity(body: Mapping[str, object], llm_router: Router | None) -> bool: + model: Final = body.get("model") + if not isinstance(model, str): + return False + return any( + _allow_model_level_clientside_configurable_parameters( + model=model, param=param, request_body_value=body[param], llm_router=llm_router + ) + for param in _CLIENT_AWS_IDENTITY_PARAMS + if _sends_a_value(body, param) + ) + + # Config dicts whose entries are spread as ``**dict`` into outbound LLM # API calls. ``litellm_embedding_config`` is consumed by the Milvus # vector store transformer. ``extra_body`` is the OpenAI-SDK passthrough diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1406b2e784c..5074a0ab22a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -100,7 +100,12 @@ from litellm.proxy.auth.auth_checks import ( request_skips_budget_checks, tag_max_budget_check_for_tags, ) -from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route +from litellm.proxy.auth.auth_utils import ( + ClientSideStorageAllowlist, + check_response_size_is_safe, + deployment_admits_a_client_aws_identity, + get_request_route, +) from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -891,6 +896,7 @@ async def refuse_s3_media_the_key_may_not_read( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], managed_files: object, + llm_router: Router | None, ) -> None: if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return @@ -903,7 +909,9 @@ async def refuse_s3_media_the_key_may_not_read( detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " f"({', '.join(managed_uris)}). Use a LiteLLM managed file id that this key's user or team created.", ) - if general_settings.get("allow_client_side_credentials") is True: + if general_settings.get("allow_client_side_credentials") is True or deployment_admits_a_client_aws_identity( + data, llm_router + ): return is_listed: Final = ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed)) @@ -2338,7 +2346,11 @@ class ProxyBaseLLMRequestProcessing: code=status.HTTP_400_BAD_REQUEST, ) await refuse_s3_media_the_key_may_not_read( - self.data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + self.data, + user_api_key_dict, + general_settings, + proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router, ) self.data = await add_litellm_data_to_request( data=self.data, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e44af2aaf6d..00b896a7f1e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17681,7 +17681,7 @@ async def async_queue_request( data = await request.json() data.pop("_litellm_strip_stream_usage", None) await refuse_s3_media_the_key_may_not_read( - data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files"), llm_router ) # Include original request and headers in the data 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 13a9d129a7b..39e2d4364bd 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 @@ -2916,6 +2916,42 @@ def test_bedrock_converse_tool_result_s3_media_keeps_its_bucket_owner(part: Mapp ] +@pytest.mark.parametrize( + "file, expected_media", + [ + pytest.param( + {"file_id": "s3://media-bucket/clips/demo.mp4"}, + [{"video": {"format": "mp4", "source": _s3_source("s3://media-bucket/clips/demo.mp4")}}], + id="s3-video-is-forwarded", + ), + pytest.param( + {"file_data": "data:video/mp4;base64,AAAAGGZ0eXBtcDQy"}, + [], + id="inline-video-is-dropped", + ), + ], +) +def test_bedrock_converse_tool_result_video_is_forwarded_only_from_s3( + file: Mapping[str, object], expected_media: list[Mapping[str, object]] +): + messages: Final = [ + {"role": "user", "content": "Describe the clip."}, + { + "tool_call_id": "tooluse_video_1", + "role": "tool", + "name": "fetch_clip", + "content": [{"type": "text", "text": "clip:"}, {"type": "file", "file": file}], + }, + ] + + with respx.mock: + converted: Final = _bedrock_converse_messages_pt( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + + assert converted[-1]["content"][-1]["toolResult"]["content"] == [{"text": "clip:"}, *expected_media] + + @pytest.mark.parametrize("bucket_owner", [111122223333, ""], ids=["int", "empty-string"]) @pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) def test_bedrock_converse_tool_result_s3_media_with_a_bucket_owner_that_is_not_an_account_id_is_a_bad_request( 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 882f4983d5a..fee717f9c6e 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -1,6 +1,7 @@ +import sys from base64 import urlsafe_b64encode from functools import reduce -from typing import Final +from typing import Final, cast from urllib.parse import quote import pytest @@ -328,6 +329,16 @@ def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_bod assert s3_media_uris(request_body) == () +def test_s3_media_uris_finds_an_s3_object_nested_deeper_than_the_recursion_limit(): + nested: Final = reduce( + lambda inner, _: [inner], + range(sys.getrecursionlimit() + 100), + cast(object, [{"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}]), + ) + + assert s3_media_uris({"messages": [{"role": "user", "content": nested}]}) == (_MANAGED_UPLOAD,) + + def test_s3_media_uris_skips_only_the_objects_a_trusted_file_id_embeds(): listed_image: Final = "s3://media-bucket/dir/cat.png" request_body: Final = _responses_turn( diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 78cafd9b709..a51deaadbd6 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -426,6 +426,7 @@ async def _send_s3_media( user_role: LitellmUserRoles, general_settings: dict[str, object], managed_files: CustomLogger | None = None, + client_auth_params: tuple[str, ...] = (), ) -> object: from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig @@ -437,7 +438,11 @@ async def _send_s3_media( model_list=[ { "model_name": "nova", - "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "mock_response": "a cat"}, + "litellm_params": { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "mock_response": "a cat", + "configurable_clientside_auth_params": list(client_auth_params), + }, } ] ) @@ -835,6 +840,72 @@ class TestProxyBaseLLMRequestProcessing: assert "a cat" in str(response) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("param", "value"), + [ + pytest.param("aws_role_name", "arn:aws:iam::111122223333:role/reader", id="role"), + pytest.param("aws_web_identity_token", "oidc-token", id="web-identity-token"), + pytest.param("aws_profile_name", "reader", id="profile"), + ], + ) + async def test_non_admin_s3_media_read_with_a_client_aws_identity_the_deployment_admits_needs_no_listed_bucket( + self, param: str, value: str + ) -> None: + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + {**_chat_image("s3://media-bucket/dir/cat.png"), param: value}, + LitellmUserRoles.INTERNAL_USER, + {}, + client_auth_params=(param,), + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("uri", "identity", "admitted", "named_in_refusal"), + [ + pytest.param( + "s3://media-bucket/dir/cat.png", + {"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, + ("aws_region_name",), + "general_settings.allowed_client_side_s3_buckets", + id="identity-the-deployment-does-not-admit", + ), + pytest.param( + "s3://media-bucket/dir/cat.png", + {"aws_role_name": ""}, + ("aws_role_name",), + "general_settings.allowed_client_side_s3_buckets", + id="empty-identity", + ), + pytest.param( + "s3://media-bucket/litellm-bedrock-files/other-team.png", + {"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, + ("aws_role_name",), + "LiteLLM-managed storage", + id="managed-key-with-an-admitted-identity", + ), + ], + ) + async def test_non_admin_s3_media_is_refused_unless_the_deployment_admits_the_client_aws_identity( + self, uri: str, identity: dict[str, object], admitted: tuple[str, ...], named_in_refusal: str + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + {**_chat_image(uri), **identity}, + LitellmUserRoles.INTERNAL_USER, + {}, + client_auth_params=admitted, + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert named_in_refusal in refused.value.detail + @pytest.mark.asyncio @pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_OUTSIDE_MEDIA_PARTS) async def test_non_admin_s3_values_outside_media_parts_pass_with_no_bucket_listed( From 5ef4feba948333bc4d1a8bcec99518253db525ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:04:59 -0800 Subject: [PATCH 10/24] fix(proxy): check s3 media on every Responses WebSocket frame on the LTS line The managed Responses WebSocket handler calls litellm.aresponses once per response.create frame, and the HTTP s3 media check only saw the first frame as a raw JSON string. A non-admin key could send an s3 object from an unlisted bucket, or one under LiteLLM-managed storage, to Bedrock over WS /v1/responses The WS route now hands the handler a frame check bound to the connection's key, general_settings, managed files hook and router. The handler runs it on the params it is about to forward for every frame, the first one included, and a refused frame gets an invalid_request_error frame with the 403 message while the socket keeps serving later frames --- litellm/llms/custom_httpx/llm_http_handler.py | 3 + .../proxy/response_api_endpoints/endpoints.py | 32 +++++++- litellm/responses/main.py | 4 +- litellm/responses/streaming_iterator.py | 13 +++- .../response_api_endpoints/test_endpoints.py | 75 +++++++++++++++++++ 5 files changed, 124 insertions(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c042264c034..646a52384d9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -119,6 +119,7 @@ from litellm.responses.streaming_iterator import ( ProjectQuotaCallback, ResponsesAPIStreamingIterator, ResponsesWebSocketStreaming, + ResponsesWsFrameCheck, SyncResponsesAPIStreamingIterator, ) from litellm.types.containers.main import ( @@ -6851,6 +6852,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_check: ResponsesWsFrameCheck | None = None, **kwargs: object, ) -> Exception | None: """ @@ -6883,6 +6885,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + frame_check=frame_check, **kwargs, ) await handler.run() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 63062569af4..9b6838dad58 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + refuse_s3_media_the_key_may_not_read, resolve_counted_tool_fields, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id @@ -71,6 +72,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.responses.streaming_iterator import ResponsesWsFrameCheck from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1904,6 +1906,25 @@ async def _enforce_responses_ws_first_frame_model_auth( ) +def _s3_media_frame_check( + admitted_model: object, + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + managed_files: object, + llm_router: Router | None, +) -> ResponsesWsFrameCheck: + async def refusal(params: Mapping[str, object], /) -> str | None: + try: + await refuse_s3_media_the_key_may_not_read( + {**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files, llm_router + ) + except HTTPException as refused: + return str(refused.detail) + return None + + return refusal + + async def _responses_websocket_session( websocket: WebSocket, model: str | None, @@ -2000,7 +2021,16 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_frame_check=_s3_media_frame_check( + data.get("model"), + user_api_key_dict, + general_settings, + proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router, + ), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index a858b91c4c8..0d04558322b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -77,7 +77,7 @@ if TYPE_CHECKING: else: MCPTool = Any -from .streaming_iterator import BaseResponsesAPIStreamingIterator +from .streaming_iterator import BaseResponsesAPIStreamingIterator, ResponsesWsFrameCheck __all__ = ( "acancel_responses", @@ -2488,6 +2488,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_frame_check: ResponsesWsFrameCheck | None = None, **kwargs, ) -> Exception | None: """ @@ -2580,5 +2581,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + frame_check=responses_ws_frame_check, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index cb745384403..61abeb1d74b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -66,6 +66,10 @@ if TYPE_CHECKING: from litellm.types.router import LiteLLM_Params +class ResponsesWsFrameCheck(Protocol): + async def __call__(self, params: Mapping[str, object], /) -> str | None: ... + + class ProjectQuotaCallback(Protocol): async def enforce_project_io_token_quota_for_frame( self, @@ -2723,6 +2727,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + frame_check: ResponsesWsFrameCheck | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -2741,6 +2746,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () + self.frame_check = frame_check # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -3144,6 +3150,12 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc) return + call_kwargs: Final = self._build_base_call_kwargs(msg_obj) + refusal: Final = None if self.frame_check is None else await self.frame_check(call_kwargs) + if refusal is not None: + await self._send_error(refusal, "invalid_request_error") + return + try: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message @@ -3152,7 +3164,6 @@ class ManagedResponsesWebSocketHandler: await self._send_error(str(e), error_type="rate_limit_exceeded") return - call_kwargs: Final = self._build_base_call_kwargs(msg_obj) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index bde570a8f61..077d7e2aa3c 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1612,6 +1612,81 @@ class TestManagedResponsesWSFirstMessage: assert processed == [subsequent] +class TestResponsesWSS3MediaFrameCheck: + @pytest.mark.asyncio + async def test_managed_socket_refuses_a_non_admin_s3_image_on_every_frame_and_forwards_a_proxy_admin_frame( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm import Router + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import _responses_websocket_session + + image_uri: Final = "s3://media-bucket/dir/cat.png" + frame: Final = json.dumps( + { + "type": "response.create", + "input": [{"role": "user", "content": [{"type": "input_image", "image_url": image_uri}]}], + } + ) + forwarded_event: Final = {"type": "response.created", "response": {"id": "resp_admin"}} + provider_inputs: Final[list[object]] = [] + + async def fake_aresponses(**kwargs: object): + provider_inputs.append(kwargs["input"]) + + async def events(): + yield forwarded_event + + return events() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + test_router: Final = Router( + model_list=[{"model_name": "nova", "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}}] + ) + + async def route_to_test_router(*, data: dict[str, Any], **_kwargs: object): + return test_router._aresponses_websocket(**data) + + async def frames_sent_back(user_role: LitellmUserRoles) -> list[object]: + ws = MagicMock() + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.receive_text = AsyncMock(side_effect=[frame, frame, Exception("disconnect")]) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + with ( + patch( # test-quality-ok: the pre-call processor needs a live proxy and is not what this test covers + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + side_effect=lambda data: MagicMock( + common_processing_pre_call_logic=AsyncMock(return_value=(data, MagicMock())) + ), + ), + patch( # test-quality-ok: routes to a router holding one Bedrock deployment instead of the proxy's + "litellm.proxy.route_llm_request.route_request", + side_effect=route_to_test_router, + ), + ): + await _responses_websocket_session( + websocket=ws, + model="nova", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="user-a", user_role=user_role), + ) + return [json.loads(call.args[0]) for call in ws.send_text.await_args_list] + + refused: Final = await frames_sent_back(LitellmUserRoles.INTERNAL_USER) + + assert provider_inputs == [] + assert [event["type"] for event in refused] == ["error", "error"] + assert all("allowed_client_side_s3_buckets" in event["error"]["message"] for event in refused) + assert all(image_uri in event["error"]["message"] for event in refused) + + admitted: Final = await frames_sent_back(LitellmUserRoles.PROXY_ADMIN) + + assert admitted == [forwarded_event, forwarded_event] + assert len(provider_inputs) == 2 + assert all(image_uri in json.dumps(provider_input) for provider_input in provider_inputs) + + class TestResponsesWSStreamingFirstMessage: @pytest.mark.asyncio async def test_client_to_backend_replays_first_message(self): From b66cf9a4354d7dc1739c07c7cb5117c83c328c8d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 20:27:31 -0800 Subject: [PATCH 11/24] fix(proxy): let a client AWS identity skip the s3 bucket list only when Bedrock signs with it An admitted aws_role_name, aws_web_identity_token or aws_profile_name now skips the bucket list only when it decides the credentials get_credentials picks, on every deployment of the model group. A web identity token without a role and a session name does not, since get_credentials then ignores it. A profile does not when the deployment or the environment sets a role, since the role branch runs first. In those cases Bedrock still reads with the proxy's credentials, so the bucket list applies BaseAWSLLM.identity_choosing_params mirrors the branch order of get_credentials, and both read the same AWS_* environment fallbacks through one helper Ported to lts/1.108.x, where the model group is read with Router.get_model_list --- litellm/llms/bedrock/base_aws_llm.py | 19 ++++- litellm/proxy/auth/auth_utils.py | 39 ++++++++-- .../proxy/test_common_request_processing.py | 75 +++++++++++++++++-- 3 files changed, 121 insertions(+), 12 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 0c4151b4b1b..eecf54eb0ca 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -232,6 +232,10 @@ class SignsRequestsWithAWS: pass +def _or_env(value: str | None, env_var: str) -> str | None: + return value if value is not None else os.getenv(env_var) + + class BaseAWSLLM(SignsRequestsWithAWS): # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request). # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static @@ -346,6 +350,19 @@ class BaseAWSLLM(SignsRequestsWithAWS): def _is_auth_with_aws_profile(aws_profile_name: str | None) -> bool: return aws_profile_name is not None + @classmethod + def identity_choosing_params(cls, params: AwsAuthParams) -> tuple[str, ...]: + token: Final = _or_env(params.aws_web_identity_token, "AWS_WEB_IDENTITY_TOKEN") + role: Final = _or_env(params.aws_role_name, "AWS_ROLE_NAME") + session: Final = _or_env(params.aws_session_name, "AWS_SESSION_NAME") + if cls._is_auth_with_web_identity_token(token, role, session): + return ("aws_web_identity_token", "aws_role_name") + if cls._is_auth_with_aws_role(role): + return ("aws_role_name",) + if cls._is_auth_with_aws_profile(_or_env(params.aws_profile_name, "AWS_PROFILE_NAME")): + return ("aws_profile_name",) + return () + @staticmethod def _is_auth_with_aws_session_token_tuple( aws_access_key_id: str | None, @@ -398,7 +415,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_sts_endpoint, aws_external_id, ) = tuple( - value if value is not None else os.getenv(env_var) + _or_env(value, env_var) for value, env_var in ( (aws_access_key_id, "AWS_ACCESS_KEY_ID"), (aws_secret_access_key, "AWS_SECRET_ACCESS_KEY"), diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a6f777d9f45..128c3c7cea1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -31,13 +31,20 @@ from litellm.litellm_core_utils.url_utils import ( validate_url, ) from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) -from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present +from litellm.types.router import ( + CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, + Deployment, + DeploymentTypedDict, + server_owned_wif_fields_present, +) from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params from litellm.types.videos.main import DecodedVideoId @@ -242,16 +249,36 @@ def _sends_a_value(body: Mapping[str, object], param: str) -> bool: return isinstance(value, str) and bool(value) +def _client_aws_fields(body: Mapping[str, object]) -> Mapping[str, object]: + return {key: body[key] for key in AWS_AUTH_PARAM_KEYS if _sends_a_value(body, key)} + + +def _identity_choosing_params(client_fields: Mapping[str, object], deployment: DeploymentTypedDict) -> tuple[str, ...]: + return BaseAWSLLM.identity_choosing_params( + AwsAuthParams.model_validate({**deployment["litellm_params"], **client_fields}) + ) + + def deployment_admits_a_client_aws_identity(body: Mapping[str, object], llm_router: Router | None) -> bool: model: Final = body.get("model") - if not isinstance(model, str): + if not isinstance(model, str) or llm_router is None: return False - return any( - _allow_model_level_clientside_configurable_parameters( - model=model, param=param, request_body_value=body[param], llm_router=llm_router - ) + admitted: Final = frozenset( + param for param in _CLIENT_AWS_IDENTITY_PARAMS if _sends_a_value(body, param) + and _allow_model_level_clientside_configurable_parameters( + model=model, param=param, request_body_value=body[param], llm_router=llm_router + ) + ) + deployments: Final = llm_router.get_model_list(model_name=model) or () + client_fields: Final = _client_aws_fields(body) + return ( + bool(admitted) + and bool(deployments) + and all( + not admitted.isdisjoint(_identity_choosing_params(client_fields, deployment)) for deployment in deployments + ) ) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index a51deaadbd6..fc5497f4293 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -419,6 +419,9 @@ _S3_VALUES_ON_NON_MEDIA_ROUTES: Final = ( ) +_AWS_IDENTITY_ENV_VARS: Final = ("AWS_WEB_IDENTITY_TOKEN", "AWS_ROLE_NAME", "AWS_SESSION_NAME", "AWS_PROFILE_NAME") + + async def _send_s3_media( route_type: _S3MediaRoute, path: str, @@ -427,6 +430,7 @@ async def _send_s3_media( general_settings: dict[str, object], managed_files: CustomLogger | None = None, client_auth_params: tuple[str, ...] = (), + deployments_aws_params: tuple[Mapping[str, str], ...] = (MappingProxyType({}),), ) -> object: from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig @@ -442,8 +446,10 @@ async def _send_s3_media( "model": "bedrock/us.amazon.nova-2-lite-v1:0", "mock_response": "a cat", "configurable_clientside_auth_params": list(client_auth_params), + **deployment_aws_params, }, } + for deployment_aws_params in deployments_aws_params ] ) return await ProxyBaseLLMRequestProcessing(data={"model": "nova", **body}).base_process_llm_request( @@ -842,16 +848,27 @@ class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio @pytest.mark.parametrize( - ("param", "value"), + ("param", "value", "deployments_aws_params"), [ - pytest.param("aws_role_name", "arn:aws:iam::111122223333:role/reader", id="role"), - pytest.param("aws_web_identity_token", "oidc-token", id="web-identity-token"), - pytest.param("aws_profile_name", "reader", id="profile"), + pytest.param("aws_role_name", "arn:aws:iam::111122223333:role/reader", ({},), id="role"), + pytest.param( + "aws_web_identity_token", + "oidc-token", + ({"aws_role_name": "arn:aws:iam::111122223333:role/web", "aws_session_name": "proxy"},), + id="web-identity-token-with-the-deployment-role-and-session", + ), + pytest.param("aws_profile_name", "reader", ({},), id="profile"), ], ) async def test_non_admin_s3_media_read_with_a_client_aws_identity_the_deployment_admits_needs_no_listed_bucket( - self, param: str, value: str + self, + param: str, + value: str, + deployments_aws_params: tuple[dict[str, str], ...], + monkeypatch: pytest.MonkeyPatch, ) -> None: + for env_var in _AWS_IDENTITY_ENV_VARS: + monkeypatch.delenv(env_var, raising=False) response: Final = await _send_s3_media( "acompletion", "/v1/chat/completions", @@ -859,10 +876,58 @@ class TestProxyBaseLLMRequestProcessing: LitellmUserRoles.INTERNAL_USER, {}, client_auth_params=(param,), + deployments_aws_params=deployments_aws_params, ) assert "a cat" in str(response) + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("identity", "admitted", "deployments_aws_params"), + [ + pytest.param( + {"aws_web_identity_token": "oidc-token"}, + ("aws_web_identity_token",), + ({},), + id="web-identity-token-without-a-role-and-session", + ), + pytest.param( + {"aws_profile_name": "reader"}, + ("aws_profile_name",), + ({"aws_role_name": "arn:aws:iam::111122223333:role/deployment"},), + id="profile-under-a-deployment-role", + ), + pytest.param( + {"aws_profile_name": "reader"}, + ("aws_profile_name",), + ({}, {"aws_role_name": "arn:aws:iam::111122223333:role/deployment"}), + id="profile-under-a-role-on-one-deployment-of-the-group", + ), + ], + ) + async def test_non_admin_s3_media_is_refused_when_the_admitted_client_aws_field_does_not_pick_the_credentials( + self, + identity: dict[str, str], + admitted: tuple[str, ...], + deployments_aws_params: tuple[dict[str, str], ...], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + for env_var in _AWS_IDENTITY_ENV_VARS: + monkeypatch.delenv(env_var, raising=False) + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + {**_chat_image("s3://media-bucket/dir/cat.png"), **identity}, + LitellmUserRoles.INTERNAL_USER, + {}, + client_auth_params=admitted, + deployments_aws_params=deployments_aws_params, + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert "general_settings.allowed_client_side_s3_buckets" in str(refused.value.detail) + @pytest.mark.asyncio @pytest.mark.parametrize( ("uri", "identity", "admitted", "named_in_refusal"), From c40e4fb2f30390eb2df3377f992aa3059ddad539 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 21:25:58 -0800 Subject: [PATCH 12/24] fix(proxy): keep the s3 bucket list for non-admin keys that send a client AWS identity A client AWS identity field does not prove that the caller can read the bucket. A role or profile name picks credentials the proxy holds, and Bedrock reads the object with them. So the bucket list now applies to every non-admin s3 media part, whatever AWS identity fields the request sends. allow_client_side_credentials: true still turns it off This removes BaseAWSLLM.identity_choosing_params and the router lookup that only the skip needed Ported to lts/1.108.x --- litellm/llms/bedrock/base_aws_llm.py | 19 +-- litellm/proxy/auth/auth_utils.py | 44 ------- litellm/proxy/common_request_processing.py | 13 +- litellm/proxy/proxy_server.py | 2 +- .../proxy/response_api_endpoints/endpoints.py | 4 +- .../proxy/test_common_request_processing.py | 119 ++---------------- 6 files changed, 15 insertions(+), 186 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index eecf54eb0ca..0c4151b4b1b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -232,10 +232,6 @@ class SignsRequestsWithAWS: pass -def _or_env(value: str | None, env_var: str) -> str | None: - return value if value is not None else os.getenv(env_var) - - class BaseAWSLLM(SignsRequestsWithAWS): # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request). # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static @@ -350,19 +346,6 @@ class BaseAWSLLM(SignsRequestsWithAWS): def _is_auth_with_aws_profile(aws_profile_name: str | None) -> bool: return aws_profile_name is not None - @classmethod - def identity_choosing_params(cls, params: AwsAuthParams) -> tuple[str, ...]: - token: Final = _or_env(params.aws_web_identity_token, "AWS_WEB_IDENTITY_TOKEN") - role: Final = _or_env(params.aws_role_name, "AWS_ROLE_NAME") - session: Final = _or_env(params.aws_session_name, "AWS_SESSION_NAME") - if cls._is_auth_with_web_identity_token(token, role, session): - return ("aws_web_identity_token", "aws_role_name") - if cls._is_auth_with_aws_role(role): - return ("aws_role_name",) - if cls._is_auth_with_aws_profile(_or_env(params.aws_profile_name, "AWS_PROFILE_NAME")): - return ("aws_profile_name",) - return () - @staticmethod def _is_auth_with_aws_session_token_tuple( aws_access_key_id: str | None, @@ -415,7 +398,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_sts_endpoint, aws_external_id, ) = tuple( - _or_env(value, env_var) + value if value is not None else os.getenv(env_var) for value, env_var in ( (aws_access_key_id, "AWS_ACCESS_KEY_ID"), (aws_secret_access_key, "AWS_SECRET_ACCESS_KEY"), diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 128c3c7cea1..22001e29267 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -31,18 +31,15 @@ from litellm.litellm_core_utils.url_utils import ( validate_url, ) from litellm.llms.azure.passthrough.transformation import azure_router_model_in_endpoint -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path from litellm.proxy._types import * from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_metadata -from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) from litellm.types.router import ( CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, - DeploymentTypedDict, server_owned_wif_fields_present, ) from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params @@ -241,47 +238,6 @@ def _allow_model_level_clientside_configurable_parameters( ) -_CLIENT_AWS_IDENTITY_PARAMS: Final = ("aws_role_name", "aws_web_identity_token", "aws_profile_name") - - -def _sends_a_value(body: Mapping[str, object], param: str) -> bool: - value: Final = body.get(param) - return isinstance(value, str) and bool(value) - - -def _client_aws_fields(body: Mapping[str, object]) -> Mapping[str, object]: - return {key: body[key] for key in AWS_AUTH_PARAM_KEYS if _sends_a_value(body, key)} - - -def _identity_choosing_params(client_fields: Mapping[str, object], deployment: DeploymentTypedDict) -> tuple[str, ...]: - return BaseAWSLLM.identity_choosing_params( - AwsAuthParams.model_validate({**deployment["litellm_params"], **client_fields}) - ) - - -def deployment_admits_a_client_aws_identity(body: Mapping[str, object], llm_router: Router | None) -> bool: - model: Final = body.get("model") - if not isinstance(model, str) or llm_router is None: - return False - admitted: Final = frozenset( - param - for param in _CLIENT_AWS_IDENTITY_PARAMS - if _sends_a_value(body, param) - and _allow_model_level_clientside_configurable_parameters( - model=model, param=param, request_body_value=body[param], llm_router=llm_router - ) - ) - deployments: Final = llm_router.get_model_list(model_name=model) or () - client_fields: Final = _client_aws_fields(body) - return ( - bool(admitted) - and bool(deployments) - and all( - not admitted.isdisjoint(_identity_choosing_params(client_fields, deployment)) for deployment in deployments - ) - ) - - # Config dicts whose entries are spread as ``**dict`` into outbound LLM # API calls. ``litellm_embedding_config`` is consumed by the Milvus # vector store transformer. ``extra_body`` is the OpenAI-SDK passthrough diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5074a0ab22a..4c87390bb5c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -100,12 +100,7 @@ from litellm.proxy.auth.auth_checks import ( request_skips_budget_checks, tag_max_budget_check_for_tags, ) -from litellm.proxy.auth.auth_utils import ( - ClientSideStorageAllowlist, - check_response_size_is_safe, - deployment_admits_a_client_aws_identity, - get_request_route, -) +from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -896,7 +891,6 @@ async def refuse_s3_media_the_key_may_not_read( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], managed_files: object, - llm_router: Router | None, ) -> None: if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return @@ -909,9 +903,7 @@ async def refuse_s3_media_the_key_may_not_read( detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " f"({', '.join(managed_uris)}). Use a LiteLLM managed file id that this key's user or team created.", ) - if general_settings.get("allow_client_side_credentials") is True or deployment_admits_a_client_aws_identity( - data, llm_router - ): + if general_settings.get("allow_client_side_credentials") is True: return is_listed: Final = ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed)) @@ -2350,7 +2342,6 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files"), - llm_router, ) self.data = await add_litellm_data_to_request( data=self.data, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 00b896a7f1e..e44af2aaf6d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17681,7 +17681,7 @@ async def async_queue_request( data = await request.json() data.pop("_litellm_strip_stream_usage", None) await refuse_s3_media_the_key_may_not_read( - data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files"), llm_router + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") ) # Include original request and headers in the data diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9b6838dad58..0d77c9750d2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1911,12 +1911,11 @@ def _s3_media_frame_check( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], managed_files: object, - llm_router: Router | None, ) -> ResponsesWsFrameCheck: async def refusal(params: Mapping[str, object], /) -> str | None: try: await refuse_s3_media_the_key_may_not_read( - {**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files, llm_router + {**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files ) except HTTPException as refused: return str(refused.detail) @@ -2028,7 +2027,6 @@ async def _responses_websocket_session( user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files"), - llm_router, ), **_routing_hints_from_first_ws_frame(first_message), ) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fc5497f4293..2922fca970c 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -419,9 +419,6 @@ _S3_VALUES_ON_NON_MEDIA_ROUTES: Final = ( ) -_AWS_IDENTITY_ENV_VARS: Final = ("AWS_WEB_IDENTITY_TOKEN", "AWS_ROLE_NAME", "AWS_SESSION_NAME", "AWS_PROFILE_NAME") - - async def _send_s3_media( route_type: _S3MediaRoute, path: str, @@ -430,7 +427,7 @@ async def _send_s3_media( general_settings: dict[str, object], managed_files: CustomLogger | None = None, client_auth_params: tuple[str, ...] = (), - deployments_aws_params: tuple[Mapping[str, str], ...] = (MappingProxyType({}),), + deployment_aws_params: Mapping[str, str] = MappingProxyType({}), ) -> object: from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig @@ -449,7 +446,6 @@ async def _send_s3_media( **deployment_aws_params, }, } - for deployment_aws_params in deployments_aws_params ] ) return await ProxyBaseLLMRequestProcessing(data={"model": "nova", **body}).base_process_llm_request( @@ -848,72 +844,20 @@ class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio @pytest.mark.parametrize( - ("param", "value", "deployments_aws_params"), - [ - pytest.param("aws_role_name", "arn:aws:iam::111122223333:role/reader", ({},), id="role"), - pytest.param( - "aws_web_identity_token", - "oidc-token", - ({"aws_role_name": "arn:aws:iam::111122223333:role/web", "aws_session_name": "proxy"},), - id="web-identity-token-with-the-deployment-role-and-session", - ), - pytest.param("aws_profile_name", "reader", ({},), id="profile"), - ], - ) - async def test_non_admin_s3_media_read_with_a_client_aws_identity_the_deployment_admits_needs_no_listed_bucket( - self, - param: str, - value: str, - deployments_aws_params: tuple[dict[str, str], ...], - monkeypatch: pytest.MonkeyPatch, - ) -> None: - for env_var in _AWS_IDENTITY_ENV_VARS: - monkeypatch.delenv(env_var, raising=False) - response: Final = await _send_s3_media( - "acompletion", - "/v1/chat/completions", - {**_chat_image("s3://media-bucket/dir/cat.png"), param: value}, - LitellmUserRoles.INTERNAL_USER, - {}, - client_auth_params=(param,), - deployments_aws_params=deployments_aws_params, - ) - - assert "a cat" in str(response) - - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("identity", "admitted", "deployments_aws_params"), + ("identity", "deployment_aws_params"), [ + pytest.param({"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, {}, id="role"), pytest.param( {"aws_web_identity_token": "oidc-token"}, - ("aws_web_identity_token",), - ({},), - id="web-identity-token-without-a-role-and-session", - ), - pytest.param( - {"aws_profile_name": "reader"}, - ("aws_profile_name",), - ({"aws_role_name": "arn:aws:iam::111122223333:role/deployment"},), - id="profile-under-a-deployment-role", - ), - pytest.param( - {"aws_profile_name": "reader"}, - ("aws_profile_name",), - ({}, {"aws_role_name": "arn:aws:iam::111122223333:role/deployment"}), - id="profile-under-a-role-on-one-deployment-of-the-group", + {"aws_role_name": "arn:aws:iam::111122223333:role/web", "aws_session_name": "proxy"}, + id="web-identity-token-with-the-deployment-role-and-session", ), + pytest.param({"aws_profile_name": "reader"}, {}, id="profile"), ], ) - async def test_non_admin_s3_media_is_refused_when_the_admitted_client_aws_field_does_not_pick_the_credentials( - self, - identity: dict[str, str], - admitted: tuple[str, ...], - deployments_aws_params: tuple[dict[str, str], ...], - monkeypatch: pytest.MonkeyPatch, + async def test_non_admin_s3_media_needs_a_listed_bucket_even_with_a_client_aws_identity_the_deployment_admits( + self, identity: dict[str, str], deployment_aws_params: dict[str, str] ) -> None: - for env_var in _AWS_IDENTITY_ENV_VARS: - monkeypatch.delenv(env_var, raising=False) with pytest.raises(HTTPException) as refused: await _send_s3_media( "acompletion", @@ -921,56 +865,13 @@ class TestProxyBaseLLMRequestProcessing: {**_chat_image("s3://media-bucket/dir/cat.png"), **identity}, LitellmUserRoles.INTERNAL_USER, {}, - client_auth_params=admitted, - deployments_aws_params=deployments_aws_params, + client_auth_params=tuple(identity), + deployment_aws_params=deployment_aws_params, ) assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail assert "general_settings.allowed_client_side_s3_buckets" in str(refused.value.detail) - @pytest.mark.asyncio - @pytest.mark.parametrize( - ("uri", "identity", "admitted", "named_in_refusal"), - [ - pytest.param( - "s3://media-bucket/dir/cat.png", - {"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, - ("aws_region_name",), - "general_settings.allowed_client_side_s3_buckets", - id="identity-the-deployment-does-not-admit", - ), - pytest.param( - "s3://media-bucket/dir/cat.png", - {"aws_role_name": ""}, - ("aws_role_name",), - "general_settings.allowed_client_side_s3_buckets", - id="empty-identity", - ), - pytest.param( - "s3://media-bucket/litellm-bedrock-files/other-team.png", - {"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, - ("aws_role_name",), - "LiteLLM-managed storage", - id="managed-key-with-an-admitted-identity", - ), - ], - ) - async def test_non_admin_s3_media_is_refused_unless_the_deployment_admits_the_client_aws_identity( - self, uri: str, identity: dict[str, object], admitted: tuple[str, ...], named_in_refusal: str - ) -> None: - with pytest.raises(HTTPException) as refused: - await _send_s3_media( - "acompletion", - "/v1/chat/completions", - {**_chat_image(uri), **identity}, - LitellmUserRoles.INTERNAL_USER, - {}, - client_auth_params=admitted, - ) - - assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail - assert named_in_refusal in refused.value.detail - @pytest.mark.asyncio @pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_OUTSIDE_MEDIA_PARTS) async def test_non_admin_s3_values_outside_media_parts_pass_with_no_bucket_listed( From 9faf74ef8ee45224bbdf4c8718ef4536e499ec45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:32:29 -0800 Subject: [PATCH 13/24] fix(proxy): look up managed file ids only when they can change the s3 refusal (cherry picked from commit 0043e3020f5258d6542cad2031f70b2928b89101) --- .../cloud_storage_security.py | 4 + litellm/proxy/auth/auth_utils.py | 10 +- litellm/proxy/common_request_processing.py | 54 ++++++++--- .../proxy/test_common_request_processing.py | 94 +++++++++++++++++++ 4 files changed, 145 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index edf54fd375f..b88ce55ffa3 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -177,6 +177,10 @@ def media_file_ids(request_body: Mapping[str, object]) -> tuple[str, ...]: return tuple(dict.fromkeys(filter(None, map(_media_file_id, _body_nodes(request_body))))) +def file_id_embeds_s3_uri_matching(file_id: str, matches: Callable[[str], bool]) -> bool: + return any(map(matches, _s3_uris_embedded_in_file_id(file_id))) + + def s3_media_uris( request_body: Mapping[str, object], trusted_file_ids: frozenset[str] = frozenset() ) -> tuple[str, ...]: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 96450b7fd72..b2467457eb4 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,8 +2,8 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping -from functools import lru_cache +from collections.abc import Callable, Collection, Iterator, Mapping +from functools import lru_cache, partial from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType @@ -532,7 +532,7 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( ) -class ClientSideStorageAllowlist(BaseModel): +class _ClientSideStorageAllowlist(BaseModel): model_config = ConfigDict(frozen=True) allowed_client_side_s3_buckets: tuple[str, ...] | None = None @@ -541,6 +541,10 @@ class ClientSideStorageAllowlist(BaseModel): return bucket in (self.allowed_client_side_s3_buckets or ()) +def client_side_s3_bucket_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]: + return _ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket + + def _check_banned_params( body: Mapping[str, object], general_settings: dict, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4c87390bb5c..b7b80858ba2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,10 +3,10 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence from datetime import datetime -from functools import lru_cache -from itertools import chain, compress +from functools import lru_cache, partial +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -59,6 +59,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.cloud_storage_security import ( + file_id_embeds_s3_uri_matching, is_managed_s3_uri, media_file_ids, names_an_unlisted_s3_bucket, @@ -100,7 +101,7 @@ from litellm.proxy.auth.auth_checks import ( request_skips_budget_checks, tag_max_budget_check_for_tags, ) -from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route +from litellm.proxy.auth.auth_utils import check_response_size_is_safe, client_side_s3_bucket_check, get_request_route from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -874,16 +875,41 @@ async def _key_may_read_unified_file_id( return False +async def _leading_readable_file_ids( + managed_files: ManagedResourceAccessChecker, user_api_key_dict: UserAPIKeyAuth, unified_file_ids: Iterable[str] +) -> AsyncIterator[str]: + for file_id in unified_file_ids: + if not await _key_may_read_unified_file_id(managed_files, user_api_key_dict, file_id): + return + yield file_id + + async def _unified_file_ids_the_key_may_read( - data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, managed_files: object + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + managed_files: object, + is_refused: Callable[[str], bool], ) -> frozenset[str]: if not isinstance(managed_files, ManagedResourceAccessChecker): return frozenset() - unified_file_ids: Final = tuple(filter(is_base64_encoded_unified_file_id, media_file_ids(data))) - readable: Final = await asyncio.gather( - *(_key_may_read_unified_file_id(managed_files, user_api_key_dict, file_id) for file_id in unified_file_ids) + file_ids_hiding_a_refused_uri: Final = filter( + partial(file_id_embeds_s3_uri_matching, matches=is_refused), + filter(is_base64_encoded_unified_file_id, media_file_ids(data)), ) - return frozenset(compress(unified_file_ids, readable)) + return frozenset( + [ + file_id + async for file_id in _leading_readable_file_ids( + managed_files, user_api_key_dict, file_ids_hiding_a_refused_uri + ) + ] + ) + + +def _unlisted_s3_uri_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]: + if general_settings.get("allow_client_side_credentials") is True: + return lambda _: False + return partial(names_an_unlisted_s3_bucket, is_listed=client_side_s3_bucket_check(general_settings)) async def refuse_s3_media_the_key_may_not_read( @@ -894,7 +920,10 @@ async def refuse_s3_media_the_key_may_not_read( ) -> None: if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return - readable_file_ids: Final = await _unified_file_ids_the_key_may_read(data, user_api_key_dict, managed_files) + is_unlisted: Final = _unlisted_s3_uri_check(general_settings) + readable_file_ids: Final = await _unified_file_ids_the_key_may_read( + data, user_api_key_dict, managed_files, lambda uri: is_managed_s3_uri(uri) or is_unlisted(uri) + ) s3_uris: Final = s3_media_uris(data, readable_file_ids) managed_uris: Final = tuple(filter(is_managed_s3_uri, s3_uris)) if managed_uris: @@ -903,10 +932,7 @@ async def refuse_s3_media_the_key_may_not_read( detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " f"({', '.join(managed_uris)}). Use a LiteLLM managed file id that this key's user or team created.", ) - if general_settings.get("allow_client_side_credentials") is True: - return - is_listed: Final = ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket - unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed)) + unlisted_uris: Final = tuple(filter(is_unlisted, s3_uris)) if unlisted_uris: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2922fca970c..27d56ade859 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -283,6 +283,30 @@ class _ManagedFilesCreatedBy(CustomLogger): return False +class _ManagedFilesCountingLookups(_ManagedFilesCreatedBy): + def __init__(self, creators: Mapping[str, str]) -> None: + super().__init__(creators) + self.looked_up: Final[list[str]] = [] # mutable-ok: the fake records the lookups the gate makes + + async def can_user_call_unified_file_id( + self, unified_file_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + self.looked_up.append(unified_file_id) + return await super().can_user_call_unified_file_id(unified_file_id, user_api_key_dict) + + +def _unified_file_id_for(s3_uri: str) -> str: + managed_id: Final = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/pdf", "u1", "nova", s3_uri, "deployment-1" + ) + return urlsafe_b64encode(managed_id.encode()).decode().rstrip("=") + + +def _chat_files(file_ids: Sequence[str]) -> dict[str, object]: + parts: Final = [{"type": "file", "file": {"file_id": file_id}} for file_id in file_ids] + return {"messages": [{"role": "user", "content": parts}]} + + _S3_OBJECT: Final = "s3://media-bucket/dir/cat.png" _HELLO: Final = {"role": "user", "content": "hello"} _S3_VALUES_OUTSIDE_MEDIA_PARTS: Final = ( @@ -941,6 +965,76 @@ class TestProxyBaseLLMRequestProcessing: assert _UPLOADED_OBJECT in refused.value.detail assert "LiteLLM-managed storage" in refused.value.detail + @pytest.mark.asyncio + async def test_non_admin_forged_managed_file_ids_stop_at_the_first_failed_lookup(self) -> None: + forged_ids: Final = tuple( + _unified_file_id_for(f"s3://media-bucket/litellm-bedrock-files/forged-{index}.pdf") for index in range(2000) + ) + managed_files: Final = _ManagedFilesCountingLookups({}) + + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files(forged_ids), + LitellmUserRoles.INTERNAL_USER, + {}, + managed_files, + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert managed_files.looked_up == [forged_ids[0]] + + @pytest.mark.asyncio + async def test_non_admin_file_ids_that_cannot_change_the_outcome_cost_no_lookup(self) -> None: + listed_object_id: Final = _unified_file_id_for("s3://media-bucket/dir/report.pdf") + id_without_s3_uri: Final = _unified_file_id_for("gs://other/report.pdf") + managed_files: Final = _ManagedFilesCountingLookups({_UPLOADED_FILE_ID: "user-a"}) + + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files((listed_object_id, id_without_s3_uri, _UPLOADED_FILE_ID)), + LitellmUserRoles.INTERNAL_USER, + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + managed_files, + ) + + assert "a cat" in str(response) + assert managed_files.looked_up == [_UPLOADED_FILE_ID] + + @pytest.mark.asyncio + async def test_non_admin_managed_file_id_its_user_created_passes_from_an_unlisted_bucket(self) -> None: + owned_id: Final = _unified_file_id_for("s3://other-bucket/report.pdf") + + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files((owned_id,)), + LitellmUserRoles.INTERNAL_USER, + {}, + _ManagedFilesCreatedBy({owned_id: "user-a"}), + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + async def test_non_admin_unlisted_file_id_costs_no_lookup_when_client_side_credentials_are_allowed(self) -> None: + unlisted_object_id: Final = _unified_file_id_for("s3://other-bucket/report.pdf") + managed_files: Final = _ManagedFilesCountingLookups({}) + + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files((unlisted_object_id,)), + LitellmUserRoles.INTERNAL_USER, + {"allow_client_side_credentials": True}, + managed_files, + ) + + assert "a cat" in str(response) + assert managed_files.looked_up == [] + @pytest.mark.asyncio @pytest.mark.parametrize( ("mapped_source", "named_in_refusal"), From 810608d74e75ccc32b4bce52b1bfe7092deafa06 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 14/24] fix(bedrock): check every file id field, skip MCP tool calls, accept registered S3 media types (cherry picked from commit 2b13898be4959f51255d7a40978cc6515a522a8a) --- .../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 b88ce55ffa3..b2f429a748a 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -54,7 +54,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: @@ -116,13 +117,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 6da1738dcc8..2e11ef2c775 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2193,6 +2193,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]" @@ -3855,7 +3859,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)) @@ -3876,14 +3884,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}. " @@ -3892,7 +3901,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 b7b80858ba2..9c0ab11459e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2363,12 +2363,13 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) - await refuse_s3_media_the_key_may_not_read( - self.data, - user_api_key_dict, - 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( + self.data, + user_api_key_dict, + 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 39e2d4364bd..ddc37b80705 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 @@ -2786,6 +2786,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 fee717f9c6e..64f86f185c7 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -88,6 +88,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"}}), @@ -366,6 +378,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 27d56ade859..12143241e08 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -249,9 +249,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 = ( @@ -492,6 +504,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( @@ -768,6 +784,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( From afd373f11250a147dc00a10ef0b3bb68e59cb0e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 21:20:35 -0800 Subject: [PATCH 15/24] fix(proxy): refuse unlisted s3 media on the rag query and adapter routes (cherry picked from commit 1d85ba2c1de440543a38b1ca58f0636d8c1bb631) --- .../pass_through_endpoints.py | 5 ++ litellm/proxy/rag_endpoints/endpoints.py | 6 ++ .../test_pass_through_endpoints.py | 83 +++++++++++++++++++ .../proxy/rag_endpoints/test_rag_endpoints.py | 42 ++++++++++ 4 files changed, 136 insertions(+) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a41e2ea068..16001d5bf43 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -93,6 +93,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, log_llm_api_exception, open_sse_before_first_byte, + refuse_s3_media_the_key_may_not_read, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id @@ -253,6 +254,10 @@ async def chat_completion_pass_through_endpoint( data["adapter_id"] = adapter_id + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) + verbose_proxy_logger.debug("Request received by LiteLLM:\n%s", data) data["model"] = ( general_settings.get("completion_model", None) # server default diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a1a7d62899f..4b06c430755 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_au from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, open_sse_before_first_byte, + refuse_s3_media_the_key_may_not_read, ttft_keepalive_interval, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -767,6 +768,7 @@ async def rag_query( llm_router, proxy_config, select_data_generator, + proxy_logging_obj, version, ) @@ -828,6 +830,10 @@ async def rag_query( **store_data, } + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) + # Add litellm data request_data: dict[str, object] = {} request_data = await add_litellm_data_to_request( diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 53bff9dd30c..29c37e3fa89 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -10328,3 +10328,86 @@ def test_init_kwargs_no_authenticated_identity_field_is_client_settable(mock_req metadata: Final = result["litellm_params"]["metadata"] survived: Final = {key: metadata.get(key) for key in expected if metadata.get(key) != expected[key]} assert survived == {}, f"client-supplied values survived for: {sorted(survived)}" + + +def _stub_adapter_proxy(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + proxy_logging: Final = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + return proxy_logging + +def _request_with_body(content_type: str, body: bytes) -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/pass-through", + "query_string": b"", + "headers": [(b"content-type", content_type.encode())], + }, + receive, + ) + +_ADAPTER_BODY_WITH_S3_IMAGE: Final = json.dumps( + { + "model": "unknown-model", + "messages": [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}}]} + ], + } +).encode() + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_refuses_an_s3_image_in_an_unlisted_bucket( + monkeypatch: pytest.MonkeyPatch, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock(return_value=_ADAPTER_BODY_WITH_S3_IMAGE) + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.code == "403" + proxy_logging.pre_call_hook.assert_not_called() + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_forwards_an_s3_image_in_a_listed_bucket( + monkeypatch: pytest.MonkeyPatch, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", {"allowed_client_side_s3_buckets": ["private-bucket"]} + ) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock(return_value=_ADAPTER_BODY_WITH_S3_IMAGE) + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.code != "403" + proxy_logging.pre_call_hook.assert_awaited_once() diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 2d654ea28ec..79a2b950c31 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -10,6 +10,7 @@ import json from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx from fastapi.testclient import TestClient import litellm @@ -1174,3 +1175,44 @@ class TestVectorStoreUploadControls: assert server_filename.endswith(".txt") assert secured_content_type == "text/plain" assert content_bytes == b"benign document text\n" + + + +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize("listed, expected_status", [(False, 403), (True, 200)]) +def test_rag_query_checks_the_s3_image_bucket_before_search( + client_internal_user: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + listed: bool, expected_status: int, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + search = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond( + json={"object": "vector_store.search_results.page", "search_query": ["hello"], "data": [], + "has_more": False, "next_page": None} + ) + completion = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + json={"id": "chatcmpl-test", "object": "chat.completion", "created": 1, "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + ) + with ( + patch("litellm.proxy.proxy_server.general_settings", + {"allowed_client_side_s3_buckets": ["private-bucket"]} if listed else {}), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.vector_store_registry", None), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={"model": "gpt-4o-mini", + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}} + ]}], + "retrieval_config": {"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}}, + ) + assert response.status_code == expected_status, response.text + assert search.call_count == int(listed) + assert completion.call_count == int(listed) + if not listed: + assert "allowed_client_side_s3_buckets" in response.json()["detail"] From 032f1f9ae97d7e3092d24e5fa0995d3de685f9ac Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 11 Nov 2026 15:52:40 -0800 Subject: [PATCH 16/24] style(proxy): format the merged collections.abc import (cherry picked from commit 91bb8946180b52a708944a7c939e337704834d0e) --- litellm/proxy/common_request_processing.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9c0ab11459e..ddf2917d881 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,16 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Iterable, + Mapping, + Sequence, +) from datetime import datetime from functools import lru_cache, partial from itertools import chain From 339c2bc346b10991c654a8d8f3089e807263847e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 11 Nov 2026 17:52:09 -0800 Subject: [PATCH 17/24] fix(proxy): walk tuples when checking s3 media and skip plain number lists (cherry picked from commit 5d8f25696341df2acf559ada53ef66ff53dd448d) --- .../litellm_core_utils/cloud_storage_security.py | 16 +++++++++++++--- .../test_pass_through_endpoints.py | 16 +++++++++++++++- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index b2f429a748a..87b3fccb48d 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -9,6 +9,8 @@ from types import MappingProxyType from typing import Any, Final, cast from urllib.parse import quote, unquote +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping @@ -56,6 +58,8 @@ _EMBEDDED_S3_URI: Final = re.compile(r"s3://[^;]*") _MODEL_FILE_ID_MAPPING_KEY: Final = "model_file_id_mapping" _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} +_SCALAR_TYPES: Final = frozenset({int, float, bool, type(None)}) +_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) def is_s3_uri(value: str) -> bool: @@ -82,6 +86,12 @@ def _with_embedded_s3_uris(file_id: str) -> Iterator[str]: return chain((file_id,), _s3_uris_embedded_in_file_id(file_id)) +def _sequence_items(value: object) -> Sequence[object] | None: + if isinstance(value, tuple): + return _OBJECT_TUPLE.validate_python(value) + return as_object_list(value) + + @dataclass(frozen=True, slots=True) class _BodyNode: value: object @@ -99,7 +109,7 @@ def _child_nodes(node: _BodyNode) -> Iterator[_BodyNode]: if fields is not None: inside_media_part: Final = node.inside_media_part or _is_media_part(fields, node.key) return (_BodyNode(child, child_key, inside_media_part) for child_key, child in fields.items()) - items: Final = as_object_list(node.value) + items: Final = _sequence_items(node.value) return (_BodyNode(item, None, node.inside_media_part) for item in items) if items is not None else iter(()) @@ -163,8 +173,8 @@ def _may_name_an_s3_uri(value: object) -> bool: fields: Final = as_str_mapping(value) if fields is not None: return not _KEYS_THAT_MAY_HIDE_AN_S3_URI.isdisjoint(fields) or any(map(_may_name_an_s3_uri, fields.values())) - items: Final = as_object_list(value) - return items is not None and any(map(_may_name_an_s3_uri, items)) + items: Final = _sequence_items(value) + return items is not None and not _SCALAR_TYPES.issuperset(map(type, items)) and any(map(_may_name_an_s3_uri, items)) def _may_carry_s3_media(request_body: Mapping[str, object]) -> bool: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 29c37e3fa89..574d393e633 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -10368,15 +10368,29 @@ _ADAPTER_BODY_WITH_S3_IMAGE: Final = json.dumps( ], } ).encode() +_ADAPTER_PYTHON_LITERAL_BODY_WITH_S3_IMAGE_IN_TUPLES: Final = repr( + { + "model": "unknown-model", + "messages": ( + {"role": "user", "content": ({"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}},)}, + ), + } +).encode() @pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [_ADAPTER_BODY_WITH_S3_IMAGE, _ADAPTER_PYTHON_LITERAL_BODY_WITH_S3_IMAGE_IN_TUPLES], + ids=["json", "python_literal_with_tuples"], +) async def test_chat_completion_pass_through_endpoint_refuses_an_s3_image_in_an_unlisted_bucket( monkeypatch: pytest.MonkeyPatch, + body: bytes, ): proxy_logging: Final = _stub_adapter_proxy(monkeypatch) request: Final = MagicMock(spec=Request) - request.body = AsyncMock(return_value=_ADAPTER_BODY_WITH_S3_IMAGE) + request.body = AsyncMock(return_value=body) with pytest.raises(ProxyException) as raised: await chat_completion_pass_through_endpoint( From cc109d6d8ca84b3b9aea1dac7455373e87c9a069 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 11 Nov 2026 18:57:14 -0800 Subject: [PATCH 18/24] fix(proxy): skip numbers when walking a request body for s3 media A body with large token-id lists and any s3:// string or file_id key made the full walk build a node per number, taking seconds per request. The walk and its prefilter now skip int, float, bool and None values. (cherry picked from commit 57e73a8f7de294abbc90bef55bdd58acb30f24a6) --- .../cloud_storage_security.py | 22 ++++++++++++++----- .../test_cloud_storage_security.py | 12 ++++++++++ .../proxy/test_common_request_processing.py | 6 +++++ 3 files changed, 35 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 87b3fccb48d..411759059fe 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,7 +1,7 @@ import base64 import posixpath import re -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial from itertools import accumulate, chain, repeat, takewhile @@ -86,6 +86,10 @@ def _with_embedded_s3_uris(file_id: str) -> Iterator[str]: return chain((file_id,), _s3_uris_embedded_in_file_id(file_id)) +def _non_scalar(values: Iterable[object]) -> Iterator[object]: + return (value for value in values if type(value) not in _SCALAR_TYPES) + + def _sequence_items(value: object) -> Sequence[object] | None: if isinstance(value, tuple): return _OBJECT_TUPLE.validate_python(value) @@ -108,9 +112,15 @@ def _child_nodes(node: _BodyNode) -> Iterator[_BodyNode]: fields: Final = as_str_mapping(node.value) if fields is not None: inside_media_part: Final = node.inside_media_part or _is_media_part(fields, node.key) - return (_BodyNode(child, child_key, inside_media_part) for child_key, child in fields.items()) + return ( + _BodyNode(child, child_key, inside_media_part) + for child_key, child in fields.items() + if type(child) not in _SCALAR_TYPES + ) items: Final = _sequence_items(node.value) - return (_BodyNode(item, None, node.inside_media_part) for item in items) if items is not None else iter(()) + if items is None: + return iter(()) + return (_BodyNode(item, None, node.inside_media_part) for item in _non_scalar(items)) def _next_level(level: tuple[_BodyNode, ...], _: None) -> tuple[_BodyNode, ...]: @@ -172,9 +182,11 @@ def _may_name_an_s3_uri(value: object) -> bool: return is_s3_uri(value) fields: Final = as_str_mapping(value) if fields is not None: - return not _KEYS_THAT_MAY_HIDE_AN_S3_URI.isdisjoint(fields) or any(map(_may_name_an_s3_uri, fields.values())) + return not _KEYS_THAT_MAY_HIDE_AN_S3_URI.isdisjoint(fields) or any( + map(_may_name_an_s3_uri, _non_scalar(fields.values())) + ) items: Final = _sequence_items(value) - return items is not None and not _SCALAR_TYPES.issuperset(map(type, items)) and any(map(_may_name_an_s3_uri, items)) + return items is not None and any(map(_may_name_an_s3_uri, _non_scalar(items))) def _may_carry_s3_media(request_body: Mapping[str, object]) -> bool: 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 64f86f185c7..dd9c16025ac 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -164,6 +164,18 @@ def _gemini_turn(*parts: object) -> dict[str, object]: }, id="responses-computer-screenshot", ), + pytest.param( + { + "messages": ( + {"role": "user", "content": ({"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}},)}, + ) + }, + id="tuples", + ), + pytest.param( + _user_turn(1, 2.5, True, None, {"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}), + id="numbers-before-the-media-part", + ), ), ) def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 12143241e08..17462abd3b0 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -211,6 +211,11 @@ def _chat_image(uri: str) -> dict[str, object]: return {"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": uri}}]}]} +def _chat_image_after_numbers(uri: str) -> dict[str, object]: + image_part: Final = {"type": "image_url", "image_url": {"url": uri}} + return {"messages": [{"role": "user", "content": [1, 2.5, True, None, image_part]}]} + + def _responses_image(uri: str) -> dict[str, object]: return {"input": [{"role": "user", "content": [{"type": "input_image", "image_url": uri}]}]} @@ -234,6 +239,7 @@ _S3MediaRoute: TypeAlias = Literal["acompletion", "aresponses", "anthropic_messa _GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/nova:generateContent" _S3_MEDIA_SURFACES: Final = ( pytest.param("acompletion", "/v1/chat/completions", _chat_image, id="chat-completions"), + pytest.param("acompletion", "/v1/chat/completions", _chat_image_after_numbers, id="chat-completions-numbers-first"), pytest.param("aresponses", "/v1/responses", _responses_image, id="responses"), pytest.param("anthropic_messages", "/v1/messages", _anthropic_image, id="anthropic-messages"), pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_snake_case_image, id="gemini-snake-case"), From 5ed90e79b11abe4761cd6ecd089cbcc04ed9c665 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 11 Nov 2026 21:49:32 -0800 Subject: [PATCH 19/24] test(proxy): check the s3 refusal through the real request setup (cherry picked from commit 03624164af4ba7dc35112041a201775436ddd4a6) --- .../proxy/test_common_request_processing.py | 37 +++++++++---------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 17462abd3b0..e48624cd6b2 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -510,10 +510,6 @@ 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( @@ -795,28 +791,31 @@ class TestProxyBaseLLMRequestProcessing: ("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 + self, 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 + + pre_call: Final = 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, ) - 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, - ) + if not refused: + await pre_call + return + with pytest.raises(HTTPException) as refusal: + await pre_call + assert refusal.value.status_code == 403 + assert "s3://other-bucket/cat.png" in str(refusal.value.detail) @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", [{"bad": "value"}, ["gpt-5.2"], 1]) From 4f6cc9275d88f81a8076a10455b3fb6fc1a679d3 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:27:01 +0000 Subject: [PATCH 20/24] test(proxy): exercise real setup for allowed s3 media Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit 5f11068f76d77e05519fcc90a7087a2d9ef7fe95) --- .../test_pass_through_endpoints.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 574d393e633..e8c4161c6b1 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -30,6 +30,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger @@ -61,6 +62,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.utils import ProxyLogging from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -10405,15 +10407,16 @@ async def test_chat_completion_pass_through_endpoint_refuses_an_s3_image_in_an_u @pytest.mark.asyncio -async def test_chat_completion_pass_through_endpoint_forwards_an_s3_image_in_a_listed_bucket( +async def test_chat_completion_pass_through_endpoint_routes_an_s3_image_in_a_listed_bucket( monkeypatch: pytest.MonkeyPatch, ): - proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) monkeypatch.setattr( "litellm.proxy.proxy_server.general_settings", {"allowed_client_side_s3_buckets": ["private-bucket"]} ) - request: Final = MagicMock(spec=Request) - request.body = AsyncMock(return_value=_ADAPTER_BODY_WITH_S3_IMAGE) + request: Final = _request_with_body("application/json", _ADAPTER_BODY_WITH_S3_IMAGE) with pytest.raises(ProxyException) as raised: await chat_completion_pass_through_endpoint( @@ -10423,5 +10426,5 @@ async def test_chat_completion_pass_through_endpoint_forwards_an_s3_image_in_a_l user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) - assert raised.value.code != "403" - proxy_logging.pre_call_hook.assert_awaited_once() + assert raised.value.code == "400" + assert raised.value.message == str(ProxyModelNotFoundError(route="completion", model_name="unknown-model")) From 6a6d974c682c30d289bbe21bf9ff41ab51e6c636 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:38:49 +0000 Subject: [PATCH 21/24] fix(proxy): preserve lts allowlist contract after backport Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 6 +- litellm/proxy/rag_endpoints/endpoints.py | 2 +- .../test_cloud_storage_security.py | 2 +- .../proxy/rag_endpoints/test_rag_endpoints.py | 61 +++++++++++++------ 4 files changed, 46 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index b2467457eb4..dd57fffc6de 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -3,7 +3,7 @@ import os import re import sys from collections.abc import Callable, Collection, Iterator, Mapping -from functools import lru_cache, partial +from functools import lru_cache from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType @@ -532,7 +532,7 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( ) -class _ClientSideStorageAllowlist(BaseModel): +class ClientSideStorageAllowlist(BaseModel): model_config = ConfigDict(frozen=True) allowed_client_side_s3_buckets: tuple[str, ...] | None = None @@ -542,7 +542,7 @@ class _ClientSideStorageAllowlist(BaseModel): def client_side_s3_bucket_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]: - return _ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket + return ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket def _check_banned_params( diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 4b06c430755..ea3c9036d20 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -767,8 +767,8 @@ async def rag_query( general_settings, llm_router, proxy_config, - select_data_generator, proxy_logging_obj, + select_data_generator, version, ) 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 dd9c16025ac..5e711d58d28 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -10,8 +10,8 @@ from litellm.litellm_core_utils.cloud_storage_security import ( MAX_FILE_ID_DECODE_PASSES, is_managed_cloud_storage_uri, is_managed_s3_uri, - names_an_unlisted_s3_bucket, media_file_ids, + names_an_unlisted_s3_bucket, s3_media_uris, ) diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 79a2b950c31..680813bcc3e 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -7,6 +7,7 @@ Covers: import io import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1177,39 +1178,59 @@ class TestVectorStoreUploadControls: assert content_bytes == b"benign document text\n" - @pytest.mark.usefixtures("httpx_transport") @pytest.mark.parametrize("listed, expected_status", [(False, 403), (True, 200)]) def test_rag_query_checks_the_s3_image_bucket_before_search( - client_internal_user: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, - listed: bool, expected_status: int, + client_internal_user: TestClient, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + listed: bool, + expected_status: int, ) -> None: monkeypatch.setenv("OPENAI_API_KEY", "sk-test") - search = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond( - json={"object": "vector_store.search_results.page", "search_query": ["hello"], "data": [], - "has_more": False, "next_page": None} + search: Final = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond( + json={ + "object": "vector_store.search_results.page", + "search_query": ["hello"], + "data": [], + "has_more": False, + "next_page": None, + } ) - completion = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( - json={"id": "chatcmpl-test", "object": "chat.completion", "created": 1, "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, - "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} + completion: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } ) with ( - patch("litellm.proxy.proxy_server.general_settings", - {"allowed_client_side_s3_buckets": ["private-bucket"]} if listed else {}), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allowed_client_side_s3_buckets": ["private-bucket"]} if listed else {}, + ), patch("litellm.proxy.proxy_server.llm_router", None), patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.vector_store_registry", None), ): - response = client_internal_user.post( + response: Final = client_internal_user.post( "/v1/rag/query", - json={"model": "gpt-4o-mini", - "messages": [{"role": "user", "content": [ - {"type": "text", "text": "hello"}, - {"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}} - ]}], - "retrieval_config": {"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}}, + json={ + "model": "gpt-4o-mini", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}}, + ], + } + ], + "retrieval_config": {"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + }, ) assert response.status_code == expected_status, response.text assert search.call_count == int(listed) From ac7ad3ff7bd47ea70ccf26b9bcbe6ca02135b20e Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 12 Nov 2026 13:54:26 +0000 Subject: [PATCH 22/24] fix(responses): preserve explicit media formats on lts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 6 ++ .../test_litellm_completion_responses.py | 62 +++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 8f5cd72cbcd..b8e60d40e29 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2259,6 +2259,9 @@ class LiteLLMCompletionResponsesConfig: file_dict["file_data"] = item["file_data"] if item.get("filename"): file_dict["filename"] = item["filename"] + media_format: Final = item.get("format") + if media_format is not None: + file_dict["format"] = media_format bucket_owner: Final = item.get("bucket_owner") if bucket_owner is not None and _is_s3_url(file_id or item.get("file_data")): file_dict["bucket_owner"] = bucket_owner @@ -2283,6 +2286,9 @@ class LiteLLMCompletionResponsesConfig: if bucket_owner is not None and _is_s3_url(url) else ChatCompletionImageUrlObject(url=url, detail=detail) ) + media_format: Final = item.get("format") + if media_format is not None: + image_url_obj["format"] = media_format return ChatCompletionImageObject(type="image_url", image_url=image_url_obj) diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index c3a74465139..c17946a2f6b 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -44,6 +44,68 @@ from litellm.types.utils import ( class TestLiteLLMCompletionResponsesConfig: + @pytest.mark.parametrize("part_type", ["input_image", "input_file"]) + @pytest.mark.parametrize( + ("key", "media_format", "expected_format"), + [("key", "image/png", "png"), ("key.jpeg", "image/png", "png"), ("key.png", None, "png")], + ) + def test_explicit_responses_media_format_reaches_the_converse_s3_image( + self, part_type: str, key: str, media_format: str | None, expected_format: str + ) -> None: + uri: Final = f"s3://bucket/{key}" + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + { + "role": "user", + "content": [ + { + "type": part_type, + "image_url" if part_type == "input_image" else "file_id": uri, + **({"format": media_format} if media_format is not None else {}), + } + ], + } + ], + responses_api_request={}, + ) + request: Final = litellm.AmazonConverseConfig()._transform_request( + model="us.amazon.nova-2-lite-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [ + { + "role": "user", + "content": [{"image": {"format": expected_format, "source": {"s3Location": {"uri": uri}}}}], + } + ] + + @pytest.mark.parametrize("part_type", ["input_image", "input_file"]) + def test_extensionless_responses_s3_media_without_format_is_a_bad_request(self, part_type: str) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + { + "role": "user", + "content": [ + {"type": part_type, "image_url" if part_type == "input_image" else "file_id": "s3://bucket/key"} + ], + } + ], + responses_api_request={}, + ) + + with pytest.raises(litellm.BadRequestError, match="Cannot tell the media type"): + litellm.AmazonConverseConfig()._transform_request( + model="us.amazon.nova-2-lite-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + def test_transform_input_file_item_to_file_item_with_file_id(self): """Test transformation of input_file item with file_id to Chat Completion file format""" # Setup From 349120a9f3553afddf4a3a9c313243ebfc95cf56 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 12 Nov 2026 14:09:37 +0000 Subject: [PATCH 23/24] fix(model-armor): fail closed on remote s3 documents Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit f9ad9f8aad6e9df611d8b4f470aeacee7ac0ff85) --- .../model_armor/file_scanning.py | 4 +- .../guardrail_hooks/test_model_armor.py | 51 +++++++++++++++++-- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py index 366ea7326ae..931aacbe8be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py @@ -7,7 +7,7 @@ maps its MIME type to a Model Armor ``byteDataType``, and returns the decoded by guardrail hooks can submit them. ``plan_file_scans`` classifies each block: blocks with no inline bytes (``file_id`` or remote -``gs://`` / ``http(s)`` references) and supported documents whose base64 will not decode are +``gs://`` / ``s3://`` / ``http(s)`` references) and supported documents whose base64 will not decode are reported as unscannable so the guardrail hook can fail closed (blocking unless ``fail_on_error`` is false) rather than letting an unscanned document reach the model. """ @@ -27,7 +27,7 @@ from litellm.types.llms.openai import AllMessageValues MODEL_ARMOR_MAX_FILE_SIZE_BYTES: Final = 4 * 1024 * 1024 -_REMOTE_URI_SCHEMES: Final = ("gs://", "http://", "https://") +_REMOTE_URI_SCHEMES: Final = ("gs://", "s3://", "http://", "https://") ModelArmorByteDataType = Literal["PDF", "WORD_DOCUMENT", "EXCEL_DOCUMENT", "POWERPOINT_DOCUMENT", "CSV", "TXT"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..c0901f83853 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3,13 +3,11 @@ import base64 import io import json from collections.abc import Iterator, Sequence -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch -import pytest - - import httpx +import pytest from fastapi import HTTPException import litellm @@ -17,13 +15,18 @@ import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError -from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail +from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( + FileScanPlan, + plan_file_scans, +) from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( ModelArmorAPIError, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues @pytest.mark.asyncio @@ -2670,6 +2673,44 @@ async def test_pre_call_blocks_unresolvable_file_id_reference(): assert "could not scan" in str(exc_info.value.detail) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "uri", + [ + "s3://listed-bucket/report.pdf", + " S3://listed-bucket/report.pdf ", + "gs://listed-bucket/report.pdf", + "https://example.com/report.pdf", + "http://example.com/report.pdf", + ], +) +async def test_pre_call_blocks_remote_document_without_format_or_filename(uri: str) -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": [{"type": "file", "file": {"file_data": uri}}]} + ] + assert plan_file_scans(messages) == FileScanPlan(attachments=(), unscannable_count=1) + + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=True, + ) + data: Final = {"messages": messages, "metadata": {"guardrails": ["model-armor-test"]}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "could not scan" in str(exc_info.value.detail) + assert data["metadata"]["_model_armor_status"] == "blocked" + + @pytest.mark.asyncio async def test_pre_call_blocks_remote_url_document_reference(): """A remote (https) document reference cannot be fetched here, so it fails closed by default.""" From 6a4d59ba629afae7490567f48b77f46beaa78cfe Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 14 Nov 2026 01:40:08 +0000 Subject: [PATCH 24/24] fix(proxy): suppress cyclic fallback payload inference Limit the suppression to reportGeneralTypeIssues on the legacy payload assignment. The fallback payload is constructed as a dict and the runtime AST is unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit d5a24bcc15d2ebd4bd5ac6ad044c083ab599e1f0)