Support enabled/disabled tools

This commit is contained in:
Jeremiah Lowin 2025-06-10 10:32:50 -04:00
commit 6e92335510
4 changed files with 100 additions and 0 deletions

View file

@ -625,6 +625,7 @@ class FastMCP(Generic[LifespanResultT]):
tags: set[str] | None = None,
annotations: ToolAnnotations | dict[str, Any] | None = None,
exclude_args: list[str] | None = None,
enabled: bool | None = None,
) -> FunctionTool: ...
@overload
@ -637,6 +638,7 @@ class FastMCP(Generic[LifespanResultT]):
tags: set[str] | None = None,
annotations: ToolAnnotations | dict[str, Any] | None = None,
exclude_args: list[str] | None = None,
enabled: bool | None = None,
) -> Callable[[AnyFunction], FunctionTool]: ...
def tool(
@ -648,6 +650,7 @@ class FastMCP(Generic[LifespanResultT]):
tags: set[str] | None = None,
annotations: ToolAnnotations | dict[str, Any] | None = None,
exclude_args: list[str] | None = None,
enabled: bool | None = None,
) -> Callable[[AnyFunction], FunctionTool] | FunctionTool:
"""Decorator to register a tool.
@ -721,6 +724,7 @@ class FastMCP(Generic[LifespanResultT]):
annotations=annotations,
exclude_args=exclude_args,
serializer=self._tool_serializer,
enabled=enabled,
)
self.add_tool(tool)
return tool
@ -749,6 +753,7 @@ class FastMCP(Generic[LifespanResultT]):
tags=tags,
annotations=annotations,
exclude_args=exclude_args,
enabled=enabled,
)
def add_resource(self, resource: Resource, key: str | None = None) -> None:

View file

@ -62,6 +62,7 @@ class Tool(FastMCPComponent, ABC):
annotations: ToolAnnotations | None = None,
exclude_args: list[str] | None = None,
serializer: Callable[[Any], str] | None = None,
enabled: bool | None = None,
) -> FunctionTool:
"""Create a Tool from a function."""
return FunctionTool.from_function(
@ -72,6 +73,7 @@ class Tool(FastMCPComponent, ABC):
annotations=annotations,
exclude_args=exclude_args,
serializer=serializer,
enabled=enabled,
)
@abstractmethod
@ -92,6 +94,7 @@ class Tool(FastMCPComponent, ABC):
tags: set[str] | None = None,
annotations: ToolAnnotations | None = None,
serializer: Callable[[Any], str] | None = None,
enabled: bool | None = None,
) -> TransformedTool:
from fastmcp.tools.tool_transform import TransformedTool
@ -104,6 +107,7 @@ class Tool(FastMCPComponent, ABC):
tags=tags,
annotations=annotations,
serializer=serializer,
enabled=enabled,
)
@ -120,6 +124,7 @@ class FunctionTool(Tool):
annotations: ToolAnnotations | None = None,
exclude_args: list[str] | None = None,
serializer: Callable[[Any], str] | None = None,
enabled: bool | None = None,
) -> FunctionTool:
"""Create a Tool from a function."""
@ -136,6 +141,7 @@ class FunctionTool(Tool):
tags=tags or set(),
annotations=annotations,
serializer=serializer,
enabled=enabled if enabled is not None else True,
)
async def run(

View file

@ -267,6 +267,7 @@ class TransformedTool(Tool):
transform_args: dict[str, ArgTransform] | None = None,
annotations: ToolAnnotations | None = None,
serializer: Callable[[Any], str] | None = None,
enabled: bool | None = None,
) -> TransformedTool:
"""Create a transformed tool from a parent tool.
@ -399,6 +400,7 @@ class TransformedTool(Tool):
annotations=annotations or tool.annotations,
serializer=serializer or tool.serializer,
transform_args=transform_args,
enabled=enabled if enabled is not None else True,
)
return transformed_tool

View file

@ -699,6 +699,93 @@ class TestToolContextInjection:
assert result[0].text == "3" # type: ignore[attr-defined]
class TestToolEnabled:
async def test_toggle_enabled(self):
mcp = FastMCP()
@mcp.tool
def sample_tool(x: int) -> int:
return x * 2
assert sample_tool.enabled
tool = await mcp.get_tool("sample_tool")
assert tool.enabled
tool.disable()
assert not tool.enabled
assert not sample_tool.enabled
tool.enable()
assert tool.enabled
assert sample_tool.enabled
async def test_tool_disabled_in_decorator(self):
mcp = FastMCP()
@mcp.tool(enabled=False)
def sample_tool(x: int) -> int:
return x * 2
async with Client(mcp) as client:
tools = await client.list_tools()
assert len(tools) == 0
async def test_tool_toggle_enabled(self):
mcp = FastMCP()
@mcp.tool(enabled=False)
def sample_tool(x: int) -> int:
return x * 2
sample_tool.enable()
async with Client(mcp) as client:
tools = await client.list_tools()
assert len(tools) == 1
async def test_tool_toggle_disabled(self):
mcp = FastMCP()
@mcp.tool
def sample_tool(x: int) -> int:
return x * 2
sample_tool.disable()
async with Client(mcp) as client:
tools = await client.list_tools()
assert len(tools) == 0
async def test_get_tool_and_disable(self):
mcp = FastMCP()
@mcp.tool
def sample_tool(x: int) -> int:
return x * 2
tool = await mcp.get_tool("sample_tool")
assert tool.enabled
sample_tool.disable()
async with Client(mcp) as client:
result = await client.list_tools()
assert len(result) == 0
async def test_cant_call_disabled_tool(self):
mcp = FastMCP()
@mcp.tool(enabled=False)
def sample_tool(x: int) -> int:
return x * 2
with pytest.raises(Exception, match="Unknown tool"):
async with Client(mcp) as client:
await client.call_tool("sample_tool", {"x": 5})
class TestResource:
async def test_text_resource(self):
mcp = FastMCP()