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

This commit is contained in:
Jeremiah Lowin 2025-07-25 15:46:48 -07:00 committed by GitHub
commit 07655e6620
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 71 additions and 24 deletions

View file

@ -202,6 +202,40 @@ class ListingFilterMiddleware(Middleware):
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.
<Tip>
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.
</Tip>
### Tool Call Modification
For execution operations like tool calls, you can modify arguments before execution or transform results afterward:
```python
from fastmcp.server.middleware import Middleware, MiddlewareContext
class ToolCallMiddleware(Middleware):
async def on_call_tool(self, context: MiddlewareContext, call_next):
# Modify arguments before execution
if context.message.name == "calculate":
# Ensure positive inputs
if context.message.arguments.get("value", 0) < 0:
context.message.arguments["value"] = abs(context.message.arguments["value"])
result = await call_next(context)
# Transform result after execution
if context.message.name == "get_data":
# Add metadata to result
if result.structured_content:
result.structured_content["processed_at"] = "2024-01-01T00:00:00Z"
return result
```
<Tip>
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.
</Tip>
### Anatomy of a Hook
Every middleware hook follows the same pattern. Let's examine the `on_message` hook to understand the structure:

View file

@ -20,7 +20,7 @@ import mcp.types as mt
from fastmcp.prompts.prompt import Prompt
from fastmcp.resources.resource import Resource
from fastmcp.resources.template import ResourceTemplate
from fastmcp.tools.tool import Tool
from fastmcp.tools.tool import Tool, ToolResult
if TYPE_CHECKING:
from fastmcp.server.context import Context
@ -43,26 +43,6 @@ class CallNext(Protocol[T, R]):
def __call__(self, context: MiddlewareContext[T]) -> Awaitable[R]: ...
ServerResultT = TypeVar(
"ServerResultT",
bound=mt.EmptyResult
| mt.InitializeResult
| mt.CompleteResult
| mt.GetPromptResult
| mt.ListPromptsResult
| mt.ListResourcesResult
| mt.ListResourceTemplatesResult
| mt.ReadResourceResult
| mt.CallToolResult
| mt.ListToolsResult,
)
@runtime_checkable
class ServerResultProtocol(Protocol[ServerResultT]):
root: ServerResultT
@dataclass(kw_only=True, frozen=True)
class MiddlewareContext(Generic[T]):
"""
@ -167,8 +147,8 @@ class Middleware:
async def on_call_tool(
self,
context: MiddlewareContext[mt.CallToolRequestParams],
call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult],
) -> mt.CallToolResult:
call_next: CallNext[mt.CallToolRequestParams, ToolResult],
) -> ToolResult:
return await call_next(context)
async def on_read_resource(

View file

@ -7,7 +7,8 @@ import pytest
from fastmcp import Client, FastMCP
from fastmcp.server.context import Context
from fastmcp.server.middleware import Middleware, MiddlewareContext
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from fastmcp.tools.tool import ToolResult
@dataclass
@ -411,6 +412,38 @@ class TestMiddlewareHooks:
assert len(prompts) == 1
assert prompts[0].name == "public_prompt"
async def test_call_tool_middleware(self):
server = FastMCP()
@server.tool
def add(a: int, b: int) -> int:
return a + b
class CallToolMiddleware(Middleware):
async def on_call_tool(
self,
context: MiddlewareContext[mcp.types.CallToolRequestParams],
call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult],
):
# modify argument
if context.message.name == "add":
context.message.arguments["a"] += 100 # type: ignore
result = await call_next(context)
# modify result
if context.message.name == "add":
result.structured_content["result"] += 5 # type: ignore
return result
server.add_middleware(CallToolMiddleware())
async with Client(server) as client:
result = await client.call_tool("add", {"a": 1, "b": 2})
assert result.structured_content["result"] == 108 # type: ignore
class TestNestedMiddlewareHooks:
@pytest.fixture