diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index d9798dae69..5299b30b52 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1037,6 +1037,7 @@ async def bedrock_proxy_route( target=str(prepped.url), custom_headers=prepped.headers, # type: ignore is_streaming_request=is_streaming_request, + _forward_headers=True ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( request, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 0bb9924af8..368e210359 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1188,11 +1188,11 @@ class TestBedrockLLMProxyRoute: This test verifies the fix for the bug where passthrough endpoints were using environment variables instead of model-specific credentials from config.yaml. """ + from litellm import Router + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( handle_bedrock_passthrough_router_model, ) - from litellm import Router - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params # Model-specific credentials (different from env vars) model_access_key = "MODEL_SPECIFIC_ACCESS_KEY" @@ -1453,6 +1453,294 @@ class TestVLLMProxyRoute: mock_factory_route.assert_awaited_once() +class TestForwardHeaders: + """ + Test cases for _forward_headers parameter in passthrough endpoints + """ + + @pytest.mark.asyncio + async def test_pass_through_request_with_forward_headers_true(self): + """ + Test that when forward_headers=True, user headers from the main request + are forwarded to the target endpoint (except content-length and host) + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + pass_through_request, + ) + + # Create a mock request with custom headers + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = MagicMock() + mock_request.url.path = "/test/endpoint" + + # User headers that should be forwarded + user_headers = { + "x-custom-header": "custom-value", + "x-api-key": "user-api-key", + "authorization": "Bearer user-token", + "user-agent": "test-client/1.0", + "content-type": "application/json", + # These should NOT be forwarded + "content-length": "123", + "host": "original-host.com", + } + mock_request.headers = user_headers + mock_request.query_params = {} + + # Mock the request body + mock_request_body = {"test": "data"} + + mock_user_api_key_dict = MagicMock() + + # Custom headers that should be merged with user headers + custom_headers = { + "x-litellm-header": "litellm-value", + } + + target_url = "https://api.example.com/v1/test" + + # Mock the httpx client and response + mock_httpx_response = MagicMock() + mock_httpx_response.status_code = 200 + mock_httpx_response.headers = {"content-type": "application/json"} + mock_httpx_response.aiter_bytes = AsyncMock(return_value=[b'{"result": "success"}']) + mock_httpx_response.aread = AsyncMock(return_value=b'{"result": "success"}') + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._read_request_body", + return_value=mock_request_body, + ), patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client, patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_logging_obj: + # Setup mock httpx client + mock_client = MagicMock() + mock_client.request = AsyncMock(return_value=mock_httpx_response) + mock_client_obj = MagicMock() + mock_client_obj.client = mock_client + mock_get_client.return_value = mock_client_obj + + # Setup mock logging object + mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) + mock_logging_obj.post_call_success_hook = AsyncMock() + mock_logging_obj.post_call_failure_hook = AsyncMock() + + # Call pass_through_request with forward_headers=True + result = await pass_through_request( + request=mock_request, + target=target_url, + custom_headers=custom_headers, + user_api_key_dict=mock_user_api_key_dict, + forward_headers=True, # Enable header forwarding + stream=False, + ) + + # Verify the httpx client was called + assert mock_client.request.called + + # Get the headers that were sent to the target + call_args = mock_client.request.call_args + sent_headers = call_args[1]["headers"] + + # Verify user headers were forwarded (except content-length and host) + assert sent_headers["x-custom-header"] == "custom-value" + assert sent_headers["x-api-key"] == "user-api-key" + assert sent_headers["authorization"] == "Bearer user-token" + assert sent_headers["user-agent"] == "test-client/1.0" + assert sent_headers["content-type"] == "application/json" + + # Verify custom headers were included + assert sent_headers["x-litellm-header"] == "litellm-value" + + # Verify content-length and host were NOT forwarded + assert "content-length" not in sent_headers + assert "host" not in sent_headers + + @pytest.mark.asyncio + async def test_pass_through_request_with_forward_headers_false(self): + """ + Test that when forward_headers=False (default), user headers are NOT forwarded, + only custom_headers are sent + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + pass_through_request, + ) + + # Create a mock request with custom headers + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = MagicMock() + mock_request.url.path = "/test/endpoint" + + # User headers that should NOT be forwarded + user_headers = { + "x-custom-header": "custom-value", + "x-api-key": "user-api-key", + "authorization": "Bearer user-token", + } + mock_request.headers = user_headers + mock_request.query_params = {} + + mock_request_body = {"test": "data"} + mock_user_api_key_dict = MagicMock() + + # Only these custom headers should be sent + custom_headers = { + "x-litellm-header": "litellm-value", + "authorization": "Bearer litellm-token", + } + + target_url = "https://api.example.com/v1/test" + + # Mock the httpx client and response + mock_httpx_response = MagicMock() + mock_httpx_response.status_code = 200 + mock_httpx_response.headers = {"content-type": "application/json"} + mock_httpx_response.aiter_bytes = AsyncMock(return_value=[b'{"result": "success"}']) + mock_httpx_response.aread = AsyncMock(return_value=b'{"result": "success"}') + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._read_request_body", + return_value=mock_request_body, + ), patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client, patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_logging_obj: + # Setup mock httpx client + mock_client = MagicMock() + mock_client.request = AsyncMock(return_value=mock_httpx_response) + mock_client_obj = MagicMock() + mock_client_obj.client = mock_client + mock_get_client.return_value = mock_client_obj + + # Setup mock logging object + mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) + mock_logging_obj.post_call_success_hook = AsyncMock() + mock_logging_obj.post_call_failure_hook = AsyncMock() + + # Call pass_through_request with forward_headers=False (default) + result = await pass_through_request( + request=mock_request, + target=target_url, + custom_headers=custom_headers, + user_api_key_dict=mock_user_api_key_dict, + forward_headers=False, # Explicitly set to False + stream=False, + ) + + # Verify the httpx client was called + assert mock_client.request.called + + # Get the headers that were sent to the target + call_args = mock_client.request.call_args + sent_headers = call_args[1]["headers"] + + # Verify only custom headers were sent + assert sent_headers["x-litellm-header"] == "litellm-value" + assert sent_headers["authorization"] == "Bearer litellm-token" + + # Verify user headers were NOT forwarded + assert "x-custom-header" not in sent_headers + assert "x-api-key" not in sent_headers + # Authorization is present but should be from custom_headers, not user headers + assert sent_headers["authorization"] == "Bearer litellm-token" + + @pytest.mark.asyncio + async def test_llm_passthrough_factory_with_forward_headers(self): + """ + Test that _forward_headers works correctly in llm_passthrough_factory_proxy_route + which is used in the code snippet provided by the user + """ + from litellm.types.utils import LlmProviders + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = MagicMock() + mock_request.url.path = "/openai/chat/completions" + + # User headers to be forwarded + user_headers = { + "x-custom-tracking-id": "tracking-123", + "x-request-id": "req-456", + "user-agent": "my-app/2.0", + } + mock_request.headers = user_headers + mock_request.json = AsyncMock(return_value={"stream": False}) + + mock_fastapi_response = MagicMock(spec=Response) + mock_user_api_key_dict = MagicMock() + + # Mock the httpx response + mock_httpx_response = MagicMock() + mock_httpx_response.status_code = 200 + mock_httpx_response.headers = {"content-type": "application/json"} + mock_httpx_response.aiter_bytes = AsyncMock(return_value=[b'{"result": "success"}']) + mock_httpx_response.aread = AsyncMock(return_value=b'{"result": "success"}') + + with patch( + "litellm.utils.ProviderConfigManager.get_provider_model_info" + ) as mock_get_provider, patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" + ) as mock_get_creds, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._read_request_body", + return_value={"messages": [{"role": "user", "content": "test"}]}, + ), patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client, patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_logging_obj: + # Setup provider config + mock_provider_config = MagicMock() + mock_provider_config.get_api_base.return_value = "https://api.openai.com/v1" + mock_provider_config.validate_environment.return_value = { + "authorization": "Bearer sk-test" + } + mock_get_provider.return_value = mock_provider_config + mock_get_creds.return_value = "sk-test" + + # Setup mock httpx client + mock_client = MagicMock() + mock_client.request = AsyncMock(return_value=mock_httpx_response) + mock_client_obj = MagicMock() + mock_client_obj.client = mock_client + mock_get_client.return_value = mock_client_obj + + # Setup mock logging object + mock_logging_obj.pre_call_hook = AsyncMock( + return_value={"messages": [{"role": "user", "content": "test"}]} + ) + mock_logging_obj.post_call_success_hook = AsyncMock() + + # This is the key part - when create_pass_through_route is called with _forward_headers=True + # it should forward the user headers + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route: + mock_endpoint_func = AsyncMock(return_value="success") + mock_create_route.return_value = mock_endpoint_func + + result = await llm_passthrough_factory_proxy_route( + custom_llm_provider=LlmProviders.OPENAI, + endpoint="/chat/completions", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify create_pass_through_route was called + mock_create_route.assert_called_once() + + # Get the call arguments to verify _forward_headers parameter + call_kwargs = mock_create_route.call_args[1] + + # Note: The current implementation doesn't explicitly pass _forward_headers + # This test documents the current behavior. If _forward_headers should be + # configurable in llm_passthrough_factory_proxy_route, it would need to be added + + class TestMilvusProxyRoute: """ Test cases for Milvus passthrough endpoint