davenpi commited on
Commit
c44bbc1
·
1 Parent(s): d2d17d2

feat: add support for removing tools from server

Browse files
src/fastmcp/server/server.py CHANGED
@@ -457,6 +457,18 @@ class FastMCP(Generic[LifespanResultT]):
457
  )
458
  self._cache.clear()
459
 
 
 
 
 
 
 
 
 
 
 
 
 
460
  def tool(
461
  self,
462
  name: str | None = None,
 
457
  )
458
  self._cache.clear()
459
 
460
+ def remove_tool(self, name: str) -> None:
461
+ """Remove a tool from the server.
462
+
463
+ Args:
464
+ name: The name of the tool to remove
465
+
466
+ Raises:
467
+ NotFoundError: If the tool is not found
468
+ """
469
+ self._tool_manager.remove_tool(name)
470
+ self._cache.clear()
471
+
472
  def tool(
473
  self,
474
  name: str | None = None,
src/fastmcp/tools/tool_manager.py CHANGED
@@ -94,6 +94,20 @@ class ToolManager:
94
  self._tools[key] = tool
95
  return tool
96
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
97
  async def call_tool(
98
  self, key: str, arguments: dict[str, Any]
99
  ) -> list[TextContent | ImageContent | EmbeddedResource]:
 
94
  self._tools[key] = tool
95
  return tool
96
 
97
+ def remove_tool(self, key: str) -> None:
98
+ """Remove a tool from the server.
99
+
100
+ Args:
101
+ key: The key of the tool to remove
102
+
103
+ Raises:
104
+ NotFoundError: If the tool is not found
105
+ """
106
+ if key in self._tools:
107
+ del self._tools[key]
108
+ else:
109
+ raise NotFoundError(f"Tool {key!r} not found.")
110
+
111
  async def call_tool(
112
  self, key: str, arguments: dict[str, Any]
113
  ) -> list[TextContent | ImageContent | EmbeddedResource]:
tests/server/test_server.py CHANGED
@@ -72,6 +72,25 @@ class TestTools:
72
  assert len(mcp_tools) == 1
73
  assert mcp_tools[0].name == "custom_name"
74
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75
 
76
  class TestToolDecorator:
77
  async def test_no_tools_before_decorator(self):
 
72
  assert len(mcp_tools) == 1
73
  assert mcp_tools[0].name == "custom_name"
74
 
75
+ async def test_remove_tool_successfully(self):
76
+ """Test that FastMCP.remove_tool removes the tool from the registry."""
77
+
78
+ mcp = FastMCP()
79
+
80
+ @mcp.tool(name="adder")
81
+ def add(a: int, b: int) -> int:
82
+ return a + b
83
+
84
+ mcp_tools = await mcp.get_tools()
85
+ assert "adder" in mcp_tools
86
+
87
+ mcp.remove_tool("adder")
88
+ mcp_tools = await mcp.get_tools()
89
+ assert "adder" not in mcp_tools
90
+
91
+ with pytest.raises(NotFoundError, match="Unknown tool: adder"):
92
+ await mcp._mcp_call_tool("adder", {"a": 1, "b": 2})
93
+
94
 
95
  class TestToolDecorator:
96
  async def test_no_tools_before_decorator(self):
tests/tools/test_tool_manager.py CHANGED
@@ -100,6 +100,26 @@ class TestAddTools:
100
  ):
101
  manager.add_tool_from_fn(lambda x: x)
102
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
103
  def test_warn_on_duplicate_tools(self, caplog):
104
  """Test warning on duplicate tools."""
105
  manager = ToolManager(duplicate_behavior="warn")
 
100
  ):
101
  manager.add_tool_from_fn(lambda x: x)
102
 
103
+ def test_remove_tool_successfully(self):
104
+ """Test removing an added tool by key."""
105
+ manager = ToolManager()
106
+
107
+ def add(a: int, b: int) -> int:
108
+ return a + b
109
+
110
+ manager.add_tool_from_fn(add)
111
+ assert manager.get_tool("add") is not None
112
+
113
+ manager.remove_tool("add")
114
+ with pytest.raises(NotFoundError):
115
+ manager.get_tool("add")
116
+
117
+ def test_remove_tool_missing_key(self):
118
+ """Test removing a tool that does not exist raises NotFoundError."""
119
+ manager = ToolManager()
120
+ with pytest.raises(NotFoundError, match="Tool 'missing' not found"):
121
+ manager.remove_tool("missing")
122
+
123
  def test_warn_on_duplicate_tools(self, caplog):
124
  """Test warning on duplicate tools."""
125
  manager = ToolManager(duplicate_behavior="warn")