From 93f2ae537dd7b336e599eaf02a8ae9615532dc14 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Tue, 25 Aug 2026 06:42:02 +0000 Subject: [PATCH 01/21] fix(guardrails): send image parts to bedrock ApplyGuardrail The Bedrock guardrail built its ApplyGuardrail payload from text content items only. An OpenAI-format image part has no "text" key, so it was forwarded to the model but silently dropped from the guardrail payload: the proxy reported the guardrail as having run while the guardrail never saw the image (#35332 measured contentPolicyImageUnits: 0 on a request that demonstrably carried one). BedrockContentItem had no image field either, and masking dropped image parts out of the request entirely. - Send image content items in the ApplyGuardrail INPUT payload, decoded through the existing BedrockImageProcessor. - Add the image types ApplyGuardrail accepts (png/jpeg only). - Keep non-text parts when masking rewrites message content. An image ApplyGuardrail cannot scan (anything but png/jpeg, or a remote url we refuse to fetch) still reaches the model, so skipping it silently would let a caller defeat an IMAGE-modality guardrail just by picking a format the API rejects. on_unscannable_image controls that and defaults to "block"; set it to "allow" to keep serving such requests unscanned. Remote urls are only fetched while litellm.user_url_validation is on. With validation disabled async_safe_get degrades to an unrestricted, redirect-following GET and the url comes straight from the caller, so fetching there would turn the guardrail into an SSRF primitive. Fixes #35332 Supersedes #35338 Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 166 ++++++++++++++-- litellm/types/guardrails.py | 7 + .../guardrail_hooks/bedrock_guardrails.py | 15 +- .../test_bedrock_guardrails.py | 4 +- .../test_bedrock_guardrails.py | 178 +++++++++++++++++- .../guardrails/test_guardrail_endpoints.py | 4 +- 6 files changed, 349 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index c70a2ee8a74..62454b89532 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -10,6 +10,7 @@ import sys sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import asyncio +import binascii import copy import json import re @@ -32,6 +33,8 @@ from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost +from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor +from litellm.litellm_core_utils.url_utils import SSRFError from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, @@ -59,10 +62,13 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksViolation, BedrockContentItem, BedrockGuardrailChecksResponse, + BedrockGuardrailImageFormat, + BedrockGuardrailImageSource, BedrockGuardrailOutput, BedrockGuardrailQualifier, BedrockGuardrailResponse, BedrockGuardrailUsage, + BedrockImageContent, BedrockRequest, BedrockTextContent, ) @@ -133,6 +139,16 @@ _CONTENT_TYPE_TO_QUALIFIER: Final[dict[str, BedrockGuardrailQualifier]] = { # must not be graded against as if it were the application's own source material. _GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"}) +# ApplyGuardrail only accepts png/jpeg image blocks. Anything else (gif, webp, ...) +# has no representation in the payload, so it cannot be scanned at all. +_APPLY_GUARDRAIL_IMAGE_FORMATS: Final[ + dict[str, BedrockGuardrailImageFormat] +] = { # mutable-ok: module-level lookup table, never mutated + "png": "png", + "jpeg": "jpeg", + "jpg": "jpeg", +} + class QualifiedTextBlock(NamedTuple): """A piece of message text paired with its Bedrock grounding qualifier (if any).""" @@ -221,6 +237,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + on_unscannable_image: Literal["block", "allow"] = "block", **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -229,6 +246,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.guardrail_provider = "bedrock" self.chunk_budget_chars = chunk_budget_chars self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) + # What to do with an image part ApplyGuardrail cannot scan (non png/jpeg, or a + # remote url we refuse to fetch). Defaults to blocking: the image reaches the + # model regardless, so allowing it would be a silent guardrail bypass. + self.on_unscannable_image: Literal["block", "allow"] = on_unscannable_image # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` # routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail. @@ -316,28 +337,139 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) return cleaned or None - def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest: + async def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest: """ Create a bedrock request for the input content - the LLM request. + + Text and image parts are both sent, so a guardrail with the IMAGE modality + enabled inspects the image the caller actually sent instead of only the text + that happened to sit next to it. """ bedrock_request: Final[BedrockRequest] = BedrockRequest(source="INPUT") - bedrock_request_content: Final[list[BedrockContentItem]] = [] if messages is None: return bedrock_request - for message in messages: - blocks = self.get_content_items_for_message(message=message) - if blocks is None: - continue - for block in blocks: - # INPUT scans send plain text only. Grounding qualifiers are attached - # exclusively when assembling the OUTPUT request, so a caller cannot use - # a grounding_source/query tag to change how input-safety policies treat - # their content (which would be an input-guardrail bypass). - bedrock_request_content.append(BedrockContentItem(text=BedrockTextContent(text=block.text))) - bedrock_request["content"] = bedrock_request_content + per_message = await asyncio.gather(*(self._build_input_content_items(message=message) for message in messages)) + # mutable-ok: BedrockRequest["content"] is a list in the AWS wire format + bedrock_request["content"] = [item for items in per_message for item in items] return bedrock_request + async def _build_input_content_items(self, message: AllMessageValues) -> tuple[BedrockContentItem, ...]: + """Flatten one request message into ApplyGuardrail INPUT content items. + + INPUT scans send text and image parts. Grounding qualifiers are attached + exclusively when assembling the OUTPUT request, so a caller cannot use a + grounding_source/query tag to change how input-safety policies treat their + content (which would be an input-guardrail bypass). + """ + content = message.get("content") + if content is None: + return () + if isinstance(content, str): + return (BedrockContentItem(text=BedrockTextContent(text=content)),) + if not isinstance(content, list): + return () + parts: Final = cast( # cast-ok: AllMessageValues content is a union of part TypedDicts + tuple[object, ...], tuple(content) + ) + items: Final = await asyncio.gather(*(self._build_input_content_item(item=item) for item in parts)) + return tuple(item for item in items if item is not None) + + async def _build_input_content_item(self, item: object) -> BedrockContentItem | None: + if isinstance(item, str): + return BedrockContentItem(text=BedrockTextContent(text=item)) + if not isinstance(item, dict): + return None + part: Final = cast(Mapping[str, object], item) # cast-ok: narrowed to dict on the line above + text: Final = part.get("text") + if isinstance(text, str): + return BedrockContentItem(text=BedrockTextContent(text=text)) + image_url: Final = self._get_image_url(item=part) + if image_url is None: + return None + return await self._build_image_content_item(image_url=image_url) + + @staticmethod + def _get_image_url(item: Mapping[str, object]) -> str | None: + if item.get("type") != "image_url": + return None + image_url: Final = item.get("image_url") + if isinstance(image_url, str): + return image_url + if isinstance(image_url, dict): + url: Final = cast(Mapping[str, object], image_url).get("url") # cast-ok: narrowed to dict on the line above + return url if isinstance(url, str) else None + return None + + def _handle_unscannable_image(self, reason: str) -> None: + """Decide what to do with an image part ApplyGuardrail cannot scan. + + The image still reaches the model either way, so skipping it silently would + let a caller defeat an IMAGE-modality guardrail just by picking a format the + API does not accept. `on_unscannable_image` defaults to "block" for that + reason; "allow" restores the permissive behavior for deployments that would + rather serve the request than fail it. + + Returns None so callers can `return self._handle_unscannable_image(...)`. + """ + if self.on_unscannable_image == "block": + raise HTTPException( + status_code=400, + detail={ # mutable-ok: HTTPException detail payload, serialized immediately + "error": "Violated guardrail policy", + "bedrock_guardrail_response": ( + f"Request contains an image the guardrail cannot scan ({reason}). " + "ApplyGuardrail accepts png/jpeg images only. Set " + "'on_unscannable_image: allow' on this guardrail to send such " + "requests to the model unscanned." + ), + "guardrail_name": self.guardrail_name, + }, + ) + verbose_proxy_logger.warning( + "Bedrock Guardrail %s: image part will not be scanned (%s); on_unscannable_image=allow, forwarding it to the model anyway", + self.guardrail_name, + reason, + ) + + async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: + """Decode/fetch an image part into an ApplyGuardrail image block. + + Remote URLs are only fetched while LiteLLM's URL validation is on. With + validation disabled `async_safe_get` degrades to an unrestricted, redirect + following GET, and the URL comes straight from the caller, so fetching here + would turn the guardrail into an SSRF primitive. Such an image is treated as + unscannable instead. + """ + if not image_url.startswith("data:") and not getattr(litellm, "user_url_validation", True): + self._handle_unscannable_image( + reason=f"remote image url not fetched because litellm.user_url_validation is disabled: {image_url}" + ) + return None + + try: + block: Final = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None) + except (httpx.HTTPError, SSRFError, ValueError, TypeError, KeyError, binascii.Error) as e: + self._handle_unscannable_image(reason=f"image content could not be read: {e}") + return None + + image_block: Final = block.get("image") + image_format: Final = ( + _APPLY_GUARDRAIL_IMAGE_FORMATS.get(str(image_block.get("format"))) if image_block else None + ) + # mutable-ok: {} is an empty default for .get, never mutated + image_bytes: Final = image_block.get("source", {}).get("bytes") if image_block else None + if image_format is None or not image_bytes: + self._handle_unscannable_image(reason="attachment is not a png/jpeg image") + return None + + return BedrockContentItem( + image=BedrockImageContent( + format=image_format, + source=BedrockGuardrailImageSource(bytes=image_bytes), + ) + ) + def _create_bedrock_output_content_request( self, response: Any | ModelResponse, @@ -386,7 +518,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): items.append(self._build_content_item(block)) return items - def convert_to_bedrock_format( + async def convert_to_bedrock_format( self, source: Literal["INPUT", "OUTPUT"], messages: list[AllMessageValues] | None = None, @@ -403,7 +535,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ bedrock_request: BedrockRequest = BedrockRequest(source=source) if source == "INPUT": - bedrock_request = self._create_bedrock_input_content_request(messages=messages) + bedrock_request = await self._create_bedrock_input_content_request(messages=messages) elif source == "OUTPUT": bedrock_request = self._create_bedrock_output_content_request(response=response, messages=messages) return bedrock_request @@ -839,7 +971,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ start_time: Final = datetime.now(timezone.utc) bedrock_request_data: Final[dict] = dict( - self.convert_to_bedrock_format(source=source, messages=messages, response=response) + await self.convert_to_bedrock_format(source=source, messages=messages, response=response) ) api_key: str | None = None if request_data: @@ -2880,6 +3012,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masking_index += 1 if item is not None: new_content.append(item) + else: + new_content.append(item) return new_content, masking_index diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c7cdfaad780..0223d651d4a 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -535,6 +535,13 @@ class BedrockGuardrailConfigModel(BaseModel): "still rejects is bisected automatically, so this value only trades round trips against " "batch size and cannot fail a request on its own.", ) + on_unscannable_image: Literal["block", "allow"] = Field( + default="block", + description="What to do with an image the guardrail cannot scan - ApplyGuardrail " + "accepts png/jpeg only, and remote image URLs are not fetched while " + "litellm.user_url_validation is disabled. 'block' (default) rejects the request; " + "'allow' logs a warning and sends the image to the model unscanned.", + ) class LakeraV2GuardrailConfigModel(BaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8d66b624341..e194f4f1132 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -1,6 +1,6 @@ from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict # Bedrock contextual grounding tags each content block so the guardrail knows # which text is the reference source, the user question, and the content to grade. @@ -12,8 +12,21 @@ class BedrockTextContent(TypedDict, total=False): qualifiers: list[BedrockGuardrailQualifier] +BedrockGuardrailImageFormat = Literal["png", "jpeg"] + + +class BedrockGuardrailImageSource(TypedDict, total=False): + bytes: ReadOnly[str] + + +class BedrockImageContent(TypedDict, total=False): + format: ReadOnly[BedrockGuardrailImageFormat] + source: ReadOnly[BedrockGuardrailImageSource] + + class BedrockContentItem(TypedDict, total=False): text: BedrockTextContent + image: ReadOnly[BedrockImageContent] class BedrockRequest(TypedDict, total=False): diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 43d088268eb..bb958bcc304 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -945,7 +945,7 @@ async def test_convert_to_bedrock_format_input_source(): ] # Call the method - result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) + result = await guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) # Verify the result structure assert isinstance(result, dict) @@ -1007,7 +1007,7 @@ async def test_convert_to_bedrock_format_output_source(): ) # Call the method - result = guardrail.convert_to_bedrock_format( + result = await guardrail.convert_to_bedrock_format( source="OUTPUT", response=mock_response ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index dd339d4e51f..5eb7ff867ee 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2,6 +2,8 @@ Unit tests for Bedrock Guardrails """ +import asyncio +import base64 import json import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -2408,12 +2410,14 @@ _GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + return asyncio.run(_grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) + return asyncio.run( + _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) + ) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2961,9 +2965,7 @@ async def test_streaming_hook_reraises_guardrail_service_failures(): guardrail = _sse_guardrail() with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.side_effect = HTTPException( - status_code=500, detail="Bedrock guardrail throttle retries exhausted" - ) + mock_api.side_effect = HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") with pytest.raises(HTTPException) as exc: await _drain_streaming_hook(guardrail) @@ -5274,3 +5276,169 @@ async def test_terminal_failure_logs_usage_and_cost_of_prior_passed_chunks(monke assert logged["guardrail_cost"] == pytest.approx(0.0003) assert logged["guardrail_response"]["usage"] == {"contentPolicyUnits": 2, "wordPolicyUnits": 1} assert "error" in logged["guardrail_response"] + + +class TestBedrockGuardrailImageInput: + """Image parts must reach ApplyGuardrail, not just the text sitting next to them.""" + + _PNG_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg==" + _GIF_DATA_URI = "data:image/gif;base64,R0lGODlhAQABAAAAACw=" + _JPEG_BYTES = b"\xff\xd8\xff\xdb" + # A literal, globally-routable IP keeps validate_url's getaddrinfo off DNS; the + # transport is faked in every test below, so no request leaves the process. + _REMOTE_IMAGE_URL = "https://93.184.216.34/a.jpg" + + def _jpeg_response(self, url: str) -> httpx.Response: + return httpx.Response( + 200, + content=self._JPEG_BYTES, + headers={"content-type": "image/jpeg"}, + request=httpx.Request("GET", url), + ) + + def _guardrail(self, **kwargs) -> BedrockGuardrail: + return BedrockGuardrail( + guardrail_name="bedrock-image", + guardrailIdentifier="gr-image", + guardrailVersion="DRAFT", + **kwargs, + ) + + @pytest.mark.asyncio + async def test_inline_image_is_sent_for_scanning(self): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what does this say?"}, + {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}}, + ], + } + ] + + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert request["content"] == [ + {"text": {"text": "what does this say?"}}, + {"image": {"format": "png", "source": {"bytes": self._PNG_DATA_URI.split(",")[1]}}}, + ] + + @pytest.mark.asyncio + async def test_unscannable_image_blocks_the_request_by_default(self): + """ApplyGuardrail takes png/jpeg only, and the image reaches the model either way. + + Skipping it silently would let a caller defeat an IMAGE-modality guardrail by + sending a gif, so the default is to reject the request. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}}, + ], + } + ] + + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated guardrail policy" + + @pytest.mark.asyncio + async def test_undecodable_image_blocks_the_request_by_default(self): + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "not-an-image"}}], + } + ] + + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_unscannable_image_is_skipped_when_explicitly_allowed(self): + """on_unscannable_image: allow restores the permissive behavior, opt-in only.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}}, + {"type": "image_url", "image_url": {"url": "not-an-image"}}, + ], + } + ] + + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + assert request["content"] == [{"text": {"text": "hello"}}] + + @pytest.mark.asyncio + async def test_remote_image_is_fetched_and_scanned(self): + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}], + } + ] + get = AsyncMock(return_value=self._jpeg_response(self._REMOTE_IMAGE_URL)) + + with patch.object(httpx.AsyncClient, "get", new=get): + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert request["content"] == [ + { + "image": { + "format": "jpeg", + "source": {"bytes": base64.b64encode(self._JPEG_BYTES).decode()}, + } + } + ] + + @pytest.mark.asyncio + async def test_remote_image_is_not_fetched_when_url_validation_is_disabled(self, monkeypatch): + """With validation off, async_safe_get is an unrestricted redirect-following GET. + + The url comes straight from the caller, so fetching it here would make the + guardrail an SSRF primitive. Treat the image as unscannable instead, and make + no request at all. + """ + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + url = "http://169.254.169.254/latest/meta-data/" + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}}]}] + get = AsyncMock(return_value=self._jpeg_response(url)) + + with patch.object(httpx.AsyncClient, "get", new=get): + with pytest.raises(HTTPException): + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + get.assert_not_awaited() + assert request["content"] == [] + + def test_masking_keeps_image_parts_in_the_request(self): + """Masking rewrites text in place; the image must survive to reach the model.""" + image_part = {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "my ssn is 123-45-6789"}, image_part], + } + ] + + updated = self._guardrail()._apply_masking_to_messages(messages=messages, masked_texts=["my ssn is {SSN}"]) + + assert updated[0]["content"] == [ + {"type": "text", "text": "my ssn is {SSN}"}, + image_part, + ] diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 45f5afef1bc..d6f7d5305be 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -896,7 +896,9 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response) ), patch.object(guardrail_hook, "_load_credentials") as mock_load_creds, - patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert, + patch.object( + guardrail_hook, "convert_to_bedrock_format", new_callable=AsyncMock + ) as mock_convert, patch.object( guardrail_hook, "get_guardrail_dynamic_request_body_params" ) as mock_get_params, From 4c98cf95b5a46d1204a721ffcb9ab9d47640a824 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:13:00 +0800 Subject: [PATCH 02/21] fix(guardrails): classify content parts by declared type before reading text A part tagged `type: "image_url"` that also carries a `text` key was scanned as text and returned before the image branch ran, while the provider transformations branch on `type` and send it to the model as an image. That let a caller defeat an IMAGE-modality guardrail, and defeat on_unscannable_image, by pairing the image with a benign decoy string. Classify by the declared type first, so an image_url part always takes the image path regardless of what other fields it carries. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 15 +++++-- .../test_bedrock_guardrails.py | 41 +++++++++++++++++++ 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 62454b89532..a5c49fc2ebc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -381,13 +381,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if not isinstance(item, dict): return None part: Final = cast(Mapping[str, object], item) # cast-ok: narrowed to dict on the line above + # Classify by the declared type before reading any field. Provider + # transformations branch on `type`, so a part tagged image_url reaches the + # model as an image even when it also carries a `text` key. Reading `text` + # first would scan that decoy and forward the image unscanned, which is the + # bypass this whole extractor exists to close. + if part.get("type") == "image_url": + image_url: Final = self._get_image_url(item=part) + if image_url is None: + return None + return await self._build_image_content_item(image_url=image_url) text: Final = part.get("text") if isinstance(text, str): return BedrockContentItem(text=BedrockTextContent(text=text)) - image_url: Final = self._get_image_url(item=part) - if image_url is None: - return None - return await self._build_image_content_item(image_url=image_url) + return None @staticmethod def _get_image_url(item: Mapping[str, object]) -> str | None: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 5eb7ff867ee..542cdab64ff 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5323,6 +5323,47 @@ class TestBedrockGuardrailImageInput: {"image": {"format": "png", "source": {"bytes": self._PNG_DATA_URI.split(",")[1]}}}, ] + @pytest.mark.asyncio + async def test_image_part_carrying_a_text_field_is_still_scanned_as_an_image(self): + """A part tagged image_url reaches the model as an image, text field or not. + + Provider transformations branch on `type`, so reading `text` first would scan + the decoy and forward the image unscanned - the exact bypass this change closes. + """ + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": self._PNG_DATA_URI}, + "text": "just a friendly note", + } + ], + } + ] + + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert request["content"] == [ + {"image": {"format": "png", "source": {"bytes": self._PNG_DATA_URI.split(",")[1]}}} + ] + + @pytest.mark.asyncio + async def test_unscannable_image_part_carrying_a_text_field_still_blocks(self): + """The decoy text must not turn an unscannable image into a scanned request.""" + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": self._GIF_DATA_URI}, "text": "hello"}], + } + ] + + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert exc_info.value.status_code == 400 + @pytest.mark.asyncio async def test_unscannable_image_blocks_the_request_by_default(self): """ApplyGuardrail takes png/jpeg only, and the image reaches the model either way. From a19dd2d92e0df7513400567515a217ffaf71d2f9 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:27:51 +0800 Subject: [PATCH 03/21] fix(guardrails): forward on_unscannable_image through initialize_bedrock initialize_bedrock enumerates its kwargs explicitly, so the new setting parsed and rendered but never reached the guardrail. An operator who opted into "allow" kept getting 400s on unscannable images with nothing to explain why. Same shape as the chunk_budget_chars regression, so the test follows that one and asserts through initialize_guardrail rather than the constructor. Also trims the docstrings added by this branch down to the rationale that is not already obvious from the code, and drops a stale line describing a return convention these helpers no longer use. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 42 ++++++------------- .../guardrails/guardrail_initializers.py | 1 + .../proxy/guardrails/test_init_guardrails.py | 34 +++++++++++++++ 3 files changed, 47 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index a5c49fc2ebc..3f02d512e08 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -338,13 +338,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return cleaned or None async def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest: - """ - Create a bedrock request for the input content - the LLM request. - - Text and image parts are both sent, so a guardrail with the IMAGE modality - enabled inspects the image the caller actually sent instead of only the text - that happened to sit next to it. - """ + """Create a bedrock request for the input content - the LLM request.""" bedrock_request: Final[BedrockRequest] = BedrockRequest(source="INPUT") if messages is None: return bedrock_request @@ -357,10 +351,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def _build_input_content_items(self, message: AllMessageValues) -> tuple[BedrockContentItem, ...]: """Flatten one request message into ApplyGuardrail INPUT content items. - INPUT scans send text and image parts. Grounding qualifiers are attached - exclusively when assembling the OUTPUT request, so a caller cannot use a - grounding_source/query tag to change how input-safety policies treat their - content (which would be an input-guardrail bypass). + Grounding qualifiers are attached only when assembling the OUTPUT request, so a + grounding_source/query tag cannot change how input-safety policies treat content """ content = message.get("content") if content is None: @@ -381,11 +373,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if not isinstance(item, dict): return None part: Final = cast(Mapping[str, object], item) # cast-ok: narrowed to dict on the line above - # Classify by the declared type before reading any field. Provider - # transformations branch on `type`, so a part tagged image_url reaches the - # model as an image even when it also carries a `text` key. Reading `text` - # first would scan that decoy and forward the image unscanned, which is the - # bypass this whole extractor exists to close. + # Provider transformations branch on `type`, so an image_url part reaches the + # model as an image even when it also carries `text`. Reading `text` first would + # scan that decoy and forward the image unscanned if part.get("type") == "image_url": image_url: Final = self._get_image_url(item=part) if image_url is None: @@ -409,15 +399,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return None def _handle_unscannable_image(self, reason: str) -> None: - """Decide what to do with an image part ApplyGuardrail cannot scan. + """Block or warn for an image part ApplyGuardrail cannot scan. - The image still reaches the model either way, so skipping it silently would - let a caller defeat an IMAGE-modality guardrail just by picking a format the - API does not accept. `on_unscannable_image` defaults to "block" for that - reason; "allow" restores the permissive behavior for deployments that would - rather serve the request than fail it. - - Returns None so callers can `return self._handle_unscannable_image(...)`. + The image reaches the model either way, so skipping it silently would let a + caller defeat an IMAGE-modality guardrail by picking a format the API rejects """ if self.on_unscannable_image == "block": raise HTTPException( @@ -440,13 +425,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: - """Decode/fetch an image part into an ApplyGuardrail image block. + """Decode or fetch an image part into an ApplyGuardrail image block. - Remote URLs are only fetched while LiteLLM's URL validation is on. With - validation disabled `async_safe_get` degrades to an unrestricted, redirect - following GET, and the URL comes straight from the caller, so fetching here - would turn the guardrail into an SSRF primitive. Such an image is treated as - unscannable instead. + With `user_url_validation` off, `async_safe_get` degrades to an unrestricted, + redirect-following GET on a caller-supplied URL, so it is not fetched at all """ if not image_url.startswith("data:") and not getattr(litellm, "user_url_validation", True): self._handle_unscannable_image( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 0d23e19f88d..4ba23474d6c 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -21,6 +21,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): prompt_attack_threshold=litellm_params.prompt_attack_threshold, pii_confidence_threshold=litellm_params.pii_confidence_threshold, chunk_budget_chars=litellm_params.chunk_budget_chars, + on_unscannable_image=litellm_params.on_unscannable_image, default_on=litellm_params.default_on, disable_exception_on_block=litellm_params.disable_exception_on_block, mask_request_content=litellm_params.mask_request_content, diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 82363302d2e..70004847d59 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -118,3 +118,37 @@ def test_initialize_guardrail_sets_run_in_parallel(config_value, expected): custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]] assert custom_guardrail.run_in_parallel is expected + + +def test_initialize_bedrock_forwards_on_unscannable_image(): + """Regression: `on_unscannable_image` set in config.yaml must reach the guardrail. + + Same shape as chunk_budget_chars above: the field lives on + BedrockGuardrailConfigModel so LitellmParams parses it, but initialize_bedrock + enumerates its kwargs explicitly. Dropped here, an operator who opted into + `allow` would keep getting 400s on unscannable images with no indication why. + """ + import litellm + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + test_guardrail = { + "guardrail_name": "test_bedrock_unscannable_image", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.BEDROCK.value, + "mode": "pre_call", + "guardrailIdentifier": "test-guardrail", + "guardrailVersion": "DRAFT", + "on_unscannable_image": "allow", + }, + } + + guardrail_handler = InMemoryGuardrailHandler() + guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_unscannable_image" + ] + assert initialized, "bedrock guardrail was not registered as a callback" + assert initialized[-1].on_unscannable_image == "allow" From 735e560c1160523a3bfb93d9f066adfb761e0dee Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:03:44 +0800 Subject: [PATCH 04/21] fix(guardrails): scan images in apply_guardrail, the path the proxy actually uses The payload fix alone never runs on a real request. ProxyLogging._execute_guardrail_hook (proxy/utils.py:1216) routes any guardrail that defines `apply_guardrail` through unified_guardrail unless it sets `use_native_lifecycle_hooks`: has_apply_guardrail = "apply_guardrail" in type(callback).__dict__ and not getattr( callback, "use_native_lifecycle_hooks", False) target = unified_guardrail if has_apply_guardrail else callback BedrockGuardrail defines apply_guardrail and does not set that flag, so a /v1/chat/completions request reaches apply_guardrail, which read inputs["texts"] only. Its own docstring said "images unchanged". Measured against a live guardrail with the IMAGE modality enabled, same messages, same account: guardrail.async_pre_call_hook png imageUnits=1 gif blocked 400 ProxyLogging.pre_call_hook png imageUnits=0 gif passed through The second row is what a proxy user gets, and it matches the imageUnits: 0 the reporter measured in #35332. Nothing new is needed on the extraction side. The endpoint translations already populate inputs["images"]: OpenAIChatCompletionsHandler from `image_url` parts (openai/chat/guardrail_translation/handler.py:253), AnthropicMessagesHandler from `image`/`source` blocks. Five guardrails already consume that field (vigil_guard, custom_code, deepkeep, straiker, generic_guardrail_api), so the contract exists and Bedrock was the one dropping it. apply_guardrail now appends the images as one extra user message and lets the normal message path build the payload, so the unified and native routes share the decoding, the png/jpeg check and on_unscannable_image instead of drifting apart. Requests that carry an image but no text are no longer skipped. _normalize_image_input handles the two shapes that field carries. The OpenAI translation appends the caller's image_url verbatim, already a data: URI or an https URL. The Anthropic one returns source["data"] only, dropping media_type, so the entry is bare base64 that the decoder would reject as unreadable and, under on_unscannable_image=block, turn a legitimate /v1/messages call into a 400. The format is sniffed back from the base64 prefix. Images are only read for input_type == "request"; the OUTPUT source scans model-generated text. Three tests, all failing before this change with "image never reached the payload: ['text']". Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 72 +++++++++++++++++-- .../test_bedrock_guardrails.py | 71 ++++++++++++++++++ 2 files changed, 137 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 3f02d512e08..fa9b3567638 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -56,7 +56,12 @@ from litellm.proxy.guardrails.anthropic_sse import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks -from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionImageObject, + ChatCompletionImageUrlObject, + ChatCompletionUserMessage, +) from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksMessage, BedrockChecksViolation, @@ -424,6 +429,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): reason, ) + #: base64 magic-byte prefixes for the formats ApplyGuardrail accepts. + _BASE64_IMAGE_PREFIXES: ClassVar[tuple[tuple[str, str], ...]] = ( + ("iVBORw0KGgo", "image/png"), + ("/9j/", "image/jpeg"), + ) + + @staticmethod + def _image_content_part(url: str) -> ChatCompletionImageObject: + """One OpenAI-format image content part, so the payload builder handles it.""" + return ChatCompletionImageObject(type="image_url", image_url=ChatCompletionImageUrlObject(url=url)) + + @classmethod + def _normalize_image_input(cls, value: str) -> str: + """Return a data URI or URL that `_build_image_content_item` can consume. + + `GenericGuardrailAPIInputs["images"]` is not a single shape. The OpenAI chat + translation appends the caller's `image_url` verbatim, so entries are already a + `data:` URI or an `https://` URL. The Anthropic translation's `_image_sources` + returns `source["data"]` only, which is bare base64 with the `media_type` + dropped. Sniff the format back from the base64 prefix so both shapes reach the + same decoder instead of the bare-base64 one failing as unreadable. + """ + if value.startswith(("data:", "http://", "https://")): + return value + for prefix, media_type in cls._BASE64_IMAGE_PREFIXES: + if value.startswith(prefix): + return f"data:{media_type};base64,{value}" + # Unrecognized: hand it over as-is and let the decoder reject it, so the + # on_unscannable_image policy decides rather than this helper. + return value + async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: """Decode or fetch an image part into an ApplyGuardrail image block. @@ -3170,7 +3206,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): logging_obj: Optional logging object Returns: - GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged + GenericGuardrailAPIInputs - processed_texts may be masked, images are + scanned but returned unchanged (ApplyGuardrail does not rewrite images) Raises: Exception: If content is blocked by Bedrock guardrail @@ -3178,8 +3215,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None. # dict.get("texts", []) would return None if the key exists with a None value. texts: Final = inputs.get("texts") or [] + # Images only exist on the request side; ApplyGuardrail's OUTPUT source takes + # model-generated text. The endpoint translations already extract them: + # OpenAIChatCompletionsHandler from `image_url` parts, AnthropicMessagesHandler + # from `image`/`source` blocks. Five other guardrails already consume this + # field; Bedrock was the one that dropped it on the floor. + image_urls: Final = tuple(inputs.get("images") or ()) if input_type == "request" else () try: - verbose_proxy_logger.debug("Bedrock Guardrail: Applying guardrail to %s text(s)", len(texts)) + verbose_proxy_logger.debug( + "Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls) + ) if input_type == "request": incremental_result: Final = await self._apply_incremental_request_scan( @@ -3204,8 +3249,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): scanned_slice: Final = selection.scanned_slice scanned_role_subset: Final = selection.scanned_role_subset - # Bedrock will throw an error if there is no text to process - if filtered_messages: + # Bedrock rejects an empty content list, so only skip when there is + # neither text nor an image to scan. + if filtered_messages or image_urls: _log_hook = GuardrailEventHooks.pre_call if input_type == "request" else GuardrailEventHooks.post_call # Map the abstract input_type to the Bedrock source parameter. # "request" -> INPUT (scan user-supplied content) @@ -3238,9 +3284,23 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): logging_event_type=_log_hook, ) else: + # Append the images as one extra user message. Reusing the normal + # message path means `_create_bedrock_input_content_request` does the + # decoding, format check and on_unscannable_image handling, so the + # unified and native lifecycle paths cannot drift apart. + image_parts: Final = [ # mutable-ok: OpenAI message content is a list in the wire format + self._image_content_part(self._normalize_image_input(url)) for url in image_urls + ] + image_message: Final = ( + (ChatCompletionUserMessage(role="user", content=image_parts),) if image_parts else () + ) + scan_messages: Final = [ # mutable-ok: make_bedrock_api_request takes a list of messages + *filtered_messages, + *image_message, + ] bedrock_response = await self.make_bedrock_api_request( source="INPUT", - messages=filtered_messages, + messages=scan_messages, request_data=request_data, logging_event_type=_log_hook, ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 542cdab64ff..931a4f30d24 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5467,6 +5467,77 @@ class TestBedrockGuardrailImageInput: get.assert_not_awaited() assert request["content"] == [] + @pytest.mark.asyncio + async def test_apply_guardrail_scans_images_from_inputs(self): + """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. + + ProxyLogging._execute_guardrail_hook routes any guardrail that defines + `apply_guardrail` through unified_guardrail, so a fix that only touches + async_pre_call_hook never runs on a real request. The endpoint translations + already put images in inputs["images"]; this asserts they reach the payload. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={"texts": ["what does this say?"], "images": [self._PNG_DATA_URI]}, + request_data={}, + input_type="request", + ) + + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds, f"image never reached the payload: {kinds}" + + @pytest.mark.asyncio + async def test_apply_guardrail_scans_bare_base64_images(self): + """Anthropic's translation drops media_type and passes bare base64. + + `_image_sources` returns source["data"] only, so the entry is not a data URI. + Without sniffing the format back it would be rejected as unreadable and, under + on_unscannable_image=block, turn a legitimate /v1/messages call into a 400. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + bare_base64 = self._PNG_DATA_URI.split(",")[1] + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={"texts": [], "images": [bare_base64]}, + request_data={}, + input_type="request", + ) + + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds, f"bare base64 image never reached the payload: {kinds}" + + @pytest.mark.asyncio + async def test_apply_guardrail_ignores_images_on_the_response_side(self): + """ApplyGuardrail's OUTPUT source scans model-generated text, not input images.""" + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={"texts": ["some model output"], "images": [self._PNG_DATA_URI]}, + request_data={}, + input_type="response", + ) + + assert sent and sent[0]["source"] == "OUTPUT" + def test_masking_keeps_image_parts_in_the_request(self): """Masking rewrites text in place; the image must survive to reach the model.""" image_part = {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} From 589086dc5b812fdd34f2620ae2d653a31b3edab0 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:57:27 +0800 Subject: [PATCH 05/21] fix(guardrails): carry Anthropic url and file image sources through to guardrails _image_sources returned source["data"] only. An Anthropic image block has three shapes (types/llms/anthropic.py:259) and only the base64 one carries "data", so {"type": "url", "url": ...} yielded nothing and the image never reached any guardrail at all. This is not Bedrock-specific. Five guardrails consume GenericGuardrailAPIInputs["images"] (vigil_guard, custom_code, deepkeep, straiker, generic_guardrail_api) and every one of them was blind to url sources on /v1/messages. base64 now returns a data URI rather than the bare payload. A consumer otherwise has no way to recover media_type, and an API like Bedrock's ApplyGuardrail needs the format to build its request. The file shape stays unresolvable here: the bytes live behind the Files API and this extractor has no client to fetch them. Documented rather than silently dropped, so a consumer treating a missing entry as "no image to scan" is a known gap and not a surprise. Co-Authored-By: Claude Opus 5 (1M context) --- .../chat/guardrail_translation/handler.py | 32 +++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 721a6653597..a09048be684 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -788,12 +788,40 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: + """Normalize an Anthropic image block into strings a guardrail can read. + + `source` is one of three shapes (types/llms/anthropic.py:259): + + {"type": "base64", "media_type": "image/png", "data": ""} + {"type": "url", "url": "https://..."} + {"type": "file", "file_id": "..."} + + base64 is returned as a data URI rather than the bare payload: consumers of + ``GenericGuardrailAPIInputs["images"]`` otherwise have no way to know the + format, and an API like Bedrock's ApplyGuardrail requires it. url is passed + through so the consumer can fetch it under its own SSRF policy. + + file is not resolvable here (the bytes live behind the Files API), so it + yields nothing. That is a silent gap for any consumer that treats a missing + entry as "no image to scan"; scanning a file_id needs a fetch this extractor + has no client for. + """ source: Final = block.get("source") if not isinstance(source, Mapping): return () - # Could be base64 or url + + source_type: Final = source.get("type") + if source_type == "url": + url: Final = source.get("url") + return (url,) if isinstance(url, str) and url else () + data: Final = source.get("data") - return (data,) if data else () + if not isinstance(data, str) or not data: + return () + media_type: Final = source.get("media_type") + if isinstance(media_type, str) and media_type: + return (f"data:{media_type};base64,{data}",) + return (data,) async def _apply_guardrail_responses_to_input( self, From 9f1f9ebe9e5fcad2e8a51a7f86637a07977c3ae5 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:57:44 +0800 Subject: [PATCH 06/21] fix(guardrails): enforce ApplyGuardrail's image limits and pack batches by image count AWS caps images at 4 MB each and 20 per request (https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-mmfilter.html). Nothing checked either; BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS is the text-unit quota and says nothing about images. An oversized image is now reported through on_unscannable_image with the measured size instead of surfacing as an opaque AWS error. _bin_pack_bedrock_content measured only text, so an image counted as 0 and `used + 0 <= budget` always held: 45 images packed into a single batch of 45. That measurement was complete when a content item could only be text; putting images in the payload is what invalidated it. Packing now carries a second dimension for the image count. Image bytes are deliberately not charged against `budget`, which is a different quota. _apply_guardrail_content_with_chunking splits up front when the image count is over the limit. Chunking is otherwise reactive, and the substrings _is_input_too_large_error matches ("text unit", "too long", ...) are all text-shaped, so an image-count rejection may never reach that fallback. Bisection could not rescue it either: _split_bedrock_content halves a lone item by its "text", which is empty for an image, so it gives up and re-raises the original error. Recursion terminates because every batch _bin_pack_bedrock_content returns is already within the image limit. Nested images needed no work here: _extract_tool_result already collects them out of tool_result blocks, so with apply_guardrail reading inputs["images"] an image inside a tool result is scanned. Four tests. The oversized case builds a real 5 MB data URI rather than patching the decoder, so the size check runs against what the decoder actually produces. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 92 +++++++++++++++++-- .../test_bedrock_guardrails.py | 64 +++++++++++++ 2 files changed, 148 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index fa9b3567638..2f98c172e04 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -146,6 +146,17 @@ _GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"}) # ApplyGuardrail only accepts png/jpeg image blocks. Anything else (gif, webp, ...) # has no representation in the payload, so it cannot be scanned at all. +# AWS's hard limits for image content filters: +# https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-mmfilter.html +# 4 MB per image, 20 images per request, 8000x8000, PNG/JPEG only, 25 images/second, +# and only the first 100 words of text inside an image are evaluated. +# Nothing in litellm checked these. BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS is the +# text-side quota and says nothing about images, and the strings +# _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS matches are all text-shaped, so an image-count +# rejection would not even reach the chunking fallback. +_MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024 +_MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: Final = 20 + _APPLY_GUARDRAIL_IMAGE_FORMATS: Final[ dict[str, BedrockGuardrailImageFormat] ] = { # mutable-ok: module-level lookup table, never mutated @@ -488,6 +499,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self._handle_unscannable_image(reason="attachment is not a png/jpeg image") return None + # base64 decodes to roughly 3/4 of its length; estimate rather than decode the + # whole image a second time just to measure it. + decoded_size: Final = len(image_bytes) * 3 // 4 + if decoded_size > _MAX_IMAGE_BYTES: + self._handle_unscannable_image( + reason=f"image is {decoded_size / 1024 / 1024:.1f} MB, over ApplyGuardrail's 4 MB limit" + ) + return None + return BedrockContentItem( image=BedrockImageContent( format=image_format, @@ -1112,7 +1132,48 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): guardrail block on any (sub-)chunk raises immediately -- callers must not lose that signal by continuing to post the remaining chunks. + + Images are the one case split up front rather than reactively. Chunking here + is a recovery path: the content goes out as a single call and is only + re-batched once AWS rejects it AND `_is_input_too_large_error` matches. Those + substrings ("text unit", "too long", ...) are all text-shaped, so a rejection + for exceeding 20 images per request may never reach this fallback at all, and + bisection cannot rescue it either -- `_split_bedrock_content` reads + `item["text"]` to halve a lone item, which is empty for an image, so it gives + up and re-raises the original error. Splitting before the first call keeps a + many-image request inside the documented limit instead. """ + image_count: Final = sum(1 for item in content if "image" in item) + if allow_chunking and image_count > _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: + preemptive_batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars) + if len(preemptive_batches) > 1: + verbose_proxy_logger.warning( + "Bedrock Guardrail: %d image(s) exceeds ApplyGuardrail's limit of %d per request; " + "splitting into %d calls before sending", + image_count, + _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL, + len(preemptive_batches), + ) + preemptive_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below + await self._apply_guardrail_content_with_chunking( + content=batch, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + # Safe to keep enabled: every batch _bin_pack_bedrock_content + # returns holds at most 20 images, so the recursive call falls + # straight through this branch and text chunking still applies. + allow_chunking=allow_chunking, + completed_chunk_usages=completed_chunk_usages, + ) + for batch in preemptive_batches + ] + return tuple(result for results in preemptive_results for result in results) + try: response: Final = await self._post_apply_guardrail_content_with_retry( content=content, @@ -1504,15 +1565,27 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if not content: return (tuple(content),) - lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content) + # An image item has no `text`, so it measures 0 against `budget` and + # `used + 0 <= budget` always holds: without a second dimension every image + # lands in whichever batch is open, however many there are. That measurement + # was complete when a content item could only be text. Packing also has to + # respect ApplyGuardrail's limit of 20 images per request, which is a count + # rather than a character budget, so the two are carried separately: image + # bytes are deliberately not charged against `budget`, which is the text-unit + # quota. + measured: Final = tuple( + (len((item.get("text") or BedrockTextContent()).get("text") or ""), 1 if "image" in item else 0) + for item in content + ) - def assign(carried: tuple[int, int], length: int) -> tuple[int, int]: - batch_index, used = carried - if used + length <= budget: - return batch_index, used + length - return batch_index + 1, length + def assign(carried: tuple[int, int, int], item: tuple[int, int]) -> tuple[int, int, int]: + batch_index, used, images = carried + length, is_image = item + if used + length <= budget and images + is_image <= _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: + return batch_index, used + length, images + is_image + return batch_index + 1, length, is_image - batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:]) + batch_numbers: Final = (index for index, _, _ in tuple(accumulate(measured, assign, initial=(0, 0, 0)))[1:]) return tuple( tuple(item for _, item in group) for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0]) @@ -3294,8 +3367,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): image_message: Final = ( (ChatCompletionUserMessage(role="user", content=image_parts),) if image_parts else () ) + # `filtered_messages` is Optional; the guard above only proves one of + # it and `image_urls` is truthy, so it can still be None here when the + # request carries an image and no text. scan_messages: Final = [ # mutable-ok: make_bedrock_api_request takes a list of messages - *filtered_messages, + *(filtered_messages or ()), *image_message, ] bedrock_response = await self.make_bedrock_api_request( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 931a4f30d24..4ec507c922e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -6,6 +6,7 @@ import asyncio import base64 import json import sys +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5538,6 +5539,69 @@ class TestBedrockGuardrailImageInput: assert sent and sent[0]["source"] == "OUTPUT" + @pytest.mark.asyncio + async def test_oversized_image_is_rejected_before_sending(self): + """ApplyGuardrail caps images at 4 MB; AWS's rejection is not worth a round trip. + + Built as a real data URI rather than a patched decoder so the size check runs + against what the decoder actually produces. + """ + oversized_png = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * (5 * 1024 * 1024)).decode() + + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().convert_to_bedrock_format( + source="INPUT", + messages=[ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{oversized_png}"}} + ], + } + ], + ) + + assert "4 MB limit" in str(exc_info.value.detail) + + def test_bin_packing_respects_the_twenty_image_limit(self): + """An image measures 0 against the character budget, so count them separately.""" + image = {"image": {"format": "png", "source": {"bytes": "AAAA"}}} + batches = BedrockGuardrail._bin_pack_bedrock_content([image] * 45, budget=25_000) + sizes = [len(batch) for batch in batches] + assert all(size <= 20 for size in sizes), sizes + assert sum(sizes) == 45, sizes + + def test_bin_packing_still_splits_on_the_text_budget(self): + text = {"text": {"text": "x" * 20_000}} + batches = BedrockGuardrail._bin_pack_bedrock_content([text] * 3, budget=25_000) + assert [len(batch) for batch in batches] == [1, 1, 1] + + @pytest.mark.asyncio + async def test_many_images_are_split_before_the_first_call(self): + """Chunking is reactive; an image-count rejection may never match the too-large check.""" + g = self._guardrail() + sent: list = [] + + async def fake_post(content, **kwargs): + sent.append(sum(1 for item in content if "image" in item)) + return {"action": "NONE", "outputs": []} + + image = {"image": {"format": "png", "source": {"bytes": "AAAA"}}} + with patch.object(g, "_post_apply_guardrail_content_with_retry", new=fake_post): + await g._apply_guardrail_content_with_chunking( + content=[image] * 45, + base_request_data={}, + credentials=None, + aws_region_name="us-west-2", + api_key=None, + request_data=None, + event_type=GuardrailEventHooks.pre_call, + start_time=datetime.now(timezone.utc), + allow_chunking=True, + completed_chunk_usages=[], + ) + assert sent == [20, 20, 5], sent + def test_masking_keeps_image_parts_in_the_request(self): """Masking rewrites text in place; the image must survive to reach the model.""" image_part = {"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} From 762a61a83034651f0f789eba0e60d7309ec2769e Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:12:31 +0800 Subject: [PATCH 07/21] fix(guardrails): stop the text-only scan shortcuts from skipping images apply_guardrail reads inputs["images"], but two optimizations return before that point and both decide what to scan from `texts` alone: _apply_incremental_request_scan (only_scan_new_messages) sends just the new text segments and returns. Benign text plus a policy-violating image, on a session whose text is already cached, is never scanned -- and the proxy still reports the guardrail as having run. Nothing caches images either, so "already seen" cannot be established for them in the first place. _select_messages_for_apply_guardrail (experimental_use_latest_role_message_only) marks a latest user message with no text as skip_scan. An image-only message is exactly that shape, so the whole request was dropped from the scan. An image now forces the full image-aware path in both cases. The optimization is lost for image-carrying requests -- an incremental turn with an image rescans its text rather than skipping it -- which is the safe direction: correctness over the optimization, failing closed rather than silently open. Falling through skip_scan leaves filtered_messages None; the existing `*(filtered_messages or ())` unpack and _merge_masked_texts's empty-input guard already handle that, so the image-only scan needs no other change. Both tests fail on the code without this change, for the reason named in each. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 12 +++- .../test_bedrock_guardrails.py | 63 +++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 2f98c172e04..61f3003a0ea 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -3299,7 +3299,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): "Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls) ) - if input_type == "request": + # Both of the optimizations below decide what to scan by looking at + # `texts` alone, so an image rides along unscanned and the proxy still + # reports the guardrail as run. Neither tracks images in its session + # cache either, so "already seen" cannot be established for them. + # Presence of an image therefore forces the full image-aware path: + # correctness over the optimization, and it fails closed. + if input_type == "request" and not image_urls: incremental_result: Final = await self._apply_incremental_request_scan( texts=texts, inputs=inputs, @@ -3316,7 +3322,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, input_type=input_type, ) - if selection.skip_scan: + # `experimental_use_latest_role_message_only` marks a latest message with + # no text as skip_scan; an image-only message is exactly that shape. + if selection.skip_scan and not image_urls: return inputs filtered_messages: Final = selection.filtered_messages scanned_slice: Final = selection.scanned_slice diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 4ec507c922e..3b4064f4385 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5468,6 +5468,69 @@ class TestBedrockGuardrailImageInput: get.assert_not_awaited() assert request["content"] == [] + @pytest.mark.asyncio + async def test_incremental_scan_does_not_skip_a_request_carrying_an_image(self): + """`only_scan_new_messages` decides what to scan from `texts` alone. + + With a session id and every text segment already seen, the incremental path + returns before the image-aware code runs, so benign text plus a violating + image is never scanned -- and the proxy still reports the guardrail as run. + Nothing caches images either, so "already seen" cannot cover them. + """ + g = self._guardrail(only_scan_new_messages=True) + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + # Every text is already in the session cache: the exact state that made the + # incremental path return early. + with ( + patch.object(g, "filter_new_texts_for_session", new=AsyncMock(return_value=[])), + patch.object(g, "make_bedrock_api_request", new=spy), + ): + await g.apply_guardrail( + inputs={"texts": ["look at this"], "images": [self._PNG_DATA_URI]}, + request_data={"litellm_session_id": "sess-1"}, + input_type="request", + ) + + assert sent, "image-carrying request was skipped entirely by the incremental path" + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds, f"image never reached the payload: {kinds}" + + @pytest.mark.asyncio + async def test_latest_message_only_does_not_skip_an_image_only_message(self): + """`experimental_use_latest_role_message_only` skips a latest message with no text. + + An image-only user message is exactly that shape, so the whole request was + dropped from the scan. + """ + g = self._guardrail(experimental_use_latest_role_message_only=True) + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={ + "texts": [], + "images": [self._PNG_DATA_URI], + "structured_messages": [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}}]} + ], + }, + request_data={}, + input_type="request", + ) + + assert sent, "image-only latest message was skipped entirely" + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds, f"image never reached the payload: {kinds}" + @pytest.mark.asyncio async def test_apply_guardrail_scans_images_from_inputs(self): """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. From 657dbce543d50591b00a10b1e4db35ded6af7465 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:13:09 +0800 Subject: [PATCH 08/21] fix(guardrails): cap the image fetch instead of buffering whatever the url serves _build_image_content_item hands a caller-supplied url to BedrockImageProcessor.process_image_async, which buffers the whole body through async_safe_get before returning. The 4 MB check runs on bytes that are already resident, so it rejects an oversized image without preventing the allocation -- and _build_input_content_items gathers these concurrently, so one request with several urls multiplies it. An arbitrarily large or indefinitely chunked response is enough to exhaust proxy memory. async_safe_get takes an optional max_bytes and, when given one, streams the body and aborts past the cap with PayloadTooLargeError. Omitted, it keeps the previous buffering, so every existing caller -- including the model-call image paths in factory.py -- is byte-for-byte unchanged. Only the guardrail passes it. The rebuilt response drops content-encoding and content-length: aiter_bytes yields decoded bytes, so carrying those over would describe the body wrongly. PayloadTooLargeError subclasses ValueError, like SSRFError, so callers already treating a bad remote response as a rejected fetch need no new except arm. The guardrail names it explicitly anyway, so an operator reading the log sees "too large" rather than "could not be read". The two existing remote-url tests now stub `stream` rather than `get`, which is the transport the capped path uses. The new test asserts on how many bytes were pulled -- without that, an unstubbed transport would raise for the wrong reason and the test would pass against the unfixed code. Co-Authored-By: Claude Opus 5 (1M context) --- .../prompt_templates/factory.py | 17 +++- litellm/litellm_core_utils/url_utils.py | 91 +++++++++++++++++-- .../guardrail_hooks/bedrock_guardrails.py | 16 +++- .../test_bedrock_guardrails.py | 68 +++++++++++++- 4 files changed, 173 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b676077ab0e..dafdf986c33 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3407,14 +3407,14 @@ class BedrockImageProcessor: return base64_bytes, content_type @staticmethod - async def get_image_details_async(image_url) -> tuple[str, str]: + async def get_image_details_async(image_url, max_bytes: int | None = None) -> tuple[str, str]: try: client: Final = get_async_httpx_client( llm_provider=httpxSpecialProvider.PromptFactory, params={"concurrent_limit": 1}, ) # Send a GET request to the image URL - response: Final[httpx.Response] = await async_safe_get(client, image_url) + response: Final[httpx.Response] = await async_safe_get(client, image_url, max_bytes=max_bytes) response.raise_for_status() # Raise an exception for HTTP errors return BedrockImageProcessor._post_call_image_processing(response, image_url) @@ -3595,13 +3595,20 @@ class BedrockImageProcessor: return cls._create_bedrock_block(img_bytes, mime_type, image_format) @classmethod - async def process_image_async(cls, image_url: str, format: str | None) -> BedrockContentBlock: - """Asynchronous image processing.""" + async def process_image_async( + cls, image_url: str, format: str | None, max_bytes: int | None = None + ) -> BedrockContentBlock: + """Asynchronous image processing. + + ``max_bytes`` caps a remote fetch and is ignored for a base64 data URI, + whose size the caller already knows before calling. Omitting it keeps the + previous unbounded fetch. + """ 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: - img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url) + img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url, max_bytes=max_bytes) image_format = mime_type.split("/")[1] else: raise ValueError("Unsupported image type. Expected either image url or base64 encoded string") diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..0b26eaa869e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -44,6 +44,14 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" +class PayloadTooLargeError(ValueError): + """Raised when a fetched body exceeds the caller's byte cap. + + A ``ValueError`` subclass, like :class:`SSRFError`, so callers that already + treat a malformed remote response as a rejected fetch need no new except arm. + """ + + def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. @@ -412,21 +420,88 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: raise SSRFError("Too many redirects") -async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: - """Async version of safe_get.""" +# Headers that describe the wire encoding of the body. `aiter_bytes` yields +# decoded bytes, so carrying these onto the rebuilt response would describe it +# wrongly (a gzip label on already-inflated bytes, a length from before decoding). +_TRANSFER_ENCODING_HEADERS: Final = frozenset({"content-encoding", "content-length"}) + + +def _underlying_httpx_client(client: object) -> httpx.AsyncClient: + """Return the object exposing ``stream``. + + ``AsyncHTTPHandler``/``HTTPHandler`` wrap an httpx client and forward only + ``get``/``post``/..., so streaming has to go through the wrapped ``.client``. + A raw httpx client is returned unchanged. + """ + inner: Final = getattr(client, "client", client) + if not isinstance(inner, httpx.AsyncClient): + raise TypeError(f"cannot stream from {type(client).__name__}: no httpx client to stream with") + return inner + + +async def _async_get_capped( + client: object, + url: str, + max_bytes: int, + request_kwargs: dict[str, object], # mutable-ok: forwarded straight to httpx as **kwargs +) -> httpx.Response: + """GET ``url``, aborting the transfer once the body exceeds ``max_bytes``. + + ``client.get`` buffers the whole body before returning, so a caller-supplied + URL serving an arbitrarily large or indefinitely chunked response is an + unbounded allocation. Streaming makes the cap effective during the transfer + rather than after it. + + Returns a fully-read response so callers keep using ``.content`` as before. + """ + async with _underlying_httpx_client(client).stream("GET", url, **request_kwargs) as response: + chunks: list[bytes] = [] # mutable-ok: accumulator for the capped body + total = 0 + async for chunk in response.aiter_bytes(): + total += len(chunk) + if total > max_bytes: + raise PayloadTooLargeError(f"remote body exceeds {max_bytes} bytes") + chunks.append(chunk) + kept_headers: Final = [ # mutable-ok: httpx.Response takes the header pairs as a list + (k, v) for k, v in response.headers.multi_items() if k.lower() not in _TRANSFER_ENCODING_HEADERS + ] + return httpx.Response( + status_code=response.status_code, + headers=kept_headers, + content=b"".join(chunks), + request=response.request, + ) + + +async def async_safe_get(client: Any, url: str, max_bytes: int | None = None, **kwargs: Any) -> Any: + """Async version of safe_get. + + ``max_bytes`` caps the response body, rejecting an oversized transfer with + :class:`PayloadTooLargeError` while it is still in flight. Omitting it keeps + the previous unbounded buffering, so existing callers are unaffected. + """ + + async def _issue( + target_url: str, + request_kwargs: dict[str, object], # mutable-ok: forwarded straight to httpx as **kwargs + ) -> httpx.Response: + if max_bytes is None: + return await client.get(target_url, **request_kwargs) + return await _async_get_capped(client, target_url, max_bytes, request_kwargs) + if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) - return await client.get(url, **kwargs) + return await _issue(url, kwargs) kwargs.pop("follow_redirects", None) caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = await client.get( - validated_url, - headers={**caller_headers, "Host": original_host}, - follow_redirects=False, + hop_kwargs: dict[str, object] = { # mutable-ok: a fresh per-hop kwargs dict, consumed by this call **kwargs, - ) + "headers": {**caller_headers, "Host": original_host}, # mutable-ok: httpx takes headers as a dict + "follow_redirects": False, + } + response = await _issue(validated_url, hop_kwargs) if not response.is_redirect: return response # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 61f3003a0ea..139674b8aff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -34,7 +34,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor -from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.litellm_core_utils.url_utils import PayloadTooLargeError, SSRFError from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, @@ -484,7 +484,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return None try: - block: Final = await BedrockImageProcessor.process_image_async(image_url=image_url, format=None) + # Cap the fetch itself. The decoded-size check below runs after the + # bytes are already resident, so on its own it does not stop a caller + # from pointing the proxy at an arbitrarily large or chunked response + # -- and `_build_input_content_items` gathers these concurrently, so + # one request with several URLs multiplies the allocation. + block: Final = await BedrockImageProcessor.process_image_async( + image_url=image_url, format=None, max_bytes=_MAX_IMAGE_BYTES + ) + except PayloadTooLargeError as e: + # Named before the ValueError arm it subclasses, so the operator sees + # "too large" rather than "could not be read" for a size rejection. + self._handle_unscannable_image(reason=f"remote image over ApplyGuardrail's 4 MB limit: {e}") + return None except (httpx.HTTPError, SSRFError, ValueError, TypeError, KeyError, binascii.Error) as e: self._handle_unscannable_image(reason=f"image content could not be read: {e}") return None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 3b4064f4385..a849cfbfd71 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -4,6 +4,7 @@ Unit tests for Bedrock Guardrails import asyncio import base64 +import contextlib import json import sys from datetime import datetime, timezone @@ -5297,6 +5298,32 @@ class TestBedrockGuardrailImageInput: request=httpx.Request("GET", url), ) + @staticmethod + def _fake_stream(chunks: list[bytes], served: list[int] | None = None): + """Stand in for httpx.AsyncClient.stream, serving `chunks` one at a time. + + The guardrail caps the transfer, so the fetch goes through `stream` rather + than `get`. `served` counts the chunks actually pulled, which is how a test + tells "stopped mid-transfer" apart from "read everything, then rejected". + """ + + @contextlib.asynccontextmanager + async def _stream(self, method: str, url, **kwargs): + async def _aiter_bytes(): + for chunk in chunks: + if served is not None: + served.append(len(chunk)) + yield chunk + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers({"content-type": "image/jpeg"}) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = _aiter_bytes + yield response + + return _stream + def _guardrail(self, **kwargs) -> BedrockGuardrail: return BedrockGuardrail( guardrail_name="bedrock-image", @@ -5430,9 +5457,7 @@ class TestBedrockGuardrailImageInput: "content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}], } ] - get = AsyncMock(return_value=self._jpeg_response(self._REMOTE_IMAGE_URL)) - - with patch.object(httpx.AsyncClient, "get", new=get): + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([self._JPEG_BYTES])): request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) assert request["content"] == [ @@ -5456,8 +5481,12 @@ class TestBedrockGuardrailImageInput: url = "http://169.254.169.254/latest/meta-data/" messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}}]}] get = AsyncMock(return_value=self._jpeg_response(url)) + served: list[int] = [] - with patch.object(httpx.AsyncClient, "get", new=get): + with ( + patch.object(httpx.AsyncClient, "get", new=get), + patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([self._JPEG_BYTES], served)), + ): with pytest.raises(HTTPException): await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) @@ -5465,7 +5494,9 @@ class TestBedrockGuardrailImageInput: source="INPUT", messages=messages ) + # Both transports are stubbed: neither the plain nor the capped fetch ran. get.assert_not_awaited() + assert served == [] assert request["content"] == [] @pytest.mark.asyncio @@ -5531,6 +5562,35 @@ class TestBedrockGuardrailImageInput: kinds = [k for item in sent[0]["content"] for k in item] assert "image" in kinds, f"image never reached the payload: {kinds}" + @pytest.mark.asyncio + async def test_oversized_remote_image_is_cut_off_during_the_transfer(self): + """A caller-supplied url can serve an unbounded or indefinitely chunked body. + + The decoded-size check runs once the bytes are already resident, so the cap + has to apply while the transfer is in flight. Asserting on how much was + pulled is what separates that from buffering it all and rejecting after. + """ + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}], + } + ] + # 8 MB offered one MB at a time against a 4 MB cap. + chunks: list[bytes] = [b"\0" * (1024 * 1024) for _ in range(8)] + served: list[int] = [] + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(chunks, served)): + with pytest.raises(HTTPException): + await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + # Without the cap the fetch buffers through `get` instead, so `served` stays + # empty and the HTTPException above would come from an unstubbed transport + # rather than from the size rejection. Assert the capped path actually ran. + assert served, "the fetch did not go through the capped stream path" + assert sum(served) <= 5 * 1024 * 1024, f"read {sum(served)} bytes past a 4 MB cap" + assert len(served) < len(chunks), "the whole body was pulled before rejecting it" + @pytest.mark.asyncio async def test_apply_guardrail_scans_images_from_inputs(self): """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. From 0e1ec44e35bf331e4c0dece7458537a58755770f Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:22:51 +0800 Subject: [PATCH 09/21] test(guardrails): cover every Anthropic image source shape in the extractor's own suite _image_sources had no test asserting what it extracts. The existing image tests live on the Bedrock side and all use base64 without a media_type, which is the one path the fix left unchanged, so both behaviors it does change went unverified: the url shape reaching the guardrail at all, and base64 arriving as a data URI. Against the pre-fix extractor the url case sees [] and the media_type case sees ['AAAA'] instead of ['data:image/png;base64,AAAA']. The remaining three assert behavior the fix deliberately preserves -- bare base64 passed through, a file source yielding nothing, a malformed source dropped rather than handed on for a consumer to choke on. Each message carries a text block because a message with no text never reaches the guardrail, which would make every source shape look equally dropped. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_anthropic_guardrail_handler.py | 86 +++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 2b392456763..ed25ada891f 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -1472,6 +1472,92 @@ class MockCanaryMaskingGuardrail(CustomGuardrail): return inputs +class TestAnthropicMessagesImageSources: + """An Anthropic image block has three source shapes (types/llms/anthropic.py:259). + + Only the base64 one carries "data", so reading that key alone drops url images + entirely -- for every guardrail consuming GenericGuardrailAPIInputs["images"], + not just Bedrock. + """ + + def _data(self, messages): + return {"model": "claude-sonnet-4-5", "messages": messages} + + async def _images_seen(self, content) -> list[str]: + handler = AnthropicMessagesHandler() + + class ImageRecordingGuardrail(MockCanaryMaskingGuardrail): + def __init__(self): + super().__init__() + self.seen_images: list[str] = [] # mutable-ok: accumulator for the assertion + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.seen_images.extend(inputs.get("images") or []) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + guardrail = ImageRecordingGuardrail() + # The text block is what gets the guardrail invoked at all: a message with + # no text gives the handler nothing to scan, so it never reaches the + # guardrail and every source shape would look equally "dropped". + await handler.process_input_messages( + data=self._data([{"role": "user", "content": [{"type": "text", "text": "describe it"}, *content]}]), + guardrail_to_apply=guardrail, + ) + return guardrail.seen_images + + @pytest.mark.asyncio + async def test_url_source_reaches_the_guardrail(self): + """A url source has no "data" key, so it used to yield nothing at all.""" + seen = await self._images_seen( + [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}] + ) + + assert seen == ["https://example.com/a.png"] + + @pytest.mark.asyncio + async def test_base64_source_carries_its_media_type(self): + """Bare base64 leaves the consumer no way to recover the format. + + An API like Bedrock's ApplyGuardrail needs it to build the request, so the + media_type travels with the payload as a data URI. + """ + seen = await self._images_seen( + [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}}] + ) + + assert seen == ["data:image/png;base64,AAAA"] + + @pytest.mark.asyncio + async def test_base64_source_without_a_media_type_is_passed_through(self): + """There is no format to attach, so the payload goes through unchanged.""" + seen = await self._images_seen([{"type": "image", "source": {"type": "base64", "data": "AAAA"}}]) + + assert seen == ["AAAA"] + + @pytest.mark.asyncio + async def test_file_source_yields_nothing(self): + """The bytes live behind the Files API and this extractor has no client. + + Documented as a known gap rather than silently handed on as a file_id string, + which a consumer would try to decode as an image. + """ + seen = await self._images_seen([{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}]) + + assert seen == [] + + @pytest.mark.asyncio + async def test_a_malformed_source_is_dropped_rather_than_passed_on(self): + seen = await self._images_seen( + [ + {"type": "image", "source": {"type": "base64"}}, + {"type": "image", "source": {"type": "url"}}, + {"type": "image", "source": {"type": "base64", "data": ""}}, + ] + ) + + assert seen == [] + + class TestAnthropicMessagesToolResultScanning: """LIT-5251: tool_result blocks carry whatever a client's local tool fetched, so they are the request-path payload an indirect prompt injection actually arrives in. From 355a44494403584c7180fd297b4fd0682567398f Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:08:18 +0800 Subject: [PATCH 10/21] fix: forward max_bytes only when the caller set one `process_image_async` passed `max_bytes=` on every call. The parameter has a default, so the signature stayed compatible, but the *call* did not: an override or test stub written against the previous signature now gets an unexpected keyword and raises. test_url_with_format_param caught it. It stubs `get_image_details_async` with a one-parameter fake, so the model call died on TypeError and the provider mock it asserts on was never reached. Omitting the keyword when there is nothing to cap leaves the model-call image paths byte-for-byte as they were, which is what the parameter being additive was supposed to mean. Co-Authored-By: Claude Opus 5 (1M context) --- litellm/litellm_core_utils/prompt_templates/factory.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index dafdf986c33..e59982e54c2 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3608,7 +3608,12 @@ class BedrockImageProcessor: 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: - img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url, max_bytes=max_bytes) + # Forward max_bytes only when the caller set one. Passing it + # unconditionally would reach every override and test stub written + # against the previous signature, so an additive parameter would + # break them; omitting it keeps the call byte-for-byte as it was. + capped: Final = {} if max_bytes is None else {"max_bytes": max_bytes} # mutable-ok: kwargs for one call + img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url, **capped) image_format = mime_type.split("/")[1] else: raise ValueError("Unsupported image type. Expected either image url or base64 encoded string") From 66cddac50420ac502c78feff6426015af5ab61c2 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:08:35 +0800 Subject: [PATCH 11/21] fix(guardrails): bound a request's image downloads in total, not just per image Capping each image at 4 MB does not bound a request. _create_bedrock_input_content_request gathers over every message and _build_input_content_items then gathers over every part, so the fetches all start together, and _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL is not consulted until _bin_pack_bedrock_content runs on items that are already resident. 200 urls at 4 MB is 800 MB, chosen by the caller. Two request-scoped bounds. A byte budget of _MAX_IMAGE_BYTES times _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL, held for one content-request build and passed down rather than kept on the guardrail, which is a callback instance shared by every request. And a semaphore of 4, so the reserved bytes are not all in flight at once. Each fetch reserves a cap and refunds what a usable image did not take. A response that decoded to nothing is charged in full: the transfer happened, and refunding it would let a url serving megabytes of junk be repeated down the whole list for free, which is the exhaustion being guarded against. An earlier draft refunded the whole reservation and bounded only in-flight bytes; the regression test caught that the decoded images still accumulated without limit. Inline base64 draws on neither budget nor gate. Those bytes arrived in the request body the proxy already accepted, so charging them to a download quota would refuse inline images for no reason. The total is what a single ApplyGuardrail call would accept anyway (20 images at 4 MB), so no request the API would take in one call is refused. A conversation chunked across several calls can exceed it, and images past the budget are then unscannable and left to on_unscannable_image, which blocks by default. Without the budget the regression test fetches 209,715,200 bytes for one request. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 124 ++++++++++++++++-- .../test_bedrock_guardrails.py | 49 +++++++ 2 files changed, 160 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 139674b8aff..36a74f9a949 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -157,6 +157,63 @@ _GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"}) _MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024 _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: Final = 20 +# A per-image cap does not bound a request. Content items are built by gathering +# over every message and then every part, so N urls are fetched concurrently, and +# the image-count limit above is not reached until _bin_pack_bedrock_content runs +# on items that are already resident. 500 urls is 500 fetches before anything says +# stop. These two bound the request itself: how much may be fetched in total, and +# how much of it may be in flight at once. +# +# The total is what a single ApplyGuardrail call would accept anyway (20 images at +# 4 MB), so no request the API would take in one call is refused. A conversation +# chunked across several calls can exceed it, and images past the budget are then +# unscannable and left to on_unscannable_image -- blocked by default. +_MAX_TOTAL_IMAGE_FETCH_BYTES: Final = _MAX_IMAGE_BYTES * _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL +_MAX_CONCURRENT_IMAGE_FETCHES: Final = 4 + + +class _ImageFetchBudget: + """Bytes still fetchable for one guardrail request, and a concurrency gate. + + Held for the lifetime of a single content-request build and passed down rather + than kept on the guardrail, which is a callback instance shared by every + request. Claims are optimistic: a fetch reserves the largest cap it could use + and hands back what it did not, so the worst-case resident size is the budget + plus the in-flight remainder, not the sum of every url a caller listed. + """ + + def __init__(self, total: int = _MAX_TOTAL_IMAGE_FETCH_BYTES) -> None: + self._remaining = total + self._gate = asyncio.Semaphore(_MAX_CONCURRENT_IMAGE_FETCHES) + + def claim(self) -> int: + """Reserve up to one image's worth of budget. 0 means exhausted.""" + granted: Final = min(_MAX_IMAGE_BYTES, self._remaining) + self._remaining -= granted + return granted + + def give_back(self, unused: int) -> None: + self._remaining += unused + + def gate(self) -> "asyncio.Semaphore": + return self._gate + + +def _retained_image_bytes(item: "BedrockContentItem | None") -> int: + """Approximate what a built image item holds, for budget accounting. + + Measured from the base64 payload rather than decoding it a second time; the + ratio is exact enough for a quota and costs nothing. + """ + if item is None: + return 0 + image: Final = item.get("image") + if not image: + return 0 + encoded: Final = image.get("source", {}).get("bytes") # mutable-ok: {} is a .get default, never mutated + return len(encoded) * 3 // 4 if isinstance(encoded, str) else 0 + + _APPLY_GUARDRAIL_IMAGE_FORMATS: Final[ dict[str, BedrockGuardrailImageFormat] ] = { # mutable-ok: module-level lookup table, never mutated @@ -359,12 +416,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if messages is None: return bedrock_request - per_message = await asyncio.gather(*(self._build_input_content_items(message=message) for message in messages)) + budget: Final = _ImageFetchBudget() + per_message = await asyncio.gather( + *(self._build_input_content_items(message=message, budget=budget) for message in messages) + ) # mutable-ok: BedrockRequest["content"] is a list in the AWS wire format bedrock_request["content"] = [item for items in per_message for item in items] return bedrock_request - async def _build_input_content_items(self, message: AllMessageValues) -> tuple[BedrockContentItem, ...]: + async def _build_input_content_items( + self, message: AllMessageValues, budget: "_ImageFetchBudget | None" = None + ) -> tuple[BedrockContentItem, ...]: """Flatten one request message into ApplyGuardrail INPUT content items. Grounding qualifiers are attached only when assembling the OUTPUT request, so a @@ -380,10 +442,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): parts: Final = cast( # cast-ok: AllMessageValues content is a union of part TypedDicts tuple[object, ...], tuple(content) ) - items: Final = await asyncio.gather(*(self._build_input_content_item(item=item) for item in parts)) + # A direct caller gets a fresh budget rather than an unbounded fetch. + request_budget: Final = budget if budget is not None else _ImageFetchBudget() + items: Final = await asyncio.gather( + *(self._build_input_content_item(item=item, budget=request_budget) for item in parts) + ) return tuple(item for item in items if item is not None) - async def _build_input_content_item(self, item: object) -> BedrockContentItem | None: + async def _build_input_content_item( + self, item: object, budget: "_ImageFetchBudget | None" = None + ) -> BedrockContentItem | None: if isinstance(item, str): return BedrockContentItem(text=BedrockTextContent(text=item)) if not isinstance(item, dict): @@ -396,7 +464,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): image_url: Final = self._get_image_url(item=part) if image_url is None: return None - return await self._build_image_content_item(image_url=image_url) + return await self._build_image_content_item(image_url=image_url, budget=budget) text: Final = part.get("text") if isinstance(text, str): return BedrockContentItem(text=BedrockTextContent(text=text)) @@ -471,26 +539,56 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # on_unscannable_image policy decides rather than this helper. return value - async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: + async def _build_image_content_item( + self, image_url: str, budget: "_ImageFetchBudget | None" = None + ) -> BedrockContentItem | None: """Decode or fetch an image part into an ApplyGuardrail image block. With `user_url_validation` off, `async_safe_get` degrades to an unrestricted, redirect-following GET on a caller-supplied URL, so it is not fetched at all """ - if not image_url.startswith("data:") and not getattr(litellm, "user_url_validation", True): + is_remote: Final = not image_url.startswith("data:") + if is_remote and not getattr(litellm, "user_url_validation", True): self._handle_unscannable_image( reason=f"remote image url not fetched because litellm.user_url_validation is disabled: {image_url}" ) return None + if not is_remote: + # Already in the request body the proxy accepted; nothing is fetched, + # so it draws on neither the byte budget nor the concurrency gate. + return await self._decode_image_content_item(image_url=image_url, max_bytes=None) + + request_budget: Final = budget if budget is not None else _ImageFetchBudget() + granted: Final = request_budget.claim() + if granted <= 0: + self._handle_unscannable_image( + reason="remote image skipped: this request already used its image download budget" + ) + return None + async with request_budget.gate(): + item: Final = await self._decode_image_content_item(image_url=image_url, max_bytes=granted) + + # Refund only what a usable image did not take. Returning the whole + # reservation would make the budget bound in-flight bytes alone, while + # decoded images stay resident in the request being assembled. + # + # A response that produced nothing is charged in full rather than refunded: + # the transfer still happened, and refunding it would let a url serving + # megabytes of unusable bytes be repeated down the whole list for free -- + # the exact shape of the exhaustion this guards against. + # + # No try/finally: the only escape from the line above is the HTTPException + # _handle_unscannable_image raises under the block policy, which ends the + # request and takes this request-scoped budget with it. + request_budget.give_back(granted - _retained_image_bytes(item) if item is not None else 0) + return item + + async def _decode_image_content_item(self, image_url: str, max_bytes: int | None) -> BedrockContentItem | None: + """Turn a data URI or a fetched url into an ApplyGuardrail image block.""" try: - # Cap the fetch itself. The decoded-size check below runs after the - # bytes are already resident, so on its own it does not stop a caller - # from pointing the proxy at an arbitrarily large or chunked response - # -- and `_build_input_content_items` gathers these concurrently, so - # one request with several URLs multiplies the allocation. block: Final = await BedrockImageProcessor.process_image_async( - image_url=image_url, format=None, max_bytes=_MAX_IMAGE_BYTES + image_url=image_url, format=None, max_bytes=max_bytes ) except PayloadTooLargeError as e: # Named before the ValueError arm it subclasses, so the operator sees diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index a849cfbfd71..bfab8eff6de 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5591,6 +5591,55 @@ class TestBedrockGuardrailImageInput: assert sum(served) <= 5 * 1024 * 1024, f"read {sum(served)} bytes past a 4 MB cap" assert len(served) < len(chunks), "the whole body was pulled before rejecting it" + @pytest.mark.asyncio + async def test_a_request_full_of_urls_is_bounded_in_total_not_just_per_image(self): + """A per-image cap does not bound a request. + + Content items are gathered over every message and every part, so the urls + are fetched concurrently, and the 20-image limit is not applied until + _bin_pack_bedrock_content runs on items that are already resident. Without + a request-wide budget, 200 urls at 4 MB is 800 MB the caller chose. + """ + # 200 parts, each serving a 1 MB image, against a 20 x 4 MB budget. + served: list[int] = [] + one_mb: list[bytes] = [b"\0" * (1024 * 1024)] + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"{self._REMOTE_IMAGE_URL}?i={i}"}} for i in range(200) + ], + } + ] + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(one_mb, served)): + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + fetched: int = sum(served) + assert fetched <= 20 * 4 * 1024 * 1024, f"fetched {fetched} bytes for one request" + assert len(request["content"]) < 200, "every url was kept despite the budget" + assert request["content"], "the budget swallowed the whole request" + + @pytest.mark.asyncio + async def test_inline_images_do_not_draw_on_the_download_budget(self): + """Base64 arrives in the request body the proxy already accepted. + + Nothing is fetched for it, so charging it against a download quota would + refuse inline images for no reason. + """ + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} for _ in range(40)], + } + ] + + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert len(request["content"]) == 40 + @pytest.mark.asyncio async def test_apply_guardrail_scans_images_from_inputs(self): """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. From fb5c886542049d1bdb494a9d951ddd97368bd0da Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:30:23 +0800 Subject: [PATCH 12/21] test(guardrails): pin the default-deny arms this PR relies on The PR's position is that anything it cannot recognise must not reach the model unscanned, but nothing asserted that. Every unrecognised shape landed on an untested `return None`, so a refactor letting one fall through to the text branch would have been reported to the operator as scanned content. One parametrized test walks the shapes a caller can actually put on the wire -- absent content, content that is not a list, a part that is not a mapping, an image part with no url, a url that is neither string nor mapping, a part carrying neither image nor text -- and asserts none of them becomes a content item. Stating the contract once beats restating each branch. The rest cover paths that are behaviour rather than defence: - `image_url` as a bare string, which OpenAI accepts alongside `{"url": ...}` and which reaches the model as an image either way - a bare base64 payload whose format cannot be sniffed, left to on_unscannable_image instead of guessed at. It goes through apply_guardrail, the only caller that normalizes before decoding - an oversized inline image under `allow`. Under `block` that branch raises and never returns, so this is the only path reaching its fall-through _get_image_url and _retained_image_bytes are exercised directly. Both are reached only through callers that already checked the shape, so their guards are otherwise unreachable and would be dropped in a refactor without anything noticing. Added lines in bedrock_guardrails.py now measure at full coverage. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_bedrock_guardrails.py | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index bfab8eff6de..56be71885da 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -24,6 +24,7 @@ from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockContentChunkResult, BedrockGuardrail, _redact_pii_matches, + _retained_image_bytes, ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -5640,6 +5641,111 @@ class TestBedrockGuardrailImageInput: assert len(request["content"]) == 40 + @pytest.mark.parametrize( + "content", + [ + pytest.param(None, id="no content"), + pytest.param(123, id="content is not a list"), + pytest.param([123], id="part is not a mapping"), + pytest.param([{"type": "image_url"}], id="image part with no url"), + pytest.param([{"type": "image_url", "image_url": 123}], id="url is not a string or mapping"), + pytest.param([{"type": "image_url", "image_url": {"url": 123}}], id="url value is not a string"), + pytest.param([{"type": "input_audio"}], id="part carries neither image nor text"), + ], + ) + @pytest.mark.asyncio + async def test_a_malformed_part_never_becomes_scannable_content(self, content): + """Default-deny is the point of this PR, so pin it rather than trust it. + + Each shape below is one a caller can put on the wire. None of them may turn + into a content item: an unrecognised part that fell through to the text + branch would be reported to the operator as scanned when it was not. + """ + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=[{"role": "user", "content": content}] + ) + + assert request["content"] == [] + + @pytest.mark.asyncio + async def test_a_bare_string_part_is_scanned_as_text(self): + """A content list may hold plain strings, not only typed parts.""" + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", messages=[{"role": "user", "content": ["just text"]}] + ) + + assert request["content"] == [{"text": {"text": "just text"}}] + + @pytest.mark.asyncio + async def test_image_url_given_as_a_plain_string_is_accepted(self): + """OpenAI accepts `image_url` as a bare string as well as `{"url": ...}`. + + Both reach the model as an image, so both have to reach the scan. + """ + request = await self._guardrail().convert_to_bedrock_format( + source="INPUT", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": self._PNG_DATA_URI}]}], + ) + + kinds = [k for item in request["content"] for k in item] + assert kinds == ["image"] + + @pytest.mark.asyncio + async def test_an_unrecognized_payload_is_left_to_the_unscannable_policy(self): + """_normalize_image_input sniffs png and jpeg out of bare base64. + + Anything else is handed to the decoder as-is rather than guessed at, so the + rejection comes from on_unscannable_image and not from a helper deciding + quietly on its own. + """ + # Reached through apply_guardrail: bare base64 arrives in inputs["images"], + # which is the only caller that normalizes before decoding. + with pytest.raises(HTTPException) as exc_info: + await self._guardrail().apply_guardrail( + inputs={"texts": [], "images": ["R0lGODlhAQABAAAAACw="]}, + request_data={}, + input_type="request", + ) + + assert "could not be read" in str(exc_info.value.detail) or "not a png/jpeg" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_an_oversized_inline_image_is_dropped_under_the_allow_policy(self): + """The allow policy has to survive the size rejection, not just the format one. + + Under `block` the oversized branch raises and never returns, so this is the + only path that reaches its fall-through. + """ + oversized_png = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * (5 * 1024 * 1024)).decode() + + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{oversized_png}"}}, + ], + } + ], + ) + + assert request["content"] == [{"text": {"text": "look"}}] + + def test_the_url_and_budget_helpers_guard_their_own_inputs(self): + """Both are reached only through callers that already checked the shape. + + Exercised directly so the guards are not silently dropped in a refactor that + gives either one a second caller. + """ + assert BedrockGuardrail._get_image_url(item={"type": "text", "text": "hi"}) is None + + assert _retained_image_bytes(None) == 0 + assert _retained_image_bytes({"text": {"text": "not an image"}}) == 0 + assert _retained_image_bytes({"image": {"format": "png", "source": {"bytes": 123}}}) == 0 + assert _retained_image_bytes({"image": {"format": "png", "source": {"bytes": "AAAA"}}}) == 3 + @pytest.mark.asyncio async def test_apply_guardrail_scans_images_from_inputs(self): """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. From ef71da890e5e1c66cf163d3c335571f499263621 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:00:31 +0800 Subject: [PATCH 13/21] test: cover the capped fetch in the suites that own it The capped path was only ever exercised through the Bedrock guardrail, so the two shared files it lives in were the last uncovered lines on this PR. Both now have tests where the code does, not only where its first caller does. url_utils: _underlying_httpx_client's TypeError had no test at all. Every existing one goes through AsyncHTTPHandler, whose `.client` is a real httpx client, so the isinstance guard always held. It exists because `cast` is banned here, and without it the declared return type would be a lie. Also covered: aborting mid-transfer, asserting on how many bytes were pulled rather than only that it raised, and the rebuilt response dropping content-encoding and content-length -- aiter_bytes yields decoded bytes, so carrying those over would describe the body wrongly. factory: get_image_details_async's body never ran under test. The model-path tests stub the whole method out, and the guardrail tests fake the transport underneath it. One test drives the real method with a cap; the other pins the reason the previous commit exists, calling process_image_async with no cap against a stub written to the old one-parameter signature. Co-Authored-By: Claude Opus 5 (1M context) --- ...llm_core_utils_prompt_templates_factory.py | 72 ++++++++++++++ .../litellm_core_utils/test_url_utils.py | 94 +++++++++++++++++++ 2 files changed, 166 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 3a7e06d085a..c24d2e806a3 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1,11 +1,15 @@ import base64 import json import os +import socket from unittest.mock import MagicMock, patch import pytest +import httpx + import litellm +from litellm.litellm_core_utils.url_utils import PayloadTooLargeError from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, BEDROCK_DOCUMENT_PLACEHOLDER_TEXT, @@ -3516,3 +3520,71 @@ async def test_bedrock_converse_pdf_only_user_message_gets_text_block_async(): assert len(result) == 1 assert any("document" in block for block in result[0]["content"]) assert _text_blocks(result[0]) == [BEDROCK_DOCUMENT_PLACEHOLDER_TEXT] + + +def _resolve_to_public(host, port, *args, **kwargs): + """Keep validate_url's DNS lookup off the network without faking the fetch.""" + return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port or 443))] + + +class TestBedrockImageProcessorMaxBytes: + """`max_bytes` is threaded from the caller down to the fetch. + + The Bedrock guardrail is the only caller that sets it. Everything else, the + model-call image paths included, must keep the previous unbounded fetch, and the + keyword has to be absent from the call rather than merely defaulted -- an + override or stub written against the old signature would otherwise break. + """ + + _REMOTE_URL = "https://93.184.216.34/a.png" + + @staticmethod + def _fake_stream(chunks): + import contextlib + + @contextlib.asynccontextmanager + async def _stream(self, method, url, **kwargs): + async def aiter_bytes(): + for chunk in chunks: + yield chunk + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers({"content-type": "image/png"}) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = aiter_bytes + yield response + + return _stream + + @pytest.mark.asyncio + async def test_a_remote_fetch_is_capped_when_max_bytes_is_given(self, monkeypatch): + monkeypatch.setattr(socket, "getaddrinfo", _resolve_to_public, raising=False) + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream([b"\0" * 8192])): + with pytest.raises(PayloadTooLargeError): + await BedrockImageProcessor.get_image_details_async(self._REMOTE_URL, max_bytes=1024) + + @pytest.mark.asyncio + async def test_omitting_max_bytes_leaves_the_call_as_it_was(self, monkeypatch): + """A stub written against the previous one-parameter signature still works. + + This is what test_url_with_format_param asserts through the model path; here + it is pinned on the helper itself so the plumbing cannot start passing the + keyword unconditionally again. + """ + monkeypatch.setattr(socket, "getaddrinfo", _resolve_to_public, raising=False) + seen: list = [] + + async def one_parameter_stub(image_url): + seen.append(image_url) + return "ZmFrZQ==", "image/png" + + monkeypatch.setattr( + BedrockImageProcessor, "get_image_details_async", staticmethod(one_parameter_stub) + ) + + block = await BedrockImageProcessor.process_image_async(image_url=self._REMOTE_URL, format=None) + + assert seen == [self._REMOTE_URL] + assert block["image"]["source"]["bytes"] == "ZmFrZQ==" diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index aaaa43a0dc4..e285f749ebc 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -1,11 +1,17 @@ +import contextlib import socket +from types import SimpleNamespace +from unittest.mock import MagicMock, patch +import httpx import pytest import litellm from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.url_utils import ( + PayloadTooLargeError, SSRFError, + _underlying_httpx_client, _is_blocked_ip, assert_same_origin, encode_url_path_segment, @@ -535,3 +541,91 @@ def test_assert_same_origin_error_message_does_not_leak_hostnames(): detail = str(exc.value) assert "attacker.example.com" not in detail assert "api.internal-corp.example" not in detail + + +class TestCappedFetch: + """`async_safe_get(max_bytes=...)` streams and aborts past the cap. + + `client.get` buffers the whole body first, so a caller-supplied url serving an + arbitrarily large or indefinitely chunked response is an unbounded allocation. + """ + + def test_a_client_with_no_httpx_client_is_rejected(self): + """Streaming needs the wrapped httpx client. + + AsyncHTTPHandler forwards get/post but not stream, so the wrapped `.client` + is what gets used. Something with neither is a programming error and says so, + rather than failing later inside httpx with nothing pointing back here. + """ + with pytest.raises(TypeError) as exc: + _underlying_httpx_client(object()) + + assert "no httpx client" in str(exc.value) + + def test_a_raw_httpx_client_is_used_as_is(self): + client = httpx.AsyncClient() + + assert _underlying_httpx_client(client) is client + + def test_a_wrapped_client_resolves_to_the_one_it_wraps(self): + inner = httpx.AsyncClient() + wrapper = SimpleNamespace(client=inner) + + assert _underlying_httpx_client(wrapper) is inner + + @pytest.mark.asyncio + async def test_the_body_is_cut_off_once_it_passes_the_cap(self, mock_dns_public): + """Asserting on how much was pulled is what separates a streamed abort from + buffering everything and rejecting afterwards.""" + served: list[int] = [] + chunks = [b"\0" * 1024 for _ in range(100)] + + @contextlib.asynccontextmanager + async def fake_stream(self, method, url, **kwargs): + async def aiter_bytes(): + for chunk in chunks: + served.append(len(chunk)) + yield chunk + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers({"content-type": "image/png"}) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = aiter_bytes + yield response + + client = httpx.AsyncClient() + with patch.object(httpx.AsyncClient, "stream", new=fake_stream): + with pytest.raises(PayloadTooLargeError): + await url_utils.async_safe_get(client, "https://93.184.216.34/a.png", max_bytes=4096) + + assert sum(served) <= 5 * 1024, f"pulled {sum(served)} bytes past a 4096 byte cap" + assert len(served) < len(chunks), "the whole body was read before rejecting it" + + @pytest.mark.asyncio + async def test_a_body_inside_the_cap_comes_back_whole(self, mock_dns_public): + """The rebuilt response drops content-encoding and content-length: aiter_bytes + yields decoded bytes, so carrying those over would describe the body wrongly.""" + + @contextlib.asynccontextmanager + async def fake_stream(self, method, url, **kwargs): + async def aiter_bytes(): + yield b"tiny-image" + + response = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers( + {"content-type": "image/png", "content-length": "999", "content-encoding": "gzip"} + ) + response.request = httpx.Request("GET", str(url)) + response.aiter_bytes = aiter_bytes + yield response + + client = httpx.AsyncClient() + with patch.object(httpx.AsyncClient, "stream", new=fake_stream): + result = await url_utils.async_safe_get(client, "https://93.184.216.34/a.png", max_bytes=4096) + + assert result.content == b"tiny-image" + assert result.headers.get("content-type") == "image/png" + assert "content-encoding" not in result.headers + assert result.headers.get("content-length") == str(len(b"tiny-image")) From e84522118c9a58eeb0c3b57c93c7a89429978820 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:11:44 +0800 Subject: [PATCH 14/21] style: satisfy LIT010 on the lines this PR touches The base gained LIT010 (assignment without a Final declaration) while this PR was open, so code that passed the gate at the old base now trips it. Three are genuinely never rebound and take Final: `per_message`, `chunks`, and the awaited fetch result, which is now bound separately so the unpack below has room for its reason. Two are real rebinding and say so. `total` is the running byte count the cap is measured against. `mime_type` is overwritten further down when the caller passes `format`, so it cannot be Final; the reason sits on the unpack line itself, since a reason on the preceding line does not count. Co-Authored-By: Claude Opus 5 (1M context) --- litellm/litellm_core_utils/prompt_templates/factory.py | 3 ++- litellm/litellm_core_utils/url_utils.py | 4 ++-- .../proxy/guardrails/guardrail_hooks/bedrock_guardrails.py | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 624d5eaac27..ab60e827943 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3571,7 +3571,8 @@ class BedrockImageProcessor: # against the previous signature, so an additive parameter would # break them; omitting it keeps the call byte-for-byte as it was. capped: Final = {} if max_bytes is None else {"max_bytes": max_bytes} # mutable-ok: kwargs for one call - img_bytes, mime_type = await BedrockImageProcessor.get_image_details_async(image_url, **capped) + fetched: Final = await BedrockImageProcessor.get_image_details_async(image_url, **capped) + img_bytes, mime_type = fetched # rebind-ok: mime_type is overridden below by `format` image_format = mime_type.split("/")[1] else: raise ValueError("Unsupported image type. Expected either image url or base64 encoded string") diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0b26eaa869e..ec3afe7276e 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -455,8 +455,8 @@ async def _async_get_capped( Returns a fully-read response so callers keep using ``.content`` as before. """ async with _underlying_httpx_client(client).stream("GET", url, **request_kwargs) as response: - chunks: list[bytes] = [] # mutable-ok: accumulator for the capped body - total = 0 + chunks: Final[list[bytes]] = [] # mutable-ok: accumulator for the capped body + total = 0 # rebind-ok: running byte count for the cap async for chunk in response.aiter_bytes(): total += len(chunk) if total > max_bytes: diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 36a74f9a949..b49040580f7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -417,7 +417,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return bedrock_request budget: Final = _ImageFetchBudget() - per_message = await asyncio.gather( + per_message: Final = await asyncio.gather( *(self._build_input_content_items(message=message, budget=budget) for message in messages) ) # mutable-ok: BedrockRequest["content"] is a list in the AWS wire format From c830ee204c11a2c5907093833804a69c7f52ed81 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:26:31 +0800 Subject: [PATCH 15/21] fix: type the forwarded kwargs so basedpyright can check the httpx call CI's basedpyright budget moved with the base and surfaced 12 new reportArgumentType errors, ten of them on one line. `request_kwargs` was typed `dict[str, object]` to keep an earlier gate happy, but `object` does not unpack into httpx's typed parameters: `stream()` names ten of them (content, data, files, params, headers, cookies, auth, follow_redirects, timeout, extensions) and every one was an error. `dict[str, Any]` is what the values actually are -- whatever `async_safe_get`'s caller passed through -- and it lets the call be checked instead of merely tolerated. The other two are the OUTPUT branch's `enumerate(filtered_messages)`. Making skip_scan conditional on `image_urls` removed pyright's narrowing there, since skip_scan used to be the only way that name could still be None. It cannot be None at runtime -- images exist on the request side only, so a response scan returns early exactly as before -- but the branch now says so with `or ()` rather than resting on that indirection. Co-Authored-By: Claude Opus 5 (1M context) --- litellm/litellm_core_utils/url_utils.py | 6 +++--- .../proxy/guardrails/guardrail_hooks/bedrock_guardrails.py | 6 +++++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index ec3afe7276e..50cdc679f9b 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -443,7 +443,7 @@ async def _async_get_capped( client: object, url: str, max_bytes: int, - request_kwargs: dict[str, object], # mutable-ok: forwarded straight to httpx as **kwargs + request_kwargs: dict[str, Any], # mutable-ok: forwarded straight to httpx as **kwargs ) -> httpx.Response: """GET ``url``, aborting the transfer once the body exceeds ``max_bytes``. @@ -483,7 +483,7 @@ async def async_safe_get(client: Any, url: str, max_bytes: int | None = None, ** async def _issue( target_url: str, - request_kwargs: dict[str, object], # mutable-ok: forwarded straight to httpx as **kwargs + request_kwargs: dict[str, Any], # mutable-ok: forwarded straight to httpx as **kwargs ) -> httpx.Response: if max_bytes is None: return await client.get(target_url, **request_kwargs) @@ -496,7 +496,7 @@ async def async_safe_get(client: Any, url: str, max_bytes: int | None = None, ** caller_headers: Final = kwargs.pop("headers", {}) for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - hop_kwargs: dict[str, object] = { # mutable-ok: a fresh per-hop kwargs dict, consumed by this call + hop_kwargs: dict[str, Any] = { # mutable-ok: a fresh per-hop kwargs dict, consumed by this call **kwargs, "headers": {**caller_headers, "Host": original_host}, # mutable-ok: httpx takes headers as a dict "follow_redirects": False, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index b49040580f7..2e9971697cd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -3465,7 +3465,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ), finish_reason="stop", ) - for _idx, _msg in enumerate(filtered_messages) + # `or ()`: skip_scan is now bypassed when an image is present, + # and images exist on the request side only, so a response scan + # still always has messages here. Spelled out rather than left + # leaning on that indirection. + for _idx, _msg in enumerate(filtered_messages or ()) ] ) bedrock_response = await self.make_bedrock_api_request( From 43687653744f3fd112e9229f96360a43765a063a Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:40:39 +0800 Subject: [PATCH 16/21] fix(guardrails): refuse a file-backed image instead of documenting that it passes An Anthropic `{"type": "image", "source": {"type": "file", "file_id": ...}}` block carries no bytes, so the extractor yields nothing for it while the provider still forwards the file to the model. The docstring called that a known gap. Documented is not the same as safe, and a silent pass is the exact failure this whole path exists to remove, so it is now refused: blocked by default, allowed only where the operator sets on_unscannable_image. Resolving the file id would need a Files API client this guardrail does not have, which is a larger change than the one this PR is making. The check reads inputs["structured_messages"], which on /v1/messages carries the raw Anthropic blocks, so the file source is visible to the guardrail without altering GenericGuardrailAPIInputs. Putting a marker in inputs["images"] was the alternative and would have reached five other guardrails that consume that field, trading one gap for four new unknowns. It detects the file shape specifically rather than comparing an image count against inputs["images"]. structured_messages is already narrowed by the skip and scope flags, so a mismatch is not by itself evidence of a dropped image, and refusing a legitimate request would be worse than the gap being closed. A test pins that base64 and url sources are still accepted. Placed before the incremental and latest-message-only shortcuts, so a file-backed image cannot be skipped by them either. The first draft of the refusal test passed for the wrong reason: without the check the request died on absent AWS credentials rather than on the bypass. The Bedrock call is now stubbed, so removing the check makes it fail with DID NOT RAISE -- the silent pass itself. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 39 ++++++ .../test_bedrock_guardrails.py | 129 ++++++++++++++++++ 2 files changed, 168 insertions(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 2e9971697cd..5b7ee1196ce 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -539,6 +539,38 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # on_unscannable_image policy decides rather than this helper. return value + @staticmethod + def _file_backed_image_count(structured_messages: object) -> int: + """Count image parts whose bytes live behind a provider Files API. + + An Anthropic `{"type": "image", "source": {"type": "file", "file_id": ...}}` + block carries no data, so the guardrail translation yields nothing for it + while the provider still forwards the file to the model. Left alone that is + an image the policy never sees, which is the failure this whole path exists + to remove -- so it is counted here and refused rather than documented. + + Detects that one shape rather than comparing counts against + `inputs["images"]`: structured_messages is already narrowed by the + skip/scope flags, so a mismatch is not by itself evidence of a dropped + image, and blocking a legitimate request is worse than the gap. + """ + if not isinstance(structured_messages, list): + return 0 + found = 0 # rebind-ok: running count over the message list + for message in structured_messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") != "image": + continue + source = part.get("source") + if isinstance(source, dict) and source.get("type") == "file": + found += 1 + return found + async def _build_image_content_item( self, image_url: str, budget: "_ImageFetchBudget | None" = None ) -> BedrockContentItem | None: @@ -3404,6 +3436,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # from `image`/`source` blocks. Five other guardrails already consume this # field; Bedrock was the one that dropped it on the floor. image_urls: Final = tuple(inputs.get("images") or ()) if input_type == "request" else () + if input_type == "request": + file_images: Final = self._file_backed_image_count(inputs.get("structured_messages")) + if file_images: + # Before the shortcuts below, so this cannot be skipped either. + self._handle_unscannable_image( + reason=f"{file_images} image(s) reference a provider file id, whose bytes are not available here" + ) try: verbose_proxy_logger.debug( "Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 56be71885da..b29103551cd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5746,6 +5746,135 @@ class TestBedrockGuardrailImageInput: assert _retained_image_bytes({"image": {"format": "png", "source": {"bytes": 123}}}) == 0 assert _retained_image_bytes({"image": {"format": "png", "source": {"bytes": "AAAA"}}}) == 3 + @pytest.mark.asyncio + async def test_a_file_backed_image_is_refused_rather_than_ignored(self): + """`{"type": "file"}` carries no bytes, so nothing reaches inputs["images"]. + + The provider still forwards the file to the model, so ignoring it is exactly + the silent pass this path exists to remove. Documented is not the same as + safe; under the default policy the request is refused. + """ + inputs = { + "texts": ["what does this say?"], + "images": [], + "structured_messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what does this say?"}, + {"type": "image", "source": {"type": "file", "file_id": "file_abc"}}, + ], + } + ], + } + + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs["messages"]) + return {"action": "NONE", "outputs": []} + + # Stubbed so that without the refusal this request would simply succeed: + # the failure mode being pinned is a silent pass, not an AWS error. + with patch.object(g, "make_bedrock_api_request", new=spy): + with pytest.raises(HTTPException) as exc_info: + await g.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert "file id" in str(exc_info.value.detail) + assert sent == [], "refused before any scan was attempted" + + @pytest.mark.asyncio + async def test_a_file_backed_image_is_let_through_under_the_allow_policy(self): + """An operator who would rather serve it unscanned can still say so.""" + g = self._guardrail(on_unscannable_image="allow") + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs["messages"]) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={ + "texts": ["hello"], + "images": [], + "structured_messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], + } + ], + }, + request_data={}, + input_type="request", + ) + + assert sent, "the text alongside the file image still has to be scanned" + + @pytest.mark.asyncio + async def test_the_scannable_source_shapes_are_not_refused(self): + """The refusal has to be specific to the shape that cannot be read. + + structured_messages is already narrowed by the skip and scope flags, so a + count mismatch against inputs["images"] is not evidence of a dropped image. + Blocking a legitimate request would be worse than the gap being closed. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(await g.convert_to_bedrock_format(source="INPUT", messages=kwargs["messages"])) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + await g.apply_guardrail( + inputs={ + "texts": ["hello"], + "images": [self._PNG_DATA_URI], + "structured_messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image", "source": {"type": "base64", "data": "AAAA"}}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}, + ], + } + ], + }, + request_data={}, + input_type="request", + ) + + kinds = [k for item in sent[0]["content"] for k in item] + assert "image" in kinds + + @pytest.mark.asyncio + async def test_a_file_backed_image_on_the_response_side_is_not_refused(self): + """Images are a request-side concern; an OUTPUT scan takes generated text.""" + g = self._guardrail() + + async def spy(**kwargs): + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + result = await g.apply_guardrail( + inputs={ + "texts": ["the model said this"], + "structured_messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], + } + ], + }, + request_data={}, + input_type="response", + ) + + assert result is not None + @pytest.mark.asyncio async def test_apply_guardrail_scans_images_from_inputs(self): """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook. From 1930ae163418dacc9a46b0aeba47f9a084c394eb Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:51:20 +0800 Subject: [PATCH 17/21] refactor(guardrails): move the file-image refusal out of apply_guardrail The previous commit added two branches to apply_guardrail and pushed it from 14 to 16 against ruff-strict's max-complexity of 15, taking the codebase C901 total from 312 to 313 with the budget at 312. It was one branch under the ceiling, so the check had to live somewhere else. Both conditions now sit in _refuse_file_backed_images: input_type, because images are a request-side concern, and the count. apply_guardrail is back to 14 and the call site reads as one statement. Worth recording why this took two tries to see. `ruff check --select C901` uses pyproject.toml, where max-complexity is 10, and apply_guardrail was already over that both before and after -- so a before/after comparison showed no change and I read the failure as a local artifact problem. The gate runs ruff with --config ruff-strict.toml, where the ceiling is 15, and that is the only threshold the budget is counted against. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 5b7ee1196ce..ee0d057ec0e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -539,6 +539,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # on_unscannable_image policy decides rather than this helper. return value + def _refuse_file_backed_images(self, inputs: "GenericGuardrailAPIInputs", input_type: str) -> None: + """Hand a file-backed image to on_unscannable_image rather than ignoring it. + + Images are a request-side concern; an OUTPUT scan takes generated text, so a + file reference sitting in the conversation history is not this scan's problem. + """ + if input_type != "request": + return + found: Final = self._file_backed_image_count(inputs.get("structured_messages")) + if not found: + return + self._handle_unscannable_image( + reason=f"{found} image(s) reference a provider file id, whose bytes are not available here" + ) + @staticmethod def _file_backed_image_count(structured_messages: object) -> int: """Count image parts whose bytes live behind a provider Files API. @@ -3436,13 +3451,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # from `image`/`source` blocks. Five other guardrails already consume this # field; Bedrock was the one that dropped it on the floor. image_urls: Final = tuple(inputs.get("images") or ()) if input_type == "request" else () - if input_type == "request": - file_images: Final = self._file_backed_image_count(inputs.get("structured_messages")) - if file_images: - # Before the shortcuts below, so this cannot be skipped either. - self._handle_unscannable_image( - reason=f"{file_images} image(s) reference a provider file id, whose bytes are not available here" - ) + # Before the shortcuts below, so a file-backed image cannot be skipped either. + # Both conditions live in the callee: apply_guardrail sits one branch under + # ruff-strict's complexity ceiling, and two more here would cross it. + self._refuse_file_backed_images(inputs=inputs, input_type=input_type) try: verbose_proxy_logger.debug( "Bedrock Guardrail: Applying guardrail to %s text(s) and %s image(s)", len(texts), len(image_urls) From e54f69c48d24df93377278ba06db111359c83964 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:54:49 +0800 Subject: [PATCH 18/21] test(guardrails): keep the file-image scan walking past an unreadable entry Codecov left one line uncovered: the `continue` that skips a structured_messages entry which is not a mapping. Worth a test rather than a shrug. structured_messages comes from the caller, so its shape is not guaranteed, and the loop has to keep going past an entry it cannot read. Stopping or throwing there would let one junk element hide a file image sitting after it, turning the guard into the bypass it exists to prevent -- so the test puts the file image last, behind a string, an int, and a message whose content is not a list. Added lines in bedrock_guardrails.py are back to full coverage. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_bedrock_guardrails.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index b29103551cd..3eed3702720 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5850,6 +5850,44 @@ class TestBedrockGuardrailImageInput: kinds = [k for item in sent[0]["content"] for k in item] assert "image" in kinds + @pytest.mark.asyncio + async def test_a_malformed_structured_message_does_not_derail_the_file_check(self): + """structured_messages comes from the caller, so its shape is not guaranteed. + + The scan must keep walking past an entry it cannot read rather than throwing + or giving up, or a single junk element would hide a file image sitting after + it -- turning a defensive guard into the bypass it was meant to prevent. + """ + g = self._guardrail() + sent: list = [] + + async def spy(**kwargs): + sent.append(kwargs["messages"]) + return {"action": "NONE", "outputs": []} + + with patch.object(g, "make_bedrock_api_request", new=spy): + with pytest.raises(HTTPException) as exc_info: + await g.apply_guardrail( + inputs={ + "texts": ["hello"], + "images": [], + "structured_messages": [ + "not a message", + 123, + {"role": "user", "content": "a plain string, not a list"}, + { + "role": "user", + "content": [{"type": "image", "source": {"type": "file", "file_id": "file_abc"}}], + }, + ], + }, + request_data={}, + input_type="request", + ) + + assert "file id" in str(exc_info.value.detail) + assert sent == [] + @pytest.mark.asyncio async def test_a_file_backed_image_on_the_response_side_is_not_refused(self): """Images are a request-side concern; an OUTPUT scan takes generated text.""" From 05488171b31599141267aa744212b9d7a68a6759 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:13:31 +0800 Subject: [PATCH 19/21] fix(guardrails): grant a whole image's budget or none, so the rejection reads true The download budget handed out whatever was left. With a few hundred bytes remaining that became the per-image cap, and a perfectly ordinary 50 KB image was refused with remote image over ApplyGuardrail's 4 MB limit: remote body exceeds 100 bytes which blames AWS for this request having spent its own budget, and contradicts itself in the same sentence. This PR argues that a guardrail's failure modes have to be legible; that message is not. claim() is all or nothing now. "Over the per-image limit" and "this request is out of download budget" stay separately diagnosable, at the cost of up to one image's worth of the 80 MB going unused at the tail. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 17 +++++++++++++---- .../guardrail_hooks/test_bedrock_guardrails.py | 18 ++++++++++++++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index ee0d057ec0e..2e4e61d1e15 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -187,10 +187,19 @@ class _ImageFetchBudget: self._gate = asyncio.Semaphore(_MAX_CONCURRENT_IMAGE_FETCHES) def claim(self) -> int: - """Reserve up to one image's worth of budget. 0 means exhausted.""" - granted: Final = min(_MAX_IMAGE_BYTES, self._remaining) - self._remaining -= granted - return granted + """Reserve one image's worth of budget. 0 means exhausted. + + All or nothing rather than handing out whatever is left. A partial grant + would cap the fetch below the per-image limit, and the rejection then + surfaces as "over ApplyGuardrail's 4 MB limit" while naming a few hundred + bytes -- blaming AWS for this request having spent its own budget. The two + failures stay separately legible at the cost of up to one image's worth of + headroom going unused at the tail. + """ + if self._remaining < _MAX_IMAGE_BYTES: + return 0 + self._remaining -= _MAX_IMAGE_BYTES + return _MAX_IMAGE_BYTES def give_back(self, unused: int) -> None: self._remaining += unused diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 3eed3702720..1cd104aa044 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5623,6 +5623,24 @@ class TestBedrockGuardrailImageInput: assert len(request["content"]) < 200, "every url was kept despite the budget" assert request["content"], "the budget swallowed the whole request" + def test_the_budget_grants_a_whole_image_or_nothing(self): + """A partial grant would cap a fetch below the per-image limit. + + The rejection then reads "over ApplyGuardrail's 4 MB limit" while naming a + few hundred bytes, blaming AWS for this request having spent its own budget. + Keeping the two failures separately legible is worth leaving one image's + worth of headroom unused at the tail. + """ + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + _MAX_IMAGE_BYTES, + _ImageFetchBudget, + ) + + budget = _ImageFetchBudget(total=_MAX_IMAGE_BYTES + 100) + + assert budget.claim() == _MAX_IMAGE_BYTES + assert budget.claim() == 0, "100 bytes left must read as exhausted, not as a 100 byte cap" + @pytest.mark.asyncio async def test_inline_images_do_not_draw_on_the_download_budget(self): """Base64 arrives in the request body the proxy already accepted. From ed991c4f6c6e8e3e294cb6aa330ae013dadb5127 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:20:16 +0800 Subject: [PATCH 20/21] test(guardrails): an oversized remote image under `allow` drops, it does not derail The previous commit's all-or-nothing grant left one line uncovered: the return after the PayloadTooLargeError arm. Under `block` that arm raises, so nothing was reaching its fall-through. The gap is worth a test rather than a shrug. An operator who sets on_unscannable_image=allow asked for the image to go unscanned, not for the whole request to die on a url that streams past the cap. The test asserts the text alongside it still reaches the scan, and that the transfer was cut off rather than read in full first. The existing allow-policy oversize test uses an inline data URI, which is checked after decoding and never touches this arm. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_bedrock_guardrails.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 1cd104aa044..f51190ef683 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5623,6 +5623,35 @@ class TestBedrockGuardrailImageInput: assert len(request["content"]) < 200, "every url was kept despite the budget" assert request["content"], "the budget swallowed the whole request" + @pytest.mark.asyncio + async def test_an_oversized_remote_image_is_dropped_under_the_allow_policy(self): + """The transfer is cut off, and then the request has to carry on. + + `block` raises out of the size rejection, so this is the only path that + reaches its fall-through. An operator who set `allow` asked for the image to + go unscanned, not for the whole request to die on it. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": self._REMOTE_IMAGE_URL}}, + ], + } + ] + served: list[int] = [] + chunks: list[bytes] = [b"\0" * (1024 * 1024) for _ in range(8)] + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(chunks, served)): + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + assert request["content"] == [{"text": {"text": "look"}}] + assert served, "the capped stream path did not run" + assert len(served) < len(chunks), "the whole body was pulled before dropping it" + def test_the_budget_grants_a_whole_image_or_nothing(self): """A partial grant would cap a fetch below the per-image limit. From d4904ad1e7231d2539843ad555435025b9c7d004 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:30:06 +0800 Subject: [PATCH 21/21] docs(guardrails): correct two comments that outran what the code guarantees MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _ImageFetchBudget's class docstring still described claims as optimistic, taking "the largest cap it could use" — the behaviour claim() had before the previous commit made it all or nothing. Two docstrings in one class contradicting each other is worse than either being absent. _file_backed_image_count argued that structured_messages being narrowed by the skip and scope flags is why it does not compare counts against inputs["images"]. That is true and it is half the picture: the same narrowing means a file source in a message the scope excluded is not seen here at all. `images` is extracted from every message while structured_messages holds only the scoped subset -- different lists in guardrail_translation/handler.py, neither of which this PR changes. The blind spot is now named, along with why reading the unscoped list would trade it for a worse answer: refusing content the operator's skip flags deliberately took out of scanning. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 2e4e61d1e15..fc083010147 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -177,9 +177,10 @@ class _ImageFetchBudget: Held for the lifetime of a single content-request build and passed down rather than kept on the guardrail, which is a callback instance shared by every - request. Claims are optimistic: a fetch reserves the largest cap it could use - and hands back what it did not, so the worst-case resident size is the budget - plus the in-flight remainder, not the sum of every url a caller listed. + request. A fetch reserves one whole image's worth up front and hands back what + the decoded image did not take, so the worst-case resident size is the budget + plus whatever the in-flight fetches have pulled, not the sum of every url a + caller listed. See `claim` for why the reservation is all or nothing. """ def __init__(self, total: int = _MAX_TOTAL_IMAGE_FETCH_BYTES) -> None: @@ -577,6 +578,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): `inputs["images"]`: structured_messages is already narrowed by the skip/scope flags, so a mismatch is not by itself evidence of a dropped image, and blocking a legitimate request is worse than the gap. + + That narrowing cuts both ways and this check inherits it. `images` is + extracted from every message while structured_messages holds only the + scoped subset (guardrail_translation/handler.py builds them from different + lists), so a file source in a message the scope excluded is not seen here. + Reading the unscoped list instead would refuse requests for content the + operator's skip flags deliberately took out of scanning, which is a + different wrong answer. """ if not isinstance(structured_messages, list): return 0