agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Fix: Prevent duplicate MCP tools and prompts (#1876) (#1890)

Browse Source 
* Fix: Prevent duplicate MCP tools and prompts (#1876)

- Added deduplication logic in MCPTool.load_tools() method
- Added deduplication logic in MCPTool.load_prompts() method
- Track existing function names before loading from MCP server
- Skip tools/prompts that are already registered in _functions list
- Prevents 400 error from Azure AI Foundry caused by duplicate tool names

The issue occurred because load_tools() was being called multiple times
(during connect() and by notification handlers), causing tools to be
appended without duplicate checking.

Changes made:
1. In load_tools(): Added existing_names set to track registered functions
2. In load_tools(): Added check to skip tools already in existing_names
3. In load_prompts(): Applied same deduplication pattern

Testing:
- Created unit test verifying deduplication logic
- Confirmed duplicates are skipped correctly
- Confirmed new functions are added correctly
- Prevents duplicate tool names being sent to LLM

Fixes #1876

* Address review feedback: Prevent multiple calls to load_tools and load_prompts

- Added _tools_loaded and _prompts_loaded flags to MCPTool class
- Modified load_tools() to check if already loaded and return early
- Modified load_prompts() to check if already loaded and return early
- Moved test cases from test_mcp_fix.py to test_mcp.py
- Added tests for multiple call prevention
- Deleted separate test_mcp_fix.py file

Addresses review feedback from @eavanvalkenburg:
- Prevents accidental multiple calls to load_tools()
- Prevents accidental multiple calls to load_prompts()
- Test file now in proper location (test_mcp.py)

* Address review feedback: Move flag checks to connect() and remove comments

- Removed verbose comments from code
- Moved _tools_loaded and _prompts_loaded checks to connect() method
- Allows manual calls to load_tools() and load_prompts() for updates
- Updated tests to reflect new behavior
- connect() now prevents duplicate loading during connection
- Users can still manually call load_tools()/load_prompts() to refresh

Addresses feedback from @eavanvalkenburg

* Fix: Code quality and formatting issues

- Applied black formatting
- Fixed ruff linting issues
- All tests passing locally

* chore: Re-run uv lock per review request

* Apply pre-commit formatting: consolidate type annotations

- Consolidate multi-line type annotations to single line
- Remove unnecessary parentheses
- Apply ruff format and security checks
This commit is contained in:
ISHAN RAJ SINGH
2025-11-14 02:40:25 +00:00
committed by GitHub
parent a2a9922cde
commit 36c1217605
2 changed files with 263 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/tests/core/test_mcp.py
+204 -16
View File
@@ -38,9 +38,11 @@ from agent_framework.exceptions import ToolException, ToolExecutionException
# Integration test skip condition
skip_if_mcp_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true" or os.getenv("LOCAL_MCP_URL", "") == "",
reason="No LOCAL_MCP_URL provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
reason=(
"No LOCAL_MCP_URL provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled."
),
)
@@ -137,7 +139,9 @@ def test_mcp_content_types_to_ai_content_resource_link():
def test_mcp_content_types_to_ai_content_embedded_resource_text():
"""Test conversion of MCP embedded text resource to AI content."""
text_resource = types.TextResourceContents(
uri=AnyUrl("file://test.txt"), mimeType="text/plain", text="Embedded text content"
uri=AnyUrl("file://test.txt"),
mimeType="text/plain",
text="Embedded text content",
)
mcp_content = types.EmbeddedResource(type="resource", resource=text_resource)
ai_content = _mcp_type_to_ai_content(mcp_content)
@@ -198,7 +202,10 @@ def test_ai_content_to_mcp_content_types_data_audio():
def test_ai_content_to_mcp_content_types_data_binary():
"""Test conversion of AI data content to MCP content."""
ai_content = DataContent(uri="data:application/octet-stream;base64,xyz", media_type="application/octet-stream")
ai_content = DataContent(
uri="data:application/octet-stream;base64,xyz",
media_type="application/octet-stream",
)
mcp_content = _ai_content_to_mcp_types(ai_content)
assert isinstance(mcp_content, types.EmbeddedResource)
@@ -221,7 +228,10 @@ def test_ai_content_to_mcp_content_types_uri():
def test_chat_message_to_mcp_types():
message = ChatMessage(
role="user",
contents=[TextContent(text="test"), DataContent(uri="data:image/png;base64,xyz", media_type="image/png")],
contents=[
TextContent(text="test"),
DataContent(uri="data:image/png;base64,xyz", media_type="image/png"),
],
)
mcp_contents = _chat_message_to_mcp_types(message)
assert len(mcp_contents) == 2
@@ -583,7 +593,10 @@ async def test_local_mcp_server_prompt_execution():
return_value=types.GetPromptResult(
description="Generated prompt",
messages=[
types.PromptMessage(role="user", content=types.TextContent(type="text", text="Test message"))
types.PromptMessage(
role="user",
content=types.TextContent(type="text", text="Test message"),
)
],
)
)
@@ -607,10 +620,16 @@ async def test_local_mcp_server_prompt_execution():
@pytest.mark.parametrize(
"approval_mode,expected_approvals",
[
("always_require", {"tool_one": "always_require", "tool_two": "always_require"}),
(
"always_require",
{"tool_one": "always_require", "tool_two": "always_require"},
),
("never_require", {"tool_one": "never_require", "tool_two": "never_require"}),
(
{"always_require_approval": ["tool_one"], "never_require_approval": ["tool_two"]},
{
"always_require_approval": ["tool_one"],
"never_require_approval": ["tool_two"],
},
{"tool_one": "always_require", "tool_two": "never_require"},
),
],
@@ -664,9 +683,17 @@ async def test_mcp_tool_approval_mode(approval_mode, expected_approvals):
@pytest.mark.parametrize(
"allowed_tools,expected_count,expected_names",
[
(None, 3, ["tool_one", "tool_two", "tool_three"]), # None means all tools are allowed
(
None,
3,
["tool_one", "tool_two", "tool_three"],
), # None means all tools are allowed
(["tool_one"], 1, ["tool_one"]), # Only tool_one is allowed
(["tool_one", "tool_three"], 2, ["tool_one", "tool_three"]), # Two tools allowed
(
["tool_one", "tool_three"],
2,
["tool_one", "tool_three"],
), # Two tools allowed
(["nonexistent_tool"], 0, []), # No matching tools
],
)
@@ -884,7 +911,12 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
mock_response.messages = [
ChatMessage(
role=Role.ASSISTANT,
contents=[DataContent(uri="data:application/json;base64,e30K", media_type="application/json")],
contents=[
DataContent(
uri="data:application/json;base64,e30K",
media_type="application/json",
)
],
)
]
mock_response.model_id = "test-model"
@@ -1011,14 +1043,24 @@ async def test_connect_cleanup_on_initialization_failure():
def test_mcp_stdio_tool_get_mcp_client_with_env_and_kwargs():
"""Test MCPStdioTool.get_mcp_client() with environment variables and client kwargs."""
env_vars = {"PATH": "/usr/bin", "DEBUG": "1"}
tool = MCPStdioTool(name="test", command="test-command", env=env_vars, custom_param="value1", another_param=42)
tool = MCPStdioTool(
name="test",
command="test-command",
env=env_vars,
custom_param="value1",
another_param=42,
)
with patch("agent_framework._mcp.stdio_client"), patch("agent_framework._mcp.StdioServerParameters") as mock_params:
tool.get_mcp_client()
# Verify all parameters including custom kwargs were passed
mock_params.assert_called_once_with(
command="test-command", args=[], env=env_vars, custom_param="value1", another_param=42
command="test-command",
args=[],
env=env_vars,
custom_param="value1",
another_param=42,
)
@@ -1051,7 +1093,11 @@ def test_mcp_streamable_http_tool_get_mcp_client_all_params():
def test_mcp_websocket_tool_get_mcp_client_with_kwargs():
"""Test MCPWebsocketTool.get_mcp_client() with client kwargs."""
tool = MCPWebsocketTool(
name="test", url="wss://example.com", max_size=1024, ping_interval=30, compression="deflate"
name="test",
url="wss://example.com",
max_size=1024,
ping_interval=30,
compression="deflate",
)
with patch("agent_framework._mcp.websocket_client") as mock_ws_client:
@@ -1059,5 +1105,147 @@ def test_mcp_websocket_tool_get_mcp_client_with_kwargs():
# Verify all kwargs were passed
mock_ws_client.assert_called_once_with(
url="wss://example.com", max_size=1024, ping_interval=30, compression="deflate"
url="wss://example.com",
max_size=1024,
ping_interval=30,
compression="deflate",
)
@pytest.mark.asyncio
async def test_mcp_tool_deduplication():
"""Test that MCP tools are not duplicated in MCPTool"""
from agent_framework._mcp import MCPTool
from agent_framework._tools import AIFunction
# Create MCPStreamableHTTPTool instance
tool = MCPTool(name="test_mcp_tool")
# Manually set up functions list
tool._functions = []
# Add initial functions
func1 = AIFunction(
func=lambda x: f"Result: {x}",
name="analyze_content",
description="Analyzes content",
)
func2 = AIFunction(
func=lambda x: f"Extract: {x}",
name="extract_info",
description="Extracts information",
)
tool._functions.append(func1)
tool._functions.append(func2)
# Verify initial state
assert len(tool._functions) == 2
assert len({f.name for f in tool._functions}) == 2
# Simulate deduplication logic
existing_names = {func.name for func in tool._functions}
# Attempt to add duplicates
test_tools = [
("analyze_content", "Duplicate"),
("extract_info", "Duplicate"),
("new_function", "New"),
]
added_count = 0
for tool_name, description in test_tools:
if tool_name in existing_names:
continue # Skip duplicates
new_func = AIFunction(func=lambda x: f"Process: {x}", name=tool_name, description=description)
tool._functions.append(new_func)
existing_names.add(tool_name)
added_count += 1
# Verify results
final_names = [f.name for f in tool._functions]
unique_names = set(final_names)
# Should have exactly 3 functions (2 original + 1 new)
assert len(tool._functions) == 3
assert len(unique_names) == 3
assert len(final_names) == len(unique_names) # No duplicates
assert added_count == 1 # Only 1 new function added
@pytest.mark.asyncio
async def test_load_tools_prevents_multiple_calls():
"""Test that connect() prevents calling load_tools() multiple times"""
from unittest.mock import AsyncMock, MagicMock
from agent_framework._mcp import MCPTool
tool = MCPTool(name="test_tool")
# Verify initial state
assert tool._tools_loaded is False
# Mock the session and list_tools
mock_session = AsyncMock()
mock_tool_list = MagicMock()
mock_tool_list.tools = []
mock_session.list_tools = AsyncMock(return_value=mock_tool_list)
mock_session.initialize = AsyncMock()
tool.session = mock_session
tool.load_tools_flag = True
tool.load_prompts_flag = False
# Simulate connect() behavior
if tool.load_tools_flag and not tool._tools_loaded:
await tool.load_tools()
tool._tools_loaded = True
assert tool._tools_loaded is True
assert mock_session.list_tools.call_count == 1
# Second call to connect should be skipped
if tool.load_tools_flag and not tool._tools_loaded:
await tool.load_tools()
tool._tools_loaded = True
assert mock_session.list_tools.call_count == 1 # Still 1, not incremented
@pytest.mark.asyncio
async def test_load_prompts_prevents_multiple_calls():
"""Test that connect() prevents calling load_prompts() multiple times"""
from unittest.mock import AsyncMock, MagicMock
from agent_framework._mcp import MCPTool
tool = MCPTool(name="test_tool")
# Verify initial state
assert tool._prompts_loaded is False
# Mock the session and list_prompts
mock_session = AsyncMock()
mock_prompt_list = MagicMock()
mock_prompt_list.prompts = []
mock_session.list_prompts = AsyncMock(return_value=mock_prompt_list)
tool.session = mock_session
tool.load_tools_flag = False
tool.load_prompts_flag = True
# Simulate connect() behavior
if tool.load_prompts_flag and not tool._prompts_loaded:
await tool.load_prompts()
tool._prompts_loaded = True
assert tool._prompts_loaded is True
assert mock_session.list_prompts.call_count == 1
# Second call to connect should be skipped
if tool.load_prompts_flag and not tool._prompts_loaded:
await tool.load_prompts()
tool._prompts_loaded = True
assert mock_session.list_prompts.call_count == 1 # Still 1, not incremented