fix(guardrails): never split a Bedrock text into an empty fragment

_nearest_whitespace_split_index could return len(text) when the only space at
or after the midpoint was the final character, so the first fragment came back
identical to the text AWS had just rejected as too large and the second came
back empty. AWS rejects the unchanged fragment again, and each retry re-splits
it into the same fragment, so an oversized single item shaped like a long
unbroken token with one trailing space exhausted the stack with a
RecursionError instead of scanning or surfacing Bedrock's error.

Candidate boundaries that would leave either side empty are now discarded, and
the raw midpoint is used when none remain. The midpoint is always safe because
_split_bedrock_content only calls this for text of at least two characters.
This commit is contained in:
spencer-burridge 2026-07-29 15:29:20 -05:00
parent ffe529c04c
commit d976631260
2 changed files with 77 additions and 14 deletions

View file

@ -1284,22 +1284,31 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
@staticmethod
def _nearest_whitespace_split_index(text: str) -> int:
"""Return the index nearest `text`'s midpoint that falls on a
whitespace boundary, so splitting `text[:i]` / `text[i:]` there never
severs a word. Falls back to the raw midpoint when `text` has no
whitespace at all (a single giant token) -- still a correct, lossless
split, just no longer guaranteed word-safe for that pathological case.
"""Return the index nearest `text`'s midpoint that falls on a space
boundary, so splitting `text[:i]` / `text[i:]` there never severs a word.
The returned index always leaves both sides non-empty, which is what makes
the caller's recursion terminate. A boundary that would put the split at 0
or at ``len(text)`` is discarded: it would hand back a fragment identical to
the text just rejected as too large, AWS would reject that again, and each
retry would re-split it into the same unchanged fragment until the stack ran
out. The dangerous shape is a text whose only space at or after the midpoint
is its final character.
Falls back to the raw midpoint when no usable space boundary exists, either
because `text` has none at all (a single giant token) or because the only
candidates were degenerate. That is still a correct, lossless split, just no
longer guaranteed word-safe for those cases. `text` must be at least two
characters, which `_split_bedrock_content` guarantees, so the midpoint itself
is never degenerate.
"""
midpoint = len(text) // 2
left = text.rfind(" ", 0, midpoint)
right = text.find(" ", midpoint)
if left == -1 and right == -1:
return midpoint
if left == -1:
return right + 1
if right == -1:
return left + 1
return left + 1 if midpoint - left <= right - midpoint else right + 1
boundaries = (text.rfind(" ", 0, midpoint), text.find(" ", midpoint))
candidates = sorted(
(found + 1 for found in boundaries if found != -1),
key=lambda split: abs(split - midpoint),
)
return next((split for split in candidates if 0 < split < len(text)), midpoint)
@staticmethod
def _is_input_too_large_error(detail: object) -> bool:

View file

@ -4376,3 +4376,57 @@ async def test_configured_chunk_budget_changes_how_content_is_packed():
assert await _calls_made_with_budget(_BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS) == 4
assert await _calls_made_with_budget(100_000) == 1
def test_split_index_never_produces_an_empty_fragment():
"""Both fragments must be non-empty for every splittable text, so bisection always
makes progress.
A text whose only qualifying whitespace is its final character is the dangerous
shape: taking that boundary puts the split at len(text), leaving the first fragment
identical to the input that was just rejected and the second empty. The recursion
would then resubmit the unchanged fragment forever and exhaust the stack instead of
scanning or surfacing Bedrock's error."""
for text in ("ab ", "xxxx ", ("x" * 40) + " ", " ab", "a b", "ab", " "):
split_at = BedrockGuardrail._nearest_whitespace_split_index(text)
assert 0 < split_at < len(text), f"degenerate split {split_at} for {text!r}"
assert text[:split_at] and text[split_at:], f"empty fragment for {text!r}"
assert text[:split_at] + text[split_at:] == text
@pytest.mark.asyncio
async def test_oversized_single_item_with_trailing_space_gives_up_instead_of_recursing():
"""An oversized single item whose only space is trailing must bottom out and
surface Bedrock's error, not recurse forever.
AWS is modelled the way it really behaves, rejecting every attempt, because the
danger is a fragment identical to the input that was just rejected: AWS would
reject it again, and each retry would split it into the same unchanged fragment.
A split that always shrinks the text terminates and re-raises; one that can return
the whole text raises RecursionError instead. The call-count bound is generous:
halving 41 characters down to unsplittable is a handful of attempts, nowhere near
a stack limit."""
guardrail = _bedrock_guardrail_for_chunk_tests()
messages = [{"role": "user", "content": ("x" * 40) + " "}]
mock_credentials = MagicMock()
mock_credentials.access_key = "k"
mock_credentials.secret_key = "s"
mock_credentials.token = None
with (
patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post,
patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")),
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
):
mock_post.side_effect = lambda *_a, **_k: _too_large_validation_httpx_response()
with pytest.raises(HTTPException) as excinfo:
await guardrail.make_bedrock_api_request(
source="INPUT",
messages=messages,
request_data={"model": "bedrock-nova-micro"},
)
assert excinfo.value.status_code == 400
assert mock_post.await_count < 200