scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(bedrock): resolve AWS region and secret lookups off the event loop (#47331)

Browse source 
* fix(bedrock): resolve AWS region and secret lookups off the event loop

Bedrock request paths that resolve their region, endpoint or bearer token
inside the async handler now do it in the AWS signing executor, so a slow
secret manager read no longer stalls every other request on the worker

* fix(bedrock): type the STS client kwargs and claude platform environment override

Keeps the basedpyright counts at or below the merge base after the off-loop wraps, by typing the STS client kwargs as a TypedDict, the claude platform environment override with concrete dicts, and the realtime region call through the typed getter

* test(bedrock): detect event loop stalls without wall-clock timing

The Vault fake now answers a read only after the test loop runs a callback, so a read made on the loop thread is caught deterministically instead of by a sleep and a heartbeat threshold
This commit is contained in:
Mateo Wang 2026-10-20 17:04:41 +00:00 • committed by GitHub
parent 2c6b280755
commit f3513bca9e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 496 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
litellm/llms/bedrock/base_aws_llm.py
View file

@ -80,6 +80,12 @@ class _AssumeRoleParams(TypedDict):
Tags: ReadOnly[NotRequired[tuple[AwsSessionTag, ...]]]
class _StsClientKwargs(TypedDict):
verify: ReadOnly[bool | str]
endpoint_url: ReadOnly[NotRequired[str]]
region_name: ReadOnly[NotRequired[str]]
def _assume_role_params(
aws_role_name: str,
aws_session_name: str,
@ -275,7 +281,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"aws_session_tags",
]
def _get_ssl_verify(self, ssl_verify: bool | str | None = None):
def _get_ssl_verify(self, ssl_verify: bool | str | None = None) -> bool | str:
"""
Get SSL verification setting for boto3 clients.
@ -875,15 +881,18 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_sts_endpoint: str | None = None,
ssl_verify: bool | str | None = None,
aws_region_name: str | None = None,
) -> dict:
) -> _StsClientKwargs:
"""STS client kwargs with aligned endpoint_url and region_name (SigV4)."""
kwargs: Final[dict] = {"verify": self._get_ssl_verify(ssl_verify)}
if aws_sts_endpoint is not None:
kwargs["endpoint_url"] = aws_sts_endpoint
sts_region: Final = self._resolve_sts_region(aws_sts_endpoint, aws_region_name)
if sts_region is not None:
kwargs["region_name"] = sts_region
return kwargs
verify: Final = self._get_ssl_verify(ssl_verify)
match (aws_sts_endpoint, self._resolve_sts_region(aws_sts_endpoint, aws_region_name)):
case (None, None):
return _StsClientKwargs(verify=verify)
case (None, region):
return _StsClientKwargs(verify=verify, region_name=region)
case (endpoint, None):
return _StsClientKwargs(verify=verify, endpoint_url=endpoint)
case (endpoint, region):
return _StsClientKwargs(verify=verify, endpoint_url=endpoint, region_name=region)
def get_aws_region_name_for_non_llm_api_calls(
self,
@ -1813,13 +1822,13 @@ async def run_aws_signing(
)
async def sign_request_off_loop_if_aws(
async def run_off_loop_if_aws(
provider_config: object,
sign_request: Callable[_SignParams, _SignedRequest],
call: Callable[_SignParams, _SignedRequest],
/,
*args: _SignParams.args,
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped call's signature
) -> _SignedRequest:
if isinstance(provider_config, SignsRequestsWithAWS):
return await run_aws_signing(sign_request, *args, **kwargs)
return sign_request(*args, **kwargs)
return await run_aws_signing(call, *args, **kwargs)
return call(*args, **kwargs)

34
litellm/llms/bedrock/claude_platform/messages_transformation.py
View file

@ -1,10 +1,11 @@
from typing import Any, Final
from typing import Final
import litellm
from litellm.llms.anthropic.pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
AnthropicMessagesConfig,
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
@ -22,14 +23,14 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
def validate_anthropic_messages_environment(
self,
headers: dict,
headers: dict[str, str], # mutable-ok: overrides the base signature
model: str,
messages: list[Any],
optional_params: dict,
litellm_params: dict,
messages: list[object], # mutable-ok: overrides the base signature
optional_params: dict[str, object], # mutable-ok: overrides the base signature
litellm_params: dict[str, object], # mutable-ok: overrides the base signature
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]:
) -> tuple[dict[str, str], str | None]: # mutable-ok: overrides the base signature
workspace_id: Final = self._get_workspace_id(optional_params, litellm_params)
if workspace_id is None:
raise litellm.AuthenticationError(
@ -66,6 +67,27 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
api_base,
)
async def avalidate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: mirrors validate_anthropic_messages_environment
model: str,
messages: list[object], # mutable-ok: mirrors validate_anthropic_messages_environment
optional_params: dict[str, object], # mutable-ok: mirrors validate_anthropic_messages_environment
litellm_params: dict[str, object], # mutable-ok: mirrors validate_anthropic_messages_environment
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors validate_anthropic_messages_environment
return await run_aws_signing(
self.validate_anthropic_messages_environment,
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
api_base=api_base,
)
def transform_anthropic_messages_request(
self,
model: str,

6
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -48,7 +48,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model)
# Get AWS region using existing LiteLLM function
aws_region_name: Final = self.get_aws_region_name(
aws_region_name: Final = await run_aws_signing(
self.get_aws_region_name,
optional_params=litellm_params,
model=resolved_model,
model_id=None,
@ -64,7 +65,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
# Get endpoint URL using simplified function
api_base: Final = litellm_params.get("api_base", None)
aws_bedrock_runtime_endpoint: Final = litellm_params.get("aws_bedrock_runtime_endpoint", None)
endpoint_url: Final = self.get_bedrock_count_tokens_endpoint(
endpoint_url: Final = await run_aws_signing(
self.get_bedrock_count_tokens_endpoint,
model=resolved_model,
aws_region_name=aws_region_name,
api_base=api_base,

6
litellm/llms/bedrock/count_tokens/mantle_handler.py
View file

@ -54,10 +54,10 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
) -> dict[str, JsonValue]:
try:
request: Final = _COUNT_REQUEST.validate_python(request_data)
aws_region_name: Final = self.get_aws_region_name(
optional_params=litellm_params, model=resolved_model, model_id=None
aws_region_name: Final = await run_aws_signing(
self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None
)
endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)
endpoint_url: Final = await run_aws_signing(self.get_mantle_count_tokens_endpoint, aws_region_name)
body: Final = self._anthropic_config.transform_request_to_count_tokens(
model=resolved_model,
messages=request["messages"],

5
litellm/llms/bedrock/realtime/handler.py
View file

@ -337,10 +337,7 @@ class BedrockRealtime(BaseAWSLLM):
# Get AWS region
if aws_region_name is None:
optional_params: Final = {
"aws_region_name": aws_region_name,
}
aws_region_name = self._get_aws_region_name(optional_params, model)
aws_region_name = await run_aws_signing(self.get_aws_region_name, {}, model)
# Get endpoint URL
if api_base is not None:

74
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -104,7 +104,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 SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws
from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, run_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,
@ -2085,14 +2085,18 @@ class BaseLLMHTTPHandler:
raise ValueError(f"No provider config found for provider: {custom_llm_provider}")
# Validate environment and get headers
headers = provider_config.validate_environment(
headers = await run_off_loop_if_aws(
provider_config,
provider_config.validate_environment,
api_key=api_key,
api_base=api_base,
headers=headers or {},
)
# Transform the request first
data: Final = provider_config.transform_search_request(
data: Final = await run_off_loop_if_aws(
provider_config,
provider_config.transform_search_request,
query=query,
optional_params=optional_params,
api_key=api_key,
@ -2108,7 +2112,7 @@ class BaseLLMHTTPHandler:
api_key=api_key,
)
signed_headers, signed_json_body = await sign_request_off_loop_if_aws(
signed_headers, signed_json_body = await run_off_loop_if_aws(
provider_config,
provider_config.sign_request,
headers=headers,
@ -2236,7 +2240,7 @@ class BaseLLMHTTPHandler:
)
provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body)
recovered_errors = (*recovered_errors, e)
headers, signed_json_body = await sign_request_off_loop_if_aws(
headers, signed_json_body = await run_off_loop_if_aws(
provider_config,
provider_config.sign_request,
headers=headers,
@ -2405,7 +2409,9 @@ class BaseLLMHTTPHandler:
logging_obj.model_call_details.update(request_body)
# Make the request
request_url: Final = anthropic_messages_provider_config.get_complete_url(
request_url: Final = await run_off_loop_if_aws(
anthropic_messages_provider_config,
anthropic_messages_provider_config.get_complete_url,
api_base=api_base,
api_key=api_key,
model=model,
@ -2416,7 +2422,7 @@ class BaseLLMHTTPHandler:
stream=stream,
)
headers, signed_json_body = await sign_request_off_loop_if_aws(
headers, signed_json_body = await run_off_loop_if_aws(
anthropic_messages_provider_config,
anthropic_messages_provider_config.sign_request,
headers=headers,
@ -2967,7 +2973,9 @@ class BaseLLMHTTPHandler:
stream = response_api_optional_request_params.get("stream", False)
caller_requested_stream: Final = bool(stream or (extra_body or {}).get("stream"))
api_base: Final = responses_api_provider_config.get_complete_url(
api_base: Final = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.get_complete_url,
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
@ -3009,7 +3017,7 @@ class BaseLLMHTTPHandler:
fake_stream=fake_stream,
)
headers, signed_body = await sign_request_off_loop_if_aws(
headers, signed_body = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.sign_request,
headers=validated_headers,
@ -3856,7 +3864,9 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
api_key=api_key,
)
complete_api_base: Final = provider_config.get_complete_file_url(
complete_api_base: Final = await run_off_loop_if_aws(
provider_config,
provider_config.get_complete_file_url,
api_base=api_base,
api_key=api_key,
model="",
@ -3867,7 +3877,9 @@ class BaseLLMHTTPHandler:
if not complete_api_base:
raise ValueError("api_base is required for create_file")
return await self.async_create_file(
transformed_request=provider_config.transform_create_file_request(
transformed_request=await run_off_loop_if_aws(
provider_config,
provider_config.transform_create_file_request,
model="",
create_file_data=create_file_data,
litellm_params=litellm_params,
@ -4334,7 +4346,9 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
api_key=api_key,
)
complete_api_base: Final = provider_config.get_complete_batch_url(
complete_api_base: Final = await run_off_loop_if_aws(
provider_config,
provider_config.get_complete_batch_url,
api_base=api_base,
api_key=api_key,
model=model,
@ -4345,7 +4359,9 @@ class BaseLLMHTTPHandler:
if not complete_api_base:
raise ValueError("api_base is required for create_batch")
return await self.async_create_batch(
transformed_request=provider_config.transform_create_batch_request(
transformed_request=await run_off_loop_if_aws(
provider_config,
provider_config.transform_create_batch_request,
model=model,
create_batch_data=create_batch_data,
litellm_params=litellm_params,
@ -4818,7 +4834,9 @@ class BaseLLMHTTPHandler:
**(extra_headers or {}),
}
api_base: Final = responses_api_provider_config.get_complete_url(
api_base: Final = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.get_complete_url,
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
@ -4836,7 +4854,7 @@ class BaseLLMHTTPHandler:
)
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
headers, signed_body = await sign_request_off_loop_if_aws(
headers, signed_body = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.sign_request,
headers=validated_headers,
@ -4962,7 +4980,9 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
# Get URL and params from provider config
url, params = provider_config.transform_retrieve_file_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_retrieve_file_request,
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
@ -5089,7 +5109,9 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
# Get URL and params from provider config
url, params = provider_config.transform_delete_file_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_delete_file_request,
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
@ -5214,7 +5236,9 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
# Get URL and params from provider config
url, params = provider_config.transform_list_files_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_list_files_request,
purpose=purpose,
optional_params={},
litellm_params=litellm_params,
@ -5298,7 +5322,9 @@ class BaseLLMHTTPHandler:
)
yield page_files[: MAX_FILE_LIST_LIMIT - listed_count]
listed_count += len(page_files)
next_request = self._next_listing_request(latest_page, provider_config, litellm_params, listed_count)
next_request = await run_off_loop_if_aws(
provider_config, self._next_listing_request, latest_page, provider_config, litellm_params, listed_count
)
if next_request is None:
return
url, params = next_request
@ -5524,7 +5550,9 @@ class BaseLLMHTTPHandler:
headers: dict,
logging_obj: LiteLLMLoggingObj,
) -> "_PreparedFileContentRequest":
url, params = provider_config.transform_file_content_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_file_content_request,
file_content_request=file_content_request,
optional_params={},
litellm_params=litellm_params,
@ -10332,7 +10360,9 @@ class BaseLLMHTTPHandler:
if extra_headers:
headers.update(extra_headers)
api_base: Final = vector_store_provider_config.get_complete_url(
api_base: Final = await run_off_loop_if_aws(
vector_store_provider_config,
vector_store_provider_config.get_complete_url,
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
@ -10366,7 +10396,7 @@ 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 = await sign_request_off_loop_if_aws(
headers, signed_json_body = await run_off_loop_if_aws(
vector_store_provider_config,
vector_store_provider_config.sign_request,
headers=headers,

12
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -796,7 +796,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# logic becomes shared across providers.
#### CALL HOOKS - proxy only ####
def _load_credentials(self, bearer_token: str | None = None):
def _load_credentials(self, api_key: str | None = None):
try:
from botocore.credentials import Credentials
except ImportError:
@ -820,7 +820,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
credentials: Final[Credentials | None] = (
None
if bearer_token is not None
if bedrock_bearer_token(api_key) is not None
else self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
@ -1017,9 +1017,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
start_time: "datetime",
allow_chunking: bool,
) -> BedrockGuardrailResponse:
credentials, aws_region_name = await run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)
credentials, aws_region_name = await run_aws_signing(self._load_credentials, api_key=api_key)
completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
try:
@ -1960,9 +1958,7 @@ 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 run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)
credentials, aws_region_name = await run_aws_signing(self._load_credentials, api_key=api_key)
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
prepared_request: Final = await run_aws_signing(

9
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -1459,7 +1459,6 @@ async def bedrock_proxy_route(
"""
create_request_copy(request)
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(
endpoint=endpoint,
@ -1474,6 +1473,10 @@ async def bedrock_proxy_route(
detail="bedrock-agent-runtime pass-through is disabled on this proxy.",
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.chat import BedrockConverseLLM
aws_region_name: Final = await run_aws_signing(get_secret_str, secret_name="AWS_REGION_NAME")
base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
encoded_endpoint = httpx.URL(endpoint).path
@ -1488,10 +1491,6 @@ async def bedrock_proxy_route(
path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint)
)
# Add or update query parameters
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()
# Assuming the body contains JSON data, parse it
try:

29
tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py
View file

@ -10,6 +10,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
class _ProbedCountTokensHandler(BedrockCountTokensHandler):
@ -56,6 +57,34 @@ async def test_handle_count_tokens_request_signs_off_the_event_loop(monkeypatch)
assert probe.served_during_refresh is True
def _count_tokens_in_us_west_2(request: httpx.Request) -> httpx.Response:
if request.url.host != "bedrock-runtime.us-west-2.amazonaws.com":
return httpx.Response(404, request=request)
return httpx.Response(200, json={"inputTokens": 7}, request=request)
@pytest.mark.asyncio
async def test_handle_count_tokens_request_looks_up_its_region_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret")
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
vault: Final = host_in_loop_probing_vault(monkeypatch, "AWS_REGION_NAME", "AWS_BEDROCK_RUNTIME_ENDPOINT")
result = await BedrockCountTokensHandler().handle_count_tokens_request(
request_data={"model": "us.anthropic.claude-sonnet-5-5", "messages": [{"role": "user", "content": "hi"}]},
litellm_params={},
resolved_model="us.anthropic.claude-sonnet-5-5",
client=AsyncHTTPHandler(transport=httpx.MockTransport(_count_tokens_in_us_west_2)),
)
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert result == {"input_tokens": 7}
class TestBedrockCountTokensEndpoint:
def _make_handler(self) -> BedrockCountTokensConfig:
return BedrockCountTokensConfig()

20
tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py
View file

@ -8,6 +8,7 @@ import pytest
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
MANTLE_COUNT_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages/count_tokens"
LITELLM_PARAMS: Final = {
@ -107,3 +108,22 @@ async def test_non_200_answers_raise_bedrock_error_with_mantle_status() -> None:
assert raised.value.status_code == 404
assert "does not exist" in raised.value.message
@pytest.mark.asyncio
async def test_region_and_endpoint_lookups_leave_the_event_loop_free(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
vault: Final = host_in_loop_probing_vault(monkeypatch, "AWS_REGION_NAME", "BEDROCK_MANTLE_API_BASE")
mantle: Final = _MantleEndpoint(200, {"input_tokens": 2177})
result = await BedrockMantleCountTokensHandler().handle_count_tokens_request(
request_data={**REQUEST, "model": "global.anthropic.claude-opus-5-5"},
litellm_params={key: value for key, value in LITELLM_PARAMS.items() if key != "aws_region_name"},
resolved_model="anthropic.claude-opus-5-5",
client=_client(mantle),
)
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert result == {"input_tokens": 2177}
assert str(mantle.only_request().url) == MANTLE_COUNT_URL

38
tests/unit/llms/bedrock/event_loop_probe.py
View file

@ -1,4 +1,4 @@
"""Refreshable credentials whose refresh only completes while the event loop keeps serving."""
"""Credential refreshes and Vault reads that only complete while the event loop keeps serving."""
from __future__ import annotations
@ -8,10 +8,15 @@ import time
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from botocore.credentials import RefreshableCredentials
import litellm
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
REFRESH_RELEASE_TIMEOUT_SECONDS: Final = 2.0
REFRESH_START_TIMEOUT_SECONDS: Final = 10.0
LOOP_SERVE_TIMEOUT_SECONDS: Final = 2.0
class EventLoopProbe:
@ -55,3 +60,34 @@ class EventLoopProbe:
raise TimeoutError("signing finished without ever starting a credential refresh")
await asyncio.sleep(0.005)
self.loop_served.set()
class LoopProbingVault:
"""A Vault holding none of its hosted secrets, whose read waits for the event loop to run a callback.
A read made on the loop thread can never see that callback, so `stalled_the_loop` is set only then.
"""
def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
self._loop: Final = loop
self.read: Final = threading.Event()
self.stalled_the_loop: Final = threading.Event()
def sync_read_secret(self, secret_name: str) -> None:
self.read.set()
served: Final = threading.Event()
self._loop.call_soon_threadsafe(served.set)
if not served.wait(timeout=LOOP_SERVE_TIMEOUT_SECONDS):
self.stalled_the_loop.set()
def host_in_loop_probing_vault(monkeypatch: pytest.MonkeyPatch, *secret_names: str) -> LoopProbingVault:
vault: Final = LoopProbingVault(asyncio.get_running_loop())
monkeypatch.setattr(litellm, "secret_manager_client", vault)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT)
monkeypatch.setattr(
litellm,
"_key_management_settings",
KeyManagementSettings(access_mode="read_only", hosted_keys=list(secret_names)),
)
return vault

22
tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py
View file

@ -18,6 +18,7 @@ from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.realtime.handler import BedrockRealtime
from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
from litellm.realtime_api import main as realtime_main
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
@pytest.fixture(autouse=True)
@ -625,6 +626,27 @@ class TestBedrockRealtimeSessionLifecycle:
assert first_event["session"]["id"] == "trace-nova-sonic"
assert first_event["session"]["model"] == "amazon.nova-sonic-v1:0"
@pytest.mark.asyncio
async def test_region_lookup_on_connect_leaves_the_event_loop_free(
self, stub_aws_sdk_client, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
vault: Final = host_in_loop_probing_vault(monkeypatch, "AWS_REGION_NAME")
websocket: Final = RealtimeClientWS()
_ = await BedrockRealtime().async_realtime(
model="amazon.nova-2-sonic-v1:0",
websocket=websocket,
logging_obj=FakeLogging(),
aws_access_key_id="k",
aws_secret_access_key="s",
)
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert stub_aws_sdk_client["config_kwargs"]["region"] == "us-west-2"
assert stub_aws_sdk_client["config_kwargs"]["endpoint_uri"] == "https://bedrock-runtime.us-west-2.amazonaws.com"
@pytest.mark.asyncio
async def test_session_update_is_acked_with_session_updated(self, stub_aws_models):
handler = BedrockRealtime()

6
tests/unit/llms/bedrock/test_base_aws_llm.py
View file

@ -30,7 +30,7 @@ from litellm.llms.bedrock.base_aws_llm import (
BaseAWSLLM,
Boto3CredentialsInfo,
run_aws_signing,
sign_request_off_loop_if_aws,
run_off_loop_if_aws,
)
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
from litellm.llms.bedrock.common_utils import BedrockModelInfo
@ -3443,7 +3443,7 @@ class TestGetRequestHeadersResign:
@pytest.mark.asyncio
async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh():
async def test_run_off_loop_if_aws_keeps_the_loop_serving_while_credentials_refresh() -> None:
"""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()
@ -3456,7 +3456,7 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent
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"})
signed = await run_off_loop_if_aws(BaseAWSLLM(), sign, headers={"Content-Type": "application/json"})
await release
assert "Authorization" in signed

27
tests/unit/llms/bedrock/test_claude_platform_provider.py
View file

@ -10,6 +10,7 @@ from pydantic import BaseModel, ConfigDict, ValidationError
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.router import GenericLiteLLMParams
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
WORKSPACE_ALIASES: Final = ("workspace_id", "aws_workspace_id", "anthropic_workspace_id", "anthropic-workspace-id")
@ -683,6 +684,32 @@ async def test_anthropic_messages_claude_platform_sends_exact_body_through_stric
_assert_gateway_accepted(requests[0], auth_kwargs)
@pytest.mark.asyncio
async def test_anthropic_messages_claude_platform_workspace_and_region_lookups_leave_the_event_loop_free(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import litellm
monkeypatch.setenv("ANTHROPIC_AWS_WORKSPACE_ID", "wrkspc_test")
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
vault: Final = host_in_loop_probing_vault(monkeypatch, "ANTHROPIC_AWS_WORKSPACE_ID", "AWS_REGION_NAME")
sigv4_without_region: Final = {key: value for key, value in SIGV4_KWARGS.items() if key != "aws_region_name"}
requests: list[dict] = []
response = await litellm.anthropic_messages(
model="bedrock/claude_platform/claude-sonnet-5-5",
messages=[{"role": "user", "content": "hello"}],
max_tokens=10,
client=_async_gateway_client(requests),
**sigv4_without_region,
)
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert response["content"][0]["text"] == "ok"
_assert_gateway_accepted(requests[0], sigv4_without_region)
def test_chat_completion_claude_platform_drops_context_management_and_gateway_accepts():
import litellm

6
tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
View file

@ -20,7 +20,7 @@ from botocore.awsrequest import AWSRequest
import litellm
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
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.llms.bedrock.base_aws_llm import run_off_loop_if_aws
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.types.utils import LlmProviders
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
@ -1234,9 +1234,7 @@ async def test_mantle_signing_runs_off_the_event_loop():
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"}
)
signed = await run_off_loop_if_aws(BedrockMantleChatConfig(), sign, headers={"Content-Type": "application/json"})
await release
assert "Authorization" in signed

153
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -5,13 +5,14 @@ import json
import logging
import threading
import time
from collections.abc import Mapping
from collections.abc import Awaitable, Callable, Mapping
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest
import respx
from botocore.credentials import RefreshableCredentials
import litellm
@ -66,6 +67,7 @@ from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIRe
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
_ACTIVE_KEY = "_code_interpreter_interception_active"
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
@ -1245,6 +1247,155 @@ async def test_async_anthropic_messages_handler_signs_bedrock_off_the_event_loop
assert probe.served_during_refresh is True
_BEDROCK_MESSAGE: Final = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-5-5",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 1},
}
_BEDROCK_RESPONSE: Final = {
"id": "resp_1",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "openai.gpt-5.6-sol",
"output": [],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
_EMPTY_S3_LISTING: Final = (
b'<?xml version="1.0" encoding="UTF-8"?>'
b'<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">'
b"<Name>batch-inputs</Name><KeyCount>0</KeyCount><IsTruncated>false</IsTruncated></ListBucketResult>"
)
_BATCH_LINE: Final = json.dumps(
{
"custom_id": "1",
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": "us.anthropic.claude-sonnet-5-5", "messages": [{"role": "user", "content": "hi"}]},
}
).encode()
_BATCH_FILE_ID: Final = "s3://batch-inputs/litellm-bedrock-files/batch.jsonl"
_BEDROCK_ENTRY_POINTS: Final[tuple[tuple[str, Callable[[], Awaitable[object]], httpx.Response], ...]] = (
(
"anthropic_messages",
lambda: litellm.anthropic.messages.acreate(
model="bedrock/invoke/us.anthropic.claude-sonnet-5-5",
messages=[{"role": "user", "content": "hi"}],
max_tokens=16,
),
httpx.Response(200, json=_BEDROCK_MESSAGE),
),
(
"responses",
lambda: litellm.aresponses(model="bedrock/us.openai.gpt-5.6-sol", input="hi"),
httpx.Response(200, json=_BEDROCK_RESPONSE),
),
(
"compact_responses",
lambda: litellm.acompact_responses(model="bedrock/us.openai.gpt-5.6-sol", input="hi"),
httpx.Response(200, json=_BEDROCK_RESPONSE),
),
(
"search",
lambda: litellm.asearch(
query="hi",
search_provider="agentcore",
api_base="https://gw.gateway.bedrock-agentcore.us-west-2.amazonaws.com/mcp",
),
httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": {"content": []}}),
),
(
"create_file",
lambda: litellm.acreate_file(
file=("batch.jsonl", _BATCH_LINE, "application/jsonl"),
purpose="batch",
custom_llm_provider="bedrock",
model="bedrock/us.anthropic.claude-sonnet-5-5",
),
httpx.Response(200),
),
(
"retrieve_file",
lambda: litellm.afile_retrieve(file_id=_BATCH_FILE_ID, custom_llm_provider="bedrock"),
httpx.Response(200, content=_BATCH_LINE),
),
(
"delete_file",
lambda: litellm.afile_delete(file_id=_BATCH_FILE_ID, custom_llm_provider="bedrock"),
httpx.Response(204),
),
(
"list_files",
lambda: litellm.afile_list(custom_llm_provider="bedrock"),
httpx.Response(200, content=_EMPTY_S3_LISTING),
),
(
"file_content",
lambda: litellm.afile_content(file_id=_BATCH_FILE_ID, custom_llm_provider="bedrock"),
httpx.Response(200, content=_BATCH_LINE),
),
(
"create_batch",
lambda: litellm.acreate_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=_BATCH_FILE_ID,
custom_llm_provider="bedrock",
model="bedrock/us.anthropic.claude-sonnet-5-5",
aws_batch_role_arn="arn:aws:iam::123456789012:role/batch",
s3_output_bucket_name="batch-outputs",
),
httpx.Response(200, json={"jobArn": "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/batch"}),
),
(
"vector_store_search",
lambda: litellm.vector_stores.asearch(vector_store_id="KBID12345", query="hi", custom_llm_provider="bedrock"),
httpx.Response(200, json={"retrievalResults": []}),
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call", "upstream_answer"),
[pytest.param(call, answer, id=name) for name, call, answer in _BEDROCK_ENTRY_POINTS],
)
async def test_bedrock_secret_lookups_leave_the_event_loop_free(
call: Callable[[], Awaitable[object]],
upstream_answer: httpx.Response,
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTTESTTESTTEST")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret")
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "batch-inputs")
monkeypatch.setenv("AWS_S3_OUTPUT_BUCKET_NAME", "batch-outputs")
vault: Final = host_in_loop_probing_vault(
monkeypatch,
"AWS_REGION_NAME",
"AWS_BEDROCK_RUNTIME_ENDPOINT",
"AWS_BEARER_TOKEN_BEDROCK",
"AWS_S3_ENCRYPTION_KEY_ID",
"AWS_S3_BUCKET_OWNER",
"AGENTCORE_MCP_PROTOCOL_VERSION",
"AGENTCORE_SEARCH_TOOL_NAME",
"AGENTCORE_GATEWAY_TOKEN",
)
upstream: Final = respx_mock.route().mock(return_value=upstream_answer)
_ = await call()
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert ".us-west-2." in upstream.calls.last.request.url.host
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_passes_litellm_metadata():
"""Ensure litellm_metadata from kwargs is forwarded via update_from_kwargs.

42
tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -35,6 +35,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
)
from litellm.types.utils import CallTypes, ModelResponse
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
@pytest.fixture
@ -8063,3 +8064,44 @@ async def test_during_call_masking_rewrites_only_pii_block_in_content_list(
]
assert messages[1]["content"] == "Hello, how can I help you today?"
assert messages[2]["content"] == "who is the president of the united states?"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("guardrail_config", "upstream_answer"),
[
pytest.param(
{"guardrailIdentifier": "wf0hkdb5x07f", "guardrailVersion": "DRAFT"},
{"action": "NONE", "outputs": [], "assessments": []},
id="apply_guardrail",
),
pytest.param(
{"checks": {"contentFilter": {"categories": [{"category": "VIOLENCE"}]}}, "content_filter_threshold": 0.5},
{"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.1}]}}},
id="invoke_guardrail_checks",
),
],
)
async def test_guardrail_secret_lookups_leave_the_event_loop_free(
guardrail_config: dict[str, object],
upstream_answer: dict[str, object],
respx_mock: respx.MockRouter,
httpx_transport: None,
aws_test_credentials: None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
vault: Final = host_in_loop_probing_vault(monkeypatch, "AWS_BEARER_TOKEN_BEDROCK", "AWS_REGION_NAME")
guardrail: Final = BedrockGuardrail(**guardrail_config)
upstream: Final = respx_mock.route().mock(return_value=httpx.Response(200, json=upstream_answer))
_ = await guardrail.make_bedrock_api_request(
source="INPUT",
messages=[{"role": "user", "content": "hello"}],
request_data={"messages": []},
)
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert upstream.calls.last.request.url.host == "bedrock-runtime.us-west-2.amazonaws.com"

43
tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -35,6 +35,7 @@ from litellm.proxy.utils import ProxyLogging
from litellm.passthrough.main import raise_for_provider_status
from litellm.llms.bedrock.common_utils import BedrockError
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
BaseOpenAIPassThroughHandler,
@ -2519,6 +2520,48 @@ class TestBedrockAgentRuntimePassthroughVirtualKeyLeak:
assert upstream["x-request-id"] == "trace-2"
@pytest.mark.asyncio
async def test_bedrock_agent_runtime_region_lookup_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", SHARED_MASTER_KEY)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
for ambient in TestBedrockAgentRuntimePassthroughVirtualKeyLeak.AMBIENT_AWS_ENV:
monkeypatch.delenv(ambient, raising=False)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak")
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk")
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
vault: Final = host_in_loop_probing_vault(monkeypatch, "AWS_REGION_NAME")
async def receive() -> dict[str, object]:
return {"type": "http.request", "body": b'{"retrievalQuery": {"text": "hi"}}', "more_body": False}
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": "/bedrock/knowledgebases/KB1234567/retrieve",
"headers": [(b"content-type", b"application/json")],
"query_string": b"",
},
receive=receive,
)
with respx.mock(assert_all_called=True) as upstream:
retrieve: Final = upstream.post(
"https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/KB1234567/retrieve"
).mock(return_value=httpx.Response(200, json={"retrievalResults": []}))
_ = await bedrock_proxy_route(
endpoint="knowledgebases/KB1234567/retrieve",
request=request,
fastapi_response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-caller"),
)
assert vault.read.is_set()
assert not vault.stalled_the_loop.is_set()
assert retrieve.calls.last.request.headers["authorization"].startswith("AWS4-HMAC-SHA256")
class TestLLMPassthroughFactoryProxyRoute:
@pytest.mark.asyncio
async def test_llm_passthrough_factory_proxy_route_success(self):