From ab4b6197effe73b2e38dfe5b214b9c83cad90173 Mon Sep 17 00:00:00 2001 From: The Mavik <179817126+themavik@users.noreply.github.com> Date: Fri, 13 Feb 2026 19:44:40 -0500 Subject: [PATCH] fix: add custom_body parameter to endpoint_func in create_pass_through_route (#20849) * fix: add custom_body parameter to endpoint_func in create_pass_through_route The bedrock_proxy_route calls `endpoint_func(custom_body=data)` to pass a pre-parsed, SigV4-signed request body. However, the `endpoint_func` closure created by `create_pass_through_route` does not accept a `custom_body` keyword argument, causing: TypeError: endpoint_func() got an unexpected keyword argument 'custom_body' Add `custom_body: Optional[dict] = None` to both `endpoint_func` definitions (adapter-based and URL-based). In the URL-based path, when `custom_body` is provided by the caller, use it instead of re-parsing the body from the raw request. Fixes #16999 * Add tests for custom_body handling in create_pass_through_route Address reviewer feedback on PR #20849: - Document why the adapter-based endpoint_func accepts custom_body for signature compatibility but does not forward it (the underlying chat_completion_pass_through_endpoint does not support it). - Add test_create_pass_through_route_custom_body_url_target: verifies that when a caller (e.g. bedrock_proxy_route) supplies custom_body, it takes precedence over the body parsed from the raw request. - Add test_create_pass_through_route_no_custom_body_falls_back: verifies that the default path (no custom_body) correctly uses the request-parsed body, preserving existing behavior. Both tests are fully mocked following the project's CONTRIBUTING.md guidelines and the patterns established in the existing test file. Co-authored-by: Cursor --------- Co-authored-by: themavik Co-authored-by: Cursor --- .../pass_through_endpoints.py | 17 ++- .../test_pass_through_endpoints.py | 137 ++++++++++++++++++ 2 files changed, 149 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a7b60c8b185..fb587975ed5 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1099,6 +1099,7 @@ def create_pass_through_route( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), subpath: str = "", # captures sub-paths when include_subpath=True + custom_body: Optional[dict] = None, # accepted for signature compatibility with URL-based path; not forwarded because chat_completion_pass_through_endpoint does not support it ): return await chat_completion_pass_through_endpoint( fastapi_response=fastapi_response, @@ -1115,6 +1116,7 @@ def create_pass_through_route( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), subpath: str = "", # captures sub-paths when include_subpath=True + custom_body: Optional[dict] = None, ): from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, @@ -1189,11 +1191,16 @@ def create_pass_through_route( ) if query_params: final_query_params.update(query_params) - final_custom_body = ( - custom_body_data - if isinstance(custom_body_data, dict) or custom_body_data is None - else None - ) + # When a caller (e.g. bedrock_proxy_route) supplies a pre-built + # body, use it instead of the body parsed from the raw request. + if custom_body is not None: + final_custom_body = custom_body + else: + final_custom_body = ( + custom_body_data + if isinstance(custom_body_data, dict) or custom_body_data is None + else None + ) return await pass_through_request( # type: ignore request=request, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e50e10352e2..7ec97ddc185 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2087,6 +2087,143 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata(): assert "headers" in result["proxy_server_request"] +@pytest.mark.asyncio +async def test_create_pass_through_route_custom_body_url_target(): + """ + Test that the URL-based endpoint_func created by create_pass_through_route + accepts a custom_body parameter and forwards it to pass_through_request, + taking precedence over the request-parsed body. + + This verifies the fix for issue #16999 where bedrock_proxy_route passes + custom_body=data to the endpoint function, which previously crashed with: + TypeError: endpoint_func() got an unexpected keyword argument 'custom_body' + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + create_pass_through_route, + ) + + unique_path = "/test/path/unique/custom_body_url" + endpoint_func = create_pass_through_route( + endpoint=unique_path, + target="https://bedrock-agent-runtime.us-east-1.amazonaws.com", + custom_headers={"Content-Type": "application/json"}, + _forward_headers=True, + ) + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" + ) as mock_pass_through, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" + ) as mock_is_registered, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.get_registered_pass_through_route" + ) as mock_get_registered, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._parse_request_data_by_content_type" + ) as mock_parse_request: + mock_pass_through.return_value = MagicMock() + mock_is_registered.return_value = True + mock_get_registered.return_value = None + # Simulate the request parser returning a different body + mock_parse_request.return_value = ( + {}, # query_params_data + {"parsed_from_request": True}, # custom_body_data (from request) + None, # file_data + False, # stream + ) + + mock_request = MagicMock(spec=Request) + mock_request.url = MagicMock() + mock_request.url.path = unique_path + mock_request.path_params = {} + mock_request.query_params = QueryParams({}) + + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.api_key = "test-key" + + # The caller-supplied body (e.g. from bedrock_proxy_route) + bedrock_body = { + "retrievalQuery": {"text": "What is in the knowledge base?"}, + } + + # Call endpoint_func with custom_body — this is the call that + # used to crash with TypeError before the fix + await endpoint_func( + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=mock_user_api_key_dict, + custom_body=bedrock_body, + ) + + mock_pass_through.assert_called_once() + call_kwargs = mock_pass_through.call_args[1] + + # The critical assertion: custom_body takes precedence over + # the body parsed from the raw request + assert call_kwargs["custom_body"] == bedrock_body + + +@pytest.mark.asyncio +async def test_create_pass_through_route_no_custom_body_falls_back(): + """ + Test that the URL-based endpoint_func falls back to the request-parsed body + when custom_body is not provided. + + This ensures the default pass-through behavior is preserved — only the + Bedrock proxy route (and similar callers) supply a pre-built body. + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + create_pass_through_route, + ) + + unique_path = "/test/path/unique/no_custom_body" + endpoint_func = create_pass_through_route( + endpoint=unique_path, + target="http://example.com/api", + custom_headers={}, + ) + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_request" + ) as mock_pass_through, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.is_registered_pass_through_route" + ) as mock_is_registered, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.InitPassThroughEndpointHelpers.get_registered_pass_through_route" + ) as mock_get_registered, patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._parse_request_data_by_content_type" + ) as mock_parse_request: + mock_pass_through.return_value = MagicMock() + mock_is_registered.return_value = True + mock_get_registered.return_value = None + request_parsed_body = {"key": "from_request"} + mock_parse_request.return_value = ( + {}, # query_params_data + request_parsed_body, # custom_body_data + None, # file_data + False, # stream + ) + + mock_request = MagicMock(spec=Request) + mock_request.url = MagicMock() + mock_request.url.path = unique_path + mock_request.path_params = {} + mock_request.query_params = QueryParams({}) + + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.api_key = "test-key" + + # Call without custom_body — should use the request-parsed body + await endpoint_func( + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=mock_user_api_key_dict, + ) + + mock_pass_through.assert_called_once() + call_kwargs = mock_pass_through.call_args[1] + + # Should fall back to the body parsed from the request + assert call_kwargs["custom_body"] == request_parsed_body + + def test_build_full_path_with_root_default(): """ Test _build_full_path_with_root with default root path (/)