feat: implement tool filter interface for MCP servers

Author: devtalker

src/agents/mcp/server.py

@@ -17,6 +17,7 @@
 
 from ..exceptions import UserError
 from ..logger import logger
+from ..tool import ToolFilter, ToolFilterContext, ToolFilterStatic
 
 
 class MCPServer(abc.ABC):
@@ -61,8 +62,7 @@ def __init__(
         self,
         cache_tools_list: bool,
         client_session_timeout_seconds: float | None,
-        allowed_tools: list[str] | None = None,
-        excluded_tools: list[str] | None = None,
+        tool_filter: ToolFilter = None,
     ):
         """
         Args:
@@ -74,10 +74,7 @@ def __init__(
             (by avoiding a round-trip to the server every time).
 
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
-            allowed_tools: Optional list of tool names to allow (whitelist).
-                If set, only these tools will be available.
-            excluded_tools: Optional list of tool names to exclude (blacklist).
-                If set, these tools will be filtered out.
+            tool_filter: The tool filter to use for filtering tools.
         """
         self.session: ClientSession | None = None
         self.exit_stack: AsyncExitStack = AsyncExitStack()
@@ -91,8 +88,39 @@ def __init__(
         self._cache_dirty = True
         self._tools_list: list[MCPTool] | None = None
 
-        self.allowed_tools = allowed_tools
-        self.excluded_tools = excluded_tools
+        self.tool_filter = tool_filter
+
+    def _apply_tool_filter(self, tools: list[MCPTool]) -> list[MCPTool]:
+        """Apply the tool filter to the list of tools."""
+        if self.tool_filter is None:
+            return tools
+
+        # Handle static tool filter
+        if isinstance(self.tool_filter, dict):
+            static_filter: ToolFilterStatic = self.tool_filter
+            filtered_tools = tools
+
+            # Apply allowed_tool_names filter (whitelist)
+            if "allowed_tool_names" in static_filter:
+                allowed_names = static_filter["allowed_tool_names"]
+                filtered_tools = [t for t in filtered_tools if t.name in allowed_names]
+
+            # Apply blocked_tool_names filter (blacklist)
+            if "blocked_tool_names" in static_filter:
+                blocked_names = static_filter["blocked_tool_names"]
+                filtered_tools = [t for t in filtered_tools if t.name not in blocked_names]
+
+            return filtered_tools
+
+        # Handle callable tool filter
+        # For now, we can't support callable filters because we don't have access to
+        # run context and agent in the current list_tools signature.
+        # This could be enhanced in the future by modifying the call chain.
+        else:
+            raise NotImplementedError(
+                "Callable tool filters are not yet supported. Please use ToolFilterStatic "
+                "with 'allowed_tool_names' and/or 'blocked_tool_names' for now."
+            )
 
     @abc.abstractmethod
     def create_streams(
@@ -159,12 +187,10 @@ async def list_tools(self) -> list[MCPTool]:
             self._tools_list = (await self.session.list_tools()).tools
             tools = self._tools_list
 
-        # Filter tools based on allowed and excluded tools
+        # Filter tools based on tool_filter
         filtered_tools = tools
-        if self.allowed_tools is not None:
-            filtered_tools = [t for t in filtered_tools if t.name in self.allowed_tools]
-        if self.excluded_tools is not None:
-            filtered_tools = [t for t in filtered_tools if t.name not in self.excluded_tools]
+        if self.tool_filter is not None:
+            filtered_tools = self._apply_tool_filter(filtered_tools)
         return filtered_tools
 
     async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
@@ -226,8 +252,7 @@ def __init__(
         cache_tools_list: bool = False,
         name: str | None = None,
         client_session_timeout_seconds: float | None = 5,
-        allowed_tools: list[str] | None = None,
-        excluded_tools: list[str] | None = None,
+        tool_filter: ToolFilter = None,
     ):
         """Create a new MCP server based on the stdio transport.
 
@@ -245,14 +270,12 @@ def __init__(
             name: A readable name for the server. If not provided, we'll create one from the
                 command.
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
-            allowed_tools: Optional list of tool names to allow (whitelist).
-            excluded_tools: Optional list of tool names to exclude (blacklist).
+            tool_filter: The tool filter to use for filtering tools.
         """
         super().__init__(
             cache_tools_list,
             client_session_timeout_seconds,
-            allowed_tools,
-            excluded_tools,
+            tool_filter,
         )
 
         self.params = StdioServerParameters(
@@ -312,8 +335,7 @@ def __init__(
         cache_tools_list: bool = False,
         name: str | None = None,
         client_session_timeout_seconds: float | None = 5,
-        allowed_tools: list[str] | None = None,
-        excluded_tools: list[str] | None = None,
+        tool_filter: ToolFilter = None,
     ):
         """Create a new MCP server based on the HTTP with SSE transport.
 
@@ -333,14 +355,12 @@ def __init__(
                 URL.
 
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
-            allowed_tools: Optional list of tool names to allow (whitelist).
-            excluded_tools: Optional list of tool names to exclude (blacklist).
+            tool_filter: The tool filter to use for filtering tools.
         """
         super().__init__(
             cache_tools_list,
             client_session_timeout_seconds,
-            allowed_tools,
-            excluded_tools,
+            tool_filter,
         )
 
         self.params = params
@@ -400,8 +420,7 @@ def __init__(
         cache_tools_list: bool = False,
         name: str | None = None,
         client_session_timeout_seconds: float | None = 5,
-        allowed_tools: list[str] | None = None,
-        excluded_tools: list[str] | None = None,
+        tool_filter: ToolFilter = None,
     ):
         """Create a new MCP server based on the Streamable HTTP transport.
 
@@ -422,14 +441,12 @@ def __init__(
                 URL.
 
             client_session_timeout_seconds: the read timeout passed to the MCP ClientSession.
-            allowed_tools: Optional list of tool names to allow (whitelist).
-            excluded_tools: Optional list of tool names to exclude (blacklist).
+            tool_filter: The tool filter to use for filtering tools.
         """
         super().__init__(
             cache_tools_list,
             client_session_timeout_seconds,
-            allowed_tools,
-            excluded_tools,
+            tool_filter,
         )
 
         self.params = params

src/agents/tool.py

@@ -27,6 +27,7 @@
 
 if TYPE_CHECKING:
     from .agent import Agent
+    from mcp.types import Tool as MCPTool
 
 ToolParams = ParamSpec("ToolParams")
 
@@ -41,6 +42,73 @@
 ]
 
 
+@dataclass
+class ToolFilterContext:
+    """Context information available to tool filter functions."""
+    
+    run_context: RunContextWrapper[Any]
+    """The current run context."""
+    
+    agent: "Agent[Any]"
+    """The agent that is requesting the tool list."""
+    
+    server_name: str
+    """The name of the MCP server."""
+
+
+ToolFilterCallable = Callable[["ToolFilterContext", "MCPTool"], MaybeAwaitable[bool]]
+"""A function that determines whether a tool should be available.
+
+Args:
+    context: The context information including run context, agent, and server name.
+    tool: The MCP tool to filter.
+
+Returns:
+    Whether the tool should be available (True) or filtered out (False).
+"""
+
+
+class ToolFilterStatic(TypedDict):
+    """Static tool filter configuration using allowlists and blocklists."""
+    
+    allowed_tool_names: NotRequired[list[str]]
+    """Optional list of tool names to allow (whitelist). If set, only these tools will be available."""
+    
+    blocked_tool_names: NotRequired[list[str]]
+    """Optional list of tool names to exclude (blacklist). If set, these tools will be filtered out."""
+
+
+ToolFilter = Union[ToolFilterCallable, ToolFilterStatic, None]
+"""A tool filter that can be either a function, static configuration, or None (no filtering)."""
+
+
+def create_static_tool_filter(
+    allowed_tool_names: list[str] | None = None,
+    blocked_tool_names: list[str] | None = None,
+) -> ToolFilterStatic | None:
+    """Create a static tool filter from allowlist and blocklist parameters.
+    
+    This is a convenience function for creating a ToolFilterStatic.
+    
+    Args:
+        allowed_tool_names: Optional list of tool names to allow (whitelist).
+        blocked_tool_names: Optional list of tool names to exclude (blacklist).
+    
+    Returns:
+        A ToolFilterStatic if any filtering is specified, None otherwise.
+    """
+    if allowed_tool_names is None and blocked_tool_names is None:
+        return None
+    
+    filter_dict: ToolFilterStatic = {}
+    if allowed_tool_names is not None:
+        filter_dict["allowed_tool_names"] = allowed_tool_names
+    if blocked_tool_names is not None:
+        filter_dict["blocked_tool_names"] = blocked_tool_names
+    
+    return filter_dict
+
+
 @dataclass
 class FunctionToolResult:
     tool: FunctionTool

tests/mcp/test_tool_filtering.py

@@ -1,43 +1,65 @@
 import pytest
 
+from agents.tool import ToolFilterStatic
 from .helpers import FakeMCPServer
 
 
 class FilterableFakeMCPServer(FakeMCPServer):
     """Extended FakeMCPServer that supports tool filtering"""
 
-    def __init__(self, tools=None, allowed_tools=None, excluded_tools=None, server_name=None):
+    def __init__(self, tools=None, tool_filter=None, server_name=None):
         super().__init__(tools)
-        self.allowed_tools = allowed_tools
-        self.excluded_tools = excluded_tools
+        self.tool_filter = tool_filter
         self._server_name = server_name
 
     async def list_tools(self):
         tools = await super().list_tools()
 
         # Apply filtering logic similar to _MCPServerWithClientSession
         filtered_tools = tools
-        if self.allowed_tools is not None:
-            filtered_tools = [t for t in filtered_tools if t.name in self.allowed_tools]
-        if self.excluded_tools is not None:
-            filtered_tools = [t for t in filtered_tools if t.name not in self.excluded_tools]
+        if self.tool_filter is not None:
+            filtered_tools = self._apply_tool_filter(filtered_tools)
         return filtered_tools
 
+    def _apply_tool_filter(self, tools):
+        """Apply the tool filter to the list of tools."""
+        if self.tool_filter is None:
+            return tools
+        
+        # Handle static tool filter
+        if isinstance(self.tool_filter, dict):
+            static_filter: ToolFilterStatic = self.tool_filter
+            filtered_tools = tools
+            
+            # Apply allowed_tool_names filter (whitelist)
+            if "allowed_tool_names" in static_filter:
+                allowed_names = static_filter["allowed_tool_names"]
+                filtered_tools = [t for t in filtered_tools if t.name in allowed_names]
+            
+            # Apply blocked_tool_names filter (blacklist)
+            if "blocked_tool_names" in static_filter:
+                blocked_names = static_filter["blocked_tool_names"]
+                filtered_tools = [t for t in filtered_tools if t.name not in blocked_names]
+            
+            return filtered_tools
+        
+        return tools
+
     @property
     def name(self) -> str:
         return self._server_name or "filterable_fake_server"
 
 
 @pytest.mark.asyncio
-async def test_server_allowed_tools():
-    """Test that server-level allowed_tools filters tools correctly"""
+async def test_server_allowed_tool_names():
+    """Test that server-level allowed_tool_names filters tools correctly"""
     server = FilterableFakeMCPServer(server_name="test_server")
     server.add_tool("tool1", {})
     server.add_tool("tool2", {})
     server.add_tool("tool3", {})
 
-    # Set allowed_tools to only include tool1 and tool2
-    server.allowed_tools = ["tool1", "tool2"]
+    # Set tool_filter to only include tool1 and tool2
+    server.tool_filter = {"allowed_tool_names": ["tool1", "tool2"]}
 
     # Get tools and verify filtering
     tools = await server.list_tools()
@@ -46,15 +68,15 @@ async def test_server_allowed_tools():
 
 
 @pytest.mark.asyncio
-async def test_server_excluded_tools():
-    """Test that server-level excluded_tools filters tools correctly"""
+async def test_server_blocked_tool_names():
+    """Test that server-level blocked_tool_names filters tools correctly"""
     server = FilterableFakeMCPServer(server_name="test_server")
     server.add_tool("tool1", {})
     server.add_tool("tool2", {})
     server.add_tool("tool3", {})
 
-    # Set excluded_tools to exclude tool3
-    server.excluded_tools = ["tool3"]
+    # Set tool_filter to exclude tool3
+    server.tool_filter = {"blocked_tool_names": ["tool3"]}
 
     # Get tools and verify filtering
     tools = await server.list_tools()
@@ -64,18 +86,37 @@ async def test_server_excluded_tools():
 
 @pytest.mark.asyncio
 async def test_server_both_filters():
-    """Test that server-level allowed_tools and excluded_tools work together correctly"""
+    """Test that server-level allowed_tool_names and blocked_tool_names work together correctly"""
     server = FilterableFakeMCPServer(server_name="test_server")
     server.add_tool("tool1", {})
     server.add_tool("tool2", {})
     server.add_tool("tool3", {})
     server.add_tool("tool4", {})
 
     # Set both filters
-    server.allowed_tools = ["tool1", "tool2", "tool3"]
-    server.excluded_tools = ["tool3"]
+    server.tool_filter = {
+        "allowed_tool_names": ["tool1", "tool2", "tool3"],
+        "blocked_tool_names": ["tool3"]
+    }
 
-    # Get tools and verify filtering (allowed_tools applied first, then excluded_tools)
+    # Get tools and verify filtering (allowed_tool_names applied first, then blocked_tool_names)
     tools = await server.list_tools()
     assert len(tools) == 2
     assert {t.name for t in tools} == {"tool1", "tool2"}
+
+
+@pytest.mark.asyncio
+async def test_server_no_filter():
+    """Test that when no filter is set, all tools are returned"""
+    server = FilterableFakeMCPServer(server_name="test_server")
+    server.add_tool("tool1", {})
+    server.add_tool("tool2", {})
+    server.add_tool("tool3", {})
+
+    # No filter set (None)
+    server.tool_filter = None
+
+    # Get tools and verify no filtering
+    tools = await server.list_tools()
+    assert len(tools) == 3
+    assert {t.name for t in tools} == {"tool1", "tool2", "tool3"}