From 9a944cc11921bc9efdec57a4d2d78c4f69c06671 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 3 Aug 2026 19:38:42 -0700 Subject: [PATCH] refactor(bedrock): capture Converse signing inputs once instead of threading six params 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). --- litellm/llms/bedrock/chat/converse_handler.py | 187 ++++++------------ ...test_converse_rejected_tool_field_retry.py | 40 ++-- 2 files changed, 78 insertions(+), 149 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7fa77b1dba1..69647f28088 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -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( diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py b/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py index 9505650134d..b978cf2647f 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_rejected_tool_field_retry.py @@ -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