From c1369a07baa33dc59ec065788bca7c9e07a85673 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Fri, 31 Oct 2025 08:40:17 +0530
Subject: [PATCH] Add Add per model group header forwarding for Bedrock Invoke
 API (#16042)

---
 litellm/llms/custom_httpx/llm_http_handler.py |  7 +-
 .../test_anthropic_messages_passthrough.py    | 83 +++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index b67e8823d0..0a01e9cd68 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -1774,11 +1774,16 @@ class BaseLLMHTTPHandler:
             provider_specific_header=provider_specific_header,
             custom_llm_provider=custom_llm_provider,
         )
+        forwarded_headers = kwargs.get("headers", None)
+        if forwarded_headers and extra_headers:
+            merged_headers = {**forwarded_headers, **extra_headers}
+        else:
+            merged_headers = forwarded_headers or extra_headers
         (
             headers,
             api_base,
         ) = anthropic_messages_provider_config.validate_anthropic_messages_environment(
-            headers=extra_headers or {},
+            headers=merged_headers or {},
             model=model,
             messages=messages,
             optional_params=anthropic_messages_optional_request_params,
diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py
index ee73efb4a3..17e72f2915 100644
--- a/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py
+++ b/tests/pass_through_unit_tests/test_anthropic_messages_passthrough.py
@@ -581,6 +581,89 @@ async def test_anthropic_messages_with_extra_headers():
     return response
 
 
+@pytest.mark.asyncio
+async def test_bedrock_messages_api_header_forwarding():
+    """
+    Test that headers from kwargs (set by proxy's add_headers_to_llm_call_by_model_group)
+    are correctly passed to validate_anthropic_messages_environment for Bedrock Invoke API.
+    
+    This verifies that forward_client_headers_to_llm_api works for Bedrock Invoke API (Messages API).
+    
+    Issue: When calling Anthropic models via the Messages API, LiteLLM makes a call to 
+    Bedrock's Invoke API, and custom headers were not being forwarded, even though
+    they worked correctly for Chat Completions API with Bedrock's Converse API.
+    """
+    from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.types.router import GenericLiteLLMParams
+    
+    handler = BaseLLMHTTPHandler()
+    
+    # Headers that would be set by the proxy when forward_client_headers_to_llm_api is configured
+    custom_headers = {
+        "X-Custom-Header": "CustomValue",
+        "X-Request-ID": "req-123",
+    }
+    
+    # Mock the provider config
+    mock_provider_config = MagicMock()
+    
+    # We'll check what headers are passed to this method
+    mock_provider_config.validate_anthropic_messages_environment.return_value = (
+        {"Authorization": "Bearer test"},
+        "https://bedrock-runtime.us-east-1.amazonaws.com/invoke"
+    )
+    mock_provider_config.transform_anthropic_messages_request.return_value = {"model": "test"}
+    mock_provider_config.get_complete_url.return_value = "https://test.com"
+    mock_provider_config.sign_request.return_value = ({}, None)
+    mock_provider_config.transform_anthropic_messages_response.return_value = {"id": "test"}
+    
+    # Mock HTTP client to prevent actual network calls
+    with unittest.mock.patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client") as mock_get_client:
+        mock_http_client = AsyncMock()
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"id": "test", "content": []}
+        mock_response.text = "{}"
+        mock_http_client.post.return_value = mock_response
+        mock_get_client.return_value = mock_http_client
+        
+        # Mock logging object
+        mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
+        mock_logging_obj.model_call_details = {}
+        
+        # Call the handler with headers in kwargs
+        try:
+            await handler.async_anthropic_messages_handler(
+                model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
+                messages=[{"role": "user", "content": "Hello"}],
+                anthropic_messages_provider_config=mock_provider_config,
+                anthropic_messages_optional_request_params={"max_tokens": 100},
+                custom_llm_provider="bedrock",
+                litellm_params=GenericLiteLLMParams(
+                    api_key="test-key",
+                    aws_region_name="us-east-1"
+                ),
+                logging_obj=mock_logging_obj,
+                api_key="test-key",
+                stream=False,
+                kwargs={"headers": custom_headers}  # Headers set by proxy
+            )
+        except Exception:
+            pass  # Ignore errors, we're only checking if headers were passed
+        
+        # Verify that validate_anthropic_messages_environment was called
+        assert mock_provider_config.validate_anthropic_messages_environment.called
+        
+        # Get the headers that were passed
+        call_args = mock_provider_config.validate_anthropic_messages_environment.call_args
+        passed_headers = call_args[1]["headers"]
+        
+        # The custom headers from kwargs should be in the passed headers
+        assert "X-Custom-Header" in passed_headers or "x-custom-header" in passed_headers
+        assert "X-Request-ID" in passed_headers or "x-request-id" in passed_headers
+
+
 @pytest.mark.asyncio
 async def test_anthropic_messages_with_thinking():
     """