From 04207c4fca724e10129c66cd37bd3770f526c754 Mon Sep 17 00:00:00 2001 From: Dani Date: Thu, 15 May 2025 14:58:00 -0400 Subject: [PATCH 1/4] Add tool enable/disable functionality with client notifications --- src/mcp/server/fastmcp/tools/base.py | 29 +++++++++++- src/mcp/server/fastmcp/tools/tool_manager.py | 7 ++- tests/server/fastmcp/test_tool_manager.py | 49 ++++++++++++++++++++ 3 files changed, 82 insertions(+), 3 deletions(-) diff --git a/src/mcp/server/fastmcp/tools/base.py b/src/mcp/server/fastmcp/tools/base.py index 21eb1841d..af696f3d8 100644 --- a/src/mcp/server/fastmcp/tools/base.py +++ b/src/mcp/server/fastmcp/tools/base.py @@ -8,7 +8,7 @@ from mcp.server.fastmcp.exceptions import ToolError from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata, func_metadata -from mcp.types import ToolAnnotations +from mcp.types import ServerNotification, ToolAnnotations, ToolListChangedNotification if TYPE_CHECKING: from mcp.server.fastmcp.server import Context @@ -34,6 +34,7 @@ class Tool(BaseModel): annotations: ToolAnnotations | None = Field( None, description="Optional annotations for the tool" ) + enabled: bool = Field(default=True, description="Whether the tool is enabled") @classmethod def from_function( @@ -98,3 +99,29 @@ async def run( ) except Exception as e: raise ToolError(f"Error executing tool {self.name}: {e}") from e + + async def enable( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> None: + """Enable the tool and notify clients.""" + if not self.enabled: + self.enabled = True + if context and context.session: + notification = ToolListChangedNotification( + method="notifications/tools/list_changed" + ) + server_notification = ServerNotification.model_validate(notification) + await context.session.send_notification(server_notification) + + async def disable( + self, context: Context[ServerSessionT, LifespanContextT] | None = None + ) -> None: + """Disable the tool and notify clients.""" + if self.enabled: + self.enabled = False + if context and context.session: + notification = ToolListChangedNotification( + method="notifications/tools/list_changed" + ) + server_notification = ServerNotification.model_validate(notification) + await context.session.send_notification(server_notification) diff --git a/src/mcp/server/fastmcp/tools/tool_manager.py b/src/mcp/server/fastmcp/tools/tool_manager.py index cfdaeb350..454db10fd 100644 --- a/src/mcp/server/fastmcp/tools/tool_manager.py +++ b/src/mcp/server/fastmcp/tools/tool_manager.py @@ -28,8 +28,8 @@ def get_tool(self, name: str) -> Tool | None: return self._tools.get(name) def list_tools(self) -> list[Tool]: - """List all registered tools.""" - return list(self._tools.values()) + """List all enabled registered tools.""" + return [tool for tool in self._tools.values() if tool.enabled] def add_tool( self, @@ -61,4 +61,7 @@ async def call_tool( if not tool: raise ToolError(f"Unknown tool: {name}") + if not tool.enabled: + raise ToolError(f"Tool is disabled: {name}") + return await tool.run(arguments, context=context) diff --git a/tests/server/fastmcp/test_tool_manager.py b/tests/server/fastmcp/test_tool_manager.py index e36a09d54..bacde2a6e 100644 --- a/tests/server/fastmcp/test_tool_manager.py +++ b/tests/server/fastmcp/test_tool_manager.py @@ -362,3 +362,52 @@ def echo(message: str) -> str: assert tools[0].annotations is not None assert tools[0].annotations.title == "Echo Tool" assert tools[0].annotations.readOnlyHint is True + + +class TestToolEnableDisable: + """Test enabling and disabling tools.""" + + @pytest.mark.anyio + async def test_enable_disable_tool(self): + """Test enabling and disabling a tool.""" + + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + manager = ToolManager() + tool = manager.add_tool(add) + + # Tool should be enabled by default + assert tool.enabled is True + + # Disable the tool + await tool.disable() + assert tool.enabled is False + + # Enable the tool + await tool.enable() + assert tool.enabled is True + + @pytest.mark.anyio + async def test_enable_disable_no_change(self): + """Test enabling and disabling a tool when there's no state change.""" + + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + manager = ToolManager() + tool = manager.add_tool(add) + + # Enable an already enabled tool (should not change state) + await tool.enable() + assert tool.enabled is True + + # Disable the tool + await tool.disable() + assert tool.enabled is False + + # Disable an already disabled tool (should not change state) + await tool.disable() + assert tool.enabled is False From 086cc543f2f3dc388fdde65f13504a6a3fbeb895 Mon Sep 17 00:00:00 2001 From: Dani Date: Thu, 15 May 2025 19:35:53 -0400 Subject: [PATCH 2/4] add enable and disable tool functionality to README --- README.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/README.md b/README.md index b281b5513..a4b0859fb 100644 --- a/README.md +++ b/README.md @@ -30,10 +30,12 @@ - [Prompts](#prompts) - [Images](#images) - [Context](#context) + - [Authentication](#authentication) - [Running Your Server](#running-your-server) - [Development Mode](#development-mode) - [Claude Desktop Integration](#claude-desktop-integration) - [Direct Execution](#direct-execution) + - [Streamable HTTP Transport](#streamable-http-transport) - [Mounting to an Existing ASGI Server](#mounting-to-an-existing-asgi-server) - [Examples](#examples) - [Echo Server](#echo-server) @@ -243,6 +245,15 @@ async def fetch_weather(city: str) -> str: async with httpx.AsyncClient() as client: response = await client.get(f"https://api.weather.com/{city}") return response.text + +# Get a reference to the tool +tool = mcp._tool_manager.get_tool("fetch_weather") + +# Disable the tool temporarily +await tool.disable(mcp.get_context()) + +# Later, re-enable the tool +await tool.enable(mcp.get_context()) ``` ### Prompts From 213b1c3c5edc46b01a1b7c0306f9b92c88664c34 Mon Sep 17 00:00:00 2001 From: Dani Date: Fri, 16 May 2025 10:51:53 -0400 Subject: [PATCH 3/4] Fix spaces in README.md --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index a4b0859fb..fc35d1ea1 100644 --- a/README.md +++ b/README.md @@ -246,12 +246,15 @@ async def fetch_weather(city: str) -> str: response = await client.get(f"https://api.weather.com/{city}") return response.text + # Get a reference to the tool tool = mcp._tool_manager.get_tool("fetch_weather") + # Disable the tool temporarily await tool.disable(mcp.get_context()) + # Later, re-enable the tool await tool.enable(mcp.get_context()) ``` From fd71115ed5ad04202ad3c8882a74851ee8781460 Mon Sep 17 00:00:00 2001 From: Dani Date: Fri, 16 May 2025 11:26:13 -0400 Subject: [PATCH 4/4] fix function disable_tool and enable_tool. Pytest executed successfully --- README.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 85a7dfb15..5428c03e4 100644 --- a/README.md +++ b/README.md @@ -247,16 +247,17 @@ async def fetch_weather(city: str) -> str: return response.text -# Get a reference to the tool tool = mcp._tool_manager.get_tool("fetch_weather") -# Disable the tool temporarily -await tool.disable(mcp.get_context()) +async def disable_tool(): + # Disable the tool temporarily + await tool.disable(mcp.get_context()) -# Later, re-enable the tool -await tool.enable(mcp.get_context()) +async def enable_tool(): + # Re-enable the tool when needed + await tool.enable(mcp.get_context()) ``` ### Prompts