From 29c23dbe2df7f69da447a9eb1e13461e96860bbb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 3 Jul 2026 21:28:45 -0700 Subject: [PATCH] fix(responses/mcp): execute follow-up tool calls instead of dropping the stream MCPEnhancedStreamingIterator only auto-executed one round of MCP tool calls. When a model retried a tool (e.g. after an error) in its follow-up turn, that second tool call was streamed but never executed, and the response ended with no final text. Route follow-up calls back through the same completion-check phase as the initial response, so further tool-call rounds are handled the same way, capped at MAX_MCP_TOOL_CALL_ROUNDS to avoid an unbounded loop. --- .../responses/mcp/mcp_streaming_iterator.py | 80 +++++--- .../mcp/test_mcp_streaming_iterator.py | 171 ++++++++++++++++++ 2 files changed, 226 insertions(+), 25 deletions(-) create mode 100644 tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 3c24a703b68..c705a04963c 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -24,6 +24,8 @@ if TYPE_CHECKING: else: MCPTool = Any +MAX_MCP_TOOL_CALL_ROUNDS = 5 + async def create_mcp_list_tools_events( mcp_tools_with_litellm_proxy: List[ToolParam], @@ -265,9 +267,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.should_auto_execute = self._should_auto_execute_tools() # Streaming state management - self.phase = ( - "initial_response" # initial_response -> mcp_discovery -> tool_execution -> follow_up_response -> finished - ) + self.phase = "initial_response" # initial_response -> mcp_discovery -> (continue_initial_response <-> tool_execution) -> finished self.finished = False # Event queues and generation flags @@ -281,11 +281,23 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Iterator references self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = base_iterator # Will be created when needed - self.follow_up_iterator: Optional[Any] = None # Response collection for tool execution self.collected_response: Optional[ResponsesAPIResponse] = None + # Counts completed rounds of tool execution, so a model that keeps + # calling tools (e.g. retrying after an error) can't loop forever. + # Capped in _create_follow_up_iterator, which drops "tools" from the + # request once the cap is hit so the model must answer in text. + self.tool_call_round = 0 + # The collected_response that self.tool_results was computed from. + # _create_follow_up_iterator only builds a follow-up when this is + # still the current collected_response — otherwise a round whose + # response had no tool calls (e.g. the model finally answered in + # text) would incorrectly reuse tool_results left over from an + # earlier round and keep looping instead of finishing. + self._tool_results_for_response: Optional[ResponsesAPIResponse] = None + # Set up model metadata (will be updated when we get the real iterator) self.model = self.original_request_params.get("model", "unknown") self.litellm_metadata = {} @@ -376,10 +388,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): Phase-based streaming: 1. initial_response - Stream the first LLM response (includes response.created, response.in_progress, response.output_item.added) 2. mcp_discovery - Emit MCP discovery events (after response.output_item.added) - 3. continue_initial_response - Continue streaming the initial response content + 3. continue_initial_response - Stream the current round's response (initial or follow-up). + On completion, if auto-execute is on and the response contains tool calls, loops back + through tool_execution/follow-up instead of ending, up to MAX_MCP_TOOL_CALL_ROUNDS. 4. tool_execution - Emit tool execution events - 5. follow_up_response - Stream the follow-up response - 6. finished - End iteration + 5. finished - End iteration """ # Phase 1: Initial Response Stream (emit standard OpenAI events first) @@ -415,21 +428,17 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self.tool_execution_events: return self.tool_execution_events.pop(0) - # Move to follow-up response phase - self.phase = "follow_up_response" + # Route the follow-up call back through continue_initial_response so + # its completion is checked for further tool calls the same way the + # initial response is — otherwise a model that needs a second round + # of tool calls (e.g. retrying after an error) would have that round + # silently dropped and the stream would end with no final text. await self._create_follow_up_iterator() - - # Phase 5: Follow-up Response Stream - if self.phase == "follow_up_response": - if self.follow_up_iterator: - try: - return await cast(Any, self.follow_up_iterator).__anext__() # type: ignore[attr-defined] - except StopAsyncIteration: - self.phase = "finished" - raise - else: - self.phase = "finished" - raise StopAsyncIteration + if self.base_iterator is not None: + self.phase = "continue_initial_response" + return await self.__anext__() + self.phase = "finished" + raise StopAsyncIteration # Phase 6: Finished if self.phase == "finished": @@ -599,6 +608,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_calls = [] if not tool_calls: return + self.tool_call_round += 1 for tool_call in tool_calls: ( @@ -686,6 +696,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Store tool results for follow-up call self.tool_results = tool_results + self._tool_results_for_response = self.collected_response except Exception as e: verbose_logger.error(f"Error in tool execution: {e}") @@ -693,10 +704,16 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): traceback.print_exc() self.tool_results = [] + self._tool_results_for_response = self.collected_response async def _create_follow_up_iterator(self) -> None: """Create the follow-up response iterator with tool results""" - if not self.collected_response or not hasattr(self, "tool_results"): + if self.collected_response is None or self.collected_response is not self._tool_results_for_response: + # Either no response to follow up on, or the current round's + # response had no tool calls (self.tool_results is stale from an + # earlier round) — there is nothing to follow up with, so end + # the stream instead of reusing stale tool results. + self.base_iterator = None return from litellm.responses.main import aresponses @@ -726,18 +743,31 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Remove tool_choice to avoid forcing more tool calls follow_up_params.pop("tool_choice", None) + if self.tool_call_round >= MAX_MCP_TOOL_CALL_ROUNDS: + # Hit the round cap: drop tools entirely so the model must + # answer in text instead of emitting another (unexecuted) + # tool call that would otherwise end the stream in silence. + follow_up_params.pop("tools", None) + verbose_logger.warning( + "MCP auto-execute hit MAX_MCP_TOOL_CALL_ROUNDS=%s; forcing a text-only follow-up.", + MAX_MCP_TOOL_CALL_ROUNDS, + ) + follow_up_response = await aresponses(**follow_up_params) - # Set up the follow-up iterator + # Route the follow-up through the same base_iterator machinery as + # the initial call so its completion is checked for further tool + # calls (see phase 4 in __anext__). if hasattr(follow_up_response, "__aiter__"): - self.follow_up_iterator = follow_up_response + self.base_iterator = follow_up_response + self.collected_response = None except Exception as e: verbose_logger.error(f"Error creating follow-up iterator: {e}") import traceback traceback.print_exc() - self.follow_up_iterator = None + self.base_iterator = None def __iter__(self): return self diff --git a/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py new file mode 100644 index 00000000000..ecf89ce18d4 --- /dev/null +++ b/tests/test_litellm/responses/mcp/test_mcp_streaming_iterator.py @@ -0,0 +1,171 @@ +import sys +import types +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from mcp.types import CallToolResult, TextContent + +import litellm # noqa: F401 - ensures litellm.responses.main is registered in sys.modules +from litellm.responses.mcp.mcp_streaming_iterator import ( + MAX_MCP_TOOL_CALL_ROUNDS, + MCPEnhancedStreamingIterator, +) +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamEvents + +# `litellm.__init__` re-exports a function named `responses`, which shadows the +# `litellm.responses` subpackage as an attribute — `import litellm.responses.main` +# can resolve to the unrelated third-party `responses` package instead. Look the +# real submodule up in sys.modules directly to sidestep the shadowing. +responses_main_module = sys.modules["litellm.responses.main"] + + +class _FakeAsyncStream: + """Minimal async iterator yielding pre-built chunks, one per __anext__ call.""" + + def __init__(self, chunks): + self._chunks = list(chunks) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _output_item_added_chunk(): + return SimpleNamespace(type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED) + + +def _completed_chunk(output): + response = ResponsesAPIResponse(id="resp-1", created_at=0, output=output) + return SimpleNamespace(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=response) + + +def _function_call(call_id: str, name: str, arguments: str = "{}"): + return {"type": "function_call", "call_id": call_id, "name": name, "arguments": arguments} + + +def _text_message(text: str): + return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} + + +def _tool_call_stream(call_id: str, tool_name: str) -> _FakeAsyncStream: + return _FakeAsyncStream([_completed_chunk([_function_call(call_id, tool_name)])]) + + +def _text_only_stream(text: str) -> _FakeAsyncStream: + return _FakeAsyncStream([_completed_chunk([_text_message(text)])]) + + +def _mock_mcp_environment(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: + """Patch the MCP tool-call plumbing so _execute_tool_calls can run in tests.""" + call_tool = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)) + fake_manager = types.SimpleNamespace( + call_tool=call_tool, + _get_mcp_server_from_tool_name=MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + types.SimpleNamespace(proxy_logging_obj=MagicMock()), + ) + return call_tool + + +def _make_iterator(initial_chunks) -> MCPEnhancedStreamingIterator: + return MCPEnhancedStreamingIterator( + base_iterator=_FakeAsyncStream(initial_chunks), + mcp_events=[], + tool_server_map={"read_wiki_contents": "deepwiki"}, + mcp_tools_with_litellm_proxy=[{"require_approval": "never"}], + user_api_key_auth=None, + original_request_params={ + "model": "gpt-4", + "input": "what is berriai/litellm?", + "tools": [{"type": "mcp"}], + }, + ) + + +@pytest.mark.asyncio +async def test_second_round_tool_call_is_executed_and_reaches_final_text(monkeypatch): + """ + Regression test: a tool call that errors, followed by the model retrying + the tool in its follow-up turn, must have that second tool call executed + too — the stream should not silently end after round 1 with no final + text (see PR discussion: deepwiki tool call errors, model retries once, + reply used to just stop with no explanation). + """ + call_tool = _mock_mcp_environment(monkeypatch) + + aresponses_mock = AsyncMock( + side_effect=[ + _tool_call_stream("call_2", "read_wiki_contents"), # round 2: model retries + _text_only_stream("Here's what I found after retrying."), # round 3: final answer + ] + ) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator( + [ + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")]), # round 1: errors + ] + ) + + chunks = [chunk async for chunk in iterator] + + # Both rounds' tool calls were actually executed, not just streamed unexecuted. + assert call_tool.call_count == 2 + assert iterator.tool_call_round == 2 + + # The stream reached round 3 and produced the final text response instead + # of stopping after round 1 or round 2. + completed_chunks = [c for c in chunks if getattr(c, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED] + assert len(completed_chunks) == 3 + final_output = completed_chunks[-1].response.output + assert final_output[0]["content"][0]["text"] == "Here's what I found after retrying." + + +@pytest.mark.asyncio +async def test_tool_call_rounds_are_capped(monkeypatch): + """ + A model that keeps calling tools every round must not loop forever — + auto-execution stops at MAX_MCP_TOOL_CALL_ROUNDS, and the follow-up made + at the cap drops "tools" from the request so the model is forced to + answer in text instead of the stream just hanging. + """ + call_tool = _mock_mcp_environment(monkeypatch) + + # Rounds 2..MAX_MCP_TOOL_CALL_ROUNDS keep calling the tool; the call made + # once the cap is hit returns a text-only response. + tool_call_streams = [ + _tool_call_stream(f"call_{i}", "read_wiki_contents") for i in range(2, MAX_MCP_TOOL_CALL_ROUNDS + 1) + ] + aresponses_mock = AsyncMock(side_effect=[*tool_call_streams, _text_only_stream("giving up on tools")]) + monkeypatch.setattr(responses_main_module, "aresponses", aresponses_mock) + + iterator = _make_iterator( + [ + _output_item_added_chunk(), + _completed_chunk([_function_call("call_1", "read_wiki_contents")]), + ] + ) + + _ = [chunk async for chunk in iterator] + + assert iterator.tool_call_round == MAX_MCP_TOOL_CALL_ROUNDS + assert call_tool.call_count == MAX_MCP_TOOL_CALL_ROUNDS + assert aresponses_mock.call_count == MAX_MCP_TOOL_CALL_ROUNDS + + # Every follow-up before the cap still offered tools; only the capped one drops them. + for call in aresponses_mock.call_args_list[:-1]: + assert "tools" in call.kwargs + assert "tools" not in aresponses_mock.call_args_list[-1].kwargs