diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index b914bb1e4aa..764fe003f2c 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -453,7 +453,10 @@ async def acompletion_with_mcp( ) # Make follow-up call with streaming - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(self.base_call_args) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + self.base_call_args, + original_stream_setting=True, + ) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = True # Ensure follow-up call doesn't trigger MCP handler again @@ -625,7 +628,10 @@ async def acompletion_with_mcp( ) # Make follow-up call with original stream setting - follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(base_call_args) + follow_up_call_args: Final = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params( + base_call_args, + original_stream_setting=stream, + ) follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index d72f335e89c..07e192e68b4 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -72,7 +72,9 @@ class LiteLLM_Proxy_MCP_Handler: """ @staticmethod - def _prepare_chained_call_params(params: Mapping[str, Any]) -> dict[str, Any]: + def _prepare_chained_call_params( + params: Mapping[str, Any], + ) -> dict[str, Any]: # mutable-ok: returns a sanitized copy for the next provider call """Copy request params without state owned by the previous LLM call. MCP auto-execution keeps the trace identifier so chained rounds remain @@ -1264,7 +1266,7 @@ class LiteLLM_Proxy_MCP_Handler: Restores the original streaming setting and removes tool_choice since we're now providing tool results, not requesting tool calls. """ - follow_up_params: Final = call_params.copy() + follow_up_params: Final = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(call_params) # Restore original streaming setting for follow-up call follow_up_params["stream"] = original_stream_setting diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index c552170a599..de889421e1c 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -444,12 +444,13 @@ async def test_execute_tool_calls_uses_unique_call_ids_and_preserves_parent_cont assert all(item["metadata"]["parent_litellm_call_id"] == "cid" for item in captured) -def test_prepare_chained_call_params_resets_per_call_logging_state(): +def test_prepare_follow_up_call_params_resets_per_call_logging_state(): original = { "model": "gpt-4", "litellm_call_id": "parent-call", "litellm_logging_obj": object(), "litellm_trace_id": "trace-1", + "tool_choice": "auto", "litellm_params": { "litellm_call_id": "nested-parent-call", "litellm_logging_obj": object(), @@ -457,13 +458,15 @@ def test_prepare_chained_call_params_resets_per_call_logging_state(): }, } - follow_up = LiteLLM_Proxy_MCP_Handler._prepare_chained_call_params(original) + follow_up = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(original, original_stream_setting=True) assert "litellm_call_id" not in follow_up assert "litellm_logging_obj" not in follow_up assert "litellm_call_id" not in follow_up["litellm_params"] assert "litellm_logging_obj" not in follow_up["litellm_params"] assert follow_up["litellm_trace_id"] == "trace-1" + assert follow_up["stream"] is True + assert "tool_choice" not in follow_up assert follow_up["litellm_params"]["metadata"] == {"team": "legal"} assert original["litellm_call_id"] == "parent-call"