diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3093a37c26a..3304759f749 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1571,6 +1571,46 @@ class CustomStreamWrapper: ) 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. + + This method checks if MCP metadata is stored in _hidden_params and adds it to + the chunk's delta.provider_specific_fields, similar to how RAG adds search results. + """ + try: + # Check if MCP metadata should be added to final 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: + return chunk + + # Add MCP metadata 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 MCP metadata + if isinstance(mcp_metadata, dict): + provider_fields.update(mcp_metadata) + + # 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 metadata to final chunk: {str(e)}" + ) + + return chunk + def cache_streaming_response(self, processed_chunk, cache_hit: bool): """ Caches the streaming response @@ -1712,6 +1752,8 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) response._hidden_params["usage"] = usage + # Add MCP metadata to final chunk if present + response = self._add_mcp_metadata_to_final_chunk(response) # RETURN RESULT return response @@ -1884,6 +1926,8 @@ class CustomStreamWrapper: processed_chunk ) ) + # Add MCP metadata to final chunk if present (after hooks) + processed_chunk = self._add_mcp_metadata_to_final_chunk(processed_chunk) return processed_chunk raise StopAsyncIteration diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 26853b30596..0b0004d1548 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -15,6 +15,69 @@ from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper +def _add_mcp_metadata_to_response( + response: Union[ModelResponse, CustomStreamWrapper], + openai_tools: Optional[List], + tool_calls: Optional[List] = None, + tool_results: Optional[List] = None, +) -> None: + """ + Add MCP metadata to response's provider_specific_fields. + + This function adds MCP-related information to the response so that + clients can access which tools were available, which were called, and + what results were returned. + + For ModelResponse: adds to choices[].message.provider_specific_fields + For CustomStreamWrapper: stores in _hidden_params and automatically adds to + final chunk's delta.provider_specific_fields via CustomStreamWrapper._add_mcp_metadata_to_final_chunk() + """ + if isinstance(response, CustomStreamWrapper): + # For streaming, store MCP metadata in _hidden_params + # CustomStreamWrapper._add_mcp_metadata_to_final_chunk() will automatically + # add it to the final chunk's delta.provider_specific_fields + if not hasattr(response, "_hidden_params"): + response._hidden_params = {} + + mcp_metadata = {} + if openai_tools: + mcp_metadata["mcp_list_tools"] = openai_tools + if tool_calls: + mcp_metadata["mcp_tool_calls"] = tool_calls + if tool_results: + mcp_metadata["mcp_call_results"] = tool_results + + if mcp_metadata: + response._hidden_params["mcp_metadata"] = mcp_metadata + return + + if not isinstance(response, ModelResponse): + return + + if not hasattr(response, "choices") or not response.choices: + return + + # Add MCP metadata to all choices' messages + for choice in response.choices: + message = getattr(choice, "message", None) + if message is not None: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(message, "provider_specific_fields", None) or {} + ) + + # Add MCP metadata + if openai_tools: + provider_fields["mcp_list_tools"] = openai_tools + if tool_calls: + provider_fields["mcp_tool_calls"] = tool_calls + if tool_results: + provider_fields["mcp_call_results"] = tool_results + + # Set the provider_specific_fields + setattr(message, "provider_specific_fields", provider_fields) + + async def acompletion_with_mcp( model: str, messages: List, @@ -103,7 +166,13 @@ async def acompletion_with_mcp( # If not auto-executing, just make the call with transformed tools if not should_auto_execute: - return await litellm_acompletion(**base_call_args) + response = await litellm_acompletion(**base_call_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + ) + return response # For auto-execute: disable streaming for initial call stream = kwargs.get("stream", False) @@ -130,7 +199,17 @@ async def acompletion_with_mcp( if stream: retry_args = dict(base_call_args) retry_args["stream"] = stream - return await litellm_acompletion(**retry_args) + 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, + ) return initial_response # Execute tool calls @@ -147,6 +226,11 @@ async def acompletion_with_mcp( ) if not tool_results: + _add_mcp_metadata_to_response( + response=initial_response, + openai_tools=openai_tools, + tool_calls=tool_calls, + ) return initial_response # Create follow-up messages with tool results @@ -161,4 +245,12 @@ async def acompletion_with_mcp( follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream - return await litellm_acompletion(**follow_up_call_args) + response = await litellm_acompletion(**follow_up_call_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + tool_calls=tool_calls, + tool_results=tool_results, + ) + return response diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index 973301abfb2..8857f016df8 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -312,3 +312,209 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): # Verify acompletion was called (should be called by acompletion_with_mcp) assert len(acompletion_calls) >= 1, "acompletion should be called" + + +@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. + """ + 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 + from litellm.litellm_core_utils.litellm_logging import Logging + + 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): + return ModelResponseStream( + id="test-stream", + model="gpt-4o-mini", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + ), + finish_reason=finish_reason, + ) + ], + ) + + 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 + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = chunks + self._index = 0 + self.sent_last_chunk = False + + def __iter__(self): + return self + + def __next__(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) + return chunk + raise StopIteration + + # 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", + ) + # Second call (streaming follow-up) + return MockStreamingResponse() + + 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) + + # 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 final chunk + all_chunks = list(result) + assert len(all_chunks) > 0 + + # Find the final chunk (with finish_reason) + 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 no chunk with finish_reason, use the last chunk + if final_chunk is None and all_chunks: + final_chunk = all_chunks[-1] + + assert final_chunk is not None, "Should have a final chunk" + + # 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] + 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" 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 03a749a8083..bc1f4fb72b5 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -177,3 +177,155 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): assert first_call["stream"] is False assert second_call["messages"] == ["follow-up"] assert second_call["stream"] is True + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): + """ + Test that acompletion_with_mcp adds MCP metadata to CustomStreamWrapper + and it appears in the final chunk's delta.provider_specific_fields. + """ + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + from litellm.litellm_core_utils.litellm_logging import Logging + + 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_results = [{"tool_call_id": "call-1", "result": "executed"}] + + # 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("Hello"), + create_chunk(" world", finish_reason="stop"), # Final chunk + ] + + # 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 + self.sent_last_chunk = False + + def __iter__(self): + return self + + def __next__(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) + return chunk + raise StopIteration + + 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 **_: False), + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + with patch("litellm.acompletion", mock_acompletion): + result = await acompletion_with_mcp( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + # Verify result is CustomStreamWrapper + 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_metadata["mcp_list_tools"] == openai_tools + + # 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 + 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 no chunk with finish_reason, use the last chunk + if final_chunk is None and all_chunks: + final_chunk = all_chunks[-1] + + assert final_chunk is not None, "Should have a final chunk" + + # 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] + 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