Repository: mozilla-ai/any-agent
Pull request #152: Filter MCP Tools and write a unit test to check it
Author: login: njbrake
Created: 2025-04-24T17:53:40Z | Merged: 2025-04-25T12:36:16Z

Most of the work here went into the unit testing. I figured out how to mock the underlying part of the MCP lib that everyone uses.

I only mocked and test the stdio functionality because SSE isn't yet supported across frameworks. For instance, Agno just added support for it in a commit from 5 hours ago 😆  https://github.com/agno-agi/agno/pull/2892/files 


Llama Index doesn't yet support SSE but I think it's on their roadmap.


Diff:
diff --git a/pyproject.toml b/pyproject.toml
index d7ba7ccc..134da233 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -53,7 +53,7 @@ openai = [
 ]
 
 agno = [
-  "agno>=1.3.4",
+  "agno<=1.3.5", # https://github.com/mozilla-ai/any-agent/issues/155
   "litellm"
 ]
 
@@ -141,3 +141,9 @@ strict = true
 disallow_untyped_calls = true
 disallow_untyped_defs = true
 follow_untyped_imports = true
+
+
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "session"
diff --git a/pytest.ini b/pytest.ini
deleted file mode 100644
index 0102b0a9..00000000
--- a/pytest.ini
+++ /dev/null
@@ -1,2 +0,0 @@
-[pytest]
-asyncio_default_fixture_loop_scope = function
diff --git a/src/any_agent/tools/mcp/frameworks/google.py b/src/any_agent/tools/mcp/frameworks/google.py
index 9326d410..3390b9dc 100644
--- a/src/any_agent/tools/mcp/frameworks/google.py
+++ b/src/any_agent/tools/mcp/frameworks/google.py
@@ -39,6 +39,8 @@ async def setup_tools(self) -> None:
         await self._exit_stack.enter_async_context(self.server)
         self.tools = await self.server.load_tools()
 
+        self.tools = self.filter_tools(self.tools)
+
 
 class GoogleMCPServerStdio(GoogleMCPServerBase):
     mcp_tool: MCPStdioParams
diff --git a/src/any_agent/tools/mcp/frameworks/langchain.py b/src/any_agent/tools/mcp/frameworks/langchain.py
index d7a4c8b6..cc8749e4 100644
--- a/src/any_agent/tools/mcp/frameworks/langchain.py
+++ b/src/any_agent/tools/mcp/frameworks/langchain.py
@@ -49,6 +49,8 @@ async def setup_tools(self) -> None:
         # List available tools
         self.tools = await load_mcp_tools(session)
 
+        self.tools = self.filter_tools(self.tools)
+
 
 class LangchainMCPServerStdio(LangchainMCPServerBase):
     mcp_tool: MCPStdioParams
diff --git a/src/any_agent/tools/mcp/frameworks/openai.py b/src/any_agent/tools/mcp/frameworks/openai.py
index 1cb69769..962de0b3 100644
--- a/src/any_agent/tools/mcp/frameworks/openai.py
+++ b/src/any_agent/tools/mcp/frameworks/openai.py
@@ -3,7 +3,6 @@
 from typing import Literal
 
 from any_agent.config import AgentFramework, MCPSseParams, MCPStdioParams
-from any_agent.logging import logger
 from any_agent.tools.mcp.mcp_server import MCPServerBase
 
 mcp_available = False
@@ -38,12 +37,10 @@ async def setup_tools(self) -> None:
             raise ValueError(msg)
 
         await self._exit_stack.enter_async_context(self.server)
-        # Get tools from the server
-        logger.warning(
-            "OpenAI MCP currently does not support filtering MCP available tools",
-        )
         self.tools = await self.server.list_tools()  # type: ignore[assignment]
 
+        self.tools = self.filter_tools(self.tools)
+
 
 class OpenAIMCPServerStdio(OpenAIMCPServerBase):
     mcp_tool: MCPStdioParams
diff --git a/src/any_agent/tools/mcp/frameworks/smolagents.py b/src/any_agent/tools/mcp/frameworks/smolagents.py
index 0045c62b..9bdfd640 100644
--- a/src/any_agent/tools/mcp/frameworks/smolagents.py
+++ b/src/any_agent/tools/mcp/frameworks/smolagents.py
@@ -2,11 +2,9 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from contextlib import suppress
-from textwrap import dedent
 from typing import Literal
 
 from any_agent.config import AgentFramework, MCPSseParams, MCPStdioParams
-from any_agent.logging import logger
 from any_agent.tools.mcp.mcp_server import MCPServerBase
 
 mcp_available = False
@@ -35,29 +33,7 @@ async def setup_tools(self) -> None:
             msg = "Tool collection is not set up. Please call `setup` from a concrete class."
             raise ValueError(msg)
 
-        # Only add the tools listed in mcp_tool['tools'] if specified
-        requested_tools = list(self.mcp_tool.tools or [])
-        if not requested_tools:
-            logger.info(
-                "No specific tools requested for MCP server, using all available tools:",
-            )
-            logger.info("Tools available: %s", self.smolagent_tools)
-            self.tools = self.smolagent_tools
-            return
-
-        filtered_tools = [
-            tool for tool in self.smolagent_tools if tool.name in requested_tools
-        ]
-        if len(filtered_tools) == len(requested_tools):
-            self.tools = filtered_tools
-            return
-
-        tool_names = [tool.name for tool in filtered_tools]
-        raise ValueError(
-            dedent(f"""Could not find all requested tools in the MCP server:
-                        Requested: {requested_tools}
-                        Set:   {tool_names}"""),
-        )
+        self.tools = self.filter_tools(self.smolagent_tools)
 
 
 class SmolagentsMCPServerStdio(SmolagentsMCPServerBase):
diff --git a/src/any_agent/tools/mcp/mcp_server.py b/src/any_agent/tools/mcp/mcp_server.py
index aa9faf95..610a7948 100644
--- a/src/any_agent/tools/mcp/mcp_server.py
+++ b/src/any_agent/tools/mcp/mcp_server.py
@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from contextlib import AsyncExitStack
+from textwrap import dedent
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
@@ -32,3 +33,18 @@ def check_dependencies(self) -> None:
 
         msg = f"You need to `pip install '{self.libraries}'` to use MCP."
         raise ImportError(msg)
+
+    def filter_tools(self, tools: Sequence[Any]) -> Sequence[Any]:
+        # Only add the tools listed in mcp_tool['tools'] if specified
+        requested_tools = list(self.mcp_tool.tools or [])
+        if not requested_tools:
+            return tools
+        tools = [tool for tool in tools if tool.name in requested_tools]
+        if len(tools) != len(requested_tools):
+            tool_names = [tool.name for tool in tools]
+            raise ValueError(
+                dedent(f"""Could not find all requested tools in the MCP server:
+                            Requested: {requested_tools}
+                            Set:   {tool_names}"""),
+            )
+        return tools
diff --git a/tests/conftest.py b/tests/conftest.py
index f3993975..48979314 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,6 @@
 import json
+from collections.abc import Generator
+from unittest.mock import AsyncMock, patch
 
 import pytest
 import rich.console
@@ -76,3 +78,15 @@ def to_json(self) -> str:
 @pytest.fixture(params=list(AgentFramework), ids=lambda x: x.name)
 def agent_framework(request: pytest.FixtureRequest) -> AgentFramework:
     return request.param  # type: ignore[no-any-return]
+
+
+@pytest.fixture
+def mock_stdio_client() -> Generator[
+    tuple[AsyncMock, tuple[AsyncMock, AsyncMock]], None
+]:
+    mock_cm = AsyncMock()
+    mock_transport = (AsyncMock(), AsyncMock())
+    mock_cm.__aenter__.return_value = mock_transport
+
+    with patch("mcp.client.stdio.stdio_client", return_value=mock_cm) as patched:
+        yield patched, mock_transport
diff --git a/tests/unit/tools/test_unit_mcp.py b/tests/unit/tools/test_unit_mcp.py
index 721be6f1..c693143c 100644
--- a/tests/unit/tools/test_unit_mcp.py
+++ b/tests/unit/tools/test_unit_mcp.py
@@ -3,12 +3,14 @@
 # Disclaim
 
 import asyncio
+import shutil
 import unittest
 from contextlib import AsyncExitStack
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from mcp import Tool as MCPTool
 
 from any_agent.config import AgentConfig, AgentFramework, MCPSseParams, MCPStdioParams
 from any_agent.frameworks.any_agent import AnyAgent
@@ -348,3 +350,42 @@ async def test_agno_mcp_sse() -> None:
 
             # Check that tools instance was set as server
             assert server.server == mock_tools_instance  # type: ignore[union-attr]
+
+
+@pytest.mark.asyncio
+@patch(
+    "mcp.client.session.ClientSession.initialize",
+    new_callable=AsyncMock,
+    return_value=None,
+)
+@patch("mcp.client.session.ClientSession.list_tools")
+async def test_stdio_tool_filtering(
+    mock_list_tools: AsyncMock,
+    mock_initialize: AsyncMock,
+    agent_framework: AgentFramework,
+    mock_stdio_client: Any,
+) -> None:
+    tee = shutil.which("tee") or ""
+    assert tee, "tee not found"
+    # Mocking the command part of stdio is really tricky so instead we'll use
+    # a real command that should be available on all systems (this is what openai-agents does too)
+    mock_stdio_params = MCPStdioParams(
+        command=tee,
+        args=[],
+        tools=["write_file", "read_file"],
+    )
+
+    mock_tool_list = MagicMock()
+    mock_tool_list.tools = [
+        MCPTool(name="write_file", inputSchema={"type": "string", "properties": {}}),
+        MCPTool(name="read_file", inputSchema={"type": "string", "properties": {}}),
+        MCPTool(name="other_tool", inputSchema={"type": "string", "properties": {}}),
+    ]
+    mock_list_tools.return_value = mock_tool_list
+    server = get_mcp_server(mock_stdio_params, agent_framework)
+    await server.setup_tools()
+    if agent_framework == AgentFramework.AGNO:
+        # Check that only the specified tools are included
+        assert set(server.tools[0].functions.keys()) == {"write_file", "read_file"}  # type: ignore[union-attr]
+    else:
+        assert len(server.tools) == 2  # ignore[arg-type]


Reviews:
user: login: njbrake
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-24T17:54:25Z
user: login: daavoo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T19:57:44Z
user: login: njbrake
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-24T20:01:58Z

Review comments:
diff_hunk: @@ -348,3 +350,91 @@ async def test_agno_mcp_sse() -> None:
 
             # Check that tools instance was set as server
             assert server.server == mock_tools_instance  # type: ignore[union-attr]
+
+
+def setup_mock_server(monkeypatch: Any) -> None:
+    mock_context = AsyncMock(spec=ClientSession)
+    # calling aenter should give back a mock session
+    mock_session = AsyncMock()
+    mock_session.__aenter__.return_value = mock_context
+    mock_session.__aexit__.return_value = None
+    mock_session.initialize = AsyncMock()
+    mock_session.list_tools = AsyncMock()
+    # Create a dummy streams context manager
+    dummy_transport = (AsyncMock(), AsyncMock())
+    mock_stdio_context = MagicMock()
+    mock_stdio_context.__aenter__.return_value = dummy_transport
+
+    def mock_stdio_client_func(
+        *args: dict[str, Any], **kwargs: dict[str, Any]
+    ) -> ClientSession:
+        return mock_stdio_context
+
+    # This is fragile because it relies on the 3rd party libs not changing where they import these mcp libs
+    # I am totally open if someone has an idea about a more stable way to do this.
+    where_frameworks_import = [
+        ("agno.tools.mcp.stdio_client", "agno.tools.mcp.ClientSession"),
+        ("agents.mcp.server.stdio_client", "agents.mcp.server.ClientSession"),
+        (
+            "any_agent.tools.mcp.frameworks.langchain.stdio_client",
+            "any_agent.tools.mcp.frameworks.langchain.ClientSession",
+        ),
+        (
+            "llama_index.tools.mcp.client.stdio_client",
+            "llama_index.tools.mcp.client.ClientSession",
+        ),
+        ("mcpadapt.core.stdio_client", "mcpadapt.core.ClientSession"),
+        (
+            "google.adk.tools.mcp_tool.mcp_session_manager.stdio_client",
+            "google.adk.tools.mcp_tool.mcp_session_manager.ClientSession",
+        ),
+    ]
+    for stdio_module, client_module in where_frameworks_import:
+        monkeypatch.setattr(stdio_module, mock_stdio_client_func)
+        monkeypatch.setattr(client_module, lambda *args, **kwargs: mock_session)
path: tests/unit/tools/test_unit_mcp.py
user: login: njbrake
body: If you know of a better way to patch these two functions, I am allll ears.
created_at: 2025-04-24T17:54:25Z
author_association: MEMBER
diff_hunk: @@ -348,3 +350,91 @@ async def test_agno_mcp_sse() -> None:
 
             # Check that tools instance was set as server
             assert server.server == mock_tools_instance  # type: ignore[union-attr]
+
+
+def setup_mock_server(monkeypatch: Any) -> None:
+    mock_context = AsyncMock(spec=ClientSession)
+    # calling aenter should give back a mock session
+    mock_session = AsyncMock()
+    mock_session.__aenter__.return_value = mock_context
+    mock_session.__aexit__.return_value = None
+    mock_session.initialize = AsyncMock()
+    mock_session.list_tools = AsyncMock()
+    # Create a dummy streams context manager
+    dummy_transport = (AsyncMock(), AsyncMock())
+    mock_stdio_context = MagicMock()
+    mock_stdio_context.__aenter__.return_value = dummy_transport
+
+    def mock_stdio_client_func(
+        *args: dict[str, Any], **kwargs: dict[str, Any]
+    ) -> ClientSession:
+        return mock_stdio_context
+
+    # This is fragile because it relies on the 3rd party libs not changing where they import these mcp libs
+    # I am totally open if someone has an idea about a more stable way to do this.
+    where_frameworks_import = [
+        ("agno.tools.mcp.stdio_client", "agno.tools.mcp.ClientSession"),
+        ("agents.mcp.server.stdio_client", "agents.mcp.server.ClientSession"),
+        (
+            "any_agent.tools.mcp.frameworks.langchain.stdio_client",
+            "any_agent.tools.mcp.frameworks.langchain.ClientSession",
+        ),
+        (
+            "llama_index.tools.mcp.client.stdio_client",
+            "llama_index.tools.mcp.client.ClientSession",
+        ),
+        ("mcpadapt.core.stdio_client", "mcpadapt.core.ClientSession"),
+        (
+            "google.adk.tools.mcp_tool.mcp_session_manager.stdio_client",
+            "google.adk.tools.mcp_tool.mcp_session_manager.ClientSession",
+        ),
+    ]
+    for stdio_module, client_module in where_frameworks_import:
+        monkeypatch.setattr(stdio_module, mock_stdio_client_func)
+        monkeypatch.setattr(client_module, lambda *args, **kwargs: mock_session)
+
+    # Use different patches based on the agent framework
+    # Create three mock tools with different names
+    from mcp import Tool as MCPTool
+
+    write_tool = MCPTool(
+        name="write_file", inputSchema={"type": "string", "properties": {}}
+    )
+    read_tool = MCPTool(
+        name="read_file", inputSchema={"type": "string", "properties": {}}
+    )
+    other_tool = MCPTool(
+        name="other_tool", inputSchema={"type": "string", "properties": {}}
+    )
+
+    # Create a ToolList object or mock with the same structure
+    mock_tool_list = MagicMock()
+    mock_tool_list.tools = [
+        write_tool,
+        read_tool,
+        other_tool,
+    ]  # Set the tools attribute directly
+
+    # Mock the list_tools method to return the mock tool list
+    mock_context.list_tools.return_value = mock_tool_list
+    mock_context.initialize = AsyncMock()
+
+
+@pytest.mark.asyncio
+async def test_stdio_tool_filtering(
+    agent_framework: AgentFramework, tmp_path: Path, monkeypatch: Any
+) -> None:
+    setup_mock_server(monkeypatch)
+    mock_stdio_params = MagicMock(spec=MCPStdioParams)
+    mock_stdio_params.command = "test_command"
+    mock_stdio_params.args = ["arg1", "arg2"]
+    mock_stdio_params.tools = ["write_file", "read_file"]
+    mock_stdio_params.client_session_timeout_seconds = 10
path: tests/unit/tools/test_unit_mcp.py
user: login: daavoo
body: Why do you have to mock this? Can't you just pass an actual MCPStdioParams?
created_at: 2025-04-24T19:57:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -348,3 +350,91 @@ async def test_agno_mcp_sse() -> None:
 
             # Check that tools instance was set as server
             assert server.server == mock_tools_instance  # type: ignore[union-attr]
+
+
+def setup_mock_server(monkeypatch: Any) -> None:
+    mock_context = AsyncMock(spec=ClientSession)
+    # calling aenter should give back a mock session
+    mock_session = AsyncMock()
+    mock_session.__aenter__.return_value = mock_context
+    mock_session.__aexit__.return_value = None
+    mock_session.initialize = AsyncMock()
+    mock_session.list_tools = AsyncMock()
+    # Create a dummy streams context manager
+    dummy_transport = (AsyncMock(), AsyncMock())
+    mock_stdio_context = MagicMock()
+    mock_stdio_context.__aenter__.return_value = dummy_transport
+
+    def mock_stdio_client_func(
+        *args: dict[str, Any], **kwargs: dict[str, Any]
+    ) -> ClientSession:
+        return mock_stdio_context
+
+    # This is fragile because it relies on the 3rd party libs not changing where they import these mcp libs
+    # I am totally open if someone has an idea about a more stable way to do this.
+    where_frameworks_import = [
+        ("agno.tools.mcp.stdio_client", "agno.tools.mcp.ClientSession"),
+        ("agents.mcp.server.stdio_client", "agents.mcp.server.ClientSession"),
+        (
+            "any_agent.tools.mcp.frameworks.langchain.stdio_client",
+            "any_agent.tools.mcp.frameworks.langchain.ClientSession",
+        ),
+        (
+            "llama_index.tools.mcp.client.stdio_client",
+            "llama_index.tools.mcp.client.ClientSession",
+        ),
+        ("mcpadapt.core.stdio_client", "mcpadapt.core.ClientSession"),
+        (
+            "google.adk.tools.mcp_tool.mcp_session_manager.stdio_client",
+            "google.adk.tools.mcp_tool.mcp_session_manager.ClientSession",
+        ),
+    ]
+    for stdio_module, client_module in where_frameworks_import:
+        monkeypatch.setattr(stdio_module, mock_stdio_client_func)
+        monkeypatch.setattr(client_module, lambda *args, **kwargs: mock_session)
+
+    # Use different patches based on the agent framework
+    # Create three mock tools with different names
+    from mcp import Tool as MCPTool
+
+    write_tool = MCPTool(
+        name="write_file", inputSchema={"type": "string", "properties": {}}
+    )
+    read_tool = MCPTool(
+        name="read_file", inputSchema={"type": "string", "properties": {}}
+    )
+    other_tool = MCPTool(
+        name="other_tool", inputSchema={"type": "string", "properties": {}}
+    )
+
+    # Create a ToolList object or mock with the same structure
+    mock_tool_list = MagicMock()
+    mock_tool_list.tools = [
+        write_tool,
+        read_tool,
+        other_tool,
+    ]  # Set the tools attribute directly
+
+    # Mock the list_tools method to return the mock tool list
+    mock_context.list_tools.return_value = mock_tool_list
+    mock_context.initialize = AsyncMock()
+
+
+@pytest.mark.asyncio
+async def test_stdio_tool_filtering(
+    agent_framework: AgentFramework, tmp_path: Path, monkeypatch: Any
+) -> None:
+    setup_mock_server(monkeypatch)
+    mock_stdio_params = MagicMock(spec=MCPStdioParams)
+    mock_stdio_params.command = "test_command"
+    mock_stdio_params.args = ["arg1", "arg2"]
+    mock_stdio_params.tools = ["write_file", "read_file"]
+    mock_stdio_params.client_session_timeout_seconds = 10
path: tests/unit/tools/test_unit_mcp.py
user: login: njbrake
body: Good point, no reason
created_at: 2025-04-24T20:01:58Z
author_association: MEMBER

PR comments:
user: login: njbrake
created_at: 2025-04-24T18:06:31Z
body: The test is failing because the current latest version of Agno has a bug in it where it's not correctly including/excluding tools. It is fixed by this PR https://github.com/agno-agi/agno/commit/541a02cb1e25dc244f11230b80f178cf989700d6 which I assume will be rolled into a new version soon.
author_association: MEMBER
user: login: daavoo
created_at: 2025-04-25T10:02:57Z
body: > The test is failing because the current latest version of Agno has a bug in it where it's not correctly including/excluding tools. It is fixed by this PR [agno-agi/agno@541a02c](https://github.com/agno-agi/agno/commit/541a02cb1e25dc244f11230b80f178cf989700d6) which I assume will be rolled into a new version soon.

It looks like it was not fixed on that PR. 
I am sending a PR to Agno to properly fix the filtering for MCP
author_association: CONTRIBUTOR
user: login: daavoo
created_at: 2025-04-25T10:40:47Z
body: > > The test is failing because the current latest version of Agno has a bug in it where it's not correctly including/excluding tools. It is fixed by this PR [agno-agi/agno@541a02c](https://github.com/agno-agi/agno/commit/541a02cb1e25dc244f11230b80f178cf989700d6) which I assume will be rolled into a new version soon.
> 
> It looks like it was not fixed on that PR. I am sending a PR to Agno to properly fix the filtering for MCP

Pending of https://github.com/agno-agi/agno/pull/2985
author_association: CONTRIBUTOR