From 89c3a8216be6c42bc9d327e312c05d62c99118b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:54:43 -0700 Subject: [PATCH 1/6] fix(bedrock): sign requests off the event loop on every async path SigV4 signing resolves AWS credentials, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call. Every async Bedrock path that still signed on the event loop (/v1/messages, Converse, count tokens, the agent-runtime and Comprehend Medical pass-throughs, async-invoke status polling, realtime, AgentCore, SQS, S3) now signs on a worker thread, so one Bedrock request no longer stalls the whole worker. Fixes #40165 --- .../providers/bedrock_agentcore/handler.py | 7 +- litellm/integrations/s3_v2.py | 4 +- litellm/integrations/sqs.py | 2 +- litellm/llms/bedrock/base_aws_llm.py | 19 ++++- litellm/llms/bedrock/chat/converse_handler.py | 7 +- litellm/llms/bedrock/count_tokens/handler.py | 9 ++- litellm/llms/bedrock/embed/embedding.py | 48 ++++++------ litellm/llms/bedrock/realtime/handler.py | 5 +- litellm/llms/custom_httpx/llm_http_handler.py | 25 ++++-- .../llm_passthrough_endpoints.py | 78 +++++++++++-------- .../chat/test_bedrock_converse_handler.py | 25 +++++- .../test_bedrock_count_tokens_handler.py | 50 ++++++++++++ .../llms/bedrock/event_loop_probe.py | 52 +++++++++++++ .../llms/bedrock/test_base_aws_llm.py | 24 ++++++ .../custom_httpx/test_llm_http_handler.py | 60 ++++++++++++++ .../test_llm_pass_through_endpoints.py | 27 ++++++- 16 files changed, 363 insertions(+), 79 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py create mode 100644 tests/test_litellm/llms/bedrock/event_loop_probe.py diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index a4e6fa50901..e01e215b5f0 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -5,6 +5,7 @@ Sends JSON-RPC envelopes directly to AgentCore endpoints, bypassing the completion bridge that would otherwise strip the envelope. """ +import asyncio import json from collections.abc import AsyncIterator, Mapping from typing import Any, Final @@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await asyncio.to_thread( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, @@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + url, headers, body = await asyncio.to_thread( + BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, litellm_params=litellm_params, diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 712ce41d09e..e3297cfd27c 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) + await asyncio.to_thread(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) @@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) - S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) + await asyncio.to_thread(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 2b4c8c9928d..31cc20f5995 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth(aws_request) + await asyncio.to_thread(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index a197702921f..bbd4988b204 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,3 +1,4 @@ +import asyncio import base64 import hashlib import json @@ -7,7 +8,7 @@ import urllib.parse from collections.abc import Callable, Mapping from datetime import datetime from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload import httpx from pydantic import BaseModel, ValidationError @@ -1668,3 +1669,19 @@ class BaseAWSLLM: request_headers_dict["Authorization"] = incoming_authorization return request_headers_dict, request.body + + +_SignParams = ParamSpec("_SignParams") +_SignedRequest = TypeVar("_SignedRequest") + + +async def sign_request_off_loop_if_aws( + provider_config: object, + sign_request: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature +) -> _SignedRequest: + if isinstance(provider_config, BaseAWSLLM): + return await asyncio.to_thread(sign_request, *args, **kwargs) + return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 984ba371898..3e8e1226a55 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,3 +1,4 @@ +import asyncio import json from collections.abc import Mapping from types import MappingProxyType @@ -136,7 +137,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await asyncio.to_thread( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, @@ -206,7 +208,8 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = self.get_request_headers( + prepped: Final = await asyncio.to_thread( + self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", extra_headers=headers, diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 1fb53f6ff0a..cae2315e744 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -4,6 +4,7 @@ AWS Bedrock CountTokens API handler. Simplified handler leveraging existing LiteLLM Bedrock infrastructure. """ +import asyncio from typing import Any, Final import httpx @@ -12,7 +13,7 @@ import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client class BedrockCountTokensHandler(BedrockCountTokensConfig): @@ -27,6 +28,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): request_data: dict[str, Any], litellm_params: dict[str, Any], resolved_model: str, + client: AsyncHTTPHandler | None = None, ) -> dict[str, Any]: """ Handle a CountTokens request using existing LiteLLM patterns. @@ -75,7 +77,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Extract api_key for bearer token auth if provided api_key: Final = litellm_params.get("api_key", None) headers: Final = {"Content-Type": "application/json"} - signed_headers, signed_body = self._sign_request( + signed_headers, signed_body = await asyncio.to_thread( + self._sign_request, service_name="bedrock", headers=headers, optional_params=litellm_params, @@ -85,7 +88,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): api_key=api_key, ) - async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) + async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) response: Final = await async_client.post( endpoint_url, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index d3725434498..e5f94e51e6d 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -2,10 +2,11 @@ Handles embedding calls to Bedrock's `/invoke` endpoint """ +import asyncio import copy import json import urllib.parse -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -26,7 +27,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -41,6 +42,20 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +def _sign_get_request( + credentials: Credentials, url: str, headers: Mapping[str, str], aws_region_name: str +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + request: Final = AWSRequest(method="GET", url=url, data=None, headers=dict(headers)) + SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request) + return request.prepare() + + class BedrockEmbedding(BaseAWSLLM): @overload def _load_credentials( @@ -599,9 +614,6 @@ class BedrockEmbedding(BaseAWSLLM): dict: Status response from AWS Bedrock """ - # Get AWS credentials using the same method as other Bedrock methods - credentials, _ = self._load_credentials(kwargs) - # Get the runtime endpoint endpoint_url, _ = self.get_runtime_endpoint( api_base=None, @@ -618,27 +630,13 @@ class BedrockEmbedding(BaseAWSLLM): # Prepare headers for GET request headers: Final = {"Content-Type": "application/json"} - # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + def sign_status_request() -> AWSPreparedRequest: + credentials, _ = self._load_credentials(kwargs) + return _sign_get_request( + credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name + ) - # Create AWSRequest with GET method and encoded URL - request: Final = AWSRequest( - method="GET", - url=status_url, - data=None, # GET request, no body - headers=headers, - ) - - # Sign the request - SigV4Auth will create canonical string from request URL - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - sigv4.add_auth(request) - - # Prepare the request - prepped: Final = request.prepare() + prepped: Final = await asyncio.to_thread(sign_status_request) # LOGGING if logging_obj is not None: diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 42fe8941443..792e3169a7a 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -149,7 +149,8 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) - credentials: Final = self.get_credentials( + credentials: Final = await asyncio.to_thread( + self.get_credentials, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -169,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM): "or configure credentials in the environment" ), ) - frozen_credentials: Final = credentials.get_frozen_credentials() + frozen_credentials: Final = await asyncio.to_thread(credentials.get_frozen_credentials) # Initialize Bedrock client with aws_sdk_bedrock_runtime config: Final = Config( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..48dc009d116 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,6 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -1961,7 +1962,9 @@ class BaseLLMHTTPHandler: api_key=api_key, ) - signed_headers, signed_json_body = provider_config.sign_request( + signed_headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params, request_data=data, @@ -2062,7 +2065,9 @@ class BaseLLMHTTPHandler: max_attempts, ) provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body) - headers, signed_json_body = provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + provider_config, + provider_config.sign_request, headers=headers, optional_params=optional_params_dict, request_data=request_body, @@ -2222,7 +2227,9 @@ class BaseLLMHTTPHandler: stream=stream, ) - headers, signed_json_body = anthropic_messages_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + anthropic_messages_provider_config, + anthropic_messages_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), # dynamic aws_* params are passed under litellm_params request_data=request_body, @@ -2898,7 +2905,9 @@ class BaseLLMHTTPHandler: fake_stream=fake_stream, ) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -4606,7 +4615,9 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) - headers, signed_body = responses_api_provider_config.sign_request( + headers, signed_body = await sign_request_off_loop_if_aws( + responses_api_provider_config, + responses_api_provider_config.sign_request, headers=headers, optional_params=dict(litellm_params), request_data=data, @@ -9833,7 +9844,9 @@ class BaseLLMHTTPHandler: ) all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) - headers, signed_json_body = vector_store_provider_config.sign_request( + headers, signed_json_body = await sign_request_off_loop_if_aws( + vector_store_provider_config, + vector_store_provider_config.sign_request, headers=headers, optional_params=all_optional_params, request_data=request_body, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2ea46b740a8..6eb7b5d8541 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -8,6 +8,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. from __future__ import annotations +import asyncio import hmac import inspect import json @@ -15,6 +16,7 @@ import os import re from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast @@ -84,6 +86,9 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: + from botocore.awsrequest import AWSPreparedRequest + from botocore.credentials import Credentials + from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router @@ -1099,13 +1104,6 @@ async def bedrock_proxy_route( """ create_request_copy(request) - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME") if not _is_bedrock_agent_runtime_route(endpoint=endpoint): return await bedrock_llm_proxy_route( @@ -1139,17 +1137,20 @@ async def bedrock_proxy_route( from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() - credentials: Final[Credentials] = bedrock_llm.get_credentials() - sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name) - headers: Final = {"Content-Type": "application/json"} # Assuming the body contains JSON data, parse it try: data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - _request: Final = AWSRequest(method="POST", url=str(updated_url), data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await asyncio.to_thread( + _sign_aws_json_post, + get_credentials=bedrock_llm.get_credentials, + service_name="bedrock", + aws_region_name=aws_region_name, + url=str(updated_url), + body=json.dumps(data), + headers=MappingProxyType({"Content-Type": "application/json"}), + ) ## check for streaming is_streaming_request = False @@ -1177,6 +1178,25 @@ async def bedrock_proxy_route( return received_value +def _sign_aws_json_post( + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") + + aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=dict(headers)) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + COMPREHEND_MEDICAL_TARGET_PREFIX: Final = "ComprehendMedical_20181030" @@ -1207,13 +1227,6 @@ async def comprehend_medical_proxy_route( [Docs](https://docs.litellm.ai/docs/pass_through/comprehend_medical) """ - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - from botocore.credentials import Credentials - except ImportError: - raise ImportError("Missing boto3 to call comprehendmedical. Run 'pip install boto3'.") - from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import ( COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS, ) @@ -1246,18 +1259,21 @@ async def comprehend_medical_proxy_route( from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - credentials: Final[Credentials] = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name) - sigv4: Final = SigV4Auth(credentials, "comprehendmedical", aws_region_name) - headers: Final = MappingProxyType( - { - "Content-Type": "application/x-amz-json-1.1", - "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", - } - ) target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" - _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers) - sigv4.add_auth(_request) - prepped: Final = _request.prepare() + prepped: Final = await asyncio.to_thread( + _sign_aws_json_post, + get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + service_name="comprehendmedical", + aws_region_name=aws_region_name, + url=target_url, + body=json.dumps(data), + headers=MappingProxyType( + { + "Content-Type": "application/x-amz-json-1.1", + "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", + } + ), + ) endpoint_func: Final = create_pass_through_route( endpoint=operation, diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index f34b8eb1fb9..8831fd9061c 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -6,6 +6,7 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS. from __future__ import annotations +import asyncio from unittest.mock import MagicMock, patch import httpx @@ -16,6 +17,7 @@ from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import chat_completions as bridge from litellm.types.utils import ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe RUST_RESPONSE = { "created": 1_700_000_000, @@ -308,7 +310,9 @@ CONVERSE_RESPONSE = { } -async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): +async def _drive_async_completion( + *, skip_pre_call_logging: bool, logging_obj, credentials: Credentials = RESOLVED_CREDENTIALS +): """Run the real `async_completion` with a stubbed transport.""" import httpx as _httpx @@ -335,7 +339,7 @@ async def _drive_async_completion(*, skip_pre_call_logging: bool, logging_obj): stream=None, optional_params={"maxTokens": 16}, litellm_params={"aws_region_name": "us-west-2"}, - credentials=RESOLVED_CREDENTIALS, + credentials=credentials, headers={}, client=client, skip_pre_call_logging=skip_pre_call_logging, @@ -357,6 +361,23 @@ async def test_async_completion_logs_pre_call_by_default(): assert logging_obj.pre_call.call_count == 1 +@pytest.mark.asyncio +async def test_async_completion_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: botocore refreshes expiring credentials inside SigV4 signing with a + blocking HTTP call, so `async_completion` must sign on a worker thread to keep the loop serving.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + response = await _drive_async_completion( + skip_pre_call_logging=False, logging_obj=MagicMock(), credentials=probe.credentials() + ) + await release + + assert response.choices[0].message.content == "hi" + assert probe.served_during_refresh is True + + def _sync_client_returning_converse_response(): client = MagicMock() client.post.side_effect = lambda **_kwargs: httpx.Response( diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py new file mode 100644 index 00000000000..31e65bb6e27 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -0,0 +1,50 @@ +import asyncio +from unittest.mock import AsyncMock + +import httpx +import pytest + +from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe + + +class _ProbedCountTokensHandler(BedrockCountTokensHandler): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials(self, **kwargs): + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_handle_count_tokens_request_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the count_tokens handler signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = AsyncMock(spec=AsyncHTTPHandler) + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json={"inputTokens": 7}, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + result = await _ProbedCountTokensHandler(probe).handle_count_tokens_request( + request_data={ + "model": "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": "hi"}], + }, + litellm_params={"aws_region_name": "us-west-2"}, + resolved_model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + client=client, + ) + await release + + assert result == {"input_tokens": 7} + assert client.post.call_args.kwargs["headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/test_litellm/llms/bedrock/event_loop_probe.py new file mode 100644 index 00000000000..68acce308b8 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/event_loop_probe.py @@ -0,0 +1,52 @@ +"""Refreshable credentials whose refresh only completes while the event loop keeps serving.""" + +from __future__ import annotations + +import asyncio +import threading +from datetime import datetime, timedelta, timezone +from typing import Final + +from botocore.credentials import RefreshableCredentials + +REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 + + +class EventLoopProbe: + """Blocks inside botocore's credential refresh until a coroutine on the loop releases it. + + Signing on the event loop thread can never be released, so `served_during_refresh` reads False there + and True only when the refresh ran on another thread while the loop stayed responsive. + """ + + def __init__(self) -> None: + self.refresh_started: Final = threading.Event() + self.loop_served: Final = threading.Event() + self.served_during_refresh: bool | None = None + + def refresh(self) -> dict[str, str | None]: + self.refresh_started.set() + served: Final = self.loop_served.wait(timeout=REFRESH_RELEASE_TIMEOUT_SECONDS) + if self.served_during_refresh is None: + self.served_during_refresh = served + return { + "access_key": "AKIAREFRESHED", + "secret_key": "refreshed-secret", + "token": None, + "expiry_time": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(), + } + + def credentials(self) -> RefreshableCredentials: + return RefreshableCredentials( + access_key="AKIASTALE", + secret_key="stale-secret", + token=None, + expiry_time=datetime.now(timezone.utc) + timedelta(seconds=60), + refresh_using=self.refresh, + method="event-loop-probe", + ) + + async def release_refresh_from_the_loop(self) -> None: + while not self.refresh_started.is_set(): + await asyncio.sleep(0.005) + self.loop_served.set() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index f854d806bdc..b3fc8dbe959 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,3 +1,4 @@ +import asyncio import json import os import threading @@ -22,7 +23,9 @@ from litellm.llms.bedrock.base_aws_llm import ( AwsAuthError, BaseAWSLLM, Boto3CredentialsInfo, + sign_request_off_loop_if_aws, ) +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Global variable for the base_aws_llm.py file path @@ -3215,3 +3218,24 @@ class TestGetRequestHeadersResign: extra_headers={"Authorization": "Bearer foo"}, ) assert prepped.headers["Authorization"] == "Bearer foo" + + +@pytest.mark.asyncio +async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh(): + """Regression for issue #40165: an AWS provider's signing (and the botocore credential refresh + inside it) must run off the event loop, so other requests keep being served meanwhile.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-runtime.us-west-2.amazonaws.com/", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-west-2").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws(BaseAWSLLM(), sign, headers={"Content-Type": "application/json"}) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index e16855da8cb..6799850c4ae 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -29,10 +29,14 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" @@ -749,6 +753,62 @@ async def test_anthropic_messages_streaming_response_aclose_closes_agentic_upstr assert tracker.closed is True +class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig): + def __init__(self, probe: EventLoopProbe) -> None: + super().__init__() + self._probe = probe + + def get_credentials(self, **kwargs): + return self._probe.credentials() + + +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_signs_bedrock_off_the_event_loop(monkeypatch): + """Regression for issue #40165: /v1/messages on Bedrock signed on the loop, so botocore's blocking + credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + handler = BaseLLMHTTPHandler() + upstream_response = httpx.Response( + 200, + json={ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": "claude-haiku-4-5-20251001", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com/"), + ) + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=upstream_response) + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.dynamic_success_callbacks = None + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + + await handler.async_anthropic_messages_handler( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=_ProbedBedrockMessagesConfig(probe), + anthropic_messages_optional_request_params={"max_tokens": 16}, + custom_llm_provider="bedrock", + litellm_params=GenericLiteLLMParams(aws_region_name="us-west-2"), + logging_obj=mock_logging_obj, + client=mock_client, + stream=False, + kwargs={}, + ) + await release + + sent_headers = mock_client.post.call_args.kwargs["headers"] + assert sent_headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + + @pytest.mark.asyncio async def test_async_anthropic_messages_handler_passes_litellm_metadata(): """Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs. diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index d9969dd1dc9..a1cfd973b75 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1,3 +1,4 @@ +import asyncio import base64 import contextlib import json @@ -19,6 +20,7 @@ from starlette.datastructures import FormData import litellm from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( BaseOpenAIPassThroughHandler, @@ -1852,11 +1854,11 @@ class TestBedrockAgentRuntimePassthroughToggle: return request @contextlib.contextmanager - def _patched_dispatch(self, general_settings: Mapping[str, object]): + def _patched_dispatch(self, general_settings: Mapping[str, object], credentials: object | None = None): from botocore.credentials import Credentials bedrock_llm: Final = Mock() - bedrock_llm.get_credentials = Mock(return_value=Credentials("ak", "sk")) + bedrock_llm.get_credentials = Mock(return_value=credentials or Credentials("ak", "sk")) forwarder: Final = AsyncMock(return_value="forwarded") with ( @@ -1891,6 +1893,27 @@ class TestBedrockAgentRuntimePassthroughToggle: forwarder.assert_awaited_once() assert "bedrock-agent-runtime.us-east-1.amazonaws.com" in create_route.call_args.kwargs["target"] + @pytest.mark.asyncio + async def test_agent_runtime_dispatch_signs_off_the_event_loop(self, monkeypatch): + """Regression for issue #40165: the agent-runtime pass-through signed on the loop, so botocore's + blocking credential refresh inside SigV4 stalled every other request on the worker.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe: Final = EventLoopProbe() + release: Final = asyncio.create_task(probe.release_refresh_from_the_loop()) + + with self._patched_dispatch(MappingProxyType({}), credentials=probe.credentials()) as (create_route, forwarder): + result: Final = await bedrock_proxy_route( + endpoint=self.AGENT_RUNTIME_ENDPOINT, + request=self._mock_request(), + fastapi_response=Mock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + await release + + assert result == "forwarded" + assert create_route.call_args.kwargs["custom_headers"]["Authorization"].startswith("AWS4-HMAC-SHA256") + assert probe.served_during_refresh is True + @pytest.mark.asyncio @pytest.mark.parametrize("value", (True, "true", "True")) async def test_agent_runtime_dispatch_rejected_when_disabled(self, value: bool | str): From 71a4a6e912624c31fde159fad5f6508e6b2a217a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:20:50 -0700 Subject: [PATCH 2/6] refactor: keep the AWS JSON signing helper under llms and bound the probe's refresh wait --- litellm/llms/bedrock/base_aws_llm.py | 19 ++++++++++++ .../llm_passthrough_endpoints.py | 29 +++---------------- .../llms/bedrock/event_loop_probe.py | 5 ++++ 3 files changed, 28 insertions(+), 25 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index bbd4988b204..1c7268ffa22 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1671,6 +1671,25 @@ class BaseAWSLLM: return request_headers_dict, request.body +def sign_aws_json_post( + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") + + aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + _SignParams = ParamSpec("_SignParams") _SignedRequest = TypeVar("_SignedRequest") diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6085bc2bfce..037f663044f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -86,9 +86,6 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: - from botocore.awsrequest import AWSPreparedRequest - from botocore.credentials import Credentials - from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router @@ -1134,6 +1131,7 @@ async def bedrock_proxy_route( ) # Add or update query parameters + from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() @@ -1143,7 +1141,7 @@ async def bedrock_proxy_route( except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) prepped: Final = await asyncio.to_thread( - _sign_aws_json_post, + sign_aws_json_post, get_credentials=bedrock_llm.get_credentials, service_name="bedrock", aws_region_name=aws_region_name, @@ -1178,25 +1176,6 @@ async def bedrock_proxy_route( return received_value -def _sign_aws_json_post( - get_credentials: Callable[[], Credentials], - service_name: str, - aws_region_name: str | None, - url: str, - body: str, - headers: Mapping[str, str], -) -> AWSPreparedRequest: - try: - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - except ImportError: - raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.") - - aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) - SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) - return aws_request.prepare() - - COMPREHEND_MEDICAL_TARGET_PREFIX: Final = "ComprehendMedical_20181030" @@ -1257,11 +1236,11 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" prepped: Final = await asyncio.to_thread( - _sign_aws_json_post, + sign_aws_json_post, get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), service_name="comprehendmedical", aws_region_name=aws_region_name, diff --git a/tests/test_litellm/llms/bedrock/event_loop_probe.py b/tests/test_litellm/llms/bedrock/event_loop_probe.py index 68acce308b8..c347247ec32 100644 --- a/tests/test_litellm/llms/bedrock/event_loop_probe.py +++ b/tests/test_litellm/llms/bedrock/event_loop_probe.py @@ -4,12 +4,14 @@ from __future__ import annotations import asyncio import threading +import time from datetime import datetime, timedelta, timezone from typing import Final from botocore.credentials import RefreshableCredentials REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0 +REFRESH_START_TIMEOUT_SECONDS: Final = 10.0 class EventLoopProbe: @@ -47,6 +49,9 @@ class EventLoopProbe: ) async def release_refresh_from_the_loop(self) -> None: + deadline: Final = time.monotonic() + REFRESH_START_TIMEOUT_SECONDS while not self.refresh_started.is_set(): + if time.monotonic() > deadline: + raise TimeoutError("signing finished without ever starting a credential refresh") await asyncio.sleep(0.005) self.loop_served.set() From 0f224dd4ed1fb4bed9ff67d8d8a095e0d032fcae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:59:34 -0700 Subject: [PATCH 3/6] fix: sign Bedrock guardrail, embedding, and Mantle requests off the event loop Bedrock Guardrails resolved credentials and signed inline on the loop in its three async paths, Titan embeddings signed each batch item inline in the async loop, and Bedrock Mantle requests slipped past the off-loop gate because BedrockMantleAuthMixin composes a BaseAWSLLM instead of inheriting from it. Introduce the SignsRequestsWithAWS marker that both BaseAWSLLM and the Mantle mixin carry so sign_request_off_loop_if_aws covers Mantle, and move the guardrail and embedding signing into asyncio.to_thread. Every new test fails at the previous tip. --- litellm/llms/bedrock/base_aws_llm.py | 8 +++- litellm/llms/bedrock/embed/embedding.py | 3 +- litellm/llms/bedrock_mantle/common_utils.py | 4 +- .../guardrail_hooks/bedrock_guardrails.py | 14 +++++-- .../test_bedrock_async_invoke_embedding.py | 35 ++++++++++++++++ .../bedrock/embed/test_bedrock_embedding.py | 42 +++++++++++++++++++ .../test_bedrock_mantle_transformation.py | 28 +++++++++++++ .../test_bedrock_guardrails.py | 36 ++++++++++++++++ .../test_bedrock_invoke_guardrail_checks.py | 33 +++++++++++++++ 9 files changed, 194 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1c7268ffa22..eddd4d6f826 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -81,7 +81,11 @@ class AwsAuthError(Exception): super().__init__(self.message) # Call the base class constructor with the parameters it needs -class BaseAWSLLM: +class SignsRequestsWithAWS: + pass + + +class BaseAWSLLM(SignsRequestsWithAWS): # Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request). # Storage is in-process memory only: no Redis backend unless attached elsewhere. Entry TTL: static # access-key + secret + region use ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); ambient @@ -1701,6 +1705,6 @@ async def sign_request_off_loop_if_aws( *args: _SignParams.args, **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature ) -> _SignedRequest: - if isinstance(provider_config, BaseAWSLLM): + if isinstance(provider_config, SignsRequestsWithAWS): return await asyncio.to_thread(sign_request, *args, **kwargs) return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index d0e8d734c5e..2c66fa1ef06 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -357,7 +357,8 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = self.get_request_headers( + prepped = await asyncio.to_thread( + self.get_request_headers, credentials=credentials, aws_region_name=aws_region_name, extra_headers=extra_headers, diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 850738bc320..ac94d9cb922 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -23,7 +23,7 @@ from botocore.exceptions import ( ProfileNotFound, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" @@ -55,7 +55,7 @@ def resolve_mantle_region(params: Mapping[str, object]) -> str: ) -class BedrockMantleAuthMixin: +class BedrockMantleAuthMixin(SignsRequestsWithAWS): _aws_signer: BaseAWSLLM @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 7204839f6d3..0ae2c48a3b4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -917,7 +917,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await asyncio.to_thread( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) allow_chunking: Final = not self._content_uses_contextual_grounding(content) completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator @@ -1178,7 +1180,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): **base_request_data, "content": content, } # mutable-ok: outbound JSON request body - prepared_request: Final = self._prepare_request( + prepared_request: Final = await asyncio.to_thread( + self._prepare_request, credentials=credentials, data=bedrock_request_data, optional_params=self.optional_params, @@ -1875,10 +1878,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() api_key: Final[str | None] = request_data.get("api_key") if request_data else None - credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key)) + credentials, aws_region_name = await asyncio.to_thread( + self._load_credentials, bearer_token=bedrock_bearer_token(api_key) + ) body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} - prepared_request: Final = self._prepare_request( + prepared_request: Final = await asyncio.to_thread( + self._prepare_request, credentials=credentials, data=body, optional_params=self.optional_params, diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 74a55cc1ef2..d8dfaaeac99 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -1,11 +1,15 @@ import json +import asyncio from unittest.mock import Mock, patch +import httpx import pytest +import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.base import HiddenParams +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock async invoke responses async_invoke_response = { @@ -383,3 +387,34 @@ class TestBedrockAsyncInvokeEmbedding: async_endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/async-invoke" ) + + +@pytest.mark.asyncio +async def test_async_invoke_status_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the GetAsyncInvoke poll is a signed GET, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + embedder = BedrockEmbedding() + probe = EventLoopProbe() + + with ( + patch.object(embedder, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + respx.mock, + ): + route = respx.get(url__regex=r"https://bedrock-runtime\.us-east-1\.amazonaws\.com/async-invoke/.*").mock( + return_value=httpx.Response(200, json=async_invoke_status_response) + ) + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + status = await embedder._get_async_invoke_status( + invocation_arn=async_invoke_status_response["invocationArn"], aws_region_name="us-east-1" + ) + await release + + assert status["status"] == "InProgress" + assert "Authorization" in route.calls.last.request.headers + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 50f8bbcf584..4ebf08dc709 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1,11 +1,16 @@ import json +import asyncio import os from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, MagicMock import pytest +import httpx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.bedrock.embed.embedding import BedrockEmbedding +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe # Mock responses for different embedding models titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -1059,3 +1064,40 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo assert response.data[0]["embedding"] == titan_embedding_response["embedding"] assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_async_single_func_embeddings_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: Titan, Nova, and TwelveLabs embeddings sign one SigV4 request per + input, and botocore refreshes expiring credentials inside that signing with a blocking HTTP call, + so each signing must run on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + probe = EventLoopProbe() + client = MagicMock() + client.__class__ = AsyncHTTPHandler + client.post = AsyncMock( + return_value=httpx.Response( + 200, + json=titan_embedding_response, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + ) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await BedrockEmbedding()._async_single_func_embeddings( + client=client, + timeout=None, + batch_data=[{"inputText": test_input}], + credentials=probe.credentials(), + extra_headers=None, + endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-embed-text-v1/invoke", + aws_region_name="us-west-2", + model="amazon.titan-embed-text-v1", + logging_obj=MagicMock(), + provider="amazon", + ) + await release + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + assert "Authorization" in client.post.call_args.kwargs["headers"] + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 1be94d4daa2..e83a844c87e 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -6,15 +6,20 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht """ import json +import asyncio from unittest.mock import patch import httpx import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest import litellm from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig +from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws from litellm.types.utils import LlmProviders +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.fixture @@ -710,3 +715,26 @@ def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id): resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name) assert provider == "bedrock_mantle" assert resolved_model == model_id + + +@pytest.mark.asyncio +async def test_mantle_signing_runs_off_the_event_loop(): + """Regression for issue #40165: Mantle signs with SigV4 through a composed BaseAWSLLM, so the + off-loop gate must recognise it too, or its credential refresh blocks the loop like Bedrock's did.""" + probe = EventLoopProbe() + + def sign(headers: dict[str, str]) -> dict[str, str]: + request = AWSRequest( + method="POST", url="https://bedrock-mantle.us-east-1.api.aws/v1/responses", data="{}", headers=headers + ) + SigV4Auth(probe.credentials(), "bedrock", "us-east-1").add_auth(request) + return dict(request.headers) + + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + signed = await sign_request_off_loop_if_aws( + BedrockMantleChatConfig(), sign, headers={"Content-Type": "application/json"} + ) + await release + + assert "Authorization" in signed + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 7da55f22bda..c0762edec92 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3,6 +3,8 @@ Unit tests for Bedrock Guardrails """ import json +import asyncio +from datetime import datetime, timezone import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -28,6 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockTextContent, ) from litellm.types.utils import CallTypes, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @pytest.mark.asyncio @@ -5842,3 +5845,36 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): assert response["action"] == "NONE" assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the ApplyGuardrail request is signed with SigV4, and botocore + refreshes expiring credentials inside that signing with a blocking HTTP call, so it must run + on a worker thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"action": "NONE", "outputs": [], "assessments": []}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=allowed)): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await guardrail._post_apply_guardrail_content( + content=[{"text": {"text": "hello"}}], + base_request_data={"source": "INPUT"}, + credentials=probe.credentials(), + aws_region_name="us-east-1", + api_key=None, + request_data={}, + event_type=GuardrailEventHooks.pre_call, + start_time=datetime.now(timezone.utc), + completed_chunk_usages=[], + ) + await release + + assert response["action"] == "NONE" + assert probe.served_during_refresh is True diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index f4af77d2e40..bbc8fd539a3 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -5,10 +5,12 @@ All Bedrock HTTP calls are mocked; no real AWS calls are made. """ import json +import asyncio import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest +import httpx from fastapi import HTTPException @@ -21,6 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailResponse, ) from litellm.types.utils import Choices, Message, ModelResponse +from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe CONTENT_FILTER_CHECKS = {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}} @@ -861,3 +864,33 @@ async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeyp {"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8} ] assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +@pytest.mark.asyncio +async def test_invoke_guardrail_checks_signs_off_the_event_loop(monkeypatch): + """Regression for issue #40165: the checks request is signed with SigV4, and botocore refreshes + expiring credentials inside that signing with a blocking HTTP call, so it must run on a worker + thread to keep the loop serving other requests.""" + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + g = BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, content_filter_threshold=0.5) + probe = EventLoopProbe() + allowed = httpx.Response( + 200, + json={"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.1}]}}}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ) + + with ( + patch.object(g, "_load_credentials", return_value=(probe.credentials(), "us-east-1")), + patch.object(g.async_handler, "post", new=AsyncMock(return_value=allowed)), + ): + release = asyncio.create_task(probe.release_refresh_from_the_loop()) + response = await g.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"messages": []}, + ) + await release + + assert response == BedrockGuardrailResponse() + assert probe.served_during_refresh is True From e27a549aa919f82e590cf04513e72bf5a75f4a7d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:14:00 -0700 Subject: [PATCH 4/6] test(bedrock): type the probe credential overrides --- .../count_tokens/test_bedrock_count_tokens_handler.py | 6 +++++- .../test_litellm/llms/custom_httpx/test_llm_http_handler.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 31e65bb6e27..3622ce7f212 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -3,6 +3,7 @@ from unittest.mock import AsyncMock import httpx import pytest +from botocore.credentials import RefreshableCredentials from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -14,7 +15,10 @@ class _ProbedCountTokensHandler(BedrockCountTokensHandler): super().__init__() self._probe = probe - def get_credentials(self, **kwargs): + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: return self._probe.credentials() diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index b5eb98c4ce6..1140090f9cf 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock, patch import httpx import pytest +from botocore.credentials import RefreshableCredentials import litellm from litellm._logging import verbose_logger @@ -822,7 +823,10 @@ class _ProbedBedrockMessagesConfig(AmazonAnthropicClaudeMessagesConfig): super().__init__() self._probe = probe - def get_credentials(self, **kwargs): + def get_credentials( + self, + **kwargs: object, # kwargs-ok: mirrors the base resolver's keyword contract, which the probe ignores + ) -> RefreshableCredentials: return self._probe.credentials() From d93baa3f2c1ca3c5123ecdee8510ee3c8d827f61 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:52:00 -0700 Subject: [PATCH 5/6] fix(bedrock): sign on a dedicated executor instead of the shared default one asyncio.to_thread puts every Bedrock signing on the loop's default executor, the same pool every provider's async entry point hops through, so signings parked on botocore's refresh lock queued unrelated providers behind Bedrock. run_aws_signing runs them on a 16-thread pool only AWS signing uses --- .../providers/bedrock_agentcore/handler.py | 6 ++--- litellm/constants.py | 1 + litellm/integrations/s3_v2.py | 6 ++--- litellm/integrations/sqs.py | 4 +-- litellm/llms/bedrock/base_aws_llm.py | 20 +++++++++++++- litellm/llms/bedrock/chat/converse_handler.py | 7 +++-- litellm/llms/bedrock/count_tokens/handler.py | 4 +-- litellm/llms/bedrock/embed/embedding.py | 7 +++-- litellm/llms/bedrock/realtime/handler.py | 6 ++--- litellm/llms/custom_httpx/llm_http_handler.py | 4 +-- .../guardrail_hooks/bedrock_guardrails.py | 10 +++---- .../llm_passthrough_endpoints.py | 9 +++---- .../llms/bedrock/test_base_aws_llm.py | 26 +++++++++++++++++++ 13 files changed, 76 insertions(+), 34 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index e01e215b5f0..306a8871b12 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -5,7 +5,6 @@ Sends JSON-RPC envelopes directly to AgentCore endpoints, bypassing the completion bridge that would otherwise strip the envelope. """ -import asyncio import json from collections.abc import AsyncIterator, Mapping from typing import Any, Final @@ -14,6 +13,7 @@ from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( BedrockAgentCoreA2ATransformation, ) +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -46,7 +46,7 @@ class BedrockAgentCoreA2AHandler: Returns: A2A JSON-RPC response dict from the AgentCore agent """ - url, headers, body = await asyncio.to_thread( + url, headers, body = await run_aws_signing( BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, @@ -93,7 +93,7 @@ class BedrockAgentCoreA2AHandler: Yields: A2A streaming response events from the AgentCore agent """ - url, headers, body = await asyncio.to_thread( + url, headers, body = await run_aws_signing( BedrockAgentCoreA2ATransformation.get_url_and_signed_request, request_id=request_id, params=params, diff --git a/litellm/constants.py b/litellm/constants.py index f9389d22dea..060402cc08d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -570,6 +570,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float( LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100) LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000) LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0 +AWS_SIGNING_MAX_THREADS: Final = 16 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index e3297cfd27c..996ff9c75af 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -24,7 +24,7 @@ from litellm.integrations.s3 import ( from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) - await asyncio.to_thread(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) + await run_aws_signing(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) @@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Sign the request aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) - await asyncio.to_thread(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) + await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 31cc20f5995..d787375ca3c 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -22,7 +22,7 @@ from litellm.constants import ( SQS_SEND_MESSAGE_ACTION, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM): data=prepped.body, headers=prepped.headers, ) - await asyncio.to_thread(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) + await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request) signed_headers: Final = dict(aws_request.headers.items()) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index eddd4d6f826..96804a1fa62 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,12 +1,15 @@ import asyncio import base64 +import contextvars import hashlib import json import os import re import urllib.parse from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from datetime import datetime +from functools import partial from threading import Lock from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload @@ -17,6 +20,7 @@ from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + AWS_SIGNING_MAX_THREADS, BEDROCK_EMBEDDING_PROVIDERS_LITERAL, BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES, BEDROCK_IAM_CACHE_MAX_ENTRIES, @@ -1697,6 +1701,20 @@ def sign_aws_json_post( _SignParams = ParamSpec("_SignParams") _SignedRequest = TypeVar("_SignedRequest") +AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing") + + +async def run_aws_signing( + sign: Callable[_SignParams, _SignedRequest], + /, + *args: _SignParams.args, + **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature +) -> _SignedRequest: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor( + AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs) + ) + async def sign_request_off_loop_if_aws( provider_config: object, @@ -1706,5 +1724,5 @@ async def sign_request_off_loop_if_aws( **kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature ) -> _SignedRequest: if isinstance(provider_config, SignsRequestsWithAWS): - return await asyncio.to_thread(sign_request, *args, **kwargs) + return await run_aws_signing(sign_request, *args, **kwargs) return sign_request(*args, **kwargs) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 3e8e1226a55..6d48ff3f07c 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,4 +1,3 @@ -import asyncio import json from collections.abc import Mapping from types import MappingProxyType @@ -22,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -137,7 +136,7 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", @@ -208,7 +207,7 @@ class BedrockConverseLLM(BaseAWSLLM): ) data: Final = json.dumps(request_data) - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( self.get_request_headers, credentials=credentials, aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index cae2315e744..d7fb510f057 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -4,13 +4,13 @@ AWS Bedrock CountTokens API handler. Simplified handler leveraging existing LiteLLM Bedrock infrastructure. """ -import asyncio from typing import Any, Final import httpx import litellm from litellm._logging import verbose_logger +from litellm.llms.bedrock.base_aws_llm import run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client @@ -77,7 +77,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): # Extract api_key for bearer token auth if provided api_key: Final = litellm_params.get("api_key", None) headers: Final = {"Content-Type": "application/json"} - signed_headers, signed_body = await asyncio.to_thread( + signed_headers, signed_body = await run_aws_signing( self._sign_request, service_name="bedrock", headers=headers, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 4254d9b1d8b..be766eaedd0 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -2,7 +2,6 @@ Handles embedding calls to Bedrock's `/invoke` endpoint """ -import asyncio import copy import json import urllib.parse @@ -27,7 +26,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.types.utils import EmbeddingResponse, LlmProviders -from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token +from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing from ..common_utils import BedrockError from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config @@ -357,7 +356,7 @@ class BedrockEmbedding(BaseAWSLLM): if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} - prepped = await asyncio.to_thread( + prepped = await run_aws_signing( self.get_request_headers, credentials=credentials, aws_region_name=aws_region_name, @@ -638,7 +637,7 @@ class BedrockEmbedding(BaseAWSLLM): credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name ) - prepped: Final = await asyncio.to_thread(sign_status_request) + prepped: Final = await run_aws_signing(sign_status_request) # LOGGING if logging_obj is not None: diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 792e3169a7a..ca2370303f2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput -from ..base_aws_llm import BaseAWSLLM +from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig @@ -149,7 +149,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) - credentials: Final = await asyncio.to_thread( + credentials: Final = await run_aws_signing( self.get_credentials, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, @@ -170,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM): "or configure credentials in the environment" ), ) - frozen_credentials: Final = await asyncio.to_thread(credentials.get_frozen_credentials) + frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials) # Initialize Bedrock client with aws_sdk_bedrock_runtime config: Final = Config( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3106810b4b3..13f134fce1d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,7 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig -from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws +from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -638,7 +638,7 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed)) + return await dispatch_async(*await run_aws_signing(sign_and_log, transformed)) return transform_then_dispatch() diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 0ae2c48a3b4..6a7ac4361b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -917,7 +917,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): source, ) return BedrockGuardrailResponse() - credentials, aws_region_name = await asyncio.to_thread( + credentials, aws_region_name = await run_aws_signing( self._load_credentials, bearer_token=bedrock_bearer_token(api_key) ) allow_chunking: Final = not self._content_uses_contextual_grounding(content) @@ -1180,7 +1180,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): **base_request_data, "content": content, } # mutable-ok: outbound JSON request body - prepared_request: Final = await asyncio.to_thread( + prepared_request: Final = await run_aws_signing( self._prepare_request, credentials=credentials, data=bedrock_request_data, @@ -1878,12 +1878,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return BedrockGuardrailResponse() api_key: Final[str | None] = request_data.get("api_key") if request_data else None - credentials, aws_region_name = await asyncio.to_thread( + credentials, aws_region_name = await run_aws_signing( self._load_credentials, bearer_token=bedrock_bearer_token(api_key) ) body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks} - prepared_request: Final = await asyncio.to_thread( + prepared_request: Final = await run_aws_signing( self._prepare_request, credentials=credentials, data=body, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 037f663044f..bf6ee57b4fc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -8,7 +8,6 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. from __future__ import annotations -import asyncio import hmac import inspect import json @@ -1131,7 +1130,7 @@ async def bedrock_proxy_route( ) # Add or update query parameters - from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post from litellm.llms.bedrock.chat import BedrockConverseLLM bedrock_llm: Final = BedrockConverseLLM() @@ -1140,7 +1139,7 @@ async def bedrock_proxy_route( data: Final = await _json_request_body(request) except Exception as e: raise HTTPException(status_code=400, detail={"error": e}) - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( sign_aws_json_post, get_credentials=bedrock_llm.get_credentials, service_name="bedrock", @@ -1236,10 +1235,10 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" - prepped: Final = await asyncio.to_thread( + prepped: Final = await run_aws_signing( sign_aws_json_post, get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), service_name="comprehendmedical", diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index ca93f58f366..e24ea1ca611 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1,5 +1,6 @@ import asyncio import json +from concurrent.futures import ThreadPoolExecutor import os import threading import time @@ -23,6 +24,7 @@ from litellm.llms.bedrock.base_aws_llm import ( AwsAuthError, BaseAWSLLM, Boto3CredentialsInfo, + run_aws_signing, sign_request_off_loop_if_aws, ) from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @@ -3247,3 +3249,27 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent assert "Authorization" in signed assert probe.served_during_refresh is True + + +async def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): + """A signing parked on botocore's refresh lock must not hold a default-executor thread, since every + other provider's async entry point hops through that same executor.""" + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() + + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name + + signing = asyncio.create_task(run_aws_signing(sign)) + try: + await asyncio.wait_for(signing_parked.wait(), timeout=5) + other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) + finally: + refresh_done.set() + + assert other_provider.name != await signing + assert (await signing).startswith("aws-signing") From 3dcc09b26ccb0aa781e7bb62007e73e93c2bd9f6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:16:11 -0700 Subject: [PATCH 6/6] fix(bedrock): keep the signing pool to AWS configs on the async-transform path The async-transform path in the shared handler sent every provider's sign_and_log to the AWS pool, so Ollama, Snowflake, and watsonx queued behind Bedrock refreshes there. Only SignsRequestsWithAWS configs take run_aws_signing now, the rest keep the default-executor hop they had. The executor isolation test also runs on its own loop instead of pinning a one-thread default executor on the session-scoped pytest loop --- litellm/llms/custom_httpx/llm_http_handler.py | 9 +++- .../llms/bedrock/test_base_aws_llm.py | 41 +++++++++++-------- .../custom_httpx/test_llm_http_handler.py | 21 ++++++++++ 3 files changed, 51 insertions(+), 20 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 13f134fce1d..e720428847d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -77,7 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig -from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_request_off_loop_if_aws +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws from litellm.llms.custom_httpx.container_handler import raise_for_error_status from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -638,7 +638,12 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*await run_aws_signing(sign_and_log, transformed)) + signed_request: Final = await ( + run_aws_signing(sign_and_log, transformed) + if isinstance(provider_config, SignsRequestsWithAWS) + else asyncio.to_thread(sign_and_log, transformed) + ) + return await dispatch_async(*signed_request) return transform_then_dispatch() diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index e24ea1ca611..2c7e476e9a8 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -3251,25 +3251,30 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent assert probe.served_during_refresh is True -async def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): +def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): """A signing parked on botocore's refresh lock must not hold a default-executor thread, since every - other provider's async entry point hops through that same executor.""" - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) - signing_parked = asyncio.Event() - refresh_done = threading.Event() + other provider's async entry point hops through that same executor. The scenario runs on its own loop + so the one-thread default executor it pins never leaks into the session loop.""" - def sign() -> str: - loop.call_soon_threadsafe(signing_parked.set) - refresh_done.wait() - return threading.current_thread().name + async def scenario() -> tuple[str, str]: + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() - signing = asyncio.create_task(run_aws_signing(sign)) - try: - await asyncio.wait_for(signing_parked.wait(), timeout=5) - other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) - finally: - refresh_done.set() + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name - assert other_provider.name != await signing - assert (await signing).startswith("aws-signing") + signing = asyncio.create_task(run_aws_signing(sign)) + try: + await asyncio.wait_for(signing_parked.wait(), timeout=5) + other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5) + finally: + refresh_done.set() + return other_provider.name, await signing + + other_provider, signing_thread = asyncio.run(scenario()) + assert other_provider != signing_thread + assert signing_thread.startswith("aws-signing") diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 1140090f9cf..cea2d439198 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -20,6 +20,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( @@ -3431,6 +3432,26 @@ async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_tran assert captured["body"] == {"transformed_by": "async"} assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + assert not any(thread.name.startswith("aws-signing") for thread in config.sign_threads + pre_call_threads) + + +class _AWSTransformRecordingConfig(SignsRequestsWithAWS, _TransformRecordingConfig): + pass + + +async def test_completion_signs_aws_configs_on_the_aws_signing_pool_after_the_async_transform(): + config = _AWSTransformRecordingConfig(transform_async=True) + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread.name.startswith("aws-signing") for thread in config.sign_threads) + assert pre_call_threads and all(thread.name.startswith("aws-signing") for thread in pre_call_threads) async def test_completion_keeps_sync_transform_request_before_returning_by_default():