Jeremiah Lowin commited on
Commit
ccc2857
·
unverified ·
1 Parent(s): 2509c47

Fix typing, add tests for tool call middleware (#1269)

Browse files
docs/servers/middleware.mdx CHANGED
@@ -202,6 +202,40 @@ class ListingFilterMiddleware(Middleware):
202
 
203
  This filtering happens before the components are converted to MCP format and returned to the client, so the tags (which are FastMCP-specific) are naturally stripped in the final response.
204
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
205
  ### Anatomy of a Hook
206
 
207
  Every middleware hook follows the same pattern. Let's examine the `on_message` hook to understand the structure:
 
202
 
203
  This filtering happens before the components are converted to MCP format and returned to the client, so the tags (which are FastMCP-specific) are naturally stripped in the final response.
204
 
205
+ <Tip>
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:
212
+
213
+ ```python
214
+ from fastmcp.server.middleware import Middleware, MiddlewareContext
215
+
216
+ class ToolCallMiddleware(Middleware):
217
+ async def on_call_tool(self, context: MiddlewareContext, call_next):
218
+ # Modify arguments before execution
219
+ if context.message.name == "calculate":
220
+ # Ensure positive inputs
221
+ if context.message.arguments.get("value", 0) < 0:
222
+ context.message.arguments["value"] = abs(context.message.arguments["value"])
223
+
224
+ result = await call_next(context)
225
+
226
+ # Transform result after execution
227
+ if context.message.name == "get_data":
228
+ # Add metadata to result
229
+ if result.structured_content:
230
+ result.structured_content["processed_at"] = "2024-01-01T00:00:00Z"
231
+
232
+ return result
233
+ ```
234
+
235
+ <Tip>
236
+ For more complex tool rewriting scenarios, consider using [Tool Transformation](/patterns/tool-transformation) patterns which provide a more structured approach to creating modified tool variants.
237
+ </Tip>
238
+
239
  ### Anatomy of a Hook
240
 
241
  Every middleware hook follows the same pattern. Let's examine the `on_message` hook to understand the structure:
src/fastmcp/server/middleware/middleware.py CHANGED
@@ -20,7 +20,7 @@ import mcp.types as mt
20
  from fastmcp.prompts.prompt import Prompt
21
  from fastmcp.resources.resource import Resource
22
  from fastmcp.resources.template import ResourceTemplate
23
- from fastmcp.tools.tool import Tool
24
 
25
  if TYPE_CHECKING:
26
  from fastmcp.server.context import Context
@@ -43,26 +43,6 @@ class CallNext(Protocol[T, R]):
43
  def __call__(self, context: MiddlewareContext[T]) -> Awaitable[R]: ...
44
 
45
 
46
- ServerResultT = TypeVar(
47
- "ServerResultT",
48
- bound=mt.EmptyResult
49
- | mt.InitializeResult
50
- | mt.CompleteResult
51
- | mt.GetPromptResult
52
- | mt.ListPromptsResult
53
- | mt.ListResourcesResult
54
- | mt.ListResourceTemplatesResult
55
- | mt.ReadResourceResult
56
- | mt.CallToolResult
57
- | mt.ListToolsResult,
58
- )
59
-
60
-
61
- @runtime_checkable
62
- class ServerResultProtocol(Protocol[ServerResultT]):
63
- root: ServerResultT
64
-
65
-
66
  @dataclass(kw_only=True, frozen=True)
67
  class MiddlewareContext(Generic[T]):
68
  """
@@ -167,8 +147,8 @@ class Middleware:
167
  async def on_call_tool(
168
  self,
169
  context: MiddlewareContext[mt.CallToolRequestParams],
170
- call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult],
171
- ) -> mt.CallToolResult:
172
  return await call_next(context)
173
 
174
  async def on_read_resource(
 
20
  from fastmcp.prompts.prompt import Prompt
21
  from fastmcp.resources.resource import Resource
22
  from fastmcp.resources.template import ResourceTemplate
23
+ from fastmcp.tools.tool import Tool, ToolResult
24
 
25
  if TYPE_CHECKING:
26
  from fastmcp.server.context import Context
 
43
  def __call__(self, context: MiddlewareContext[T]) -> Awaitable[R]: ...
44
 
45
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
46
  @dataclass(kw_only=True, frozen=True)
47
  class MiddlewareContext(Generic[T]):
48
  """
 
147
  async def on_call_tool(
148
  self,
149
  context: MiddlewareContext[mt.CallToolRequestParams],
150
+ call_next: CallNext[mt.CallToolRequestParams, ToolResult],
151
+ ) -> ToolResult:
152
  return await call_next(context)
153
 
154
  async def on_read_resource(
tests/server/middleware/test_middleware.py CHANGED
@@ -7,7 +7,8 @@ import pytest
7
 
8
  from fastmcp import Client, FastMCP
9
  from fastmcp.server.context import Context
10
- from fastmcp.server.middleware import Middleware, MiddlewareContext
 
11
 
12
 
13
  @dataclass
@@ -411,6 +412,38 @@ class TestMiddlewareHooks:
411
  assert len(prompts) == 1
412
  assert prompts[0].name == "public_prompt"
413
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
414
 
415
  class TestNestedMiddlewareHooks:
416
  @pytest.fixture
 
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
12
 
13
 
14
  @dataclass
 
412
  assert len(prompts) == 1
413
  assert prompts[0].name == "public_prompt"
414
 
415
+ async def test_call_tool_middleware(self):
416
+ server = FastMCP()
417
+
418
+ @server.tool
419
+ def add(a: int, b: int) -> int:
420
+ return a + b
421
+
422
+ class CallToolMiddleware(Middleware):
423
+ async def on_call_tool(
424
+ self,
425
+ context: MiddlewareContext[mcp.types.CallToolRequestParams],
426
+ call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult],
427
+ ):
428
+ # modify argument
429
+ if context.message.name == "add":
430
+ context.message.arguments["a"] += 100 # type: ignore
431
+
432
+ result = await call_next(context)
433
+
434
+ # modify result
435
+ if context.message.name == "add":
436
+ result.structured_content["result"] += 5 # type: ignore
437
+
438
+ return result
439
+
440
+ server.add_middleware(CallToolMiddleware())
441
+
442
+ async with Client(server) as client:
443
+ result = await client.call_tool("add", {"a": 1, "b": 2})
444
+
445
+ assert result.structured_content["result"] == 108 # type: ignore
446
+
447
 
448
  class TestNestedMiddlewareHooks:
449
  @pytest.fixture