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] 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):