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

fix(bedrock): sign the retry with the same extra_headers as the first attempt

Browse source 
`extra_headers` and `headers` are not two names for one thing. `headers` is the
signing basis; `extra_headers` is consulted only to restore a caller-supplied
non-SigV4 `Authorization` after signing, so SigV4 does not overwrite a proxied
bearer token. The retry wrappers conflated them and passed the signing headers
for both, which happened to be right on the async paths (where the original
signing call passes `headers` for both) but wrong on the two sync paths, whose
original signing passes the caller's `extra_headers`.

Thread `extra_headers` through as its own argument so every retry signs exactly
the way its first attempt did, and pin the behaviour that motivated the split:
a caller-supplied bearer token must survive the retry rather than be replaced by
a SigV4 signature.

Also flattens the azure_ai field-collection to a plain comprehension. The
previous `frozenset().union(*values or (frozenset(),))` needed the guard only to
survive an empty mapping, which a flat comprehension handles without the special
case.
This commit is contained in:
Tin Chi Lo 2026-08-03 19:33:12 -07:00
parent 5ee1b09bd5
commit 438e8e2a6f
3 changed files with 32 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/llms/azure_ai/chat/transformation.py
View file

@ -294,7 +294,8 @@ class AzureAIStudioConfig(OpenAIConfig):
def _drop_tool_level_extra_fields(self, request_data: dict, error_text: str) -> dict:
from litellm.llms.base_llm.base_utils import parse_rejected_tool_fields
fields_to_drop = frozenset().union(*parse_rejected_tool_fields(error_text).values() or (frozenset(),))
rejected = parse_rejected_tool_fields(error_text)
fields_to_drop = frozenset(field for fields in rejected.values() for field in fields)
tools = request_data.get("tools")
if fields_to_drop and isinstance(tools, list):
for tool in tools:

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

@ -147,6 +147,7 @@ class BedrockConverseLLM(BaseAWSLLM):
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]:
@ -168,7 +169,7 @@ class BedrockConverseLLM(BaseAWSLLM):
error_text=_provider_error_text(err),
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=caller_headers,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
headers=caller_headers,
api_key=api_key,
@ -188,6 +189,7 @@ class BedrockConverseLLM(BaseAWSLLM):
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]:
@ -200,7 +202,7 @@ class BedrockConverseLLM(BaseAWSLLM):
error_text=_provider_error_text(err),
credentials=credentials,
aws_region_name=aws_region_name,
extra_headers=caller_headers,
extra_headers=extra_headers,
endpoint_url=endpoint_url,
headers=caller_headers,
api_key=api_key,
@ -283,6 +285,7 @@ class BedrockConverseLLM(BaseAWSLLM):
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,
)
@ -377,6 +380,7 @@ class BedrockConverseLLM(BaseAWSLLM):
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,
)
@ -619,6 +623,7 @@ class BedrockConverseLLM(BaseAWSLLM):
credentials=credentials,
aws_region_name=aws_region_name,
caller_headers=headers,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
api_key=api_key,
)
@ -656,6 +661,7 @@ class BedrockConverseLLM(BaseAWSLLM):
credentials=credentials,
aws_region_name=aws_region_name,
caller_headers=headers,
extra_headers=extra_headers,
endpoint_url=proxy_endpoint_url,
api_key=api_key,
)

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

@ -49,6 +49,7 @@ def _retry_kwargs():
"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,
}
@ -191,6 +192,27 @@ async def test_async_retry_resends_without_the_rejected_field_and_resigns() -> N
assert sent_body == attempts[1][0]
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."""
bearer = {"Authorization": "Bearer caller-supplied-token"}
attempts: list[dict] = []
def send(body: str, headers) -> str:
attempts.append(dict(headers))
if len(attempts) == 1:
raise BedrockError(status_code=400, message=_STRICT_REJECTION)
return "ok"
BedrockConverseLLM()._send_retrying_rejected_tool_fields(
send=send,
**{**_retry_kwargs(), "caller_headers": {"Content-Type": "application/json", **bearer}, "extra_headers": bearer},
)
assert attempts[1]["Authorization"] == "Bearer caller-supplied-token"
def test_reported_body_is_the_original_when_no_retry_happens() -> None:
"""A request that succeeds first time reports exactly what it sent."""