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
2026-06-16 21:04:09 +08:00 · 2025-11-14 08:10:25 +05:30
parent a2a9922cde
commit 36c1217605
2 changed files with 263 additions and 35 deletions
@@ -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,
@@ -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