Updated list_tools

This commit is contained in:
Jeremiah Lowin 2025-06-18 21:17:38 -04:00
commit c183e3a99c
3 changed files with 85 additions and 14 deletions

View file

@ -144,6 +144,14 @@ class MCPMiddleware:
handler = partial(self.on_read_resource, call_next=handler)
case "prompts/get":
handler = partial(self.on_get_prompt, call_next=handler)
case "tools/list":
handler = partial(self.on_list_tools, call_next=handler)
case "resources/list":
handler = partial(self.on_list_resources, call_next=handler)
case "resource-templates/list":
handler = partial(self.on_list_resource_templates, call_next=handler)
case "prompts/list":
handler = partial(self.on_list_prompts, call_next=handler)
match context.type:
case "request":
@ -203,3 +211,26 @@ class MCPMiddleware:
call_next: CallNext[mt.ListToolsRequest, ListToolsResult],
) -> ListToolsResult:
return await call_next(context)
async def on_list_resources(
self,
context: MiddlewareContext[mt.ListResourcesRequest],
call_next: CallNext[mt.ListResourcesRequest, ListResourcesResult],
) -> ListResourcesResult:
return await call_next(context)
async def on_list_resource_templates(
self,
context: MiddlewareContext[mt.ListResourceTemplatesRequest],
call_next: CallNext[
mt.ListResourceTemplatesRequest, ListResourceTemplatesResult
],
) -> ListResourceTemplatesResult:
return await call_next(context)
async def on_list_prompts(
self,
context: MiddlewareContext[mt.ListPromptsRequest],
call_next: CallNext[mt.ListPromptsRequest, ListPromptsResult],
) -> ListPromptsResult:
return await call_next(context)

View file

@ -341,7 +341,8 @@ class FastMCP(Generic[LifespanResultT]):
async def get_tools(self) -> dict[str, Tool]:
"""Get all registered tools, indexed by registered key."""
return await self._list_tools(apply_middleware=False)
tools = await self._list_tools(apply_middleware=False)
return {tool.key: tool for tool in tools}
async def get_tool(self, key: str) -> Tool:
tools = await self.get_tools()
@ -515,7 +516,7 @@ class FastMCP(Generic[LifespanResultT]):
tools = await self._middleware_list_tools()
return [tool.to_mcp_tool(name=tool.name) for tool in tools]
async def _middleware_list_tools(self) -> dict[str, Tool]:
async def _middleware_list_tools(self) -> list[Tool]:
"""
List all available tools, in the format expected by the low-level MCP
server.
@ -524,13 +525,13 @@ class FastMCP(Generic[LifespanResultT]):
async def _handler(
context: MiddlewareContext[mcp.types.ListToolsRequest],
) -> list[MCPTool]:
) -> list[Tool]:
tools = await self._list_tools()
mcp_tools: list[MCPTool] = []
for key, tool in tools.items():
mcp_tools: list[Tool] = []
for tool in tools:
if self._should_enable_component(tool):
mcp_tools.append(tool.to_mcp_tool(name=key))
mcp_tools.append(tool)
return mcp_tools
@ -547,13 +548,13 @@ class FastMCP(Generic[LifespanResultT]):
# Apply the middleware chain.
return await self._apply_middleware(mw_context, _handler)
async def _list_tools(self, apply_middleware: bool = True) -> dict[str, Tool]:
async def _list_tools(self, apply_middleware: bool = True) -> list[Tool]:
"""
List all available tools.
"""
if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND:
tools: dict[str, Tool] = {}
tools: list[Tool] = []
# iterate such that new mounts overwrite older ones
for mounted_server in self._mounted_servers:
@ -566,18 +567,17 @@ class FastMCP(Generic[LifespanResultT]):
server_tools = await mounted_server.server._list_tools()
# Apply prefix to each tool key if prefix exists and is not empty
if mounted_server.prefix:
for tool in server_tools.values():
for tool in server_tools:
tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}")
tools[tool.key] = tool
tools.append(tool)
else:
tools.update(server_tools)
tools.update(server_tools)
tools.extend(server_tools)
except Exception as e:
logger.warning(
f"Failed to get tools from mounted server '{mounted_server.prefix}': {e}"
)
continue
tools.update(self._tool_manager.get_tools())
tools.extend(self._tool_manager.get_tools().values())
self._cache.set("tools", tools)
return tools

View file

@ -7,7 +7,7 @@ import pytest
from fastmcp import Client, FastMCP
from fastmcp.server.context import Context
from fastmcp.server.middleware.middleware import MCPMiddleware, MiddlewareContext
from fastmcp.server.middleware import MCPMiddleware, MiddlewareContext
@dataclass
@ -194,3 +194,43 @@ class TestMiddlewareHooks:
assert recording_middleware.assert_called(hook="on_message", times=1)
assert recording_middleware.assert_called(hook="on_request", times=1)
assert recording_middleware.assert_called(hook="on_list_tools", times=1)
async def test_list_resources(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_resources()
assert recording_middleware.assert_called(times=3)
assert recording_middleware.assert_called(method="resources/list", times=3)
assert recording_middleware.assert_called(hook="on_message", times=1)
assert recording_middleware.assert_called(hook="on_request", times=1)
assert recording_middleware.assert_called(hook="on_list_resources", times=1)
async def test_list_resource_templates(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_resource_templates()
assert recording_middleware.assert_called(times=3)
assert recording_middleware.assert_called(
method="resource-templates/list", times=3
)
assert recording_middleware.assert_called(hook="on_message", times=1)
assert recording_middleware.assert_called(hook="on_request", times=1)
assert recording_middleware.assert_called(
hook="on_list_resource_templates", times=1
)
async def test_list_prompts(
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
):
async with Client(mcp_server) as client:
await client.list_prompts()
assert recording_middleware.assert_called(times=3)
assert recording_middleware.assert_called(method="prompts/list", times=3)
assert recording_middleware.assert_called(hook="on_message", times=1)
assert recording_middleware.assert_called(hook="on_request", times=1)
assert recording_middleware.assert_called(hook="on_list_prompts", times=1)