scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

refactor(bedrock): capture Converse signing inputs once instead of threading six params

Browse source 
The retry path took credentials, region, endpoint, api_key and two different
header sets as loose parameters, re-assembling the signing inputs at each retry
site. That shape caused both defects found in review: `headers` was rebound to
the signed headers mid-function so the retry signed with stale auth, and
`extra_headers` was conflated with the signing headers so a caller-supplied
bearer token would have been replaced by a SigV4 signature on the sync paths.
Neither was a logic slip; both came from the same source, which is that the
inputs to a signature were passed around loose and could disagree.

`_signer` now captures them once and returns a callable that signs a body. The
same callable produces the first signature and any retry signature, so the two
cannot diverge and there is nothing left to pass incorrectly.

This also removes the four copies of the nine-line `get_request_headers` block,
deletes `_resign_without_rejected_tool_fields` entirely, and drops the
`caller_headers` rebinding workaround. The retry wrappers lose six parameters
each and shrink to the try/except they always should have been. Net effect on
the handler is 58 fewer added lines and 15 more deleted than the shape it
replaces, with the same behaviour: all four Converse paths re-verified against
live Bedrock, and the mutation checks still fail on reusing first-attempt
headers (4 cases) and on retrying unconditionally (2 cases).
This commit is contained in:
Tin Chi Lo 2026-08-03 19:38:42 -07:00
parent 438e8e2a6f
commit 9a944cc119
2 changed files with 78 additions and 149 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

187
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -101,116 +101,81 @@ class BedrockConverseLLM(BaseAWSLLM):
def __init__(self) -> None:
super().__init__()
def _resign_without_rejected_tool_fields(
def _signer(
self,
*,
request_data: Mapping[str, Any],
error_text: str,
credentials: Credentials,
aws_region_name: str,
extra_headers: Mapping[str, str] | None,
endpoint_url: str,
headers: Mapping[str, str],
extra_headers: Mapping[str, str] | None,
api_key: str | None,
) -> tuple[str, Mapping[str, str]] | None:
) -> Callable[[str], Mapping[str, str]]:
"""
Build a re-signed payload with the ``toolSpec`` members Bedrock just rejected removed.
Capture everything needed to SigV4-sign a body for this request.
SigV4 signs a hash of the body, so a retry that edits the body has to be signed
again; reusing the original headers would fail as ``SignatureDoesNotMatch`` rather
than succeed. Returns ``None`` when the error is not a rejection of extra tool
fields, which callers treat as "surface the original error".
SigV4 commits to a hash of the body, so the first attempt and any retry of an
edited body must be signed with identical inputs. Holding those inputs in one
callable is what makes signing the retry against the wrong header set impossible.
"""
retry_data = drop_bedrock_rejected_tool_fields(request_data, error_text)
if retry_data is None:
return None
data = json.dumps(retry_data)
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=data,
headers=headers,
api_key=api_key,
)
return data, prepped.headers
def sign(data: str) -> Mapping[str, str]:
return self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
data=data,
headers=headers,
api_key=api_key,
).headers
async def _asend_retrying_rejected_tool_fields(
return sign
async def _asend_with_tool_field_retry(
self,
*,
send: Callable[[str, Mapping[str, str]], Awaitable[_SendResultT]],
sign: Callable[[str], Mapping[str, str]],
request_data: Mapping[str, Any],
data: str,
headers: Mapping[str, str],
credentials: Credentials,
aws_region_name: str,
caller_headers: Mapping[str, str],
extra_headers: Mapping[str, str] | None,
endpoint_url: str,
api_key: str | None,
) -> tuple[_SendResultT, str]:
"""
Send once, and if Bedrock rejects extra ``toolSpec`` members, drop them and send again.
Send once, and if Bedrock rejects extra ``toolSpec`` members, drop them and resend.
``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.
``send`` owns the transport and the provider error contract, so anything failing
for another reason raises exactly what it raised before. Single-shot by
construction. Returns the body that actually produced the result, so callers log
what was sent rather than what was rejected.
"""
try:
return await send(data, headers), data
except (BedrockError, httpx.HTTPStatusError) as err:
retry = self._resign_without_rejected_tool_fields(
request_data=request_data,
error_text=_provider_error_text(err),
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
headers=caller_headers,
api_key=api_key,
)
if retry is None:
retried = drop_bedrock_rejected_tool_fields(request_data, _provider_error_text(err))
if retried is None:
raise
retry_data, retry_headers = retry
return await send(retry_data, retry_headers), retry_data
body = json.dumps(retried)
return await send(body, sign(body)), body
def _send_retrying_rejected_tool_fields(
def _send_with_tool_field_retry(
self,
*,
send: Callable[[str, Mapping[str, str]], _SendResultT],
sign: Callable[[str], Mapping[str, str]],
request_data: Mapping[str, Any],
data: str,
headers: Mapping[str, str],
credentials: Credentials,
aws_region_name: str,
caller_headers: Mapping[str, str],
extra_headers: Mapping[str, str] | None,
endpoint_url: str,
api_key: str | None,
) -> tuple[_SendResultT, str]:
"""Synchronous twin of ``_asend_retrying_rejected_tool_fields``."""
"""Synchronous twin of ``_asend_with_tool_field_retry``."""
try:
return send(data, headers), data
except (BedrockError, httpx.HTTPStatusError) as err:
retry = self._resign_without_rejected_tool_fields(
request_data=request_data,
error_text=_provider_error_text(err),
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
headers=caller_headers,
api_key=api_key,
)
if retry is None:
retried = drop_bedrock_rejected_tool_fields(request_data, _provider_error_text(err))
if retried is None:
raise
retry_data, retry_headers = retry
return send(retry_data, retry_headers), retry_data
body = json.dumps(retried)
return send(body, sign(body)), body
async def async_streaming(
self,
@ -242,15 +207,15 @@ class BedrockConverseLLM(BaseAWSLLM):
)
data = json.dumps(request_data)
prepped = self.get_request_headers(
sign = self._signer(
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
extra_headers=headers,
endpoint_url=api_base,
data=data,
headers=headers,
extra_headers=headers,
api_key=api_key,
)
signed_headers = sign(data)
## LOGGING
logging_obj.pre_call(
@ -259,7 +224,7 @@ class BedrockConverseLLM(BaseAWSLLM):
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": dict(prepped.headers),
"headers": signed_headers,
},
)
@ -277,17 +242,8 @@ class BedrockConverseLLM(BaseAWSLLM):
stream_chunk_size=stream_chunk_size,
)
completion_stream, data = await self._asend_retrying_rejected_tool_fields(
send=_send,
request_data=request_data,
data=data,
headers=prepped.headers,
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
caller_headers=headers,
extra_headers=headers,
endpoint_url=api_base,
api_key=api_key,
completion_stream, data = await self._asend_with_tool_field_retry(
send=_send, sign=sign, request_data=request_data, data=data, headers=signed_headers
)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,
@ -324,15 +280,15 @@ class BedrockConverseLLM(BaseAWSLLM):
)
data = json.dumps(request_data)
prepped = self.get_request_headers(
sign = self._signer(
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
extra_headers=headers,
endpoint_url=api_base,
data=data,
headers=headers,
extra_headers=headers,
api_key=api_key,
)
signed_headers = sign(data)
## LOGGING
logging_obj.pre_call(
@ -341,12 +297,10 @@ class BedrockConverseLLM(BaseAWSLLM):
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": prepped.headers,
"headers": signed_headers,
},
)
caller_headers = headers
headers = dict(prepped.headers)
if client is None or not isinstance(client, AsyncHTTPHandler):
_params = {}
if timeout is not None:
@ -372,17 +326,8 @@ class BedrockConverseLLM(BaseAWSLLM):
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
response, data = await self._asend_retrying_rejected_tool_fields(
send=_send,
request_data=request_data,
data=data,
headers=headers,
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
caller_headers=caller_headers,
extra_headers=caller_headers,
endpoint_url=api_base,
api_key=api_key,
response, data = await self._asend_with_tool_field_retry(
send=_send, sign=sign, request_data=request_data, data=data, headers=signed_headers
)
return litellm.AmazonConverseConfig()._transform_response(
@ -569,15 +514,15 @@ class BedrockConverseLLM(BaseAWSLLM):
)
data = json.dumps(_data)
prepped = self.get_request_headers(
sign = self._signer(
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
data=data,
headers=headers,
extra_headers=extra_headers,
api_key=api_key,
)
signed_headers = sign(data)
## LOGGING
logging_obj.pre_call(
@ -586,7 +531,7 @@ class BedrockConverseLLM(BaseAWSLLM):
additional_args={
"complete_input_dict": data,
"api_base": proxy_endpoint_url,
"headers": prepped.headers,
"headers": signed_headers,
},
)
if client is None or isinstance(client, AsyncHTTPHandler):
@ -615,17 +560,8 @@ class BedrockConverseLLM(BaseAWSLLM):
stream_chunk_size=stream_chunk_size,
)
completion_stream, data = self._send_retrying_rejected_tool_fields(
send=_send_stream,
request_data=_data,
data=data,
headers=prepped.headers,
credentials=credentials,
aws_region_name=aws_region_name,
caller_headers=headers,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
api_key=api_key,
completion_stream, data = self._send_with_tool_field_retry(
send=_send_stream, sign=sign, request_data=_data, data=data, headers=signed_headers
)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,
@ -653,17 +589,8 @@ class BedrockConverseLLM(BaseAWSLLM):
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
response, data = self._send_retrying_rejected_tool_fields(
send=_send,
request_data=_data,
data=data,
headers=prepped.headers,
credentials=credentials,
aws_region_name=aws_region_name,
caller_headers=headers,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
api_key=api_key,
response, data = self._send_with_tool_field_retry(
send=_send, sign=sign, request_data=_data, data=data, headers=signed_headers
)
return litellm.AmazonConverseConfig()._transform_response(

40
tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py
View file

@ -35,23 +35,25 @@ _REQUEST_DATA = {
}
def _credentials():
def _signer(headers=None, extra_headers=None):
from botocore.credentials import Credentials
return Credentials(access_key="AKIAEXAMPLE", secret_key="secret", token=None)
return BedrockConverseLLM()._signer(
credentials=Credentials(access_key="AKIAEXAMPLE", secret_key="secret", token=None),
aws_region_name="us-east-1",
endpoint_url="https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse",
headers=headers or {"Content-Type": "application/json"},
extra_headers=extra_headers,
api_key=None,
)
def _retry_kwargs():
return {
"sign": _signer(),
"request_data": _REQUEST_DATA,
"data": "original-body",
"headers": {"Authorization": "signature-over-original"},
"credentials": _credentials(),
"aws_region_name": "us-east-1",
"caller_headers": {"Content-Type": "application/json"},
"extra_headers": None,
"endpoint_url": "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse",
"api_key": None,
}
@ -126,7 +128,7 @@ def test_sync_retry_resends_without_the_rejected_field_and_resigns(raised: Excep
raise raised
return "ok"
result, sent_body = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
result, sent_body = BedrockConverseLLM()._send_with_tool_field_retry(send=send, **_retry_kwargs())
assert result == "ok"
assert len(attempts) == 2
@ -151,7 +153,7 @@ def test_sync_retry_leaves_unrelated_errors_alone() -> None:
raise BedrockError(status_code=429, message="ThrottlingException: rate exceeded")
with pytest.raises(BedrockError) as excinfo:
BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
BedrockConverseLLM()._send_with_tool_field_retry(send=send, **_retry_kwargs())
assert excinfo.value.status_code == 429
assert len(attempts) == 1
@ -166,7 +168,7 @@ def test_sync_retry_is_single_shot() -> None:
raise BedrockError(status_code=400, message=_STRICT_REJECTION)
with pytest.raises(BedrockError):
BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
BedrockConverseLLM()._send_with_tool_field_retry(send=send, **_retry_kwargs())
assert len(attempts) == 2
@ -181,7 +183,7 @@ async def test_async_retry_resends_without_the_rejected_field_and_resigns() -> N
raise BedrockError(status_code=400, message=_STRICT_REJECTION)
return "ok"
result, sent_body = await BedrockConverseLLM()._asend_retrying_rejected_tool_fields(
result, sent_body = await BedrockConverseLLM()._asend_with_tool_field_retry(
send=send, **_retry_kwargs()
)
@ -193,9 +195,9 @@ async def test_async_retry_resends_without_the_rejected_field_and_resigns() -> N
def test_retry_preserves_a_caller_supplied_authorization_header() -> None:
"""``extra_headers`` is not a duplicate of ``caller_headers``: it is the only thing
that restores a caller's non-SigV4 ``Authorization`` after signing, so the retry has
to pass it through or a proxied bearer token is silently replaced by a SigV4 one."""
"""``extra_headers`` is the only thing that restores a caller's non-SigV4
``Authorization`` after signing, so a retry signed through the same signer keeps a
proxied bearer token instead of replacing it with a SigV4 signature."""
bearer = {"Authorization": "Bearer caller-supplied-token"}
attempts: list[dict] = []
@ -205,9 +207,9 @@ def test_retry_preserves_a_caller_supplied_authorization_header() -> None:
raise BedrockError(status_code=400, message=_STRICT_REJECTION)
return "ok"
BedrockConverseLLM()._send_retrying_rejected_tool_fields(
BedrockConverseLLM()._send_with_tool_field_retry(
send=send,
**{**_retry_kwargs(), "caller_headers": {"Content-Type": "application/json", **bearer}, "extra_headers": bearer},
**{**_retry_kwargs(), "sign": _signer(headers={"Content-Type": "application/json", **bearer}, extra_headers=bearer)},
)
assert attempts[1]["Authorization"] == "Bearer caller-supplied-token"
@ -219,7 +221,7 @@ def test_reported_body_is_the_original_when_no_retry_happens() -> None:
def send(body: str, headers) -> str:
return "ok"
_, sent_body = BedrockConverseLLM()._send_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
_, sent_body = BedrockConverseLLM()._send_with_tool_field_retry(send=send, **_retry_kwargs())
assert sent_body == "original-body"
@ -232,7 +234,7 @@ async def test_async_retry_leaves_unrelated_errors_alone() -> None:
raise BedrockError(status_code=500, message="InternalServerException")
with pytest.raises(BedrockError) as excinfo:
await BedrockConverseLLM()._asend_retrying_rejected_tool_fields(send=send, **_retry_kwargs())
await BedrockConverseLLM()._asend_with_tool_field_retry(send=send, **_retry_kwargs())
assert excinfo.value.status_code == 500
assert len(attempts) == 1