feat: implement tool filter interface for MCP servers

Author: devtalker

src/agents/mcp/__init__.py

@@ -11,7 +11,14 @@
 except ImportError:
     pass
 
-from .util import MCPUtil
+from .util import (
+    MCPUtil,
+    ToolFilter,
+    ToolFilterCallable,
+    ToolFilterContext,
+    ToolFilterStatic,
+    create_static_tool_filter,
+)
 
 __all__ = [
     "MCPServer",
@@ -22,4 +29,9 @@
     "MCPServerStreamableHttp",
     "MCPServerStreamableHttpParams",
     "MCPUtil",
+    "ToolFilter",
+    "ToolFilterCallable",
+    "ToolFilterContext",
+    "ToolFilterStatic",
+    "create_static_tool_filter",
 ]

src/agents/mcp/server.py

@@ -17,6 +17,7 @@
 
 from ..exceptions import UserError
 from ..logger import logger
+from .util import ToolFilter, 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/mcp/util.py

@@ -1,6 +1,8 @@
 import functools
 import json
-from typing import TYPE_CHECKING, Any
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Union
+from typing_extensions import NotRequired, TypedDict
 
 from agents.strict_schema import ensure_strict_json_schema
 
@@ -10,13 +12,82 @@
 from ..run_context import RunContextWrapper
 from ..tool import FunctionTool, Tool
 from ..tracing import FunctionSpanData, get_current_span, mcp_tools_span
+from ..util._types import MaybeAwaitable
 
 if TYPE_CHECKING:
     from mcp.types import Tool as MCPTool
 
+    from ..agent import Agent
     from .server import MCPServer
 
 
+@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
+
+
 class MCPUtil:
     """Set of utilities for interop between MCP and Agents SDK tools."""
 

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")