diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index ae18d5f6f1b..5575dd2a2c5 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -1,6 +1,8 @@ """Helpers for handling MCP-aware `/chat/completions` requests.""" import logging +from collections.abc import Iterator +from itertools import chain from typing import TYPE_CHECKING, Final, cast from typing_extensions import TypedDict, Unpack @@ -242,6 +244,8 @@ async def acompletion_with_mcp( self.follow_up_stream = None self.follow_up_iterator = None self.follow_up_exhausted = False + self.pending_chunks: Iterator[ModelResponseStream] = iter(()) + self.any_chunk_yielded = False def __aiter__(self): return self @@ -316,9 +320,53 @@ 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: Final = 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 + return 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() + + def _flush_held_and_final(self, final_chunk: ModelResponseStream) -> ModelResponseStream: + flushed_final: Final = self._add_mcp_tool_metadata_to_final_chunk(final_chunk) + self.pending_chunks = chain( + ( + chunk + for chunk in self.collected_chunks + if chunk is not final_chunk and self._chunk_has_tool_call_delta(chunk) + ), + (flushed_final,), + ) + return self._yield_chunk(next(self.pending_chunks)) + async def __anext__(self): + pending_chunk: Final = next(self.pending_chunks, None) + if pending_chunk is not None: + return self._yield_chunk(pending_chunk) + # 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__() @@ -330,46 +378,28 @@ 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: 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 + return self._flush_held_and_final(self.collected_chunks[-1]) + + 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 + return self._flush_held_and_final(chunk) + + if self._chunk_has_tool_call_delta(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: @@ -384,7 +414,9 @@ async def acompletion_with_mcp( from litellm._logging import verbose_logger verbose_logger.debug("Follow-up chunk yielded: %s", 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 2c1845f7b92..dfeb0b201c1 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -749,8 +749,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 ( @@ -978,35 +978,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 @@ -1015,16 +996,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 @@ -1353,6 +1328,448 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti 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): # test-quality-ok: [TQ008] Supply exact model chunks to exercise MCP stream orchestration + 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): # test-quality-ok: [TQ008] Supply exact model chunks to exercise MCP stream orchestration + 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" + + +@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): # test-quality-ok: [TQ008] Supply exact model chunks to exercise MCP stream orchestration + 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" + + +@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): # test-quality-ok: [TQ008] Supply exact model chunks to exercise MCP stream orchestration + 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): # test-quality-ok: [TQ008] Supply exact model chunks to exercise MCP stream orchestration + 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" + + @pytest.mark.asyncio @respx.mock async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch):