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 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 8/9] 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 9/9] 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.