Jeremiah Lowin commited on
Commit
96709d8
·
unverified ·
1 Parent(s): d9fc88d

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