From 66cddac50420ac502c78feff6426015af5ab61c2 Mon Sep 17 00:00:00 2001 From: samtsai15 <6171228+samtsai15@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:08:35 +0800 Subject: [PATCH] fix(guardrails): bound a request's image downloads in total, not just per image Capping each image at 4 MB does not bound a request. _create_bedrock_input_content_request gathers over every message and _build_input_content_items then gathers over every part, so the fetches all start together, and _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL is not consulted until _bin_pack_bedrock_content runs on items that are already resident. 200 urls at 4 MB is 800 MB, chosen by the caller. Two request-scoped bounds. A byte budget of _MAX_IMAGE_BYTES times _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL, held for one content-request build and passed down rather than kept on the guardrail, which is a callback instance shared by every request. And a semaphore of 4, so the reserved bytes are not all in flight at once. Each fetch reserves a cap and refunds what a usable image did not take. A response that decoded to nothing is charged in full: the transfer happened, and refunding it would let a url serving megabytes of junk be repeated down the whole list for free, which is the exhaustion being guarded against. An earlier draft refunded the whole reservation and bounded only in-flight bytes; the regression test caught that the decoded images still accumulated without limit. Inline base64 draws on neither budget nor gate. Those bytes arrived in the request body the proxy already accepted, so charging them to a download quota would refuse inline images for no reason. The total is what a single ApplyGuardrail call would accept anyway (20 images at 4 MB), so no request the API would take in one call is refused. A conversation chunked across several calls can exceed it, and images past the budget are then unscannable and left to on_unscannable_image, which blocks by default. Without the budget the regression test fetches 209,715,200 bytes for one request. Co-Authored-By: Claude Opus 5 (1M context) --- .../guardrail_hooks/bedrock_guardrails.py | 124 ++++++++++++++++-- .../test_bedrock_guardrails.py | 49 +++++++ 2 files changed, 160 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 139674b8aff..36a74f9a949 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -157,6 +157,63 @@ _GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"}) _MAX_IMAGE_BYTES: Final = 4 * 1024 * 1024 _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL: Final = 20 +# A per-image cap does not bound a request. Content items are built by gathering +# over every message and then every part, so N urls are fetched concurrently, and +# the image-count limit above is not reached until _bin_pack_bedrock_content runs +# on items that are already resident. 500 urls is 500 fetches before anything says +# stop. These two bound the request itself: how much may be fetched in total, and +# how much of it may be in flight at once. +# +# The total is what a single ApplyGuardrail call would accept anyway (20 images at +# 4 MB), so no request the API would take in one call is refused. A conversation +# chunked across several calls can exceed it, and images past the budget are then +# unscannable and left to on_unscannable_image -- blocked by default. +_MAX_TOTAL_IMAGE_FETCH_BYTES: Final = _MAX_IMAGE_BYTES * _MAX_IMAGES_PER_APPLY_GUARDRAIL_CALL +_MAX_CONCURRENT_IMAGE_FETCHES: Final = 4 + + +class _ImageFetchBudget: + """Bytes still fetchable for one guardrail request, and a concurrency gate. + + Held for the lifetime of a single content-request build and passed down rather + than kept on the guardrail, which is a callback instance shared by every + request. Claims are optimistic: a fetch reserves the largest cap it could use + and hands back what it did not, so the worst-case resident size is the budget + plus the in-flight remainder, not the sum of every url a caller listed. + """ + + def __init__(self, total: int = _MAX_TOTAL_IMAGE_FETCH_BYTES) -> None: + self._remaining = total + self._gate = asyncio.Semaphore(_MAX_CONCURRENT_IMAGE_FETCHES) + + def claim(self) -> int: + """Reserve up to one image's worth of budget. 0 means exhausted.""" + granted: Final = min(_MAX_IMAGE_BYTES, self._remaining) + self._remaining -= granted + return granted + + def give_back(self, unused: int) -> None: + self._remaining += unused + + def gate(self) -> "asyncio.Semaphore": + return self._gate + + +def _retained_image_bytes(item: "BedrockContentItem | None") -> int: + """Approximate what a built image item holds, for budget accounting. + + Measured from the base64 payload rather than decoding it a second time; the + ratio is exact enough for a quota and costs nothing. + """ + if item is None: + return 0 + image: Final = item.get("image") + if not image: + return 0 + encoded: Final = image.get("source", {}).get("bytes") # mutable-ok: {} is a .get default, never mutated + return len(encoded) * 3 // 4 if isinstance(encoded, str) else 0 + + _APPLY_GUARDRAIL_IMAGE_FORMATS: Final[ dict[str, BedrockGuardrailImageFormat] ] = { # mutable-ok: module-level lookup table, never mutated @@ -359,12 +416,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if messages is None: return bedrock_request - per_message = await asyncio.gather(*(self._build_input_content_items(message=message) for message in messages)) + budget: Final = _ImageFetchBudget() + per_message = await asyncio.gather( + *(self._build_input_content_items(message=message, budget=budget) for message in messages) + ) # mutable-ok: BedrockRequest["content"] is a list in the AWS wire format bedrock_request["content"] = [item for items in per_message for item in items] return bedrock_request - async def _build_input_content_items(self, message: AllMessageValues) -> tuple[BedrockContentItem, ...]: + async def _build_input_content_items( + self, message: AllMessageValues, budget: "_ImageFetchBudget | None" = None + ) -> tuple[BedrockContentItem, ...]: """Flatten one request message into ApplyGuardrail INPUT content items. Grounding qualifiers are attached only when assembling the OUTPUT request, so a @@ -380,10 +442,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): parts: Final = cast( # cast-ok: AllMessageValues content is a union of part TypedDicts tuple[object, ...], tuple(content) ) - items: Final = await asyncio.gather(*(self._build_input_content_item(item=item) for item in parts)) + # A direct caller gets a fresh budget rather than an unbounded fetch. + request_budget: Final = budget if budget is not None else _ImageFetchBudget() + items: Final = await asyncio.gather( + *(self._build_input_content_item(item=item, budget=request_budget) for item in parts) + ) return tuple(item for item in items if item is not None) - async def _build_input_content_item(self, item: object) -> BedrockContentItem | None: + async def _build_input_content_item( + self, item: object, budget: "_ImageFetchBudget | None" = None + ) -> BedrockContentItem | None: if isinstance(item, str): return BedrockContentItem(text=BedrockTextContent(text=item)) if not isinstance(item, dict): @@ -396,7 +464,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): image_url: Final = self._get_image_url(item=part) if image_url is None: return None - return await self._build_image_content_item(image_url=image_url) + return await self._build_image_content_item(image_url=image_url, budget=budget) text: Final = part.get("text") if isinstance(text, str): return BedrockContentItem(text=BedrockTextContent(text=text)) @@ -471,26 +539,56 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # on_unscannable_image policy decides rather than this helper. return value - async def _build_image_content_item(self, image_url: str) -> BedrockContentItem | None: + async def _build_image_content_item( + self, image_url: str, budget: "_ImageFetchBudget | None" = None + ) -> BedrockContentItem | None: """Decode or fetch an image part into an ApplyGuardrail image block. With `user_url_validation` off, `async_safe_get` degrades to an unrestricted, redirect-following GET on a caller-supplied URL, so it is not fetched at all """ - if not image_url.startswith("data:") and not getattr(litellm, "user_url_validation", True): + is_remote: Final = not image_url.startswith("data:") + if is_remote and not getattr(litellm, "user_url_validation", True): self._handle_unscannable_image( reason=f"remote image url not fetched because litellm.user_url_validation is disabled: {image_url}" ) return None + if not is_remote: + # Already in the request body the proxy accepted; nothing is fetched, + # so it draws on neither the byte budget nor the concurrency gate. + return await self._decode_image_content_item(image_url=image_url, max_bytes=None) + + request_budget: Final = budget if budget is not None else _ImageFetchBudget() + granted: Final = request_budget.claim() + if granted <= 0: + self._handle_unscannable_image( + reason="remote image skipped: this request already used its image download budget" + ) + return None + async with request_budget.gate(): + item: Final = await self._decode_image_content_item(image_url=image_url, max_bytes=granted) + + # Refund only what a usable image did not take. Returning the whole + # reservation would make the budget bound in-flight bytes alone, while + # decoded images stay resident in the request being assembled. + # + # A response that produced nothing is charged in full rather than refunded: + # the transfer still happened, and refunding it would let a url serving + # megabytes of unusable bytes be repeated down the whole list for free -- + # the exact shape of the exhaustion this guards against. + # + # No try/finally: the only escape from the line above is the HTTPException + # _handle_unscannable_image raises under the block policy, which ends the + # request and takes this request-scoped budget with it. + request_budget.give_back(granted - _retained_image_bytes(item) if item is not None else 0) + return item + + async def _decode_image_content_item(self, image_url: str, max_bytes: int | None) -> BedrockContentItem | None: + """Turn a data URI or a fetched url into an ApplyGuardrail image block.""" try: - # Cap the fetch itself. The decoded-size check below runs after the - # bytes are already resident, so on its own it does not stop a caller - # from pointing the proxy at an arbitrarily large or chunked response - # -- and `_build_input_content_items` gathers these concurrently, so - # one request with several URLs multiplies the allocation. block: Final = await BedrockImageProcessor.process_image_async( - image_url=image_url, format=None, max_bytes=_MAX_IMAGE_BYTES + image_url=image_url, format=None, max_bytes=max_bytes ) except PayloadTooLargeError as e: # Named before the ValueError arm it subclasses, so the operator sees diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index a849cfbfd71..bfab8eff6de 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5591,6 +5591,55 @@ class TestBedrockGuardrailImageInput: assert sum(served) <= 5 * 1024 * 1024, f"read {sum(served)} bytes past a 4 MB cap" assert len(served) < len(chunks), "the whole body was pulled before rejecting it" + @pytest.mark.asyncio + async def test_a_request_full_of_urls_is_bounded_in_total_not_just_per_image(self): + """A per-image cap does not bound a request. + + Content items are gathered over every message and every part, so the urls + are fetched concurrently, and the 20-image limit is not applied until + _bin_pack_bedrock_content runs on items that are already resident. Without + a request-wide budget, 200 urls at 4 MB is 800 MB the caller chose. + """ + # 200 parts, each serving a 1 MB image, against a 20 x 4 MB budget. + served: list[int] = [] + one_mb: list[bytes] = [b"\0" * (1024 * 1024)] + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"{self._REMOTE_IMAGE_URL}?i={i}"}} for i in range(200) + ], + } + ] + + with patch.object(httpx.AsyncClient, "stream", new=self._fake_stream(one_mb, served)): + request = await self._guardrail(on_unscannable_image="allow").convert_to_bedrock_format( + source="INPUT", messages=messages + ) + + fetched: int = sum(served) + assert fetched <= 20 * 4 * 1024 * 1024, f"fetched {fetched} bytes for one request" + assert len(request["content"]) < 200, "every url was kept despite the budget" + assert request["content"], "the budget swallowed the whole request" + + @pytest.mark.asyncio + async def test_inline_images_do_not_draw_on_the_download_budget(self): + """Base64 arrives in the request body the proxy already accepted. + + Nothing is fetched for it, so charging it against a download quota would + refuse inline images for no reason. + """ + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": self._PNG_DATA_URI}} for _ in range(40)], + } + ] + + request = await self._guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) + + assert len(request["content"]) == 40 + @pytest.mark.asyncio async def test_apply_guardrail_scans_images_from_inputs(self): """The proxy reaches BedrockGuardrail through `apply_guardrail`, not the native hook.