From 4164c9ca18ce72d5ad0de538d6c8914212005cd6 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 23 Jan 2026 11:44:23 +0900 Subject: [PATCH 1/4] fix: completions mcp output ordering --- .../responses/mcp/chat_completions_handler.py | 317 +++++++++++++++++- .../components/playground/chat_ui/ChatUI.tsx | 16 +- .../playground/llm_calls/chat_completion.tsx | 82 +++-- 3 files changed, 365 insertions(+), 50 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 0b0004d154..d05e6c8c68 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -5,6 +5,7 @@ from typing import ( List, Optional, Union, + cast, ) from litellm.responses.mcp.litellm_proxy_mcp_handler import ( @@ -174,10 +175,313 @@ async def acompletion_with_mcp( ) return response - # For auto-execute: disable streaming for initial call + # For auto-execute: handle streaming vs non-streaming differently stream = kwargs.get("stream", False) mock_tool_calls = base_call_args.pop("mock_tool_calls", None) + if stream: + # Streaming mode: make initial call with streaming, collect chunks, detect tool calls + initial_call_args = dict(base_call_args) + initial_call_args["stream"] = True + if mock_tool_calls is not None: + initial_call_args["mock_tool_calls"] = mock_tool_calls + + # Make initial streaming call + initial_stream = await litellm_acompletion(**initial_call_args) + + if not isinstance(initial_stream, CustomStreamWrapper): + # Not a stream, return as-is + if isinstance(initial_stream, ModelResponse): + _add_mcp_metadata_to_response( + response=initial_stream, + openai_tools=openai_tools, + ) + return initial_stream + + # Create a custom async generator that collects chunks and handles tool execution + from litellm.main import stream_chunk_builder + from litellm.types.utils import ModelResponseStream + + class MCPStreamingIterator: + """Custom iterator that collects chunks, detects tool calls, and adds MCP metadata to final chunk.""" + + def __init__(self, stream_wrapper, messages, tool_server_map, user_api_key_auth, + mcp_auth_header, mcp_server_auth_headers, oauth2_headers, raw_headers, + litellm_call_id, litellm_trace_id, openai_tools, base_call_args): + self.stream_wrapper = stream_wrapper + self.messages = messages + self.tool_server_map = tool_server_map + self.user_api_key_auth = user_api_key_auth + self.mcp_auth_header = mcp_auth_header + self.mcp_server_auth_headers = mcp_server_auth_headers + self.oauth2_headers = oauth2_headers + self.raw_headers = raw_headers + self.litellm_call_id = litellm_call_id + self.litellm_trace_id = litellm_trace_id + self.openai_tools = openai_tools + self.base_call_args = base_call_args + self.collected_chunks: List[ModelResponseStream] = [] + self.tool_calls: Optional[List] = None + self.tool_results: Optional[List] = None + self.complete_response: Optional[ModelResponse] = None + self.stream_exhausted = False + self.tool_execution_done = False + self.follow_up_stream = None + self.follow_up_iterator = None + self.follow_up_exhausted = False + + async def __aiter__(self): + return self + + def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + """Add mcp_list_tools to the first chunk.""" + from litellm.types.utils import StreamingChoices + + if not self.openai_tools: + return chunk + + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(choice.delta, "provider_specific_fields", None) or {} + ) + + # Add only mcp_list_tools to first chunk + provider_fields["mcp_list_tools"] = self.openai_tools + + # Set the provider_specific_fields + setattr(choice.delta, "provider_specific_fields", provider_fields) + + return chunk + + def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + """Add mcp_tool_calls and mcp_call_results to the final chunk.""" + from litellm.types.utils import StreamingChoices + + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(choice.delta, "provider_specific_fields", None) or {} + ) + + # Add tool_calls and tool_results if available + if self.tool_calls: + provider_fields["mcp_tool_calls"] = self.tool_calls + if self.tool_results: + provider_fields["mcp_call_results"] = self.tool_results + + # Set the provider_specific_fields + setattr(choice.delta, "provider_specific_fields", provider_fields) + + return chunk + + async def __anext__(self): + # Phase 1: Collect and yield initial stream chunks + if not self.stream_exhausted: + # Get the iterator from the stream wrapper + if not hasattr(self, '_stream_iterator'): + self._stream_iterator = self.stream_wrapper.__aiter__() + # Add mcp_list_tools to the first chunk (available from the start) + _add_mcp_metadata_to_response( + response=self.stream_wrapper, + openai_tools=self.openai_tools, + ) + + 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: + # This is the final chunk, mark stream as exhausted + self.stream_exhausted = True + # Process tool calls after we've collected all chunks + await self._process_tool_calls() + # Apply MCP metadata (tool_calls and tool_results) to final chunk + chunk = self._add_mcp_tool_metadata_to_final_chunk(chunk) + # If we have tool results, prepare follow-up call immediately + if self.tool_results and self.complete_response: + await self._prepare_follow_up_call() + + 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() + return final_chunk + + # Phase 2: Yield follow-up stream chunks if available + if self.follow_up_stream and not self.follow_up_exhausted: + if not self.follow_up_iterator: + self.follow_up_iterator = self.follow_up_stream.__aiter__() + from litellm._logging import verbose_logger + verbose_logger.debug("Follow-up stream iterator created") + + try: + chunk = await self.follow_up_iterator.__anext__() + from litellm._logging import verbose_logger + verbose_logger.debug(f"Follow-up chunk yielded: {chunk}") + return chunk + except StopAsyncIteration: + self.follow_up_exhausted = True + from litellm._logging import verbose_logger + verbose_logger.debug("Follow-up stream exhausted") + # After follow-up stream is exhausted, check if we need to raise StopAsyncIteration + raise StopAsyncIteration + + # If we're here and follow_up_stream is None but we expected it, log a warning + if self.stream_exhausted and self.tool_results and self.complete_response and self.follow_up_stream is None: + from litellm._logging import verbose_logger + verbose_logger.warning( + "Follow-up stream was not created despite having tool results" + ) + + raise StopAsyncIteration + + async def _process_tool_calls(self): + """Process tool calls after streaming completes.""" + if self.tool_execution_done: + return + + self.tool_execution_done = True + + if not self.collected_chunks: + return + + # Build complete response from chunks + complete_response = stream_chunk_builder( + chunks=self.collected_chunks, + messages=self.messages, + ) + + if isinstance(complete_response, ModelResponse): + self.complete_response = complete_response + # Extract tool calls from complete response + self.tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( + response=complete_response + ) + + if self.tool_calls: + # Execute tool calls + self.tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=self.tool_server_map, + tool_calls=self.tool_calls, + user_api_key_auth=self.user_api_key_auth, + mcp_auth_header=self.mcp_auth_header, + mcp_server_auth_headers=self.mcp_server_auth_headers, + oauth2_headers=self.oauth2_headers, + raw_headers=self.raw_headers, + litellm_call_id=self.litellm_call_id, + litellm_trace_id=self.litellm_trace_id, + ) + + async def _prepare_follow_up_call(self): + """Prepare and initiate follow-up call with tool results.""" + if self.follow_up_stream is not None: + return # Already prepared + + if not self.tool_results or not self.complete_response: + return + + from litellm import acompletion as litellm_acompletion + + # Create follow-up messages with tool results + follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=self.messages, + response=self.complete_response, + tool_results=self.tool_results, + ) + + # Make follow-up call with streaming + follow_up_call_args = dict(self.base_call_args) + follow_up_call_args["messages"] = follow_up_messages + follow_up_call_args["stream"] = True + # Ensure follow-up call doesn't trigger MCP handler again + follow_up_call_args["_skip_mcp_handler"] = True + + follow_up_response = await litellm_acompletion(**follow_up_call_args) + + # Ensure follow-up response is a CustomStreamWrapper + if isinstance(follow_up_response, CustomStreamWrapper): + self.follow_up_stream = follow_up_response + from litellm._logging import verbose_logger + verbose_logger.debug("Follow-up stream created successfully") + else: + # Unexpected response type - log and set to None + from litellm._logging import verbose_logger + verbose_logger.warning( + f"Follow-up response is not a CustomStreamWrapper: {type(follow_up_response)}" + ) + self.follow_up_stream = None + + # Create the custom iterator + iterator = MCPStreamingIterator( + stream_wrapper=initial_stream, + messages=messages, + tool_server_map=tool_server_map, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + litellm_call_id=kwargs.get("litellm_call_id"), + litellm_trace_id=kwargs.get("litellm_trace_id"), + openai_tools=openai_tools, + base_call_args=base_call_args, + ) + + # Create a wrapper class that delegates to our custom iterator + # We'll use a simple approach: just replace the __aiter__ method + class MCPStreamWrapper(CustomStreamWrapper): + def __init__(self, original_wrapper, custom_iterator): + # Initialize with the same parameters as original wrapper + super().__init__( + completion_stream=None, + model=getattr(original_wrapper, "model", "unknown"), + logging_obj=getattr(original_wrapper, "logging_obj", None), + custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), + stream_options=getattr(original_wrapper, "stream_options", None), + make_call=getattr(original_wrapper, "make_call", None), + _response_headers=getattr(original_wrapper, "_response_headers", None), + ) + self._original_wrapper = original_wrapper + self._custom_iterator = custom_iterator + # Copy important attributes from original wrapper + if hasattr(original_wrapper, "_hidden_params"): + self._hidden_params = original_wrapper._hidden_params + + def __aiter__(self): + return self._custom_iterator + + def __getattr__(self, name): + # Delegate all other attributes to original wrapper + return getattr(self._original_wrapper, name) + + return cast(CustomStreamWrapper, MCPStreamWrapper(initial_stream, iterator)) + + # Non-streaming mode: use existing logic initial_call_args = dict(base_call_args) initial_call_args["stream"] = False if mock_tool_calls is not None: @@ -195,17 +499,6 @@ async def acompletion_with_mcp( ) if not tool_calls: - # No tool calls, return response or retry with streaming if needed - if stream: - retry_args = dict(base_call_args) - retry_args["stream"] = stream - response = await litellm_acompletion(**retry_args) - if isinstance(response, (ModelResponse, CustomStreamWrapper)): - _add_mcp_metadata_to_response( - response=response, - openai_tools=openai_tools, - ) - return response _add_mcp_metadata_to_response( response=initial_response, openai_tools=openai_tools, diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 1db4d3a585..3a7f8fec65 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -607,12 +607,16 @@ const ChatUI: React.FC = ({ console.log("ChatUI: Received MCP event:", event); setMCPEvents((prev) => { // Check if this is a duplicate event (same item_id and type) - const isDuplicate = prev.some( - (existingEvent) => - existingEvent.item_id === event.item_id && - existingEvent.type === event.type && - existingEvent.sequence_number === event.sequence_number, - ); + // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools") + const isDuplicate = event.item_id + ? prev.some( + (existingEvent) => + existingEvent.item_id === event.item_id && + existingEvent.type === event.type && + (existingEvent.sequence_number === event.sequence_number || + (existingEvent.sequence_number === undefined && event.sequence_number === undefined)), + ) + : false; if (isDuplicate) { console.log("ChatUI: Duplicate MCP event, skipping"); diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index efcf7504b2..c3c623c25c 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -59,12 +59,13 @@ export async function makeOpenAIChatCompletionRequest( let fullResponseContent = ""; let fullReasoningContent = ""; - // Track MCP metadata from final chunk + // Track MCP metadata cumulatively across chunks let mcpMetadata: { mcp_list_tools?: any[]; mcp_tool_calls?: any[]; mcp_call_results?: any[]; - } | null = null; + } = {}; + let mcpListToolsProcessed = false; // Build tools array const tools: any[] = []; @@ -167,16 +168,48 @@ export async function makeOpenAIChatCompletionRequest( onSearchResults(delta.provider_specific_fields.search_results); } - // Check for MCP metadata in provider_specific_fields (typically in final chunk) + // Check for MCP metadata in provider_specific_fields if (delta && delta.provider_specific_fields) { const providerFields = delta.provider_specific_fields; + + // Merge MCP metadata cumulatively (don't overwrite) + if (providerFields.mcp_list_tools && !mcpMetadata.mcp_list_tools) { + mcpMetadata.mcp_list_tools = providerFields.mcp_list_tools; + // Process mcp_list_tools immediately when found (typically in first chunk) + if (onMCPEvent && !mcpListToolsProcessed) { + mcpListToolsProcessed = true; + const toolsEvent: MCPEvent = { + type: "response.output_item.done", + item_id: "mcp_list_tools", // Add item_id to prevent duplicate detection issues + item: { + type: "mcp_list_tools", + tools: providerFields.mcp_list_tools.map((tool: any) => ({ + name: tool.function?.name || tool.name || "", + description: tool.function?.description || tool.description || "", + input_schema: tool.function?.parameters || tool.input_schema || {}, + })), + }, + timestamp: Date.now(), + }; + onMCPEvent(toolsEvent); + console.log("MCP list_tools event sent:", toolsEvent); + } + } + + if (providerFields.mcp_tool_calls) { + mcpMetadata.mcp_tool_calls = providerFields.mcp_tool_calls; + } + + if (providerFields.mcp_call_results) { + mcpMetadata.mcp_call_results = providerFields.mcp_call_results; + } + if (providerFields.mcp_list_tools || providerFields.mcp_tool_calls || providerFields.mcp_call_results) { - mcpMetadata = { - mcp_list_tools: providerFields.mcp_list_tools, - mcp_tool_calls: providerFields.mcp_tool_calls, - mcp_call_results: providerFields.mcp_call_results, - }; - console.log("MCP metadata found in chunk:", mcpMetadata); + console.log("MCP metadata found in chunk:", { + mcp_list_tools: providerFields.mcp_list_tools ? "present" : "absent", + mcp_tool_calls: providerFields.mcp_tool_calls ? "present" : "absent", + mcp_call_results: providerFields.mcp_call_results ? "present" : "absent", + }); } } @@ -204,35 +237,19 @@ export async function makeOpenAIChatCompletionRequest( } } - // Process MCP metadata from final chunk and convert to MCPEvent format - if (mcpMetadata && onMCPEvent) { - // Convert mcp_list_tools to MCPEvent - if (mcpMetadata.mcp_list_tools && mcpMetadata.mcp_list_tools.length > 0) { - const toolsEvent: MCPEvent = { - type: "response.output_item.done", - item: { - type: "mcp_list_tools", - tools: mcpMetadata.mcp_list_tools.map((tool: any) => ({ - name: tool.function?.name || tool.name || "", - description: tool.function?.description || tool.description || "", - input_schema: tool.function?.parameters || tool.input_schema || {}, - })), - }, - timestamp: Date.now(), - }; - onMCPEvent(toolsEvent); - } - + // Process remaining MCP metadata (mcp_tool_calls and mcp_call_results) after stream completes + // Note: mcp_list_tools is already processed when found in the first chunk + if (onMCPEvent && (mcpMetadata.mcp_tool_calls || mcpMetadata.mcp_call_results)) { // Convert mcp_tool_calls and mcp_call_results to MCPEvent[] - if (mcpMetadata?.mcp_tool_calls && mcpMetadata?.mcp_tool_calls.length > 0) { - mcpMetadata?.mcp_tool_calls.forEach((toolCall: any, index: number) => { + if (mcpMetadata.mcp_tool_calls && mcpMetadata.mcp_tool_calls.length > 0) { + mcpMetadata.mcp_tool_calls.forEach((toolCall: any, index: number) => { const functionName = toolCall.function?.name || toolCall.name || ""; const functionArgs = toolCall.function?.arguments || toolCall.arguments || "{}"; // Find corresponding result - const result = mcpMetadata?.mcp_call_results?.find( + const result = mcpMetadata.mcp_call_results?.find( (r: any) => r.tool_call_id === toolCall.id || r.tool_call_id === toolCall.call_id - ) || mcpMetadata?.mcp_call_results?.[index]; + ) || mcpMetadata.mcp_call_results?.[index]; const callEvent: MCPEvent = { type: "response.output_item.done", @@ -246,6 +263,7 @@ export async function makeOpenAIChatCompletionRequest( timestamp: Date.now(), }; onMCPEvent(callEvent); + console.log("MCP call event sent:", callEvent); }); } } From f8ceab94b4839e118fc2041841b010b85ace3929 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 23 Jan 2026 13:34:07 +0900 Subject: [PATCH 2/4] fix: for test --- .../litellm_core_utils/streaming_handler.py | 55 +++++++++++++++ .../responses/mcp/chat_completions_handler.py | 69 ++++++++++++++++--- .../mcp/litellm_proxy_mcp_handler.py | 6 ++ 3 files changed, 119 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3304759f74..c6f0f67976 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1571,6 +1571,50 @@ class CustomStreamWrapper: ) return chunk + def _add_mcp_list_tools_to_first_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + """ + Add mcp_list_tools from _hidden_params to the first chunk's delta.provider_specific_fields. + + This method checks if MCP metadata with mcp_list_tools is stored in _hidden_params + and adds it to the first chunk's delta.provider_specific_fields. + """ + try: + # Check if MCP metadata should be added to first chunk + if not hasattr(self, "_hidden_params") or not self._hidden_params: + return chunk + + mcp_metadata = self._hidden_params.get("mcp_metadata") + if not mcp_metadata or not isinstance(mcp_metadata, dict): + return chunk + + # Only add mcp_list_tools to first chunk (not tool_calls or tool_results) + mcp_list_tools = mcp_metadata.get("mcp_list_tools") + if not mcp_list_tools: + return chunk + + # Add mcp_list_tools to delta.provider_specific_fields + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(choice.delta, "provider_specific_fields", None) or {} + ) + + # Add only mcp_list_tools to first chunk + provider_fields["mcp_list_tools"] = mcp_list_tools + + # Set the provider_specific_fields + setattr(choice.delta, "provider_specific_fields", provider_fields) + + except Exception as e: + from litellm._logging import verbose_logger + verbose_logger.exception( + f"Error adding MCP list tools to first chunk: {str(e)}" + ) + + return chunk + def _add_mcp_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """ Add MCP metadata from _hidden_params to the final chunk's delta.provider_specific_fields. @@ -1727,6 +1771,12 @@ class CustomStreamWrapper: ) # HANDLE STREAM OPTIONS self.chunks.append(response) + + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + response = self._add_mcp_list_tools_to_first_chunk(response) + self.sent_first_chunk = True + if hasattr( response, "usage" ): # remove usage from chunk, only send on final chunk @@ -1894,6 +1944,11 @@ class CustomStreamWrapper: input=self.response_uptil_now, model=self.model ) self.chunks.append(processed_chunk) + + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + processed_chunk = self._add_mcp_list_tools_to_first_chunk(processed_chunk) + self.sent_first_chunk = True if hasattr( processed_chunk, "usage" ): # remove usage from chunk, only send on final chunk diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index d05e6c8c68..2b4f91a35c 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -244,14 +244,14 @@ async def acompletion_with_mcp( for choice in chunk.choices: if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) or {} - ) + existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} + provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = self.openai_tools - # Set the provider_specific_fields + # Set provider_specific_fields directly using setattr + # This ensures the modification is preserved setattr(choice.delta, "provider_specific_fields", provider_fields) return chunk @@ -264,9 +264,8 @@ async def acompletion_with_mcp( for choice in chunk.choices: if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - provider_fields = ( - getattr(choice.delta, "provider_specific_fields", None) or {} - ) + existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} + provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original # Add tool_calls and tool_results if available if self.tool_calls: @@ -274,7 +273,8 @@ async def acompletion_with_mcp( if self.tool_results: provider_fields["mcp_call_results"] = self.tool_results - # Set the provider_specific_fields + # Set provider_specific_fields directly using setattr + # This ensures the modification is preserved setattr(choice.delta, "provider_specific_fields", provider_fields) return chunk @@ -405,8 +405,6 @@ async def acompletion_with_mcp( if not self.tool_results or not self.complete_response: return - from litellm import acompletion as litellm_acompletion - # Create follow-up messages with tool results follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( original_messages=self.messages, @@ -421,7 +419,10 @@ async def acompletion_with_mcp( # Ensure follow-up call doesn't trigger MCP handler again follow_up_call_args["_skip_mcp_handler"] = True - follow_up_response = await litellm_acompletion(**follow_up_call_args) + # Import litellm here to ensure we get the patched version + # This ensures the patch works correctly in tests + import litellm + follow_up_response = await litellm.acompletion(**follow_up_call_args) # Ensure follow-up response is a CustomStreamWrapper if isinstance(follow_up_response, CustomStreamWrapper): @@ -471,14 +472,60 @@ async def acompletion_with_mcp( # Copy important attributes from original wrapper if hasattr(original_wrapper, "_hidden_params"): self._hidden_params = original_wrapper._hidden_params + # For synchronous iteration, we need to run the async iterator + self._sync_iterator = None + self._sync_loop = None def __aiter__(self): return self._custom_iterator + def __iter__(self): + # For synchronous iteration, create a sync wrapper + if self._sync_iterator is None: + import asyncio + try: + self._sync_loop = asyncio.get_event_loop() + except RuntimeError: + self._sync_loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._sync_loop) + self._sync_iterator = _SyncIteratorWrapper(self._custom_iterator, self._sync_loop) + return self._sync_iterator + + def __next__(self): + # Delegate to sync iterator + if self._sync_iterator is None: + self.__iter__() + return next(self._sync_iterator) + def __getattr__(self, name): # Delegate all other attributes to original wrapper return getattr(self._original_wrapper, name) + # Helper class to wrap async iterator for sync iteration + class _SyncIteratorWrapper: + def __init__(self, async_iterator, loop): + self._async_iterator = async_iterator + self._loop = loop + self._iterator = None + + def __iter__(self): + return self + + def __next__(self): + if self._iterator is None: + # __aiter__ might be async, so we need to await it + aiter_result = self._async_iterator.__aiter__() + if hasattr(aiter_result, '__await__'): + # It's a coroutine, await it + self._iterator = self._loop.run_until_complete(aiter_result) + else: + # It's already an iterator + self._iterator = aiter_result + try: + return self._loop.run_until_complete(self._iterator.__anext__()) + except StopAsyncIteration: + raise StopIteration + return cast(CustomStreamWrapper, MCPStreamWrapper(initial_stream, iterator)) # Non-streaming mode: use existing logic diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 4376e076a9..930f7261cd 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -775,6 +775,12 @@ class LiteLLM_Proxy_MCP_Handler: first_choice, "message", None ): message_to_append = first_choice.message.model_dump(exclude_none=True) + # Ensure tool_calls have arguments field (required by OpenAI API) + if message_to_append.get("tool_calls"): + for tool_call in message_to_append["tool_calls"]: + if isinstance(tool_call, dict) and "function" in tool_call: + if "arguments" not in tool_call["function"]: + tool_call["function"]["arguments"] = "{}" except Exception: verbose_logger.exception("Failed to convert assistant message for MCP flow") From 6a60b3d848339053df895b152ac7d22262956b84 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 23 Jan 2026 15:17:14 +0900 Subject: [PATCH 3/4] test: completions mcp output test --- .../responses/mcp/chat_completions_handler.py | 27 +- tests/mcp_tests/test_mcp_chat_completions.py | 717 +++++++++++++++--- .../mcp/test_chat_completions_handler.py | 577 ++++++++++++-- 3 files changed, 1169 insertions(+), 152 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 2b4f91a35c..09c2af9589 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -235,7 +235,7 @@ async def acompletion_with_mcp( def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_list_tools to the first chunk.""" - from litellm.types.utils import StreamingChoices + from litellm.types.utils import StreamingChoices, add_provider_specific_fields if not self.openai_tools: return chunk @@ -250,22 +250,29 @@ async def acompletion_with_mcp( # Add only mcp_list_tools to first chunk provider_fields["mcp_list_tools"] = self.openai_tools - # Set provider_specific_fields directly using setattr - # This ensures the modification is preserved - setattr(choice.delta, "provider_specific_fields", provider_fields) + # Use add_provider_specific_fields to ensure proper setting + # This function handles Pydantic model attribute setting correctly + add_provider_specific_fields(choice.delta, provider_fields) return chunk def _add_mcp_tool_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: """Add mcp_tool_calls and mcp_call_results to the final chunk.""" - from litellm.types.utils import StreamingChoices + from litellm.types.utils import StreamingChoices, add_provider_specific_fields if hasattr(chunk, "choices") and chunk.choices: for choice in chunk.choices: if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: # Get existing provider_specific_fields or create new dict - existing_fields = getattr(choice.delta, "provider_specific_fields", None) or {} - provider_fields = dict(existing_fields) # Create a copy to avoid mutating the original + # Access the attribute directly to handle Pydantic model attributes correctly + existing_fields = {} + if hasattr(choice.delta, "provider_specific_fields"): + attr_value = getattr(choice.delta, "provider_specific_fields", None) + if attr_value is not None: + # Create a copy to avoid mutating the original + existing_fields = dict(attr_value) if isinstance(attr_value, dict) else {} + + provider_fields = existing_fields # Add tool_calls and tool_results if available if self.tool_calls: @@ -273,9 +280,9 @@ async def acompletion_with_mcp( if self.tool_results: provider_fields["mcp_call_results"] = self.tool_results - # Set provider_specific_fields directly using setattr - # This ensures the modification is preserved - setattr(choice.delta, "provider_specific_fields", provider_fields) + # Use add_provider_specific_fields to ensure proper setting + # This function handles Pydantic model attribute setting correctly + add_provider_specific_fields(choice.delta, provider_fields) return chunk diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index 8857f016df..0617dcb2e4 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -206,8 +206,18 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): ) # Create a mock streaming response + from unittest.mock import MagicMock, AsyncMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.async_failure_handler = AsyncMock() + class MockStreamingResponse(CustomStreamWrapper): def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) self.chunks = [ type('Chunk', (), { 'choices': [type('Choice', (), { @@ -233,39 +243,150 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): if self._index < len(self.chunks): chunk = self.chunks[self._index] self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True return chunk raise StopIteration + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopAsyncIteration + # Track calls to acompletion acompletion_calls = [] + # Create mock streaming response for initial call + from unittest.mock import MagicMock, AsyncMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.async_failure_handler = AsyncMock() + + from litellm.types.utils import ( + ModelResponseStream, + StreamingChoices, + Delta, + ChatCompletionDeltaToolCall, + Function, + ) + + # Create initial streaming chunks with tool_calls + # Add tool_calls to the final chunk so stream_chunk_builder can extract them + tool_calls = [ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ] + + initial_chunks = [ + ModelResponseStream( + id="test-1", + model="gpt-4o-mini", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="", + role="assistant", + tool_calls=tool_calls, + ), + finish_reason="tool_calls", + ) + ], + ) + ] + + class InitialStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = initial_chunks + self._index = 0 + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopIteration + + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopAsyncIteration + async def mock_acompletion(**kwargs): acompletion_calls.append(kwargs) - # First call (non-streaming for tool extraction) - if not kwargs.get("stream", False): - # Return a ModelResponse with tool_calls using dict format - return ModelResponse( - id="test-1", - model="gpt-4o-mini", - choices=[{ - "message": { - "role": "assistant", - "tool_calls": [{ - "id": "call-1", - "type": "function", - "function": { - "name": "local_search", - "arguments": "{}" - } - }] - }, - "finish_reason": "tool_calls" - }], - created=0, - object="chat.completion", + # With new implementation, first call should be streaming + if kwargs.get("stream", False): + # Check if this is the follow-up call + messages = kwargs.get("messages", []) + is_follow_up = any( + msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg)) + for msg in messages ) - # Second call (streaming follow-up) - return MockStreamingResponse() + if is_follow_up: + # Follow-up call (streaming) + return MockStreamingResponse() + else: + # Initial call (streaming) + return InitialStreamingResponse() + # Non-streaming call should not happen with new implementation, but handle it + return ModelResponse( + id="test-1", + model="gpt-4o-mini", + choices=[{ + "message": { + "role": "assistant", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": { + "name": "local_search", + "arguments": "{}" + } + }] + }, + "finish_reason": "tool_calls" + }], + created=0, + object="chat.completion", + ) with patch("litellm.acompletion", side_effect=mock_acompletion): # This should not raise RuntimeError: Timeout context manager should be used inside a task @@ -303,8 +424,12 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): # Verify response is a streaming response assert isinstance(result, CustomStreamWrapper) or hasattr(result, '__iter__') - # Consume the stream to ensure it works - chunks = list(result) + # Consume the stream to ensure it works (run in separate thread to avoid event loop conflict) + from concurrent.futures import ThreadPoolExecutor + def consume_stream(): + return list(result) + with ThreadPoolExecutor(max_workers=1) as executor: + chunks = executor.submit(consume_stream).result() assert len(chunks) > 0, "Should have received streaming chunks" # Verify tool execution was called @@ -317,8 +442,10 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): @pytest.mark.asyncio async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): """ - Test that MCP metadata is added to the final streaming chunk's - delta.provider_specific_fields when using MCP tools with streaming. + 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 """ from types import SimpleNamespace from unittest.mock import patch @@ -328,7 +455,13 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.utils import CustomStreamWrapper - from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + from litellm.types.utils import ( + ModelResponseStream, + StreamingChoices, + Delta, + ChatCompletionDeltaToolCall, + Function, + ) from litellm.litellm_core_utils.litellm_logging import Logging dummy_tool = SimpleNamespace( @@ -369,7 +502,13 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): ) # Create mock streaming chunks - def create_chunk(content, finish_reason=None): + def create_chunk(content, finish_reason=None, tool_calls=None): + delta = Delta( + content=content, + role="assistant", + ) + if tool_calls: + delta.tool_calls = tool_calls return ModelResponseStream( id="test-stream", model="gpt-4o-mini", @@ -378,36 +517,48 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): choices=[ StreamingChoices( index=0, - delta=Delta( - content=content, - role="assistant", - ), + delta=delta, finish_reason=finish_reason, ) ], ) - chunks = [ + # Create initial streaming chunks with tool_calls + # Add tool_calls to the final chunk so stream_chunk_builder can extract them + tool_calls = [ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ] + initial_chunks = [ + create_chunk("", finish_reason="tool_calls", tool_calls=tool_calls), # Final chunk with tool_calls + ] + + # Create follow-up streaming chunks + follow_up_chunks = [ create_chunk("Hello"), create_chunk(" world"), create_chunk("!", finish_reason="stop"), # Final chunk ] # Create a proper CustomStreamWrapper with logging_obj - from unittest.mock import MagicMock + from unittest.mock import MagicMock, AsyncMock logging_obj = MagicMock() logging_obj.model_call_details = {} + logging_obj.async_failure_handler = AsyncMock() - class MockStreamingResponse(CustomStreamWrapper): + class InitialStreamingResponse(CustomStreamWrapper): def __init__(self): super().__init__( completion_stream=None, model="gpt-4o-mini", logging_obj=logging_obj, ) - self.chunks = chunks + self.chunks = initial_chunks self._index = 0 - self.sent_last_chunk = False def __iter__(self): return self @@ -416,42 +567,106 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): if self._index < len(self.chunks): chunk = self.chunks[self._index] self._index += 1 - if self._index == len(self.chunks): - self.sent_last_chunk = True - # Call the method that adds MCP metadata to final chunk - chunk = self._add_mcp_metadata_to_final_chunk(chunk) + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True return chunk raise StopIteration + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopAsyncIteration + + class FollowUpStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = follow_up_chunks + self._index = 0 + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopIteration + + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopAsyncIteration + # Track calls to acompletion acompletion_calls = [] async def mock_acompletion(**kwargs): acompletion_calls.append(kwargs) - # First call (non-streaming for tool extraction) - if not kwargs.get("stream", False): - return ModelResponse( - id="test-1", - model="gpt-4o-mini", - choices=[{ - "message": { - "role": "assistant", - "tool_calls": [{ - "id": "call-1", - "type": "function", - "function": { - "name": "local_search", - "arguments": "{}" - } - }] - }, - "finish_reason": "tool_calls" - }], - created=0, - object="chat.completion", + # With new implementation, first call should be streaming + if kwargs.get("stream", False): + # Check if this is the follow-up call (has tool results in messages) + messages = kwargs.get("messages", []) + is_follow_up = any( + msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg)) + for msg in messages ) - # Second call (streaming follow-up) - return MockStreamingResponse() + + if is_follow_up: + # Follow-up call - return follow-up chunks + return FollowUpStreamingResponse() + else: + # Initial streaming call - return chunks with tool_calls + return InitialStreamingResponse() + # Non-streaming call should not happen with new implementation, but handle it + return ModelResponse( + id="test-1", + model="gpt-4o-mini", + choices=[{ + "message": { + "role": "assistant", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": { + "name": "local_search", + "arguments": "{}" + } + }] + }, + "finish_reason": "tool_calls" + }], + created=0, + object="chat.completion", + ) with patch("litellm.acompletion", side_effect=mock_acompletion): response = litellm.completion( @@ -482,39 +697,363 @@ async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): assert isinstance(result, CustomStreamWrapper) - # Verify _hidden_params contains mcp_metadata - assert hasattr(result, "_hidden_params") - assert "mcp_metadata" in result._hidden_params - mcp_metadata = result._hidden_params["mcp_metadata"] - assert "mcp_list_tools" in mcp_metadata - assert "mcp_tool_calls" in mcp_metadata - assert "mcp_call_results" in mcp_metadata + # Consume the stream and check chunks (run in separate thread to avoid event loop conflict) + from concurrent.futures import ThreadPoolExecutor + def consume_stream(): + return list(result) + with ThreadPoolExecutor(max_workers=1) as executor: + all_chunks = executor.submit(consume_stream).result() + assert len(all_chunks) > 0, "Should have received streaming chunks" - # Consume the stream and check final chunk - all_chunks = list(result) - assert len(all_chunks) > 0 - - # Find the final chunk (with finish_reason) - final_chunk = None + # 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: - final_chunk = chunk + 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] - # If no chunk with finish_reason, use the last chunk - if final_chunk is None and all_chunks: - final_chunk = all_chunks[-1] + assert initial_final_chunk is not None, "Should have a final chunk from initial response" - assert final_chunk is not None, "Should have a final chunk" + # 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" - # Verify MCP metadata is in the final chunk's delta.provider_specific_fields - if hasattr(final_chunk, "choices") and final_chunk.choices: - choice = final_chunk.choices[0] + # 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_list_tools" in provider_fields, "Should have mcp_list_tools" 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" + + +@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 + """ + from types import SimpleNamespace + from unittest.mock import patch + + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ( + ModelResponseStream, + StreamingChoices, + Delta, + ChatCompletionDeltaToolCall, + Function, + ) + + dummy_tool = SimpleNamespace( + name="local_search", + description="search", + inputSchema={"type": "object", "properties": {}}, + ) + + async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy): + return [dummy_tool], {"local_search": "local"} + + async def fake_execute(**kwargs): + tool_calls = kwargs.get("tool_calls") or [] + call_entry = tool_calls[0] + call_id = call_entry.get("id") or call_entry.get("call_id") or "call" + return [ + { + "tool_call_id": call_id, + "result": "executed", + "name": call_entry.get("name", "local_search"), + } + ] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + fake_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda secret_fields, tools: (None, None, None, None)), + ) + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None, tool_calls=None): + delta = Delta( + content=content, + role="assistant", + ) + if tool_calls: + delta.tool_calls = tool_calls + return ModelResponseStream( + id="test-stream", + model="gpt-4o-mini", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=delta, + finish_reason=finish_reason, + ) + ], + ) + + # Create initial streaming chunks with tool_calls + # Add tool_calls to the final chunk so stream_chunk_builder can extract them + tool_calls = [ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ] + initial_chunks = [ + create_chunk("", finish_reason="tool_calls", tool_calls=tool_calls), # Final chunk with tool_calls + ] + + # Create follow-up streaming chunks + follow_up_chunks = [ + create_chunk("Hello"), + create_chunk(" world"), + create_chunk("!", finish_reason="stop"), # Final chunk + ] + + # Create a proper CustomStreamWrapper with logging_obj + from unittest.mock import MagicMock, AsyncMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + logging_obj.async_failure_handler = AsyncMock() + + class InitialStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = initial_chunks + self._index = 0 + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopIteration + + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopAsyncIteration + + class FollowUpStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = follow_up_chunks + self._index = 0 + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopIteration + + def __aiter__(self): + return self + + async def __anext__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True + return chunk + raise StopAsyncIteration + + # Track calls to acompletion + acompletion_calls = [] + + async def mock_acompletion(**kwargs): + acompletion_calls.append(kwargs) + # With new implementation, first call should be streaming + if kwargs.get("stream", False): + # Check if this is the follow-up call (has tool results in messages) + messages = kwargs.get("messages", []) + is_follow_up = any( + msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg)) + for msg in messages + ) + + if is_follow_up: + # Follow-up call - return follow-up chunks + return FollowUpStreamingResponse() + else: + # Initial streaming call - return chunks with tool_calls + return InitialStreamingResponse() + # Non-streaming call should not happen with new implementation + pytest.fail("Non-streaming call should not happen with new implementation") + + with patch("litellm.acompletion", side_effect=mock_acompletion): + response = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "server_label": "local", + "require_approval": "never", + } + ], + stream=True, + mock_response="Final answer", + mock_tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ], + ) + + import asyncio + assert asyncio.iscoroutine(response) + result = await response + + assert isinstance(result, CustomStreamWrapper) + + # Consume the stream and verify order (run in separate thread to avoid event loop conflict) + from concurrent.futures import ThreadPoolExecutor + def consume_stream(): + return list(result) + with ThreadPoolExecutor(max_workers=1) as executor: + 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 + + 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" + assert tool_calls_finish_reason_seen, "Should have seen tool_calls finish_reason" + assert follow_up_content_seen, "Should have seen follow-up content" 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 bc1f4fb72b..0741dece43 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -3,6 +3,7 @@ from unittest.mock import AsyncMock, patch from litellm.types.utils import ModelResponse +from litellm.responses.mcp import chat_completions_handler from litellm.responses.mcp.chat_completions_handler import ( acompletion_with_mcp, ) @@ -91,24 +92,116 @@ async def test_acompletion_with_mcp_without_auto_execution_calls_model(monkeypat @pytest.mark.asyncio async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta, ChatCompletionDeltaToolCall, Function + from unittest.mock import MagicMock + tools = [{"type": "function", "function": {"name": "tool"}}] - initial_response = ModelResponse( - id="1", - model="test", - choices=[], - created=0, - object="chat.completion", - ) - follow_up_response = ModelResponse( - id="2", - model="test", - choices=[], - created=0, - object="chat.completion", - ) - mock_acompletion = AsyncMock( - side_effect=[initial_response, follow_up_response] - ) + + # Create mock streaming chunks for initial response + def create_chunk(content, finish_reason=None, tool_calls=None): + return ModelResponseStream( + id="test-stream", + model="test", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + tool_calls=tool_calls, + ), + finish_reason=finish_reason, + ) + ], + ) + + initial_chunks = [ + create_chunk( + "", + finish_reason="tool_calls", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="tool", arguments="{}"), + index=0, + ) + ], + ), + ] + + follow_up_chunks = [ + create_chunk("Hello"), + create_chunk(" world", finish_reason="stop"), + ] + + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class InitialStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test", + logging_obj=logging_obj, + ) + self.chunks = initial_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 + + class FollowUpStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test", + logging_obj=logging_obj, + ) + self.chunks = follow_up_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 + + async def mock_acompletion(**kwargs): + if kwargs.get("stream", False): + messages = kwargs.get("messages", []) + is_follow_up = any( + msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg)) + for msg in messages + ) + if is_follow_up: + return FollowUpStreamingResponse() + else: + return InitialStreamingResponse() + # Non-streaming should not happen + return ModelResponse( + id="1", + model="test", + choices=[], + created=0, + object="chat.completion", + ) + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) monkeypatch.setattr( LiteLLM_Proxy_MCP_Handler, @@ -141,10 +234,10 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): monkeypatch.setattr( LiteLLM_Proxy_MCP_Handler, "_extract_tool_calls_from_chat_response", - staticmethod(lambda **_: ["call"]), + staticmethod(lambda **_: [{"id": "call-1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]), ) async def mock_execute(**_): - return ["result"] + return [{"tool_call_id": "call-1", "result": "executed"}] monkeypatch.setattr( LiteLLM_Proxy_MCP_Handler, @@ -154,7 +247,11 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): monkeypatch.setattr( LiteLLM_Proxy_MCP_Handler, "_create_follow_up_messages_for_chat", - staticmethod(lambda **_: ["follow-up"]), + staticmethod(lambda **_: [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "tool", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "call-1", "name": "tool", "content": "executed"} + ]), ) monkeypatch.setattr( ResponsesAPIRequestUtils, @@ -162,21 +259,40 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): staticmethod(lambda **_: (None, None, None, None)), ) - with patch("litellm.acompletion", mock_acompletion): + # Patch litellm.acompletion at module level to catch function-level imports + 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="test-model", - messages=["msg"], + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], tools=tools, stream=True, ) - assert result is follow_up_response - assert mock_acompletion.await_count == 2 - first_call = mock_acompletion.await_args_list[0].kwargs - second_call = mock_acompletion.await_args_list[1].kwargs - assert first_call["stream"] is False - assert second_call["messages"] == ["follow-up"] - assert second_call["stream"] is True + # Consume the stream to trigger the iterator and follow-up call + # The initial stream has one chunk with finish_reason="tool_calls" + # which will trigger tool execution and follow-up call + chunks = [] + async for chunk in result: + chunks.append(chunk) + # After consuming the initial chunk, the follow-up call should be made + # Break after first chunk since that's when follow-up is triggered + break + + # With new implementation, first call should be streaming + assert mock_acompletion_func.await_count >= 2 + first_call = mock_acompletion_func.await_args_list[0].kwargs + # First call should be streaming in new implementation + assert first_call["stream"] is True + # Find the follow-up call (should have tool role messages) + follow_up_call = None + for call in mock_acompletion_func.await_args_list: + messages = call.kwargs.get("messages", []) + if messages and any(msg.get("role") == "tool" for msg in messages if isinstance(msg, dict)): + follow_up_call = call.kwargs + break + assert follow_up_call is not None, "Should have a follow-up call" + assert follow_up_call["stream"] is True @pytest.mark.asyncio @@ -191,7 +307,7 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): 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"}}] + tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}] tool_results = [{"tool_call_id": "call-1", "result": "executed"}] # Create mock streaming chunks @@ -234,19 +350,21 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): self._index = 0 self.sent_last_chunk = False - def __iter__(self): + def __aiter__(self): return self - def __next__(self): + async def __anext__(self): if self._index < len(self.chunks): chunk = self.chunks[self._index] self._index += 1 if self._index == len(self.chunks): self.sent_last_chunk = True - # Call the method that adds MCP metadata to final chunk - chunk = self._add_mcp_metadata_to_final_chunk(chunk) + # Add mcp_list_tools to first chunk if present + if not self.sent_first_chunk: + chunk = self._add_mcp_list_tools_to_first_chunk(chunk) + self.sent_first_chunk = True return chunk - raise StopIteration + raise StopAsyncIteration mock_acompletion = AsyncMock(return_value=MockStreamingResponse()) @@ -286,7 +404,7 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): with patch("litellm.acompletion", mock_acompletion): result = await acompletion_with_mcp( - model="test-model", + model="gpt-4o-mini", messages=[{"role": "user", "content": "hello"}], tools=tools, stream=True, @@ -302,30 +420,383 @@ async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): assert "mcp_list_tools" in mcp_metadata assert mcp_metadata["mcp_list_tools"] == openai_tools - # Consume the stream and check final chunk - all_chunks = list(result) + # Consume the stream and check chunks + all_chunks = [] + async for chunk in result: + all_chunks.append(chunk) assert len(all_chunks) > 0 - # Find the final chunk (with finish_reason) - final_chunk = None + # Verify mcp_list_tools is in the first chunk + first_chunk = all_chunks[0] if all_chunks 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) + # mcp_list_tools should be added to the first chunk + assert provider_fields is not None, f"First chunk should have provider_specific_fields. Delta: {choice.delta}" + assert "mcp_list_tools" in provider_fields, f"First chunk should have mcp_list_tools. Fields: {provider_fields}" + assert provider_fields["mcp_list_tools"] == openai_tools + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_streaming_initial_call_is_streaming(monkeypatch): + """ + Test that acompletion_with_mcp makes the initial LLM call with streaming=True + when stream=True is requested, instead of making a non-streaming call first. + """ + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None): + return ModelResponseStream( + id="test-stream", + model="test-model", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + ), + finish_reason=finish_reason, + ) + ], + ) + + chunks = [ + create_chunk("", finish_reason="tool_calls"), # Final chunk with tool_calls + ] + + # Create a proper CustomStreamWrapper + from unittest.mock import MagicMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test-model", + logging_obj=logging_obj, + ) + self.chunks = 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 + + mock_acompletion = AsyncMock(return_value=MockStreamingResponse()) + + 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 **_: [{"id": "call-1", "type": "function", "function": {"name": "local_search", "arguments": "{}"}}]), + ) + async def mock_execute(**_): + return [{"tool_call_id": "call-1", "result": "executed"}] + + 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)), + ) + + # Patch litellm.acompletion at module level to catch function-level imports + 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=True, + ) + + # Verify result is CustomStreamWrapper + assert isinstance(result, CustomStreamWrapper) + + # Verify that the first call was made with stream=True + assert mock_acompletion.await_count >= 1 + first_call = mock_acompletion.await_args_list[0].kwargs + assert first_call["stream"] is True, "First call should be streaming with new implementation" + + +@pytest.mark.asyncio +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 + """ + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta, 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"}] + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None, tool_calls=None): + 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, + ) + ], + ) + + initial_chunks = [ + create_chunk( + "", + finish_reason="tool_calls", + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call-1", + type="function", + function=Function(name="local_search", arguments="{}"), + index=0, + ) + ], + ), # Final chunk with tool_calls + ] + + follow_up_chunks = [ + create_chunk("Hello"), + create_chunk(" world", finish_reason="stop"), + ] + + # Create a proper CustomStreamWrapper + from unittest.mock import MagicMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class InitialStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test-model", + logging_obj=logging_obj, + ) + self.chunks = initial_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 + + class FollowUpStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test-model", + logging_obj=logging_obj, + ) + self.chunks = follow_up_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 + + acompletion_calls = [] + + async def mock_acompletion(**kwargs): + acompletion_calls.append(kwargs) + if kwargs.get("stream", False): + messages = kwargs.get("messages", []) + is_follow_up = any( + msg.get("role") == "tool" or (isinstance(msg, dict) and "tool_call_id" in str(msg)) + for msg in messages + ) + if is_follow_up: + return FollowUpStreamingResponse() + else: + return InitialStreamingResponse() + pytest.fail("Non-streaming call should not happen with new implementation") + + mock_acompletion_func = AsyncMock(side_effect=mock_acompletion) + + 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)), + ) + + # Patch litellm.acompletion at module level to catch function-level imports + with patch("litellm.acompletion", mock_acompletion_func), \ + patch.object(chat_completions_handler, "litellm_acompletion", side_effect=mock_acompletion, create=True): + result = await acompletion_with_mcp( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + # Verify result is CustomStreamWrapper + assert isinstance(result, CustomStreamWrapper) + + # Consume the stream and verify metadata placement + all_chunks = [] + async for chunk in result: + 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 + 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: - final_chunk = chunk - break + if hasattr(choice, "finish_reason") and choice.finish_reason == "tool_calls": + initial_final_chunk = chunk - # If no chunk with finish_reason, use the last chunk - if final_chunk is None and all_chunks: - final_chunk = all_chunks[-1] + 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" - assert final_chunk is not None, "Should have a final chunk" + # print(first_chunk) + # Verify mcp_list_tools is in the 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" - # Verify MCP metadata is in the final chunk's delta.provider_specific_fields - if hasattr(final_chunk, "choices") and final_chunk.choices: - choice = final_chunk.choices[0] + # 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_list_tools" in provider_fields, "Should have mcp_list_tools" - assert provider_fields["mcp_list_tools"] == openai_tools + assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls" + assert "mcp_call_results" in provider_fields, "Should have mcp_call_results" From a06b5ffe243d56c8d19684c91847d412a89384b3 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 23 Jan 2026 15:32:46 +0900 Subject: [PATCH 4/4] chore: fix lint error --- litellm/responses/mcp/chat_completions_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 09c2af9589..4a640a6106 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -79,7 +79,7 @@ def _add_mcp_metadata_to_response( setattr(message, "provider_specific_fields", provider_fields) -async def acompletion_with_mcp( +async def acompletion_with_mcp( # noqa: PLR0915 model: str, messages: List, tools: Optional[List] = None,