diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index f1426aa9375..9591b962ab5 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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) diff --git a/litellm/llms/bedrock/claude_platform/messages_transformation.py b/litellm/llms/bedrock/claude_platform/messages_transformation.py index 1469f6a6935..57ebaa0ca71 100644 --- a/litellm/llms/bedrock/claude_platform/messages_transformation.py +++ b/litellm/llms/bedrock/claude_platform/messages_transformation.py @@ -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, diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 3bae1f22490..5f8f478e4cd 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -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, diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 4b1a6a75a7a..14edd39c386 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -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"], diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3e7c2f7eee9..e9b2f26e977 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -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: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 48858dfa441..b1e55375f16 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 12b8e728b4e..adda465ad95 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index d9d3b6f8c92..16fbd1de376 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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: diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index cb97a62f635..5dadeff6d7c 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -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() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py index 71b451f834c..2d21a3ddef5 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_mantle_count_tokens_handler.py @@ -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 diff --git a/tests/unit/llms/bedrock/event_loop_probe.py b/tests/unit/llms/bedrock/event_loop_probe.py index c347247ec32..cf6cd5975ed 100644 --- a/tests/unit/llms/bedrock/event_loop_probe.py +++ b/tests/unit/llms/bedrock/event_loop_probe.py @@ -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 diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index a181c5a5425..ccea4ba44d7 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -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() diff --git a/tests/unit/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py index c755a7e6bc2..4350fdb699b 100644 --- a/tests/unit/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -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 diff --git a/tests/unit/llms/bedrock/test_claude_platform_provider.py b/tests/unit/llms/bedrock/test_claude_platform_provider.py index 5d83cba5415..e43184df50e 100644 --- a/tests/unit/llms/bedrock/test_claude_platform_provider.py +++ b/tests/unit/llms/bedrock/test_claude_platform_provider.py @@ -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 diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index ee41af2110d..f252d9886d3 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -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 diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 8b46597b603..205c610a413 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -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'' + b'' + b"batch-inputs0false" +) +_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. diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 3b33ba75eb3..75d713028d4 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -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" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 4dbe1f2ef8d..66c52a76318 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -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):