diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 527c43cd11f..3e0092fd376 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -213,12 +213,13 @@ async def test_amessages_wrapper_forwards_args(): custom_llm_provider="azure_ai", extra_headers=None, timeout=12.5, + aws_region_name="us-east-1", ) assert response == FAKE_MESSAGES_RESPONSE assert bridge.calls[0]["model"] == "claude-sonnet-4-5" assert bridge.calls[0]["timeout_seconds"] == 12.5 - assert bridge.calls[0]["aws_region_name"] is None + assert bridge.calls[0]["aws_region_name"] == "us-east-1" def _gate(**overrides): @@ -257,6 +258,31 @@ async def test_gate_invokes_rust_and_marks_response_header(): assert call["aws_region_name"] is None +@pytest.mark.asyncio +async def test_gate_invokes_rust_for_bedrock_and_forwards_region(): + bridge = RecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + + response = await _gate( + custom_llm_provider="bedrock", + litellm_params=GenericLiteLLMParams( + api_key="bedrock-token", + aws_region_name="us-east-1", + rust=True, + ), + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + api_key="bedrock-token", + api_base=None, + headers={}, + ) + + assert response is not None + assert response["_hidden_params"]["additional_headers"] == {"x-litellm-rust": "true"} + assert bridge.calls[0]["model"] == "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0" + assert bridge.calls[0]["custom_llm_provider"] == "bedrock" + assert bridge.calls[0]["aws_region_name"] == "us-east-1" + + @pytest.mark.asyncio async def test_gate_falls_back_to_python_when_bridge_raises(): bridge = RaisingAsyncMessages()