From 1c2942d808571cd50efa3744f11e403bc76b83a5 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 15 Jan 2026 15:47:45 +0900 Subject: [PATCH] test: add mcp completions test --- .../responses/mcp/chat_completions_handler.py | 185 ------------------ tests/mcp_tests/test_mcp_chat_completions.py | 171 ++++++++++++++++ 2 files changed, 171 insertions(+), 185 deletions(-) diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 641dc60ebd..6ce59e3e67 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -2,203 +2,18 @@ from typing import ( Any, - Awaitable, - Callable, - Dict, - Iterable, List, Optional, Union, - cast, ) from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.llms.openai import ToolParam from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper -CompletionCallable = Callable[..., Awaitable[Union[ModelResponse, CustomStreamWrapper]]] - -_CHAT_COMPLETION_CALL_ARG_KEYS = [ - "model", - "messages", - "functions", - "function_call", - "timeout", - "temperature", - "top_p", - "n", - "stream", - "stream_options", - "stop", - "max_tokens", - "max_completion_tokens", - "modalities", - "prediction", - "audio", - "presence_penalty", - "frequency_penalty", - "logit_bias", - "user", - "response_format", - "seed", - "tools", - "tool_choice", - "parallel_tool_calls", - "logprobs", - "top_logprobs", - "deployment_id", - "reasoning_effort", - "verbosity", - "safety_identifier", - "service_tier", - "base_url", - "api_version", - "api_key", - "model_list", - "extra_headers", - "thinking", - "web_search_options", - "shared_session", -] - - -def _build_call_args_from_context(call_context: Dict[str, Any]) -> Dict[str, Any]: - """Build kwargs for `acompletion` from the `completion` call context.""" - - call_args = { - key: call_context.get(key) - for key in _CHAT_COMPLETION_CALL_ARG_KEYS - if key in call_context - } - additional_kwargs = dict(call_context.get("kwargs") or {}) - call_args.update(additional_kwargs) - return call_args - - -async def _call_acompletion_internal( - completion_callable: CompletionCallable, **call_args: Any -) -> Union[ModelResponse, CustomStreamWrapper]: - """Invoke `acompletion` while skipping MCP interception to avoid recursion.""" - - safe_args = dict(call_args) - safe_args["_skip_mcp_handler"] = True - safe_args.pop("acompletion", None) - return await completion_callable(**safe_args) - - -async def handle_chat_completion_with_mcp( - call_context: Dict[str, Any], - completion_callable: CompletionCallable, -) -> Optional[Union[ModelResponse, CustomStreamWrapper]]: - """Handle MCP-enabled tool execution for chat completion requests.""" - - call_args = _build_call_args_from_context(call_context) - - tools = call_args.get("tools") - if not tools: - return None - - tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) - - if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( - tools=tools_for_mcp - ): - return None - - mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) - if not mcp_tools: - return None - - base_call_args = dict(call_args) - - user_api_key_auth = call_args.get("user_api_key_auth") or ( - (call_args.get("metadata", {}) or {}).get("user_api_key_auth") - ) - ( - deduplicated_mcp_tools, - tool_server_map, - ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( - user_api_key_auth=user_api_key_auth, - mcp_tools_with_litellm_proxy=mcp_tools, - ) - - openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( - deduplicated_mcp_tools, - target_format="chat", - ) - - base_call_args["tools"] = openai_tools or None - - should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( - mcp_tools_with_litellm_proxy=mcp_tools - ) - - ( - mcp_auth_header, - mcp_server_auth_headers, - oauth2_headers, - raw_headers, - ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( - secret_fields=base_call_args.get("secret_fields"), - tools=tools, - ) - - if not should_auto_execute: - return await _call_acompletion_internal(completion_callable, **base_call_args) - - mock_tool_calls = base_call_args.pop("mock_tool_calls", None) - - initial_call_args = dict(base_call_args) - initial_call_args["stream"] = False - if mock_tool_calls is not None: - initial_call_args["mock_tool_calls"] = mock_tool_calls - - initial_response = await _call_acompletion_internal( - completion_callable, **initial_call_args - ) - if not isinstance(initial_response, ModelResponse): - return initial_response - - tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( - response=initial_response - ) - - if not tool_calls: - if base_call_args.get("stream"): - retry_args = dict(base_call_args) - retry_args["stream"] = call_args.get("stream") - return await _call_acompletion_internal(completion_callable, **retry_args) - return initial_response - - tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_server_map=tool_server_map, - tool_calls=tool_calls, - 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, - ) - - if not tool_results: - return initial_response - - follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( - original_messages=call_args.get("messages", []), - response=initial_response, - tool_results=tool_results, - ) - - follow_up_call_args = dict(base_call_args) - follow_up_call_args["messages"] = follow_up_messages - follow_up_call_args["stream"] = call_args.get("stream") - - return await _call_acompletion_internal(completion_callable, **follow_up_call_args) - async def acompletion_with_mcp( model: str, diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index ae13b6ca6e..973301abfb 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -141,3 +141,174 @@ async def test_acompletion_mcp_respects_manual_approval(monkeypatch): assert isinstance(response, ModelResponse) tool_calls = response.choices[0].message.tool_calls assert tool_calls is not None and len(tool_calls) == 1 + + +@pytest.mark.asyncio +async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): + """ + Test that litellm.completion with stream=True and MCP tools does not raise + RuntimeError: Timeout context manager should be used inside a task. + + This test ensures that the fix in ba43f742ab86d51b7da63077b85b39d0ac808d30 + prevents event loop nesting issues when using MCP tools with streaming. + + The fix changes completion() to return a coroutine from acompletion_with_mcp, + which acompletion() then awaits, avoiding event loop nesting. + """ + 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 + + 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): + fake_execute.called = True # type: ignore[attr-defined] + tool_calls = kwargs.get("tool_calls") or [] + assert tool_calls, "tool calls should be present during auto execution" + 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"), + } + ] + + fake_execute.called = False # type: ignore[attr-defined] + + 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 a mock streaming response + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + self.chunks = [ + type('Chunk', (), { + 'choices': [type('Choice', (), { + 'delta': type('Delta', (), { + 'content': 'Final' + })() + })()] + })(), + type('Chunk', (), { + 'choices': [type('Choice', (), { + 'delta': type('Delta', (), { + 'content': ' answer' + })() + })()] + })(), + ] + 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 + 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 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", + ) + # Second call (streaming follow-up) + return MockStreamingResponse() + + with patch("litellm.acompletion", side_effect=mock_acompletion): + # This should not raise RuntimeError: Timeout context manager should be used inside a task + # completion() returns a coroutine when MCP tools are present, which acompletion() awaits + 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": "{}"}, + } + ], + ) + + # completion() returns a coroutine when MCP tools are present + import asyncio + assert asyncio.iscoroutine(response), "completion() should return a coroutine when MCP tools are present" + + # Await the coroutine (this is what acompletion() does internally) + # This should not raise RuntimeError: Timeout context manager should be used inside a task + result = await response + + # Verify response is a streaming response + assert isinstance(result, CustomStreamWrapper) or hasattr(result, '__iter__') + + # Consume the stream to ensure it works + chunks = list(result) + assert len(chunks) > 0, "Should have received streaming chunks" + + # Verify tool execution was called + assert fake_execute.called is True # type: ignore[attr-defined] + + # Verify acompletion was called (should be called by acompletion_with_mcp) + assert len(acompletion_calls) >= 1, "acompletion should be called"