Spaces:
Running
Running
Jeremiah Lowin commited on
Add unit tests and docs for denying tool calls with middleware (#1333)
Browse files
docs/servers/middleware.mdx
CHANGED
|
@@ -206,6 +206,30 @@ This filtering happens before the components are converted to MCP format and ret
|
|
| 206 |
When filtering components in listing operations, ensure you also prevent execution of filtered components in the corresponding execution hooks (`on_call_tool`, `on_read_resource`, `on_get_prompt`) to maintain consistency.
|
| 207 |
</Tip>
|
| 208 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 209 |
### Tool Call Modification
|
| 210 |
|
| 211 |
For execution operations like tool calls, you can modify arguments before execution or transform results afterward:
|
|
|
|
| 206 |
When filtering components in listing operations, ensure you also prevent execution of filtered components in the corresponding execution hooks (`on_call_tool`, `on_read_resource`, `on_get_prompt`) to maintain consistency.
|
| 207 |
</Tip>
|
| 208 |
|
| 209 |
+
### Tool Call Denial
|
| 210 |
+
|
| 211 |
+
You can deny access to specific tools by raising a `ToolError` in your middleware. This is the correct way to block tool execution, as it integrates properly with the FastMCP error handling system.
|
| 212 |
+
|
| 213 |
+
```python
|
| 214 |
+
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
| 215 |
+
from fastmcp.exceptions import ToolError
|
| 216 |
+
|
| 217 |
+
class AuthMiddleware(Middleware):
|
| 218 |
+
async def on_call_tool(self, context: MiddlewareContext, call_next):
|
| 219 |
+
tool_name = context.message.name
|
| 220 |
+
|
| 221 |
+
# Deny access to restricted tools
|
| 222 |
+
if tool_name.lower() in ["delete", "admin_config"]:
|
| 223 |
+
raise ToolError("Access denied: tool requires admin privileges")
|
| 224 |
+
|
| 225 |
+
# Allow other tools to proceed
|
| 226 |
+
return await call_next(context)
|
| 227 |
+
```
|
| 228 |
+
|
| 229 |
+
<Warning>
|
| 230 |
+
When denying tool calls, always raise `ToolError` rather than returning `ToolResult` objects or other values. `ToolError` ensures proper error propagation through the middleware chain and converts to the correct MCP error response format.
|
| 231 |
+
</Warning>
|
| 232 |
+
|
| 233 |
### Tool Call Modification
|
| 234 |
|
| 235 |
For execution operations like tool calls, you can modify arguments before execution or transform results afterward:
|
tests/server/middleware/test_middleware.py
CHANGED
|
@@ -6,6 +6,7 @@ import mcp.types
|
|
| 6 |
import pytest
|
| 7 |
|
| 8 |
from fastmcp import Client, FastMCP
|
|
|
|
| 9 |
from fastmcp.server.context import Context
|
| 10 |
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
| 11 |
from fastmcp.tools.tool import ToolResult
|
|
@@ -785,3 +786,109 @@ class TestProxyServer:
|
|
| 785 |
await client.list_tools()
|
| 786 |
|
| 787 |
assert TAGS == [{"add-tool"}, set(), set(), set()]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 6 |
import pytest
|
| 7 |
|
| 8 |
from fastmcp import Client, FastMCP
|
| 9 |
+
from fastmcp.exceptions import ToolError
|
| 10 |
from fastmcp.server.context import Context
|
| 11 |
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
| 12 |
from fastmcp.tools.tool import ToolResult
|
|
|
|
| 786 |
await client.list_tools()
|
| 787 |
|
| 788 |
assert TAGS == [{"add-tool"}, set(), set(), set()]
|
| 789 |
+
|
| 790 |
+
|
| 791 |
+
class TestToolCallDenial:
|
| 792 |
+
"""Test denying tool calls in middleware using ToolError."""
|
| 793 |
+
|
| 794 |
+
async def test_deny_tool_call_with_tool_error(self):
|
| 795 |
+
"""Test that middleware can deny tool calls by raising ToolError."""
|
| 796 |
+
|
| 797 |
+
class AuthMiddleware(Middleware):
|
| 798 |
+
async def on_call_tool(
|
| 799 |
+
self,
|
| 800 |
+
context: MiddlewareContext[mcp.types.CallToolRequestParams],
|
| 801 |
+
call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult],
|
| 802 |
+
) -> ToolResult:
|
| 803 |
+
tool_name = context.message.name
|
| 804 |
+
if tool_name.lower() == "restricted_tool":
|
| 805 |
+
raise ToolError("Access denied: tool is disabled")
|
| 806 |
+
return await call_next(context)
|
| 807 |
+
|
| 808 |
+
server = FastMCP("TestServer")
|
| 809 |
+
|
| 810 |
+
@server.tool
|
| 811 |
+
def allowed_tool(x: int) -> int:
|
| 812 |
+
"""This tool is allowed."""
|
| 813 |
+
return x * 2
|
| 814 |
+
|
| 815 |
+
@server.tool
|
| 816 |
+
def restricted_tool(x: int) -> int:
|
| 817 |
+
"""This tool should be denied by middleware."""
|
| 818 |
+
return x * 3
|
| 819 |
+
|
| 820 |
+
server.add_middleware(AuthMiddleware())
|
| 821 |
+
|
| 822 |
+
async with Client(server) as client:
|
| 823 |
+
# Allowed tool should work normally
|
| 824 |
+
result = await client.call_tool("allowed_tool", {"x": 5})
|
| 825 |
+
assert result.structured_content is not None
|
| 826 |
+
assert result.structured_content["result"] == 10
|
| 827 |
+
|
| 828 |
+
# Restricted tool should raise ToolError
|
| 829 |
+
with pytest.raises(ToolError) as exc_info:
|
| 830 |
+
await client.call_tool("restricted_tool", {"x": 5})
|
| 831 |
+
|
| 832 |
+
# Verify the error message is preserved
|
| 833 |
+
assert "Access denied: tool is disabled" in str(exc_info.value)
|
| 834 |
+
|
| 835 |
+
async def test_middleware_can_selectively_deny_tools(self):
|
| 836 |
+
"""Test that middleware can deny specific tools while allowing others."""
|
| 837 |
+
|
| 838 |
+
denied_tools = set()
|
| 839 |
+
|
| 840 |
+
class SelectiveAuthMiddleware(Middleware):
|
| 841 |
+
async def on_call_tool(
|
| 842 |
+
self,
|
| 843 |
+
context: MiddlewareContext[mcp.types.CallToolRequestParams],
|
| 844 |
+
call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult],
|
| 845 |
+
) -> ToolResult:
|
| 846 |
+
tool_name = context.message.name
|
| 847 |
+
|
| 848 |
+
# Deny tools that start with "admin_"
|
| 849 |
+
if tool_name.startswith("admin_"):
|
| 850 |
+
denied_tools.add(tool_name)
|
| 851 |
+
raise ToolError(
|
| 852 |
+
f"Access denied: {tool_name} requires admin privileges"
|
| 853 |
+
)
|
| 854 |
+
|
| 855 |
+
return await call_next(context)
|
| 856 |
+
|
| 857 |
+
server = FastMCP("TestServer")
|
| 858 |
+
|
| 859 |
+
@server.tool
|
| 860 |
+
def public_tool(x: int) -> int:
|
| 861 |
+
"""Public tool available to all."""
|
| 862 |
+
return x + 1
|
| 863 |
+
|
| 864 |
+
@server.tool
|
| 865 |
+
def admin_delete(item_id: str) -> str:
|
| 866 |
+
"""Admin tool that should be denied."""
|
| 867 |
+
return f"Deleted {item_id}"
|
| 868 |
+
|
| 869 |
+
@server.tool
|
| 870 |
+
def admin_config(setting: str, value: str) -> str:
|
| 871 |
+
"""Another admin tool that should be denied."""
|
| 872 |
+
return f"Set {setting} to {value}"
|
| 873 |
+
|
| 874 |
+
server.add_middleware(SelectiveAuthMiddleware())
|
| 875 |
+
|
| 876 |
+
async with Client(server) as client:
|
| 877 |
+
# Public tool should work
|
| 878 |
+
result = await client.call_tool("public_tool", {"x": 10})
|
| 879 |
+
assert result.structured_content is not None
|
| 880 |
+
assert result.structured_content["result"] == 11
|
| 881 |
+
|
| 882 |
+
# Admin tools should be denied
|
| 883 |
+
with pytest.raises(ToolError) as exc_info:
|
| 884 |
+
await client.call_tool("admin_delete", {"item_id": "test123"})
|
| 885 |
+
assert "requires admin privileges" in str(exc_info.value)
|
| 886 |
+
|
| 887 |
+
with pytest.raises(ToolError) as exc_info:
|
| 888 |
+
await client.call_tool(
|
| 889 |
+
"admin_config", {"setting": "debug", "value": "true"}
|
| 890 |
+
)
|
| 891 |
+
assert "requires admin privileges" in str(exc_info.value)
|
| 892 |
+
|
| 893 |
+
# Verify both admin tools were denied
|
| 894 |
+
assert denied_tools == {"admin_delete", "admin_config"}
|