From 36c12176057894bb10f3b0d6012cfb2149ffa1d0 Mon Sep 17 00:00:00 2001 From: ISHAN RAJ SINGH Date: Fri, 14 Nov 2025 08:10:25 +0530 Subject: [PATCH] Python: Fix: Prevent duplicate MCP tools and prompts (#1876) (#1890) * 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 --- python/packages/core/agent_framework/_mcp.py | 78 +++++-- python/packages/core/tests/core/test_mcp.py | 220 +++++++++++++++++-- 2 files changed, 263 insertions(+), 35 deletions(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 873b7f04cc..66c96425c8 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -81,10 +81,16 @@ def _mcp_type_to_ai_content( case types.TextContent(): return TextContent(text=mcp_type.text, raw_representation=mcp_type) case types.ImageContent() | types.AudioContent(): - return DataContent(uri=mcp_type.data, media_type=mcp_type.mimeType, raw_representation=mcp_type) + return DataContent( + uri=mcp_type.data, + media_type=mcp_type.mimeType, + raw_representation=mcp_type, + ) case types.ResourceLink(): return UriContent( - uri=str(mcp_type.uri), media_type=mcp_type.mimeType or "application/json", raw_representation=mcp_type + uri=str(mcp_type.uri), + media_type=mcp_type.mimeType or "application/json", + raw_representation=mcp_type, ) case _: match mcp_type.resource: @@ -92,14 +98,14 @@ def _mcp_type_to_ai_content( return TextContent( text=mcp_type.resource.text, raw_representation=mcp_type, - additional_properties=mcp_type.annotations.model_dump() if mcp_type.annotations else None, + additional_properties=(mcp_type.annotations.model_dump() if mcp_type.annotations else None), ) case types.BlobResourceContents(): return DataContent( uri=mcp_type.resource.blob, media_type=mcp_type.resource.mimeType, raw_representation=mcp_type, - additional_properties=mcp_type.annotations.model_dump() if mcp_type.annotations else None, + additional_properties=(mcp_type.annotations.model_dump() if mcp_type.annotations else None), ) @@ -124,9 +130,11 @@ def _ai_content_to_mcp_types( # uri's are not limited in MCP but they have to be set. # the uri of data content, contains the data uri, which # is not the uri meant here, UriContent would match this. - uri=content.additional_properties.get("uri", "af://binary") - if content.additional_properties - else "af://binary", # type: ignore[reportArgumentType] + uri=( + content.additional_properties.get("uri", "af://binary") + if content.additional_properties + else "af://binary" + ), # type: ignore[reportArgumentType] ), ) return None @@ -135,9 +143,9 @@ def _ai_content_to_mcp_types( type="resource_link", uri=content.uri, # type: ignore[reportArgumentType] mimeType=content.media_type, - name=content.additional_properties.get("name", "Unknown") - if content.additional_properties - else "Unknown", + name=( + content.additional_properties.get("name", "Unknown") if content.additional_properties else "Unknown" + ), ) case _: return None @@ -272,7 +280,7 @@ class MCPTool: self, name: str, description: str | None = None, - approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None, + approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, load_tools: bool = True, load_prompts: bool = True, @@ -300,6 +308,8 @@ class MCPTool: self.chat_client = chat_client self._functions: list[AIFunction[Any, Any]] = [] self.is_connected: bool = False + self._tools_loaded: bool = False + self._prompts_loaded: bool = False def __str__(self) -> str: return f"MCPTool(name={self.name}, description={self.description})" @@ -336,7 +346,9 @@ class MCPTool: ClientSession( read_stream=transport[0], write_stream=transport[1], - read_timeout_seconds=timedelta(seconds=self.request_timeout) if self.request_timeout else None, + read_timeout_seconds=( + timedelta(seconds=self.request_timeout) if self.request_timeout else None + ), message_handler=self.message_handler, logging_callback=self.logging_callback, sampling_callback=self.sampling_callback, @@ -345,7 +357,8 @@ class MCPTool: except Exception as ex: await self._exit_stack.aclose() raise ToolException( - message="Failed to create MCP session. Please check your configuration.", inner_exception=ex + message="Failed to create MCP session. Please check your configuration.", + inner_exception=ex, ) from ex try: await session.initialize() @@ -368,8 +381,10 @@ class MCPTool: self.is_connected = True if self.load_tools_flag: await self.load_tools() + self._tools_loaded = True if self.load_prompts_flag: await self.load_prompts() + self._prompts_loaded = True if logger.level != logging.NOTSET: try: @@ -380,7 +395,9 @@ class MCPTool: logger.warning("Failed to set log level to %s", logger.level, exc_info=exc) async def sampling_callback( - self, context: RequestContext[ClientSession, Any], params: types.CreateMessageRequestParams + self, + context: RequestContext[ClientSession, Any], + params: types.CreateMessageRequestParams, ) -> types.CreateMessageResult | types.ErrorData: """Callback function for sampling. @@ -458,7 +475,7 @@ class MCPTool: async def message_handler( self, - message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception, + message: (RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception), ) -> None: """Handle messages from the MCP server. @@ -517,8 +534,17 @@ class MCPTool: exc_info=exc, ) prompt_list = None + + # Track existing function names to prevent duplicates + existing_names = {func.name for func in self._functions} + for prompt in prompt_list.prompts if prompt_list else []: local_name = _normalize_mcp_name(prompt.name) + + # Skip if already loaded + if local_name in existing_names: + continue + input_model = _get_input_model_from_mcp_prompt(prompt) approval_mode = self._determine_approval_mode(local_name) func: AIFunction[BaseModel, list[ChatMessage]] = AIFunction( @@ -529,6 +555,7 @@ class MCPTool: input_model=input_model, ) self._functions.append(func) + existing_names.add(local_name) async def load_tools(self) -> None: """Load tools from the MCP server. @@ -549,8 +576,17 @@ class MCPTool: exc_info=exc, ) tool_list = None + + # Track existing function names to prevent duplicates + existing_names = {func.name for func in self._functions} + for tool in tool_list.tools if tool_list else []: local_name = _normalize_mcp_name(tool.name) + + # Skip if already loaded + if local_name in existing_names: + continue + input_model = _get_input_model_from_mcp_tool(tool) approval_mode = self._determine_approval_mode(local_name) # Create AIFunctions out of each tool @@ -562,6 +598,7 @@ class MCPTool: input_model=input_model, ) self._functions.append(func) + existing_names.add(local_name) async def close(self) -> None: """Disconnect from the MCP server. @@ -662,7 +699,10 @@ class MCPTool: raise ToolExecutionException("Failed to enter context manager.", inner_exception=ex) from ex async def __aexit__( - self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: Any + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, ) -> None: """Exit the async context manager. @@ -714,7 +754,7 @@ class MCPStdioTool(MCPTool): request_timeout: int | None = None, session: ClientSession | None = None, description: str | None = None, - approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None, + approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, args: list[str] | None = None, env: dict[str, str] | None = None, @@ -824,7 +864,7 @@ class MCPStreamableHTTPTool(MCPTool): request_timeout: int | None = None, session: ClientSession | None = None, description: str | None = None, - approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None, + approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, headers: dict[str, Any] | None = None, timeout: float | None = None, @@ -939,7 +979,7 @@ class MCPWebsocketTool(MCPTool): request_timeout: int | None = None, session: ClientSession | None = None, description: str | None = None, - approval_mode: Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None = None, + approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None, allowed_tools: Collection[str] | None = None, chat_client: "ChatClientProtocol | None" = None, additional_properties: dict[str, Any] | None = None, diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index e6dd1fd8a7..865e2ef484 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -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