mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
The test_streaming_mcp_events_validation test was flaky because: 1. It didn't mock the nested aresponses() call inside the iterator's _create_initial_response_iterator(), causing real API calls that fail without credentials 2. The iterator silently swallowed exceptions and set phase="finished", discarding pre-generated MCP discovery events 3. The _execute_tool_calls mock had wrong signature (missing tool_server_map) Production fix: MCPEnhancedStreamingIterator no longer sets phase="finished" on LLM call failure — it falls through to emit MCP discovery events first. Test fix: Added mock for litellm.responses.main.aresponses returning a fake async streaming iterator, fixed mock signatures, removed try/except that masked failures. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1414 lines
56 KiB
Python
1414 lines
56 KiB
Python
import logging
|
|
import os
|
|
import sys
|
|
import pytest
|
|
from typing import List, Any, cast
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
sys.path.insert(0, os.path.abspath("../../.."))
|
|
|
|
# Import required modules
|
|
import litellm
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse, OpenAIMcpServerTool, ToolParam
|
|
|
|
|
|
class MockUserAPIKeyAuth:
|
|
"""Mock UserAPIKeyAuth for testing"""
|
|
def __init__(self):
|
|
self.api_key = "test_key"
|
|
self.user_id = "test_user"
|
|
self.team_id = "test_team"
|
|
self.user_email = "test@example.com"
|
|
self.max_budget = 100.0
|
|
self.spend = 0.0
|
|
self.models = []
|
|
self.aliases = {}
|
|
self.config = {}
|
|
self.permissions = {}
|
|
self.metadata = {}
|
|
self.object_permission_id = "test_permission_id"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_helper_methods():
|
|
"""Test the core MCP helper methods in LiteLLM_Proxy_MCP_Handler"""
|
|
|
|
# Test _should_use_litellm_mcp_gateway
|
|
mcp_tools: List[Any] = [
|
|
{
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never"
|
|
}
|
|
]
|
|
|
|
other_tools: List[Any] = [
|
|
{
|
|
"type": "function",
|
|
"name": "get_weather",
|
|
"description": "Get weather info",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
]
|
|
|
|
# Should return True for MCP tools with litellm_proxy
|
|
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(mcp_tools) == True
|
|
|
|
# Should return False for other tools
|
|
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(other_tools) == False
|
|
|
|
# Should return False for None
|
|
assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) == False
|
|
|
|
# Test _parse_mcp_tools
|
|
mixed_tools = mcp_tools + other_tools
|
|
mcp_parsed, other_parsed = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(mixed_tools)
|
|
|
|
assert len(mcp_parsed) == 1
|
|
assert len(other_parsed) == 1
|
|
assert mcp_parsed[0]["type"] == "mcp"
|
|
assert other_parsed[0]["type"] == "function"
|
|
|
|
# Test _should_auto_execute_tools
|
|
mcp_tools_never = [{"require_approval": "never"}]
|
|
mcp_tools_always = [{"require_approval": "always"}]
|
|
|
|
assert LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_never) == True
|
|
assert LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(mcp_tools_always) == False
|
|
|
|
print("✓ MCP helper methods test passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_output_elements_addition():
|
|
"""Test adding MCP output elements to response"""
|
|
|
|
# Create a mock response
|
|
mock_response = ResponsesAPIResponse(
|
|
**{ # type: ignore
|
|
"id": "test_response_id",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"status": "completed",
|
|
"error": None,
|
|
"incomplete_details": None,
|
|
"instructions": None,
|
|
"max_output_tokens": None,
|
|
"model": "gpt-4o",
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"id": "msg_123",
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": "Hello, world!",
|
|
"annotations": []
|
|
}
|
|
]
|
|
}
|
|
],
|
|
"parallel_tool_calls": True,
|
|
"previous_response_id": None,
|
|
"reasoning": {"effort": None, "summary": None},
|
|
"store": True,
|
|
"temperature": 1.0,
|
|
"text": {"format": {"type": "text"}},
|
|
"tool_choice": "auto",
|
|
"tools": [],
|
|
"top_p": 1.0,
|
|
"truncation": "disabled",
|
|
"usage": {
|
|
"input_tokens": 10,
|
|
"input_tokens_details": {"cached_tokens": 0},
|
|
"output_tokens": 5,
|
|
"output_tokens_details": {"reasoning_tokens": 0},
|
|
"total_tokens": 15
|
|
},
|
|
"user": None,
|
|
"metadata": {}
|
|
}
|
|
)
|
|
|
|
# Mock MCP tools and tool results
|
|
mock_mcp_tools = [
|
|
{
|
|
"name": "test_tool",
|
|
"description": "A test tool",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string"}
|
|
}
|
|
}
|
|
}
|
|
]
|
|
|
|
mock_tool_results = [
|
|
{
|
|
"tool_call_id": "call_123",
|
|
"result": "Tool executed successfully"
|
|
}
|
|
]
|
|
|
|
# Test adding output elements
|
|
updated_response = LiteLLM_Proxy_MCP_Handler._add_mcp_output_elements_to_response(
|
|
response=mock_response,
|
|
mcp_tools_fetched=mock_mcp_tools,
|
|
tool_results=mock_tool_results
|
|
)
|
|
|
|
# Verify output elements were added
|
|
assert len(updated_response.output) == 3 # Original + 2 new elements
|
|
|
|
# Check that MCP tools output was added - handle both dict and object cases
|
|
mcp_tools_output = updated_response.output[1]
|
|
if hasattr(mcp_tools_output, 'type'):
|
|
# Handle as object with attributes
|
|
output_obj = cast(Any, mcp_tools_output)
|
|
assert output_obj.type == "mcp_tools_fetched"
|
|
assert output_obj.role == "system"
|
|
assert output_obj.status == "completed"
|
|
elif isinstance(mcp_tools_output, dict):
|
|
# Handle as dictionary
|
|
assert mcp_tools_output["type"] == "mcp_tools_fetched"
|
|
assert mcp_tools_output["role"] == "system"
|
|
assert mcp_tools_output["status"] == "completed"
|
|
|
|
# Check that tool results output was added
|
|
tool_results_output = updated_response.output[2]
|
|
if hasattr(tool_results_output, 'type'):
|
|
# Handle as object with attributes
|
|
output_obj = cast(Any, tool_results_output)
|
|
assert output_obj.type == "tool_execution_results"
|
|
assert output_obj.role == "system"
|
|
assert output_obj.status == "completed"
|
|
elif isinstance(tool_results_output, dict):
|
|
# Handle as dictionary
|
|
assert tool_results_output["type"] == "tool_execution_results"
|
|
assert tool_results_output["role"] == "system"
|
|
assert tool_results_output["status"] == "completed"
|
|
|
|
print("✓ MCP output elements addition test passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aresponses_api_with_mcp_mock_integration():
|
|
"""
|
|
Test the core MCP integration logic without complex external mocking.
|
|
This focuses on verifying the MCP tool parsing and handling works correctly.
|
|
"""
|
|
# Define MCP tools with litellm_proxy server_url and require_approval="never"
|
|
mcp_tools: List[OpenAIMcpServerTool] = [
|
|
{
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never",
|
|
"server_label": "test_server"
|
|
}
|
|
]
|
|
|
|
# Test the helper methods that the integration relies on
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
|
|
# Test 1: Verify MCP tools are detected correctly
|
|
should_use_mcp = LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(cast(Any, mcp_tools))
|
|
assert should_use_mcp == True, "Should detect MCP tools with litellm_proxy server_url"
|
|
|
|
# Test 2: Verify auto-execution detection works
|
|
should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(cast(Any, mcp_tools))
|
|
assert should_auto_execute == True, "Should auto-execute tools with require_approval='never'"
|
|
|
|
# Test 3: Verify tool parsing works correctly
|
|
mcp_parsed, other_parsed = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(cast(Any, mcp_tools))
|
|
assert len(mcp_parsed) == 1, "Should parse one MCP tool"
|
|
assert len(other_parsed) == 0, "Should have no other tools"
|
|
assert mcp_parsed[0]["type"] == "mcp", "Parsed tool should be MCP type"
|
|
assert mcp_parsed[0]["server_url"] == "litellm_proxy", "Should preserve server_url"
|
|
assert mcp_parsed[0].get("require_approval") == "never", "Should preserve require_approval"
|
|
|
|
# Test 4: Test with mixed tools
|
|
mixed_tools = mcp_tools + [
|
|
{
|
|
"type": "function",
|
|
"name": "test_function",
|
|
"parameters": {"type": "object"}
|
|
}
|
|
]
|
|
|
|
mcp_parsed, other_parsed = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(cast(Any, mixed_tools))
|
|
assert len(mcp_parsed) == 1, "Should parse one MCP tool from mixed list"
|
|
assert len(other_parsed) == 1, "Should have one other tool from mixed list"
|
|
|
|
print("✓ MCP integration core logic test completed successfully!")
|
|
print(f"MCP tools detected: {should_use_mcp}")
|
|
print(f"Auto-execute enabled: {should_auto_execute}")
|
|
print(f"MCP tools parsed: {len(mcp_parsed)}")
|
|
print(f"Other tools parsed: {len(other_parsed)}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aresponses_api_with_mcp_passes_mcp_server_auth_headers_to_process_tools():
|
|
"""
|
|
Test that MCP auth headers from secret_fields (e.g. x-mcp-linear_config-authorization)
|
|
are passed to _process_mcp_tools_without_openai_transform when using the responses API.
|
|
"""
|
|
from litellm.responses.main import aresponses_api_with_mcp
|
|
|
|
captured_process_kwargs = {}
|
|
|
|
async def mock_process(**kwargs):
|
|
captured_process_kwargs.update(kwargs)
|
|
return ([], {})
|
|
|
|
mock_response = ResponsesAPIResponse(
|
|
**{
|
|
"id": "resp_test",
|
|
"object": "response",
|
|
"created_at": 1234567890,
|
|
"status": "completed",
|
|
"error": None,
|
|
"incomplete_details": None,
|
|
"instructions": None,
|
|
"max_output_tokens": None,
|
|
"model": "gpt-4o",
|
|
"output": [{"type": "message", "id": "msg_1", "status": "completed", "role": "assistant", "content": []}],
|
|
"parallel_tool_calls": True,
|
|
"previous_response_id": None,
|
|
"reasoning": {"effort": None, "summary": None},
|
|
"store": True,
|
|
"temperature": 1.0,
|
|
"text": {"format": {"type": "text"}},
|
|
"tool_choice": "auto",
|
|
"tools": [],
|
|
"top_p": 1.0,
|
|
"truncation": "disabled",
|
|
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
|
"user": None,
|
|
"metadata": {},
|
|
}
|
|
)
|
|
|
|
mcp_tools = [{"type": "mcp", "server_url": "litellm_proxy"}]
|
|
secret_fields = {
|
|
"raw_headers": {"x-mcp-linear_config-authorization": "Bearer linear-token"},
|
|
}
|
|
|
|
with patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_process_mcp_tools_without_openai_transform",
|
|
mock_process,
|
|
), patch(
|
|
"litellm.responses.main.aresponses",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
):
|
|
await aresponses_api_with_mcp(
|
|
input=[{"role": "user", "type": "message", "content": "hi"}],
|
|
model="gpt-4o",
|
|
tools=mcp_tools,
|
|
secret_fields=secret_fields,
|
|
)
|
|
|
|
assert "mcp_server_auth_headers" in captured_process_kwargs
|
|
mcp_server_auth_headers = captured_process_kwargs["mcp_server_auth_headers"]
|
|
assert mcp_server_auth_headers is not None
|
|
assert "linear_config" in mcp_server_auth_headers
|
|
assert mcp_server_auth_headers["linear_config"]["Authorization"] == "Bearer linear-token"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_allowed_tools_filtering():
|
|
"""
|
|
Test the allowed_tools filtering functionality for MCP tools.
|
|
This test verifies that when allowed_tools is specified in MCP tool config,
|
|
only the allowed tools are passed to the LLM.
|
|
"""
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
|
|
# Mock MCP tools returned from the server (simulating all available tools)
|
|
mock_mcp_tools_from_server = [
|
|
# Mock MCP tool object with name attribute
|
|
type('MCPTool', (), {
|
|
'name': 'search_tiktoken_documentation',
|
|
'description': 'Search tiktoken documentation',
|
|
'inputSchema': {'type': 'object', 'properties': {'query': {'type': 'string'}}}
|
|
})(),
|
|
type('MCPTool', (), {
|
|
'name': 'fetch_tiktoken_documentation',
|
|
'description': 'Fetch tiktoken documentation',
|
|
'inputSchema': {'type': 'object', 'properties': {'path': {'type': 'string'}}}
|
|
})(),
|
|
type('MCPTool', (), {
|
|
'name': 'list_tiktoken_functions',
|
|
'description': 'List tiktoken functions',
|
|
'inputSchema': {'type': 'object', 'properties': {}}
|
|
})(),
|
|
type('MCPTool', (), {
|
|
'name': 'get_tiktoken_examples',
|
|
'description': 'Get tiktoken examples',
|
|
'inputSchema': {'type': 'object', 'properties': {}}
|
|
})()
|
|
]
|
|
|
|
allowed_mcp_servers = ["gitmcp"]
|
|
|
|
# Test Case 1: MCP tool config with allowed_tools specified
|
|
mcp_tool_config_with_allowed_tools = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": ["search_tiktoken_documentation", "fetch_tiktoken_documentation"],
|
|
"require_approval": "never"
|
|
}
|
|
]
|
|
|
|
# Filter tools using the helper function
|
|
filtered_tools = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], mcp_tool_config_with_allowed_tools)
|
|
)
|
|
|
|
# Should only return the 2 allowed tools
|
|
assert len(filtered_tools) == 2, f"Expected 2 filtered tools, got {len(filtered_tools)}"
|
|
|
|
# Check that only allowed tools are included
|
|
filtered_tool_names = [tool.name for tool in filtered_tools]
|
|
expected_allowed_tools = ["search_tiktoken_documentation", "fetch_tiktoken_documentation"]
|
|
|
|
assert set(filtered_tool_names) == set(expected_allowed_tools), \
|
|
f"Expected tools {expected_allowed_tools}, got {filtered_tool_names}"
|
|
|
|
# Verify excluded tools are not present
|
|
excluded_tools = ["list_tiktoken_functions", "get_tiktoken_examples"]
|
|
for excluded_tool in excluded_tools:
|
|
assert excluded_tool not in filtered_tool_names, \
|
|
f"Tool {excluded_tool} should have been filtered out"
|
|
|
|
print("✓ Test Case 1: allowed_tools filtering works correctly")
|
|
|
|
# Test Case 2: MCP tool config without allowed_tools (should return all tools)
|
|
mcp_tool_config_without_allowed_tools = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"require_approval": "never"
|
|
}
|
|
]
|
|
|
|
filtered_tools_all = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], mcp_tool_config_without_allowed_tools)
|
|
)
|
|
|
|
# Should return all 4 tools when no allowed_tools specified
|
|
assert len(filtered_tools_all) == 4, f"Expected 4 tools when no allowed_tools specified, got {len(filtered_tools_all)}"
|
|
|
|
print("✓ Test Case 2: no allowed_tools returns all tools")
|
|
|
|
# Test Case 3: Test deduplication of duplicate tools
|
|
mock_mcp_tools_with_duplicates = [
|
|
# First instance of duplicate tool
|
|
type('MCPTool', (), {
|
|
'name': 'GitMCP-fetch_litellm_documentation',
|
|
'description': 'Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.',
|
|
'inputSchema': {'type': 'object', 'properties': {}, 'additionalProperties': False}
|
|
})(),
|
|
# Second instance of duplicate tool (should be filtered out)
|
|
type('MCPTool', (), {
|
|
'name': 'GitMCP-fetch_litellm_documentation',
|
|
'description': 'Fetch entire documentation file from GitHub repository: BerriAI/litellm. Useful for general questions. Always call this tool first if asked about BerriAI/litellm.',
|
|
'inputSchema': {'type': 'object', 'properties': {}, 'additionalProperties': False}
|
|
})(),
|
|
# Other unique tools
|
|
type('MCPTool', (), {
|
|
'name': 'GitMCP-search_litellm_documentation',
|
|
'description': 'Semantically search within the fetched documentation from GitHub repository: BerriAI/litellm. Useful for specific queries.',
|
|
'inputSchema': {'type': 'object', 'properties': {'query': {'type': 'string'}}, 'required': ['query'], 'additionalProperties': False}
|
|
})(),
|
|
]
|
|
|
|
mcp_tool_config_with_duplicates = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "litellm",
|
|
"server_url": "litellm_proxy/mcp",
|
|
"require_approval": "never",
|
|
"allowed_tools": ["GitMCP-fetch_litellm_documentation"]
|
|
}
|
|
]
|
|
|
|
# First filter by allowed tools
|
|
filtered_tools_with_duplicates = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_with_duplicates,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], mcp_tool_config_with_duplicates)
|
|
)
|
|
|
|
# Then deduplicate the filtered tools
|
|
filtered_tools_deduplicated, _ = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(
|
|
filtered_tools_with_duplicates, []
|
|
)
|
|
|
|
# Should only return 1 tool (the duplicate should be removed)
|
|
assert len(filtered_tools_deduplicated) == 1, f"Expected 1 tool after deduplication, got {len(filtered_tools_deduplicated)}"
|
|
|
|
# Check that the correct tool is present
|
|
assert filtered_tools_deduplicated[0].name == "GitMCP-fetch_litellm_documentation", \
|
|
f"Expected GitMCP-fetch_litellm_documentation, got {filtered_tools_deduplicated[0].name}"
|
|
|
|
print("✓ Test Case 3: duplicate tools are properly deduplicated")
|
|
|
|
# Test Case 3b: Test standalone deduplication method
|
|
standalone_deduplicated, _ = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools(mock_mcp_tools_with_duplicates, allowed_mcp_servers)
|
|
|
|
# Should return 2 unique tools (GitMCP-fetch_litellm_documentation and GitMCP-search_litellm_documentation)
|
|
assert len(standalone_deduplicated) == 2, f"Expected 2 unique tools after standalone deduplication, got {len(standalone_deduplicated)}"
|
|
|
|
unique_tool_names = [tool.name for tool in standalone_deduplicated]
|
|
expected_unique_names = ["GitMCP-fetch_litellm_documentation", "GitMCP-search_litellm_documentation"]
|
|
assert set(unique_tool_names) == set(expected_unique_names), \
|
|
f"Expected {expected_unique_names}, got {unique_tool_names}"
|
|
|
|
print("✓ Test Case 3b: standalone deduplication method works correctly")
|
|
|
|
# Test Case 4: Multiple MCP tool configs with different allowed_tools
|
|
multiple_mcp_configs = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp1",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": ["search_tiktoken_documentation"],
|
|
"require_approval": "never"
|
|
},
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp2",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": ["fetch_tiktoken_documentation", "get_tiktoken_examples"],
|
|
"require_approval": "never"
|
|
}
|
|
]
|
|
|
|
filtered_tools_multiple = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], multiple_mcp_configs)
|
|
)
|
|
|
|
# Should return union of all allowed tools (3 unique tools)
|
|
assert len(filtered_tools_multiple) == 3, f"Expected 3 tools from multiple configs, got {len(filtered_tools_multiple)}"
|
|
|
|
filtered_multiple_names = [tool.name for tool in filtered_tools_multiple]
|
|
expected_multiple_tools = ["search_tiktoken_documentation", "fetch_tiktoken_documentation", "get_tiktoken_examples"]
|
|
|
|
assert set(filtered_multiple_names) == set(expected_multiple_tools), \
|
|
f"Expected tools {expected_multiple_tools}, got {filtered_multiple_names}"
|
|
|
|
print("✓ Test Case 3: multiple MCP configs with different allowed_tools works correctly")
|
|
|
|
# Test Case 4: Empty allowed_tools list (should return no tools)
|
|
mcp_config_empty_allowed = [
|
|
{
|
|
"type": "mcp",
|
|
"server_label": "gitmcp",
|
|
"server_url": "https://gitmcp.io/openai/tiktoken",
|
|
"allowed_tools": [],
|
|
"require_approval": "never"
|
|
}
|
|
]
|
|
|
|
filtered_tools_empty = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools(
|
|
mcp_tools=mock_mcp_tools_from_server,
|
|
mcp_tools_with_litellm_proxy=cast(List[ToolParam], mcp_config_empty_allowed)
|
|
)
|
|
|
|
# Should return all tools when allowed_tools is empty list (no filtering)
|
|
assert len(filtered_tools_empty) == 4, f"Expected 4 tools when allowed_tools is empty list, got {len(filtered_tools_empty)}"
|
|
|
|
print("✓ Test Case 4: empty allowed_tools list returns all tools")
|
|
|
|
print("✓ MCP allowed_tools filtering test completed successfully!")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_mcp_events_validation():
|
|
"""
|
|
Test that MCP streaming events are properly emitted when using streaming with MCP tools.
|
|
|
|
This test validates:
|
|
1. MCP discovery events are emitted first
|
|
2. Regular streaming response events follow
|
|
3. Tool execution events are emitted when tools are auto-executed
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from litellm.types.llms.openai import ResponsesAPIStreamEvents
|
|
|
|
# Mock MCP tools that would be returned from the manager
|
|
mock_mcp_tools = [
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "search_repo",
|
|
"description": "Search BerriAI/litellm repository for information",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"}
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
)(),
|
|
type(
|
|
"MCPTool",
|
|
(),
|
|
{
|
|
"name": "get_repo_info",
|
|
"description": "Get repository information",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"repo_name": {
|
|
"type": "string",
|
|
"description": "Repository name",
|
|
}
|
|
},
|
|
"required": ["repo_name"],
|
|
},
|
|
},
|
|
)(),
|
|
]
|
|
|
|
# Build fake streaming chunks that the inner aresponses() call would yield
|
|
fake_response_obj = MagicMock(spec=ResponsesAPIResponse)
|
|
fake_response_obj.id = "resp_fake_123"
|
|
fake_response_obj.output = []
|
|
|
|
fake_created_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_created_chunk.type = ResponsesAPIStreamEvents.RESPONSE_CREATED
|
|
fake_created_chunk.response = fake_response_obj
|
|
|
|
fake_in_progress_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_in_progress_chunk.type = ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS
|
|
fake_in_progress_chunk.response = fake_response_obj
|
|
|
|
fake_output_item_added_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_output_item_added_chunk.type = ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
|
|
fake_output_item_added_chunk.response = fake_response_obj
|
|
|
|
fake_completed_chunk = MagicMock(spec=ResponsesAPIStreamingResponse)
|
|
fake_completed_chunk.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED
|
|
fake_completed_chunk.response = fake_response_obj
|
|
|
|
# Create a fake async iterator for the inner LLM streaming call
|
|
class FakeAsyncIterator:
|
|
def __init__(self, chunks):
|
|
self._chunks = list(chunks)
|
|
self._index = 0
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
if self._index >= len(self._chunks):
|
|
raise StopAsyncIteration
|
|
chunk = self._chunks[self._index]
|
|
self._index += 1
|
|
return chunk
|
|
|
|
fake_stream = FakeAsyncIterator(
|
|
[
|
|
fake_created_chunk,
|
|
fake_in_progress_chunk,
|
|
fake_output_item_added_chunk,
|
|
fake_completed_chunk,
|
|
]
|
|
)
|
|
|
|
# Mock the MCP operations and the inner aresponses call
|
|
with patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_get_mcp_tools_from_manager",
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools, patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
"_execute_tool_calls",
|
|
new_callable=AsyncMock,
|
|
) as mock_execute_tools, patch(
|
|
"litellm.responses.main.aresponses",
|
|
new_callable=AsyncMock,
|
|
return_value=fake_stream,
|
|
):
|
|
# Setup MCP mocks
|
|
mock_get_tools.return_value = (mock_mcp_tools, ["test_server"])
|
|
|
|
async def mock_execute_tool_calls_side_effect(
|
|
tool_server_map, tool_calls, user_api_key_auth, **kwargs
|
|
):
|
|
"""Mock tool execution with realistic results"""
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
call_id = None
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("call_id") or tool_call.get("id")
|
|
elif hasattr(tool_call, "call_id"):
|
|
call_id = tool_call.call_id
|
|
elif hasattr(tool_call, "id"):
|
|
call_id = tool_call.id
|
|
|
|
if call_id:
|
|
results.append(
|
|
{
|
|
"tool_call_id": call_id,
|
|
"result": "LiteLLM is a unified interface for 100+ LLMs.",
|
|
}
|
|
)
|
|
return results
|
|
|
|
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
|
|
|
|
# Configure MCP tool with streaming and auto-execution
|
|
mcp_tool_config = {
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy/mcp/test_server",
|
|
"require_approval": "never", # This enables auto-execution
|
|
}
|
|
|
|
# Make streaming request with MCP tools
|
|
response = await litellm.aresponses(
|
|
model="gpt-4o-mini",
|
|
tools=[mcp_tool_config],
|
|
tool_choice="required",
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "What is LiteLLM? Give me a brief overview.",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
|
|
assert hasattr(
|
|
response, "__aiter__"
|
|
), "Response should be async iterable for streaming"
|
|
|
|
# Collect all streaming events
|
|
events = []
|
|
event_types = []
|
|
mcp_discovery_events = []
|
|
regular_events = []
|
|
|
|
async for chunk in response:
|
|
events.append(chunk)
|
|
event_type = getattr(chunk, "type", "unknown")
|
|
event_types.append(event_type)
|
|
|
|
# Categorize events
|
|
if event_type in [
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
|
|
]:
|
|
mcp_discovery_events.append(chunk)
|
|
else:
|
|
regular_events.append(chunk)
|
|
|
|
# Validate that we got streaming events
|
|
assert len(events) > 0, "Should have received at least some streaming events"
|
|
|
|
# Validate MCP discovery events were emitted
|
|
assert (
|
|
len(mcp_discovery_events) > 0
|
|
), "Should have received MCP discovery events"
|
|
|
|
# Check that discovery events come before regular content events
|
|
first_discovery_idx = next(
|
|
i
|
|
for i, e in enumerate(events)
|
|
if getattr(e, "type", None)
|
|
in [
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
|
|
ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
|
|
]
|
|
)
|
|
# The output_item.added event triggers the transition to MCP discovery,
|
|
# so discovery events should appear after it in the stream
|
|
assert first_discovery_idx > 0, "MCP discovery events should follow the initial output_item.added event"
|
|
|
|
# Verify MCP mocks were called
|
|
assert mock_get_tools.called, "MCP tools should have been fetched"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"model",
|
|
[
|
|
pytest.param("gpt-4o-mini", id="openai"),
|
|
pytest.param("claude-haiku-4-5", id="anthropic"),
|
|
],
|
|
)
|
|
async def test_streaming_responses_api_with_mcp_tools(
|
|
model: str, caplog: pytest.LogCaptureFixture
|
|
):
|
|
"""
|
|
Test the streaming responses API with MCP tools when using server_url="litellm_proxy"
|
|
|
|
Under the hood the follow occurs
|
|
|
|
- MCP: responses called litellm MCP manager.list_tools (MOCKED)
|
|
- Request 1: Made to model under test with fetched tools (REAL LLM CALL)
|
|
- MCP: Execute tool call from request 1 and returns result (MOCKED)
|
|
- Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL)
|
|
|
|
Return the user the result of request 2
|
|
"""
|
|
# Skip test if API keys are not set for the respective models
|
|
if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv("ANTHROPIC_API_KEY"):
|
|
pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test")
|
|
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"):
|
|
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
print("🧪 Testing basic streaming with MCP tools...")
|
|
|
|
# Mock MCP tools that would be returned from the manager
|
|
mock_mcp_tools = [
|
|
type('MCPTool', (), {
|
|
'name': 'search_repo',
|
|
'description': 'Search BerriAI/litellm repository for information',
|
|
'inputSchema': {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
})()
|
|
]
|
|
|
|
# Only mock the MCP-specific operations, let LLM responses be real
|
|
with caplog.at_level(logging.ERROR):
|
|
with patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
'_get_mcp_tools_from_manager',
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools, patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
'_execute_tool_calls',
|
|
new_callable=AsyncMock,
|
|
) as mock_execute_tools:
|
|
# Setup MCP mocks only
|
|
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
|
|
|
|
# Create a dynamic mock that will match the actual tool call ID from the LLM response
|
|
def mock_execute_tool_calls_side_effect(
|
|
tool_calls, user_api_key_auth, **kwargs
|
|
):
|
|
"""Mock function that returns results matching the actual tool call IDs from the LLM"""
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
# Extract call_id from the tool call
|
|
call_id = None
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("call_id") or tool_call.get("id")
|
|
elif hasattr(tool_call, 'call_id'):
|
|
call_id = tool_call.call_id
|
|
elif hasattr(tool_call, 'id'):
|
|
call_id = tool_call.id
|
|
|
|
if call_id:
|
|
results.append(
|
|
{
|
|
"tool_call_id": call_id,
|
|
"result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.",
|
|
}
|
|
)
|
|
return results
|
|
|
|
mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect
|
|
|
|
# Make the actual call - LLM responses will be real
|
|
mcp_tool_config = cast(
|
|
Any,
|
|
{
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never",
|
|
},
|
|
)
|
|
response = await litellm.aresponses(
|
|
model=model,
|
|
tools=[mcp_tool_config],
|
|
tool_choice="required",
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "give me a TLDR of what BerriAI/litellm is about",
|
|
}
|
|
],
|
|
stream=True,
|
|
)
|
|
|
|
print(f"📋 Response type: {type(response)}")
|
|
assert hasattr(response, '__aiter__'), "Response should be an async streaming response"
|
|
|
|
# Collect streaming chunks
|
|
chunks = []
|
|
async for chunk in response:
|
|
chunks.append(chunk)
|
|
print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}")
|
|
|
|
print(f"📊 Total chunks received: {len(chunks)}")
|
|
|
|
# Verify MCP mocks were called (may be called multiple times in streaming)
|
|
assert (
|
|
mock_get_tools.call_count >= 1
|
|
), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}"
|
|
print(f"MCP tools fetched: {len(mock_mcp_tools)}")
|
|
|
|
# Verify we got a response
|
|
assert response is not None
|
|
assert len(chunks) > 0, "Should have received streaming chunks"
|
|
|
|
print("Basic streaming responses API with MCP tools test passed!")
|
|
|
|
lite_errors = [
|
|
record
|
|
for record in caplog.records
|
|
if record.levelno >= logging.ERROR
|
|
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
|
|
]
|
|
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
|
|
record.getMessage() for record in lite_errors
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_parameter_preparation_helpers():
|
|
"""
|
|
Test the new parameter preparation helper methods for clean MCP handling.
|
|
|
|
Tests:
|
|
1. _prepare_initial_call_params - handles stream disabling for auto-execute
|
|
2. _prepare_follow_up_call_params - restores stream and removes tool_choice
|
|
3. _build_request_params - clean parameter merging
|
|
"""
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
|
|
print("🧪 Testing MCP parameter preparation helpers...")
|
|
|
|
# Test _prepare_initial_call_params
|
|
base_call_params = {
|
|
"stream": True,
|
|
"temperature": 0.7,
|
|
"tool_choice": "required",
|
|
"max_output_tokens": 1000
|
|
}
|
|
|
|
# Test Case 1: Auto-execute scenario (should disable streaming)
|
|
initial_params_auto = LiteLLM_Proxy_MCP_Handler._prepare_initial_call_params(
|
|
call_params=base_call_params,
|
|
should_auto_execute=True
|
|
)
|
|
|
|
assert initial_params_auto["stream"] == False, "Stream should be disabled for auto-execute"
|
|
assert initial_params_auto["temperature"] == 0.7, "Other params should be preserved"
|
|
assert initial_params_auto["tool_choice"] == "required", "tool_choice should be preserved for initial call"
|
|
assert base_call_params["stream"] == True, "Original params should not be mutated"
|
|
|
|
print("✅ _prepare_initial_call_params (auto-execute) works correctly")
|
|
|
|
# Test Case 2: No auto-execute scenario (should preserve streaming)
|
|
initial_params_no_auto = LiteLLM_Proxy_MCP_Handler._prepare_initial_call_params(
|
|
call_params=base_call_params,
|
|
should_auto_execute=False
|
|
)
|
|
|
|
assert initial_params_no_auto["stream"] == True, "Stream should be preserved when not auto-executing"
|
|
assert initial_params_no_auto["temperature"] == 0.7, "Other params should be preserved"
|
|
|
|
print("✅ _prepare_initial_call_params (no auto-execute) works correctly")
|
|
|
|
# Test _prepare_follow_up_call_params
|
|
follow_up_params = LiteLLM_Proxy_MCP_Handler._prepare_follow_up_call_params(
|
|
call_params=base_call_params,
|
|
original_stream_setting=True
|
|
)
|
|
|
|
assert follow_up_params["stream"] == True, "Stream should be restored to original setting"
|
|
assert "tool_choice" not in follow_up_params, "tool_choice should be removed for follow-up call"
|
|
assert follow_up_params["temperature"] == 0.7, "Other params should be preserved"
|
|
assert base_call_params["tool_choice"] == "required", "Original params should not be mutated"
|
|
|
|
print("✅ _prepare_follow_up_call_params works correctly")
|
|
|
|
# Test _build_request_params
|
|
input_data = [{"role": "user", "content": "test", "type": "message"}]
|
|
model = "gpt-4o-mini"
|
|
tools = [{"type": "function", "name": "test_tool"}]
|
|
call_params = {"stream": True, "temperature": 0.8}
|
|
previous_response_id = "resp_123"
|
|
extra_kwargs = {"custom_param": "test_value"}
|
|
|
|
request_params = LiteLLM_Proxy_MCP_Handler._build_request_params(
|
|
input=input_data,
|
|
model=model,
|
|
all_tools=tools,
|
|
call_params=call_params,
|
|
previous_response_id=previous_response_id,
|
|
**extra_kwargs
|
|
)
|
|
|
|
# Verify core parameters
|
|
assert request_params["input"] == input_data, "Input should be included"
|
|
assert request_params["model"] == model, "Model should be included"
|
|
assert request_params["tools"] == tools, "Tools should be included"
|
|
assert request_params["previous_response_id"] == previous_response_id, "Previous response ID should be included"
|
|
|
|
# Verify call_params are merged
|
|
assert request_params["stream"] == True, "call_params should be merged"
|
|
assert request_params["temperature"] == 0.8, "call_params should be merged"
|
|
|
|
# Verify extra kwargs are merged
|
|
assert request_params["custom_param"] == "test_value", "Extra kwargs should be merged"
|
|
|
|
print("✅ _build_request_params works correctly")
|
|
|
|
# Test _build_request_params with None previous_response_id
|
|
request_params_no_prev = LiteLLM_Proxy_MCP_Handler._build_request_params(
|
|
input=input_data,
|
|
model=model,
|
|
all_tools=tools,
|
|
call_params=call_params,
|
|
previous_response_id=None
|
|
)
|
|
|
|
assert "previous_response_id" not in request_params_no_prev, "None previous_response_id should not be included"
|
|
|
|
print("✅ _build_request_params handles None previous_response_id correctly")
|
|
|
|
print("🎉 All MCP parameter preparation helper tests passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_tool_execution_events_creation():
|
|
"""
|
|
Test the _create_tool_execution_events helper method for generating streaming events.
|
|
"""
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
|
|
print("Testing MCP tool execution events creation...")
|
|
|
|
# Mock tool calls (simulating what comes from LLM response in function_call format)
|
|
mock_tool_calls = [
|
|
{
|
|
"id": "call_abc123",
|
|
"name": "search_repo",
|
|
"arguments": '{"query": "LiteLLM overview"}',
|
|
"type": "function_call"
|
|
},
|
|
{
|
|
"id": "call_def456",
|
|
"name": "get_repo_info",
|
|
"arguments": '{"repo_name": "BerriAI/litellm"}',
|
|
"type": "function_call"
|
|
}
|
|
]
|
|
|
|
# Mock tool results (simulating what comes from tool execution)
|
|
mock_tool_results = [
|
|
{
|
|
"tool_call_id": "call_abc123",
|
|
"result": "LiteLLM is a unified interface for 100+ LLMs"
|
|
},
|
|
{
|
|
"tool_call_id": "call_def456",
|
|
"result": "Repository: BerriAI/litellm - Python library for LLM integration"
|
|
}
|
|
]
|
|
|
|
# Create tool execution events
|
|
execution_events = LiteLLM_Proxy_MCP_Handler._create_tool_execution_events(
|
|
tool_calls=mock_tool_calls,
|
|
tool_results=mock_tool_results
|
|
)
|
|
|
|
# Verify events were created
|
|
assert len(execution_events) > 0, "Should create tool execution events"
|
|
print(f"Created {len(execution_events)} tool execution events")
|
|
|
|
# Verify events have proper structure
|
|
for event in execution_events:
|
|
assert hasattr(event, 'type'), "Event should have type attribute"
|
|
event_type = str(event.type)
|
|
assert 'mcp_call' in event_type.lower() or 'output_item' in event_type.lower(), f"Event should be MCP-related: {event_type}"
|
|
|
|
# Check for sequence numbers
|
|
if hasattr(event, 'sequence_number'):
|
|
assert isinstance(event.sequence_number, int), "Sequence number should be integer"
|
|
assert event.sequence_number > 0, "Sequence number should be positive"
|
|
|
|
print("Tool execution events have proper structure")
|
|
|
|
# Test with empty inputs
|
|
empty_events = LiteLLM_Proxy_MCP_Handler._create_tool_execution_events(
|
|
tool_calls=[],
|
|
tool_results=[]
|
|
)
|
|
|
|
assert len(empty_events) == 0, "Should create no events for empty inputs"
|
|
print("Handles empty inputs correctly")
|
|
|
|
print("MCP tool execution events creation test passed!")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_duplicate_mcp_tools_in_streaming_e2e():
|
|
"""
|
|
End-to-end test to validate that MCP tools are not duplicated when using streaming.
|
|
|
|
This test protects against the bug where:
|
|
1. Parent function (aresponses_api_with_mcp) processed MCP tools once
|
|
2. Streaming iterator processed MCP tools again, causing duplicates
|
|
|
|
The test mocks the MCP manager response but validates the actual tools
|
|
sent to the LLM to ensure no duplication occurs.
|
|
"""
|
|
from unittest.mock import AsyncMock, patch, call
|
|
from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler
|
|
|
|
print("Testing no duplicate MCP tools in streaming E2E...")
|
|
|
|
# Mock MCP tools that would be returned from the manager
|
|
mock_mcp_tools = [
|
|
type('MCPTool', (), {
|
|
'name': 'search_docs',
|
|
'description': 'Search documentation for information',
|
|
'inputSchema': {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Search query"}
|
|
},
|
|
"required": ["query"]
|
|
}
|
|
})(),
|
|
type('MCPTool', (), {
|
|
'name': 'get_file_content',
|
|
'description': 'Get content of a specific file',
|
|
'inputSchema': {
|
|
"type": "object",
|
|
"properties": {
|
|
"file_path": {"type": "string", "description": "Path to file"}
|
|
},
|
|
"required": ["file_path"]
|
|
}
|
|
})()
|
|
]
|
|
|
|
# Track all calls to the underlying LLM to detect duplicates
|
|
llm_call_tools = []
|
|
|
|
async def capture_llm_tools(**kwargs):
|
|
"""Capture the tools parameter from LLM calls"""
|
|
tools = kwargs.get('tools', [])
|
|
llm_call_tools.append(tools)
|
|
|
|
# Return a minimal mock async streaming response
|
|
class MockStreamingResponse:
|
|
async def __aiter__(self):
|
|
yield type('MockChunk', (), {
|
|
'type': 'response.completed',
|
|
'output': []
|
|
})()
|
|
|
|
return MockStreamingResponse()
|
|
|
|
# Mock both the MCP manager and the underlying LLM call
|
|
with patch.object(LiteLLM_Proxy_MCP_Handler, '_get_mcp_tools_from_manager', new_callable=AsyncMock) as mock_get_tools, \
|
|
patch('litellm.aresponses', side_effect=capture_llm_tools) as mock_aresponses:
|
|
|
|
# Setup MCP mock to return our test tools
|
|
mock_get_tools.return_value = mock_mcp_tools
|
|
|
|
# Configure MCP tool for streaming
|
|
mcp_tool_config = {
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy/mcp/test_server",
|
|
"require_approval": "always" # Disable auto-execution to focus on tool duplication
|
|
}
|
|
|
|
print("Making streaming request with MCP tools...")
|
|
|
|
# Make streaming request with MCP tools
|
|
try:
|
|
response = await litellm.aresponses(
|
|
model="gpt-4o-mini",
|
|
tools=[mcp_tool_config],
|
|
input=[{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "Search the documentation for information about authentication."
|
|
}],
|
|
stream=True
|
|
)
|
|
|
|
# Consume the streaming response
|
|
chunks = []
|
|
async for chunk in response:
|
|
chunks.append(chunk)
|
|
|
|
except Exception as e:
|
|
print(f"Request failed (expected for test): {e}")
|
|
# Continue with validation even if request fails
|
|
|
|
# Validate underlying LLM was called (this proves our mocking works)
|
|
assert len(llm_call_tools) > 0, "LLM should have been called at least once"
|
|
print(f"LLM called {len(llm_call_tools)} time(s)")
|
|
|
|
# If MCP tools were processed, validate they were fetched exactly once
|
|
# (This protects against duplicate fetching)
|
|
if mock_get_tools.call_count > 0:
|
|
assert mock_get_tools.call_count == 1, f"MCP tools should be fetched exactly once, got {mock_get_tools.call_count} calls"
|
|
print(f"MCP tools fetched exactly once: {mock_get_tools.call_count}")
|
|
else:
|
|
print("MCP tools not fetched (likely due to test mocking - this is OK for validation)")
|
|
|
|
# Analyze tools sent to LLM for duplicates
|
|
for call_idx, tools_in_call in enumerate(llm_call_tools):
|
|
print(f"LLM Call {call_idx + 1}: {len(tools_in_call)} tools")
|
|
|
|
if tools_in_call:
|
|
# Extract tool names to check for duplicates
|
|
tool_names = []
|
|
for tool in tools_in_call:
|
|
if isinstance(tool, dict):
|
|
tool_name = tool.get('function', {}).get('name') or tool.get('name')
|
|
else:
|
|
tool_name = getattr(tool, 'name', str(tool))
|
|
|
|
if tool_name:
|
|
tool_names.append(tool_name)
|
|
|
|
print(f" Tool names: {tool_names}")
|
|
|
|
# Check for duplicate tool names
|
|
unique_tool_names = set(tool_names)
|
|
duplicates = [name for name in tool_names if tool_names.count(name) > 1]
|
|
|
|
assert len(duplicates) == 0, f"Found duplicate tools in LLM call {call_idx + 1}: {duplicates}"
|
|
assert len(tool_names) == len(unique_tool_names), f"Tool names should be unique in call {call_idx + 1}"
|
|
|
|
print(f" No duplicate tools found in call {call_idx + 1}")
|
|
|
|
# Validate that MCP tools were properly transformed to OpenAI format
|
|
openai_format_tools = [tool for tool in tools_in_call if isinstance(tool, dict) and 'function' in tool]
|
|
if openai_format_tools:
|
|
print(f" Found {len(openai_format_tools)} OpenAI-format tools")
|
|
|
|
# Verify tools have proper OpenAI structure
|
|
for tool in openai_format_tools:
|
|
assert 'type' in tool, "Tool should have 'type' field"
|
|
assert tool['type'] == 'function', "Tool type should be 'function'"
|
|
assert 'function' in tool, "Tool should have 'function' field"
|
|
assert 'name' in tool['function'], "Function should have 'name'"
|
|
assert 'description' in tool['function'], "Function should have 'description'"
|
|
assert 'parameters' in tool['function'], "Function should have 'parameters'"
|
|
|
|
print(f" All tools have proper OpenAI format")
|
|
|
|
# The key validation: ensure no duplicate fetching occurred
|
|
# This is the main protection against the bug we fixed
|
|
if mock_get_tools.call_count > 1:
|
|
print(f"ERROR: Duplicate MCP fetching detected! Called {mock_get_tools.call_count} times")
|
|
assert False, f"MCP tools should be fetched exactly once, but were fetched {mock_get_tools.call_count} times"
|
|
|
|
# Additional validation: ensure no duplicate tools in any LLM call
|
|
total_duplicates_found = 0
|
|
for call_idx, tools_in_call in enumerate(llm_call_tools):
|
|
if tools_in_call:
|
|
tool_names = []
|
|
for tool in tools_in_call:
|
|
if isinstance(tool, dict):
|
|
tool_name = tool.get('function', {}).get('name') or tool.get('name')
|
|
if tool_name:
|
|
tool_names.append(tool_name)
|
|
|
|
duplicates = [name for name in tool_names if tool_names.count(name) > 1]
|
|
if duplicates:
|
|
total_duplicates_found += len(set(duplicates))
|
|
print(f"ERROR: Duplicate tools in call {call_idx + 1}: {set(duplicates)}")
|
|
|
|
if total_duplicates_found > 0:
|
|
assert False, f"Found {total_duplicates_found} duplicate tools across all LLM calls"
|
|
|
|
print("No duplicate MCP tools E2E test passed!")
|
|
print(f"Summary:")
|
|
print(f" - MCP manager called: {mock_get_tools.call_count} time(s)")
|
|
print(f" - LLM called: {len(llm_call_tools)} time(s)")
|
|
print(f" - Unique tools per call: {[len(set(getattr(t.get('function', {}), 'name', 'unknown') if isinstance(t, dict) else str(t) for t in tools)) for tools in llm_call_tools]}")
|
|
print(f" - No duplicate tools detected")
|
|
|
|
return {
|
|
'mcp_manager_calls': mock_get_tools.call_count,
|
|
'llm_calls': len(llm_call_tools),
|
|
'tools_per_call': [len(tools) for tools in llm_call_tools],
|
|
'duplicate_tools_found': False
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("model", ["gpt-4o-mini"])
|
|
async def test_streaming_mcp_event_order_and_response_id_consistency(
|
|
model: str, caplog: pytest.LogCaptureFixture
|
|
):
|
|
"""
|
|
Test that:
|
|
1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events)
|
|
2. All response lifecycle events share the same response ID within a cycle
|
|
"""
|
|
if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv("OPENAI_API_KEY"):
|
|
pytest.skip("OPENAI_API_KEY not set, skipping openai model test")
|
|
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
mock_mcp_tools = [
|
|
type('MCPTool', (), {
|
|
'name': 'get_weather',
|
|
'description': 'Get weather for a city',
|
|
'inputSchema': {
|
|
"type": "object",
|
|
"properties": {
|
|
"city": {"type": "string", "description": "City name"}
|
|
},
|
|
"required": ["city"]
|
|
}
|
|
})()
|
|
]
|
|
|
|
with caplog.at_level(logging.ERROR):
|
|
with patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
'_get_mcp_tools_from_manager',
|
|
new_callable=AsyncMock,
|
|
) as mock_get_tools, patch.object(
|
|
LiteLLM_Proxy_MCP_Handler,
|
|
'_execute_tool_calls',
|
|
new_callable=AsyncMock,
|
|
) as mock_execute_tools:
|
|
mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"])
|
|
|
|
def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs):
|
|
results = []
|
|
for tool_call in tool_calls:
|
|
call_id = None
|
|
if isinstance(tool_call, dict):
|
|
call_id = tool_call.get("call_id") or tool_call.get("id")
|
|
elif hasattr(tool_call, 'call_id'):
|
|
call_id = tool_call.call_id
|
|
elif hasattr(tool_call, 'id'):
|
|
call_id = tool_call.id
|
|
if call_id:
|
|
results.append({
|
|
"tool_call_id": call_id,
|
|
"result": "Sunny, 72°F",
|
|
})
|
|
return results
|
|
|
|
mock_execute_tools.side_effect = mock_execute_side_effect
|
|
|
|
mcp_tool_config = cast(Any, {
|
|
"type": "mcp",
|
|
"server_url": "litellm_proxy",
|
|
"require_approval": "never",
|
|
})
|
|
|
|
response = await litellm.aresponses(
|
|
model=model,
|
|
tools=[mcp_tool_config],
|
|
input=[{
|
|
"role": "user",
|
|
"type": "message",
|
|
"content": "What's the weather in San Francisco?"
|
|
}],
|
|
stream=True,
|
|
)
|
|
|
|
events = []
|
|
async for chunk in response:
|
|
events.append(chunk)
|
|
|
|
assert len(events) > 0, "Should receive streaming events"
|
|
|
|
created_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.created'), None)
|
|
in_progress_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.in_progress'), None)
|
|
output_item_added_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.output_item.added'), None)
|
|
mcp_in_progress_idx = next((i for i, e in enumerate(events) if 'mcp_list_tools.in_progress' in str(getattr(e, 'type', ''))), None)
|
|
completed_idx = next((i for i, e in enumerate(events) if getattr(e, 'type', None) == 'response.completed'), None)
|
|
|
|
assert created_idx is not None, "response.created event should be present"
|
|
assert in_progress_idx is not None, "response.in_progress event should be present"
|
|
assert output_item_added_idx is not None, "response.output_item.added event should be present"
|
|
|
|
assert created_idx < in_progress_idx, "response.created should come before response.in_progress"
|
|
assert in_progress_idx < output_item_added_idx, "response.in_progress should come before response.output_item.added"
|
|
|
|
if mcp_in_progress_idx is not None:
|
|
assert output_item_added_idx < mcp_in_progress_idx, "response.output_item.added should come before response.mcp_list_tools.in_progress"
|
|
|
|
response_ids = []
|
|
for i, event in enumerate(events):
|
|
event_type = getattr(event, 'type', None)
|
|
if hasattr(event, 'response'):
|
|
response_obj = getattr(event, 'response', None)
|
|
if response_obj and hasattr(response_obj, 'id'):
|
|
event_type_value = event_type.value if hasattr(event_type, 'value') else str(event_type)
|
|
if any(x in event_type_value for x in ['response.created', 'response.in_progress', 'response.completed']):
|
|
response_ids.append((i, event_type_value, response_obj.id))
|
|
|
|
assert len(response_ids) >= 2, f"Should have at least 2 response lifecycle events. Found {len(response_ids)}"
|
|
|
|
cycles = []
|
|
current_cycle = []
|
|
current_id = None
|
|
|
|
for idx, event_type, resp_id in response_ids:
|
|
if current_id is None or resp_id == current_id:
|
|
current_cycle.append((idx, event_type, resp_id))
|
|
current_id = resp_id
|
|
else:
|
|
if current_cycle:
|
|
cycles.append(current_cycle)
|
|
current_cycle = [(idx, event_type, resp_id)]
|
|
current_id = resp_id
|
|
if current_cycle:
|
|
cycles.append(current_cycle)
|
|
|
|
for cycle_num, cycle in enumerate(cycles):
|
|
cycle_ids = set(resp_id for _, _, resp_id in cycle)
|
|
assert len(cycle_ids) == 1, f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs"
|
|
|
|
assert completed_idx is not None, "response.completed event should be present"
|
|
|
|
lite_errors = [
|
|
record for record in caplog.records
|
|
if record.levelno >= logging.ERROR
|
|
and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage())
|
|
]
|
|
assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join(
|
|
record.getMessage() for record in lite_errors
|
|
)
|
|
|
|
|
|
|