From c183e3a99cc2977dabaa56d03d5ba1cc3d87ffbd Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 18 Jun 2025 21:17:38 -0400 Subject: [PATCH] Updated list_tools --- src/fastmcp/server/middleware.py | 31 ++++++++++++++++ src/fastmcp/server/server.py | 26 +++++++------- tests/server/middleware/test_middleware.py | 42 +++++++++++++++++++++- 3 files changed, 85 insertions(+), 14 deletions(-) diff --git a/src/fastmcp/server/middleware.py b/src/fastmcp/server/middleware.py index 3b0ba369e..3f73bfe07 100644 --- a/src/fastmcp/server/middleware.py +++ b/src/fastmcp/server/middleware.py @@ -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) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index c8d93f29f..86298d5e1 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -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 diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index c21591788..3a2bf7f2c 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -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)