From 8d43894b25b3a48378923442b0edea6683f2b7a0 Mon Sep 17 00:00:00 2001 From: nvr Date: Thu, 9 Apr 2026 05:09:35 -0700 Subject: [PATCH] Address MCP chat review follow-ups Expose the non-streaming MCP auto-execution round limit as litellm.max_mcp_auto_execution_rounds with an env-backed default. Keep MCP metadata consistent by only aggregating tool calls after execution returns results for that round. Add coverage for no-tool-call responses, empty execution results, and the configurable round-limit path. --- litellm/__init__.py | 3 + .../responses/mcp/chat_completions_handler.py | 21 +- .../mcp/test_chat_completions_handler.py | 330 ++++++++++++++++++ 3 files changed, 350 insertions(+), 4 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index d4418c661a3..350e3f0186a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -213,6 +213,9 @@ token: Optional[ ] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults +max_mcp_auto_execution_rounds: int = int( + os.getenv("LITELLM_MAX_MCP_AUTO_EXECUTION_ROUNDS", 6) +) drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) use_chat_completions_url_for_anthropic_messages: bool = bool( diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a8843a8ea9c..833de34c86a 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -15,7 +15,21 @@ from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -MAX_CHAT_MCP_AUTO_EXECUTION_ROUNDS = 6 +DEFAULT_CHAT_MCP_AUTO_EXECUTION_ROUNDS = 6 + + +def _get_max_chat_mcp_auto_execution_rounds() -> int: + import litellm + + configured_rounds = getattr( + litellm, + "max_mcp_auto_execution_rounds", + DEFAULT_CHAT_MCP_AUTO_EXECUTION_ROUNDS, + ) + try: + return max(1, int(configured_rounds)) + except (TypeError, ValueError): + return DEFAULT_CHAT_MCP_AUTO_EXECUTION_ROUNDS def _add_mcp_metadata_to_response( @@ -615,7 +629,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 aggregated_tool_calls: List[Any] = [] aggregated_tool_results: List[Any] = [] - for _round_idx in range(MAX_CHAT_MCP_AUTO_EXECUTION_ROUNDS): + for _round_idx in range(_get_max_chat_mcp_auto_execution_rounds()): if not isinstance(response, ModelResponse): return response @@ -632,8 +646,6 @@ async def acompletion_with_mcp( # noqa: PLR0915 ) return response - aggregated_tool_calls.extend(tool_calls) - tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( tool_server_map=tool_server_map, tool_calls=tool_calls, @@ -655,6 +667,7 @@ async def acompletion_with_mcp( # noqa: PLR0915 ) return response + aggregated_tool_calls.extend(tool_calls) aggregated_tool_results.extend(tool_results) current_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( original_messages=current_messages, diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index e1c203ecd55..485c79f0f8f 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1,6 +1,7 @@ import pytest from unittest.mock import AsyncMock, patch +import litellm from litellm.types.utils import ModelResponse from litellm.responses.mcp import chat_completions_handler @@ -531,6 +532,335 @@ async def test_acompletion_with_mcp_non_streaming_auto_exec_supports_multiple_ro assert provider_fields["mcp_call_results"][1]["tool_call_id"] == "call-2" +@pytest.mark.asyncio +async def test_acompletion_with_mcp_non_streaming_returns_response_when_no_tool_calls( + monkeypatch, +): + tools = [ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "require_approval": "never", + } + ] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + response_without_tool_calls = ModelResponse( + id="resp-1", + model="test-model", + created=1234567890, + object="chat.completion", + choices=[ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "final answer", + }, + } + ], + ) + + mock_acompletion = AsyncMock(return_value=response_without_tool_calls) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_use_litellm_mcp_gateway", + staticmethod(lambda tools: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_parse_mcp_tools", + staticmethod(lambda tools: (tools, [])), + ) + + async def mock_process(**_): + return (tools, {"local_search": "local"}) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda *_, **__: openai_tools), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: True), + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + with patch("litellm.acompletion", mock_acompletion), patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion, + create=True, + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=False, + tool_choice="required", + ) + + assert isinstance(result, ModelResponse) + provider_fields = result.choices[0].message.provider_specific_fields + assert provider_fields is not None + assert provider_fields["mcp_list_tools"] == openai_tools + assert "mcp_tool_calls" not in provider_fields + assert "mcp_call_results" not in provider_fields + mock_acompletion.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_non_streaming_omits_unexecuted_tool_calls_from_metadata( + monkeypatch, +): + tools = [ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "require_approval": "never", + } + ] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_call = { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + response_with_tool_call = ModelResponse( + id="resp-1", + model="test-model", + created=1234567890, + object="chat.completion", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [tool_call], + }, + } + ], + ) + + mock_acompletion = AsyncMock(return_value=response_with_tool_call) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_use_litellm_mcp_gateway", + staticmethod(lambda tools: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_parse_mcp_tools", + staticmethod(lambda tools: (tools, [])), + ) + + async def mock_process(**_): + return (tools, {"local_search": "local"}) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda *_, **__: openai_tools), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: True), + ) + + async def mock_execute(**_): + return [] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + mock_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + with patch("litellm.acompletion", mock_acompletion), patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion, + create=True, + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=False, + tool_choice="required", + ) + + assert isinstance(result, ModelResponse) + provider_fields = result.choices[0].message.provider_specific_fields + assert provider_fields is not None + assert provider_fields["mcp_list_tools"] == openai_tools + assert "mcp_tool_calls" not in provider_fields + assert "mcp_call_results" not in provider_fields + mock_acompletion.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_non_streaming_uses_configured_round_limit( + monkeypatch, +): + tools = [ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "require_approval": "never", + } + ] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + first_tool_call = { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + second_tool_call = { + "id": "call-2", + "type": "function", + "function": {"name": "local_search", "arguments": '{"phase":"second"}'}, + } + first_response = ModelResponse( + id="resp-1", + model="test-model", + created=1234567890, + object="chat.completion", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [first_tool_call], + }, + } + ], + ) + second_response = ModelResponse( + id="resp-2", + model="test-model", + created=1234567891, + object="chat.completion", + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": "", + "tool_calls": [second_tool_call], + }, + } + ], + ) + mock_acompletion = AsyncMock(side_effect=[first_response, second_response]) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_use_litellm_mcp_gateway", + staticmethod(lambda tools: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_parse_mcp_tools", + staticmethod(lambda tools: (tools, [])), + ) + + async def mock_process(**_): + return (tools, {"local_search": "local"}) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda *_, **__: openai_tools), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: True), + ) + + async def mock_execute(**kwargs): + tool_id = kwargs["tool_calls"][0]["id"] + return [ + { + "tool_call_id": tool_id, + "name": "local_search", + "result": f"result-for-{tool_id}", + } + ] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + mock_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + monkeypatch.setattr(litellm, "max_mcp_auto_execution_rounds", 1) + + with patch("litellm.acompletion", mock_acompletion), patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion, + create=True, + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=False, + tool_choice="required", + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.tool_calls is not None + assert result.choices[0].message.tool_calls[0].id == "call-2" + assert mock_acompletion.await_count == 2 + provider_fields = result.choices[0].message.provider_specific_fields + assert provider_fields is not None + assert len(provider_fields["mcp_tool_calls"]) == 1 + assert len(provider_fields["mcp_call_results"]) == 1 + assert provider_fields["mcp_tool_calls"][0]["id"] == "call-1" + assert provider_fields["mcp_call_results"][0]["tool_call_id"] == "call-1" + + @pytest.mark.asyncio async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): """