From 8e8e96d6a84720c632cc79491dc10bca8e5c5cce Mon Sep 17 00:00:00 2001 From: fangkangmi Date: Thu, 2 Jul 2026 22:00:45 +0100 Subject: [PATCH 1/3] fix(mcp): suppress intermediate tool-call turn in auto-exec streaming When the proxy auto-executes MCP tools (require_approval: never) with stream=true, the client received the model's first turn verbatim: raw tool_call deltas plus a mid-stream finish_reason "tool_calls", followed by the follow-up answer and a second finish_reason. OpenAI-compatible clients treat the first finish_reason as end-of-message, so chat UIs rendered an empty reply and discarded the real answer. MCPStreamingIterator now holds back tool-call delta chunks and, once the follow-up stream is created, drops them along with the intermediate finish_reason chunk. The client sees a single assistant message ending with one terminal finish_reason. Content deltas from the first turn still stream immediately, and if tool execution produces no follow-up the held chunks are flushed so the stream still terminates as before. mcp_list_tools metadata now lands on the first chunk actually yielded to the client, and mcp_tool_calls / mcp_call_results move to the terminal chunk of the follow-up stream since their previous carrier chunk is no longer surfaced. Fixes #31910 --- .../responses/mcp/chat_completions_handler.py | 107 +++--- tests/mcp_tests/test_mcp_chat_completions.py | 222 ++++-------- .../mcp/test_chat_completions_handler.py | 337 ++++++++++++++++-- 3 files changed, 445 insertions(+), 221 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index f2ccfd430ae..d440d9a93ac 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -245,6 +245,9 @@ async def acompletion_with_mcp( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self.held_tool_call_chunks: list[ModelResponseStream] = [] + self.pending_chunks: list[ModelResponseStream] = [] + self.any_chunk_yielded = False async def __aiter__(self): return self @@ -319,9 +322,42 @@ async def acompletion_with_mcp( "Error draining inner MCP stream after final chunk; spend logging may be incomplete" ) + def _is_final_chunk(self, chunk: ModelResponseStream) -> bool: + return bool( + hasattr(chunk, "choices") + and chunk.choices + and hasattr(chunk.choices[0], "finish_reason") + and chunk.choices[0].finish_reason is not None + ) + + def _chunk_has_tool_call_delta(self, chunk: ModelResponseStream) -> bool: + choices = getattr(chunk, "choices", None) or [] + return any(getattr(getattr(choice, "delta", None), "tool_calls", None) for choice in choices) + + def _yield_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + if not self.any_chunk_yielded: + self.any_chunk_yielded = True + chunk = self._add_mcp_list_tools_to_chunk(chunk) + return chunk + + async def _finish_initial_turn(self): + await self._process_tool_calls() + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + # Drain inner stream so CustomStreamWrapper fires its + # end-of-stream handler (dispatch_success_handlers → + # _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may + # yield one usage chunk before raising StopAsyncIteration. + await self._drain_inner_stream() + if self.follow_up_stream is not None: + self.held_tool_call_chunks = [] + async def __anext__(self): + if self.pending_chunks: + return self._yield_chunk(self.pending_chunks.pop(0)) + # Phase 1: Collect and yield initial stream chunks - if not self.stream_exhausted: + while not self.stream_exhausted: # Get the iterator from the stream wrapper if not hasattr(self, "_stream_iterator"): self._stream_iterator = self.stream_wrapper.__aiter__() @@ -333,46 +369,35 @@ async def acompletion_with_mcp( try: chunk = await self._stream_iterator.__anext__() - self.collected_chunks.append(chunk) - - # Add mcp_list_tools to the first chunk - if len(self.collected_chunks) == 1: - chunk = self._add_mcp_list_tools_to_chunk(chunk) - - # Check if this is the final chunk (has finish_reason) - is_final = ( - hasattr(chunk, "choices") - and chunk.choices - and hasattr(chunk.choices[0], "finish_reason") - and chunk.choices[0].finish_reason is not None - ) - - if is_final: - self.stream_exhausted = True - await self._process_tool_calls() - chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - # Drain inner stream so CustomStreamWrapper fires its - # end-of-stream handler (dispatch_success_handlers → - # _ProxyDBLogger → LiteLLM_SpendLogs). The CSW may - # yield one usage chunk before raising StopAsyncIteration. - await self._drain_inner_stream() - - return chunk except StopAsyncIteration: self.stream_exhausted = True - # Process tool calls after stream is exhausted - await self._process_tool_calls() - # If we have chunks, yield the final one with metadata - if self.collected_chunks: - final_chunk = self.collected_chunks[-1] - final_chunk = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) - # If we have tool results, prepare follow-up call - if self.tool_results and self.complete_response: - await self._prepare_follow_up_call() - await self._drain_inner_stream() - return final_chunk + if not self.collected_chunks: + break + await self._finish_initial_turn() + if self.follow_up_stream is not None: + break + final_chunk = self._add_mcp_tool_metadata_to_final_chunk(self.collected_chunks[-1]) + self.pending_chunks = self.held_tool_call_chunks + [final_chunk] + self.held_tool_call_chunks = [] + return self._yield_chunk(self.pending_chunks.pop(0)) + + self.collected_chunks.append(chunk) + + if self._is_final_chunk(chunk): + self.stream_exhausted = True + await self._finish_initial_turn() + if self.follow_up_stream is not None: + break + chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) + self.pending_chunks = self.held_tool_call_chunks + [chunk] + self.held_tool_call_chunks = [] + return self._yield_chunk(self.pending_chunks.pop(0)) + + if self._chunk_has_tool_call_delta(chunk): + self.held_tool_call_chunks.append(chunk) + continue + + return self._yield_chunk(chunk) # Phase 2: Yield follow-up stream chunks if available if self.follow_up_stream and not self.follow_up_exhausted: @@ -387,7 +412,9 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug(f"Follow-up chunk yielded: {chunk}") - return chunk + if self._is_final_chunk(chunk): + chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) + return self._yield_chunk(chunk) except StopAsyncIteration: self.follow_up_exhausted = True from litellm._logging import verbose_logger diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index fbdbf9152aa..5171605be1d 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -464,9 +464,9 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): """ Test that MCP metadata is added correctly to streaming chunks: - - mcp_list_tools should be in the first chunk - - mcp_tool_calls and mcp_call_results should be in the final chunk of initial response - - Follow-up response should be streamed after initial response + - mcp_list_tools should be in the first chunk yielded to the client + - mcp_tool_calls and mcp_call_results should be in the final chunk of the follow-up response + - The intermediate tool-call turn must not leak into the client stream (issue #31910) """ from types import SimpleNamespace from unittest.mock import patch @@ -734,95 +734,59 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): all_chunks = executor.submit(consume_stream).result() assert len(all_chunks) > 0, "Should have received streaming chunks" - # Find chunks from initial response (with tool_calls finish_reason) - initial_chunks_list = [] - follow_up_chunks_list = [] - for chunk in all_chunks: - if hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if ( - hasattr(choice, "finish_reason") - and choice.finish_reason == "tool_calls" - ): - initial_chunks_list.append(chunk) - elif ( - hasattr(choice, "finish_reason") and choice.finish_reason == "stop" - ): - follow_up_chunks_list.append(chunk) - elif ( - not hasattr(choice, "finish_reason") or choice.finish_reason is None - ): - # Chunks without finish_reason could be from either stream - # Check if we've seen tool_calls yet - if initial_chunks_list: - follow_up_chunks_list.append(chunk) - else: - initial_chunks_list.append(chunk) - - # Verify initial response chunks - assert len(initial_chunks_list) > 0, "Should have initial response chunks" - - # Find the final chunk from initial response (with tool_calls finish_reason) - initial_final_chunk = None - for chunk in initial_chunks_list: - if hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if ( - hasattr(choice, "finish_reason") - and choice.finish_reason == "tool_calls" - ): - initial_final_chunk = chunk - break - - if initial_final_chunk is None and initial_chunks_list: - initial_final_chunk = initial_chunks_list[-1] + finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in all_chunks + if chunk.choices and chunk.choices[0].finish_reason is not None + ] + assert finish_reasons == ["stop"], ( + f"Intermediate tool-call turn must not leak; expected a single stop. Got: {finish_reasons}" + ) + assert all( + not getattr(chunk.choices[0].delta, "tool_calls", None) + for chunk in all_chunks + if chunk.choices and chunk.choices[0].delta + ), "tool_call deltas must not leak into the client stream" + first_chunk = all_chunks[0] + first_provider_fields = getattr( + first_chunk.choices[0].delta, "provider_specific_fields", None + ) assert ( - initial_final_chunk is not None - ), "Should have a final chunk from initial response" + first_provider_fields is not None + ), "First chunk should have provider_specific_fields" + assert ( + "mcp_list_tools" in first_provider_fields + ), "First chunk should have mcp_list_tools" - # Verify mcp_list_tools is in the first chunk of initial response - first_chunk = initial_chunks_list[0] if initial_chunks_list else None - assert first_chunk is not None, "Should have a first chunk" - if hasattr(first_chunk, "choices") and first_chunk.choices: - choice = first_chunk.choices[0] - if hasattr(choice, "delta") and choice.delta: - provider_fields = getattr( - choice.delta, "provider_specific_fields", None - ) - assert ( - provider_fields is not None - ), "First chunk should have provider_specific_fields" - assert ( - "mcp_list_tools" in provider_fields - ), "First chunk should have mcp_list_tools" + final_chunk = all_chunks[-1] + assert final_chunk.choices[0].finish_reason == "stop" + final_provider_fields = getattr( + final_chunk.choices[0].delta, "provider_specific_fields", None + ) + assert ( + final_provider_fields is not None + ), "Final chunk should have provider_specific_fields" + assert "mcp_tool_calls" in final_provider_fields, "Should have mcp_tool_calls" + assert ( + "mcp_call_results" in final_provider_fields + ), "Should have mcp_call_results" - # Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response - if hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices: - choice = initial_final_chunk.choices[0] - if hasattr(choice, "delta") and choice.delta: - provider_fields = getattr( - choice.delta, "provider_specific_fields", None - ) - assert ( - provider_fields is not None - ), "Final chunk should have provider_specific_fields" - assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls" - assert ( - "mcp_call_results" in provider_fields - ), "Should have mcp_call_results" - - # Verify follow-up response chunks are present - assert len(follow_up_chunks_list) > 0, "Should have follow-up response chunks" + content = "".join( + chunk.choices[0].delta.content or "" + for chunk in all_chunks + if chunk.choices and chunk.choices[0].delta + ) + assert content == "Hello world!" @pytest.mark.asyncio async def test_mcp_streaming_metadata_ordering(monkeypatch): """ Test that MCP metadata appears in the correct order: - - mcp_list_tools should appear in the first chunk (before tool_calls) - - mcp_tool_calls and mcp_call_results should appear in the final chunk of initial response - - Follow-up response should be streamed after initial response completes + - mcp_list_tools should appear in the first chunk yielded to the client + - mcp_tool_calls and mcp_call_results should appear in the terminal chunk of the stream + - The client stream must contain exactly one terminal finish_reason (issue #31910) """ from types import SimpleNamespace from unittest.mock import patch @@ -1069,66 +1033,38 @@ async def test_mcp_streaming_metadata_ordering(monkeypatch): all_chunks = executor.submit(consume_stream).result() assert len(all_chunks) > 0, "Should have received streaming chunks" - # Track when we see each type of metadata - mcp_list_tools_seen = False - mcp_tool_calls_seen = False - mcp_call_results_seen = False - tool_calls_finish_reason_seen = False - follow_up_content_seen = False + finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in all_chunks + if chunk.choices and chunk.choices[0].finish_reason is not None + ] + assert finish_reasons == ["stop"], ( + f"Client stream must contain exactly one terminal finish_reason. Got: {finish_reasons}" + ) - for i, chunk in enumerate(all_chunks): - if hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if hasattr(choice, "delta") and choice.delta: - provider_fields = getattr( - choice.delta, "provider_specific_fields", None - ) - if provider_fields: - if "mcp_list_tools" in provider_fields: - mcp_list_tools_seen = True - # mcp_list_tools should appear before tool_calls finish_reason - assert ( - not tool_calls_finish_reason_seen - ), "mcp_list_tools should appear before tool_calls finish_reason" - if "mcp_tool_calls" in provider_fields: - mcp_tool_calls_seen = True - if "mcp_call_results" in provider_fields: - mcp_call_results_seen = True - - if ( - hasattr(choice, "finish_reason") - and choice.finish_reason == "tool_calls" - ): - tool_calls_finish_reason_seen = True - # mcp_tool_calls and mcp_call_results should be in the same chunk as tool_calls finish_reason - if hasattr(choice, "delta") and choice.delta: - provider_fields = getattr( - choice.delta, "provider_specific_fields", None - ) - assert provider_fields is not None - assert ( - "mcp_tool_calls" in provider_fields - ), "mcp_tool_calls should be in the chunk with tool_calls finish_reason" - assert ( - "mcp_call_results" in provider_fields - ), "mcp_call_results should be in the chunk with tool_calls finish_reason" - - if hasattr(choice, "delta") and choice.delta and choice.delta.content: - content = choice.delta.content - if content and ( - "Hello" in content or "world" in content or "!" in content - ): - follow_up_content_seen = True - # Follow-up content should appear after tool_calls finish_reason - assert ( - tool_calls_finish_reason_seen - ), "Follow-up content should appear after tool_calls finish_reason" - - # Verify all metadata was seen - assert mcp_list_tools_seen, "Should have seen mcp_list_tools" - assert mcp_tool_calls_seen, "Should have seen mcp_tool_calls" - assert mcp_call_results_seen, "Should have seen mcp_call_results" + first_provider_fields = getattr( + all_chunks[0].choices[0].delta, "provider_specific_fields", None + ) assert ( - tool_calls_finish_reason_seen - ), "Should have seen tool_calls finish_reason" - assert follow_up_content_seen, "Should have seen follow-up content" + first_provider_fields is not None and "mcp_list_tools" in first_provider_fields + ), "mcp_list_tools should be in the first chunk" + + terminal_chunk = all_chunks[-1] + assert terminal_chunk.choices[0].finish_reason == "stop" + terminal_provider_fields = getattr( + terminal_chunk.choices[0].delta, "provider_specific_fields", None + ) + assert terminal_provider_fields is not None + assert ( + "mcp_tool_calls" in terminal_provider_fields + ), "mcp_tool_calls should be in the terminal chunk" + assert ( + "mcp_call_results" in terminal_provider_fields + ), "mcp_call_results should be in the terminal chunk" + + content = "".join( + chunk.choices[0].delta.content or "" + for chunk in all_chunks + if chunk.choices and chunk.choices[0].delta + ) + assert content == "Hello world!", "Follow-up answer content must reach the client" 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 ab4c5185057..64808c529ef 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -742,8 +742,8 @@ async def test_acompletion_with_mcp_streaming_initial_call_is_streaming(monkeypa async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeypatch): """ Test that MCP metadata is added to the correct chunks: - - mcp_list_tools should be in the first chunk - - mcp_tool_calls and mcp_call_results should be in the final chunk of initial response + - mcp_list_tools should be in the first chunk yielded to the client + - mcp_tool_calls and mcp_call_results should be in the final chunk of the follow-up response """ from litellm.utils import CustomStreamWrapper from litellm.types.utils import ( @@ -971,35 +971,16 @@ async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeyp all_chunks.append(chunk) assert len(all_chunks) > 0 - # Find first chunk and final chunk from initial response - # mcp_list_tools is added to the first chunk (all_chunks[0]) - first_chunk = all_chunks[0] if all_chunks else None - initial_final_chunk = None + finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in all_chunks + if chunk.choices and chunk.choices[0].finish_reason is not None + ] + assert finish_reasons == ["stop"], f"Client stream must end with a single stop. Got: {finish_reasons}" - for chunk in all_chunks: - if hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if ( - hasattr(choice, "finish_reason") - and choice.finish_reason == "tool_calls" - ): - initial_final_chunk = chunk - - assert first_chunk is not None, "Should have a first chunk" - assert ( - initial_final_chunk is not None - ), "Should have a final chunk from initial response" - - # Verify mcp_list_tools is in the first chunk - assert ( - hasattr(first_chunk, "choices") and first_chunk.choices - ), "First chunk must have choices" - first_choice = first_chunk.choices[0] - assert ( - hasattr(first_choice, "delta") and first_choice.delta - ), "First choice must have delta" + first_chunk = all_chunks[0] first_provider_fields = getattr( - first_choice.delta, "provider_specific_fields", None + first_chunk.choices[0].delta, "provider_specific_fields", None ) assert ( first_provider_fields is not None @@ -1008,16 +989,10 @@ async def test_acompletion_with_mcp_streaming_metadata_in_correct_chunks(monkeyp "mcp_list_tools" in first_provider_fields ), "First chunk should have mcp_list_tools" - # Verify mcp_tool_calls and mcp_call_results are in the final chunk of initial response - assert ( - hasattr(initial_final_chunk, "choices") and initial_final_chunk.choices - ), "Final chunk must have choices" - final_choice = initial_final_chunk.choices[0] - assert ( - hasattr(final_choice, "delta") and final_choice.delta - ), "Final choice must have delta" + final_chunk = all_chunks[-1] + assert final_chunk.choices[0].finish_reason == "stop" final_provider_fields = getattr( - final_choice.delta, "provider_specific_fields", None + final_chunk.choices[0].delta, "provider_specific_fields", None ) assert ( final_provider_fields is not None @@ -1344,3 +1319,289 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti assert len(all_chunks) == 3 assert initial_stream.drained_after_exhaustion is True + + +def _create_mcp_stream_chunk(content, finish_reason=None, tool_calls=None): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + id="test-stream", + model="test-model", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=content, role="assistant", tool_calls=tool_calls), + finish_reason=finish_reason, + ) + ], + ) + + +def _make_mock_stream_class(stream_chunks): + from unittest.mock import MagicMock + + from litellm.utils import CustomStreamWrapper + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStream(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test-model", + logging_obj=logging_obj, + ) + self.chunks = stream_chunks + self._index = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + return chunk + raise StopAsyncIteration + + return MockStream + + +def _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results): + 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( + LiteLLM_Proxy_MCP_Handler, + "_extract_tool_calls_from_chat_response", + staticmethod(lambda **_: tool_calls), + ) + + async def mock_execute(**_): + return tool_results + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + mock_execute, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_create_follow_up_messages_for_chat", + staticmethod( + lambda **_: [ + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call-1", + "name": "local_search", + "content": "executed", + }, + ] + ), + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_streaming_suppresses_intermediate_tool_call_turn(monkeypatch): + """ + Regression test for https://github.com/BerriAI/litellm/issues/31910: + when the proxy auto-executes MCP tools with stream=True, the intermediate + tool-call turn (raw tool_call deltas + finish_reason "tool_calls") must not + leak into the client stream. The client should see a single assistant + message ending with exactly one terminal finish_reason. + """ + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ] + tool_results = [{"tool_call_id": "call-1", "result": "executed"}] + + initial_chunks = [ + _create_mcp_stream_chunk("Let me check. "), + _create_mcp_stream_chunk( + None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ], + ), + _create_mcp_stream_chunk("", finish_reason="tool_calls"), + ] + follow_up_chunks = [ + _create_mcp_stream_chunk("Hello"), + _create_mcp_stream_chunk(" world", finish_reason="stop"), + ] + + InitialStream = _make_mock_stream_class(initial_chunks) + FollowUpStream = _make_mock_stream_class(follow_up_chunks) + + async def mock_acompletion(**kwargs): + messages = kwargs.get("messages", []) + is_follow_up = any( + isinstance(msg, dict) and msg.get("role") == "tool" for msg in messages + ) + return FollowUpStream() if is_follow_up else InitialStream() + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) + _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results) + + with ( + patch("litellm.acompletion", mock_acompletion_func), + patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion_func, + create=True, + ), + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) + + assert all( + not getattr(chunk.choices[0].delta, "tool_calls", None) for chunk in all_chunks if chunk.choices + ), f"tool_call deltas must not leak into the client stream. Got: {all_chunks}" + + finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in all_chunks + if chunk.choices and chunk.choices[0].finish_reason is not None + ] + assert finish_reasons == ["stop"], f"Expected a single terminal stop. Got: {finish_reasons}" + assert all_chunks[-1].choices[0].finish_reason == "stop" + + content = "".join( + chunk.choices[0].delta.content or "" for chunk in all_chunks if chunk.choices and chunk.choices[0].delta + ) + assert content == "Let me check. Hello world" + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_streaming_flushes_tool_call_turn_when_no_follow_up(monkeypatch): + """ + When tool execution produces no results (so no follow-up stream is created), + the held tool-call chunks and the finish_reason "tool_calls" chunk must be + flushed so the client still receives a complete, terminated stream. + """ + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ] + + initial_chunks = [ + _create_mcp_stream_chunk( + None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ], + ), + _create_mcp_stream_chunk("", finish_reason="tool_calls"), + ] + + InitialStream = _make_mock_stream_class(initial_chunks) + + async def mock_acompletion(**kwargs): + return InitialStream() + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) + _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results=[]) + + with ( + patch("litellm.acompletion", mock_acompletion_func), + patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion_func, + create=True, + ), + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) + + assert any( + getattr(chunk.choices[0].delta, "tool_calls", None) for chunk in all_chunks if chunk.choices + ), "Held tool-call chunks must be flushed when no follow-up stream is created" + assert all_chunks[-1].choices[0].finish_reason == "tool_calls" From 119926f1123496336066f77000c2aeebe3e67f0d Mon Sep 17 00:00:00 2001 From: fangkangmi Date: Thu, 2 Jul 2026 22:45:28 +0100 Subject: [PATCH 2/3] fix(mcp): dedupe held tool-call chunk in abrupt-termination fallback Review follow-up: when the initial stream ends via StopAsyncIteration without a finish_reason chunk and the last collected chunk is a held tool-call delta, the chunk existed in both held_tool_call_chunks and collected_chunks[-1], so the fallback flush yielded it twice. The flush now filters the final chunk out of the held list by identity, via a shared _flush_held_and_final helper used by both fallback paths Also switch pending_chunks from list.pop(0) to collections.deque per review suggestion --- .../responses/mcp/chat_completions_handler.py | 23 +++--- .../mcp/test_chat_completions_handler.py | 73 +++++++++++++++++++ 2 files changed, 86 insertions(+), 10 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index d440d9a93ac..14d1cb641f9 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,6 +1,7 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" import logging +from collections import deque from typing import ( Any, List, @@ -246,7 +247,7 @@ async def acompletion_with_mcp( self.follow_up_iterator = None self.follow_up_exhausted = False self.held_tool_call_chunks: list[ModelResponseStream] = [] - self.pending_chunks: list[ModelResponseStream] = [] + self.pending_chunks: deque[ModelResponseStream] = deque() self.any_chunk_yielded = False async def __aiter__(self): @@ -352,9 +353,17 @@ async def acompletion_with_mcp( if self.follow_up_stream is not None: self.held_tool_call_chunks = [] + def _flush_held_and_final(self, final_chunk: ModelResponseStream) -> ModelResponseStream: + flushed_final = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) + self.pending_chunks = deque( + [chunk for chunk in self.held_tool_call_chunks if chunk is not final_chunk] + [flushed_final] + ) + self.held_tool_call_chunks = [] + return self._yield_chunk(self.pending_chunks.popleft()) + async def __anext__(self): if self.pending_chunks: - return self._yield_chunk(self.pending_chunks.pop(0)) + return self._yield_chunk(self.pending_chunks.popleft()) # Phase 1: Collect and yield initial stream chunks while not self.stream_exhausted: @@ -376,10 +385,7 @@ async def acompletion_with_mcp( await self._finish_initial_turn() if self.follow_up_stream is not None: break - final_chunk = self._add_mcp_tool_metadata_to_final_chunk(self.collected_chunks[-1]) - self.pending_chunks = self.held_tool_call_chunks + [final_chunk] - self.held_tool_call_chunks = [] - return self._yield_chunk(self.pending_chunks.pop(0)) + return self._flush_held_and_final(self.collected_chunks[-1]) self.collected_chunks.append(chunk) @@ -388,10 +394,7 @@ async def acompletion_with_mcp( await self._finish_initial_turn() if self.follow_up_stream is not None: break - chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) - self.pending_chunks = self.held_tool_call_chunks + [chunk] - self.held_tool_call_chunks = [] - return self._yield_chunk(self.pending_chunks.pop(0)) + return self._flush_held_and_final(chunk) if self._chunk_has_tool_call_delta(chunk): self.held_tool_call_chunks.append(chunk) 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 64808c529ef..780a3f7a656 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1605,3 +1605,76 @@ async def test_acompletion_with_mcp_streaming_flushes_tool_call_turn_when_no_fol getattr(chunk.choices[0].delta, "tool_calls", None) for chunk in all_chunks if chunk.choices ), "Held tool-call chunks must be flushed when no follow-up stream is created" assert all_chunks[-1].choices[0].finish_reason == "tool_calls" + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_streaming_no_duplicate_chunk_on_abrupt_termination(monkeypatch): + """ + Regression test for a duplicate-chunk bug in the abrupt-termination fallback: + when the initial stream ends via StopAsyncIteration without a finish_reason + chunk and the last collected chunk is a held tool-call delta, that chunk is + both in held_tool_call_chunks and collected_chunks[-1]. It must be yielded + to the client exactly once. + """ + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ] + + initial_chunks = [ + _create_mcp_stream_chunk("partial answer "), + _create_mcp_stream_chunk( + None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ], + ), + ] + + InitialStream = _make_mock_stream_class(initial_chunks) + + async def mock_acompletion(**kwargs): + return InitialStream() + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) + _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results=[]) + + with ( + patch("litellm.acompletion", mock_acompletion_func), + patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion_func, + create=True, + ), + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) + + tool_call_chunk_count = sum( + 1 for chunk in all_chunks if chunk.choices and getattr(chunk.choices[0].delta, "tool_calls", None) + ) + assert tool_call_chunk_count == 1, ( + f"The held tool-call chunk must be yielded exactly once on abrupt termination. Got chunks: {all_chunks}" + ) + assert len(all_chunks) == len(set(id(chunk) for chunk in all_chunks)), "No chunk object may be yielded twice" From 9a83e07894da8ddc27687301d1c55711a395193a Mon Sep 17 00:00:00 2001 From: fangkangmi Date: Thu, 2 Jul 2026 23:05:08 +0100 Subject: [PATCH 3/3] test(mcp): cover empty-stream and abrupt-termination-with-follow-up paths Codecov flagged the two break statements in the StopAsyncIteration branch as uncovered. Both are real paths: an initial stream that ends with zero chunks must terminate cleanly, and a stream that ends without a finish_reason chunk while tool execution still succeeds must suppress the tool-call turn and stream the follow-up answer, same as clean termination --- .../mcp/test_chat_completions_handler.py | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) 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 780a3f7a656..5a49ebe0ba5 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1678,3 +1678,126 @@ async def test_acompletion_with_mcp_streaming_no_duplicate_chunk_on_abrupt_termi f"The held tool-call chunk must be yielded exactly once on abrupt termination. Got chunks: {all_chunks}" ) assert len(all_chunks) == len(set(id(chunk) for chunk in all_chunks)), "No chunk object may be yielded twice" + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_streaming_abrupt_termination_with_follow_up_suppresses_tool_turn(monkeypatch): + """ + When the initial stream ends via StopAsyncIteration without ever emitting a + finish_reason chunk but tool execution still succeeds, the held tool-call + chunks must be suppressed and the follow-up answer streamed, same as the + clean-termination path. + """ + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ] + tool_results = [{"tool_call_id": "call-1", "result": "executed"}] + + initial_chunks = [ + _create_mcp_stream_chunk( + None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ], + ), + ] + follow_up_chunks = [ + _create_mcp_stream_chunk("Hello"), + _create_mcp_stream_chunk(" world", finish_reason="stop"), + ] + + InitialStream = _make_mock_stream_class(initial_chunks) + FollowUpStream = _make_mock_stream_class(follow_up_chunks) + + async def mock_acompletion(**kwargs): + messages = kwargs.get("messages", []) + is_follow_up = any( + isinstance(msg, dict) and msg.get("role") == "tool" for msg in messages + ) + return FollowUpStream() if is_follow_up else InitialStream() + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) + _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls, tool_results) + + with ( + patch("litellm.acompletion", mock_acompletion_func), + patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion_func, + create=True, + ), + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) + + assert all( + not getattr(chunk.choices[0].delta, "tool_calls", None) for chunk in all_chunks if chunk.choices + ), f"tool_call deltas must not leak even when the initial stream ends abruptly. Got: {all_chunks}" + finish_reasons = [ + chunk.choices[0].finish_reason + for chunk in all_chunks + if chunk.choices and chunk.choices[0].finish_reason is not None + ] + assert finish_reasons == ["stop"], f"Expected a single terminal stop. Got: {finish_reasons}" + content = "".join( + chunk.choices[0].delta.content or "" for chunk in all_chunks if chunk.choices and chunk.choices[0].delta + ) + assert content == "Hello world" + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_streaming_empty_initial_stream_terminates_cleanly(monkeypatch): + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + + InitialStream = _make_mock_stream_class([]) + + async def mock_acompletion(**kwargs): + return InitialStream() + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) + _patch_mcp_auto_exec_scaffolding(monkeypatch, tools, openai_tools, tool_calls=[], tool_results=[]) + + with ( + patch("litellm.acompletion", mock_acompletion_func), + patch.object( + chat_completions_handler, + "litellm_acompletion", + mock_acompletion_func, + create=True, + ), + ): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) + + assert all_chunks == [], "An empty initial stream must terminate cleanly with no chunks"