Spaces:
Running
Running
Jeremiah Lowin commited on
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,
|
| 171 |
-
) ->
|
| 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
|