mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 13:04:18 +02:00
Updated list_tools
This commit is contained in:
parent
e11417c3c2
commit
c183e3a99c
3 changed files with 85 additions and 14 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue