fix(bedrock): report the retried body, not the rejected one, after a successful retry

The retry wrappers returned only the response, so the two non-streaming Converse
paths kept passing their original `data` into `_transform_response`. A request
that succeeded on the second attempt therefore logged the payload Bedrock had
just rejected, claiming a field was sent that was in fact removed. That is
exactly backwards for the debugging this feature exists to support, and it
diverges from the established pattern in `llm_http_handler.py`, where the retry
loop rebinds `data` so everything downstream sees the corrected body.

The wrappers now return the body that actually produced the response alongside
the result, and every call site binds it. The streaming paths were already
correct, since they pass the retried body through to `make_call` themselves;
this makes all four consistent.
This commit is contained in:
Tin Chi Lo 2026-08-03 19:12:55 -07:00
parent 37c8efde0f
commit 5ee1b09bd5
2 changed files with 32 additions and 12 deletions

View file

@ -149,16 +149,19 @@ class BedrockConverseLLM(BaseAWSLLM):
caller_headers: Mapping[str, str],
endpoint_url: str,
api_key: str | None,
) -> _SendResultT:
) -> tuple[_SendResultT, str]:
"""
Send once, and if Bedrock rejects extra ``toolSpec`` members, drop them and send again.
``send`` owns the transport and the provider-error contract, so a request that
fails for any other reason raises exactly what it raised before. The retry is
single-shot: a second rejection surfaces rather than looping.
Returns the result together with the body that actually produced it, so callers
log and transform against what was sent rather than the payload that was rejected.
"""
try:
return await send(data, headers)
return await send(data, headers), data
except (BedrockError, httpx.HTTPStatusError) as err:
retry = self._resign_without_rejected_tool_fields(
request_data=request_data,
@ -172,7 +175,8 @@ class BedrockConverseLLM(BaseAWSLLM):
)
if retry is None:
raise
return await send(*retry)
retry_data, retry_headers = retry
return await send(retry_data, retry_headers), retry_data
def _send_retrying_rejected_tool_fields(
self,
@ -186,10 +190,10 @@ class BedrockConverseLLM(BaseAWSLLM):
caller_headers: Mapping[str, str],
endpoint_url: str,
api_key: str | None,
) -> _SendResultT:
) -> tuple[_SendResultT, str]:
"""Synchronous twin of ``_asend_retrying_rejected_tool_fields``."""
try:
return send(data, headers)
return send(data, headers), data
except (BedrockError, httpx.HTTPStatusError) as err:
retry = self._resign_without_rejected_tool_fields(
request_data=request_data,
@ -203,7 +207,8 @@ class BedrockConverseLLM(BaseAWSLLM):
)
if retry is None:
raise
return send(*retry)
retry_data, retry_headers = retry
return send(retry_data, retry_headers), retry_data
async def async_streaming(
self,
@ -270,7 +275,7 @@ class BedrockConverseLLM(BaseAWSLLM):
stream_chunk_size=stream_chunk_size,
)
completion_stream = await self._asend_retrying_rejected_tool_fields(
completion_stream, data = await self._asend_retrying_rejected_tool_fields(
send=_send,
request_data=request_data,
data=data,
@ -364,7 +369,7 @@ class BedrockConverseLLM(BaseAWSLLM):
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
response = await self._asend_retrying_rejected_tool_fields(
response, data = await self._asend_retrying_rejected_tool_fields(
send=_send,
request_data=request_data,
data=data,
@ -606,7 +611,7 @@ class BedrockConverseLLM(BaseAWSLLM):
stream_chunk_size=stream_chunk_size,
)
completion_stream = self._send_retrying_rejected_tool_fields(
completion_stream, data = self._send_retrying_rejected_tool_fields(
send=_send_stream,
request_data=_data,
data=data,
@ -643,7 +648,7 @@ class BedrockConverseLLM(BaseAWSLLM):
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
response = self._send_retrying_rejected_tool_fields(
response, data = self._send_retrying_rejected_tool_fields(
send=_send,
request_data=_data,
data=data,

View file

@ -125,10 +125,12 @@ def test_sync_retry_resends_without_the_rejected_field_and_resigns(raised: Excep
raise raised
return "ok"
result = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
result, sent_body = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
assert result == "ok"
assert len(attempts) == 2
assert sent_body == attempts[1][0]
assert '"strict"' not in sent_body
first_body, first_headers = attempts[0]
retry_body, retry_headers = attempts[1]
@ -178,12 +180,25 @@ async def test_async_retry_resends_without_the_rejected_field_and_resigns() -> N
raise BedrockError(status_code=400, message=_STRICT_REJECTION)
return "ok"
result = await BedrockConverseLLM()._asend_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
result, sent_body = await BedrockConverseLLM()._asend_retrying_rejected_tool_fields(
send=send, **_retry_kwargs()
)
assert result == "ok"
assert len(attempts) == 2
assert '"strict"' not in attempts[1][0]
assert attempts[1][1]["Authorization"].startswith("AWS4-HMAC-SHA256")
assert sent_body == attempts[1][0]
def test_reported_body_is_the_original_when_no_retry_happens() -> None:
"""A request that succeeds first time reports exactly what it sent."""
def send(body: str, headers) -> str:
return "ok"
_, sent_body = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
assert sent_body == "original-body"
@pytest.mark.asyncio