Update tool manager

src/fastmcp/server/server.py

@@ -434,10 +434,10 @@ class FastMCP(Generic[LifespanResultT]):
         logger.debug("Handler called: list_tools")
 
         with fastmcp.server.context.Context(fastmcp=self):
-            tools = await self._middleware_list_tools()
+            tools = await self._list_tools()
             return [tool.to_mcp_tool(name=tool.key) for tool in tools]
 
-    async def _middleware_list_tools(self) -> list[Tool]:
+    async def _list_tools(self) -> list[Tool]:
         """
         List all available tools, in the format expected by the low-level MCP
         server.
@@ -447,7 +447,7 @@ class FastMCP(Generic[LifespanResultT]):
         async def _handler(
             context: MiddlewareContext[mcp.types.ListToolsRequest],
         ) -> list[Tool]:
-            tools = await self._list_tools()
+            tools = await self._tool_manager.list_tools()
 
             mcp_tools: list[Tool] = []
             for tool in tools:
@@ -469,39 +469,6 @@ 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) -> list[Tool]:
-        """
-        List all available tools.
-        """
-
-        if (tools := self._cache.get("tools")) is self._cache.NOT_FOUND:
-            tools: dict[str, Tool] = {}
-
-            # iterate such that new mounts overwrite older ones
-            for mounted_server in self._mounted_servers:
-                try:
-                    if apply_middleware:
-                        server_tools = (
-                            await mounted_server.server._middleware_list_tools()
-                        )
-                    else:
-                        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:
-                            tool = tool.with_key(f"{mounted_server.prefix}_{tool.key}")
-                            tools[tool.key] = tool
-                    else:
-                        tools.update({tool.key: tool for tool in 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())
-            self._cache.set("tools", tools)
-        return list(tools.values())
-
     async def _mcp_list_resources(self) -> list[MCPResource]:
         logger.debug("Handler called: list_resources")
 
@@ -580,7 +547,11 @@ class FastMCP(Generic[LifespanResultT]):
                         f"Failed to get resources from mounted server '{mounted_server.prefix}': {e}"
                     )
                     continue
-            resources.update(self._resource_manager.get_resources())
+            (
+                local_resources,
+                _,
+            ) = await self._resource_manager.get_resources_and_templates()
+            resources.update(local_resources)
             self._cache.set("resources", resources)
         return list(resources.values())
 
@@ -668,7 +639,11 @@ class FastMCP(Generic[LifespanResultT]):
                         f"'{mounted_server.prefix}': {e}"
                     )
                     continue
-            templates.update(self._resource_manager.get_templates())
+            (
+                _,
+                local_templates,
+            ) = await self._resource_manager.get_resources_and_templates()
+            templates.update(local_templates)
             self._cache.set("resource_templates", templates)
         return list(templates.values())
 
@@ -744,7 +719,7 @@ class FastMCP(Generic[LifespanResultT]):
                         f"Failed to get prompts from mounted server '{mounted_server.prefix}': {e}"
                     )
                     continue
-            prompts.update(self._prompt_manager.get_prompts())
+            prompts.update(await self._prompt_manager.get_prompts())
             self._cache.set("prompts", prompts)
         return list(prompts.values())
 
@@ -767,27 +742,26 @@ class FastMCP(Generic[LifespanResultT]):
 
         with fastmcp.server.context.Context(fastmcp=self):
             try:
-                return await self._middleware_call_tool(key, arguments)
+                return await self._call_tool(key, arguments)
             except DisabledError:
                 raise NotFoundError(f"Unknown tool: {key}")
             except NotFoundError:
                 raise NotFoundError(f"Unknown tool: {key}")
 
-    async def _middleware_call_tool(
-        self,
-        key: str,
-        arguments: dict[str, Any],
-    ) -> list[MCPContent]:
+    async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]:
         """
-        Call a tool with middleware.
+        Applies this server's middleware and delegates the filtered call to the manager.
         """
 
         async def _handler(
             context: MiddlewareContext[mcp.types.CallToolRequestParams],
         ) -> list[MCPContent]:
-            return await self._call_tool(
-                key=context.message.name,
-                arguments=context.message.arguments or {},
+            tool = await self._tool_manager.get_tool(context.message.name)
+            if not self._should_enable_component(tool):
+                raise NotFoundError(f"Unknown tool: {context.message.name!r}")
+
+            return await self._tool_manager.call_tool(
+                key=context.message.name, arguments=context.message.arguments or {}
             )
 
         mw_context = MiddlewareContext(
@@ -799,46 +773,6 @@ class FastMCP(Generic[LifespanResultT]):
         )
         return await self._apply_middleware(mw_context, _handler)
 
-    async def _call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]:
-        """
-        Call a tool with raw MCP arguments. FastMCP subclasses should override
-        this method, not _mcp_call_tool.
-
-        Args:
-            key: The name of the tool to call arguments: Arguments to pass to
-            the tool
-
-        Returns:
-            List of MCP Content objects containing the tool results
-        """
-
-        # Get tool, checking first from our tools, then from the mounted servers
-        if self._tool_manager.has_tool(key):
-            tool = self._tool_manager.get_tool(key)
-            if not self._should_enable_component(tool):
-                raise DisabledError(f"Tool {key!r} is disabled")
-            return await self._tool_manager.call_tool(key, arguments)
-
-        # Check mounted servers to see if they have the tool
-        # iterate such that new mounts take precedence over older ones
-        for mounted_server in reversed(self._mounted_servers):
-            tool_key = key
-            try:
-                # If server has a prefix, check if key matches and strip prefix
-                if mounted_server.prefix:
-                    if tool_key.startswith(f"{mounted_server.prefix}_"):
-                        tool_key = tool_key.removeprefix(f"{mounted_server.prefix}_")
-                    else:
-                        continue
-                return await mounted_server.server._middleware_call_tool(
-                    tool_key, arguments
-                )
-            except NotFoundError:
-                # Tool not found on this server, try the next one
-                continue
-
-        raise NotFoundError(f"Unknown tool: {key!r}")
-
     async def _mcp_read_resource(self, uri: AnyUrl | str) -> list[ReadResourceContents]:
         """
         Handle MCP 'readResource' requests.
@@ -1853,11 +1787,12 @@ class FastMCP(Generic[LifespanResultT]):
         if as_proxy and not isinstance(server, FastMCPProxy):
             server = FastMCPProxy(Client(transport=FastMCPTransport(server)))
 
-        mounted_server = MountedServer(
-            server=server,
-            prefix=prefix,
-        )
-        self._mounted_servers.append(mounted_server)
+        # Delegate mounting to all three managers
+        mounted_server = MountedServer(prefix=prefix, server=server)
+        self._tool_manager.mount(mounted_server)
+        self._resource_manager.mount(mounted_server)
+        self._prompt_manager.mount(mounted_server)
+
         self._cache.clear()
 
     async def import_server(

src/fastmcp/tools/tool_manager.py

@@ -1,8 +1,8 @@
-from __future__ import annotations as _annotations
+from __future__ import annotations
 
 import warnings
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from mcp.types import ToolAnnotations
 
@@ -14,7 +14,7 @@ from fastmcp.utilities.logging import get_logger
 from fastmcp.utilities.types import MCPContent
 
 if TYPE_CHECKING:
-    pass
+    from fastmcp.server.server import MountedServer
 
 logger = get_logger(__name__)
 
@@ -28,6 +28,7 @@ class ToolManager:
         mask_error_details: bool | None = None,
     ):
         self._tools: dict[str, Tool] = {}
+        self._mounted_sources: list[MountedServer] = []
         self.mask_error_details = mask_error_details or settings.mask_error_details
 
         # Default to "warn" if None is provided
@@ -42,23 +43,78 @@ class ToolManager:
 
         self.duplicate_behavior = duplicate_behavior
 
-    def has_tool(self, key: str) -> bool:
+    def mount(self, server: MountedServer) -> None:
+        """Adds a mounted server as a source for tools."""
+        self._mounted_sources.append(server)
+
+    async def _load_tools(
+        self, *, mode: Literal["inventory", "protocol"]
+    ) -> dict[str, Tool]:
+        """
+        The single, consolidated recursive method for fetching tools. The 'mode'
+        parameter determines the communication path.
+
+        - mode="inventory": Manager-to-manager path for complete, unfiltered inventory
+        - mode="protocol": Server-to-server path for filtered MCP requests
+        """
+        all_tools: dict[str, Tool] = {}
+
+        for mounted in self._mounted_sources:
+            try:
+                if mode == "protocol":
+                    # PATH 2: Use the server-to-server filtered path
+                    child_results = await mounted.server._list_tools()
+                else:  # mode == "inventory"
+                    # PATH 1: Use the manager-to-manager unfiltered path
+                    child_results = await mounted.server._tool_manager.get_tools()
+
+                # The combination logic is the same for both paths
+                child_dict = (
+                    {t.key: t for t in child_results}
+                    if isinstance(child_results, list)
+                    else child_results
+                )
+                if mounted.prefix:
+                    for tool in child_dict.values():
+                        prefixed_tool = tool.with_key(f"{mounted.prefix}_{tool.key}")
+                        all_tools[prefixed_tool.key] = prefixed_tool
+                else:
+                    all_tools.update(child_dict)
+            except Exception as e:
+                # Skip failed mounts silently, matches existing behavior
+                logger.warning(
+                    f"Failed to get tools from mounted server '{mounted.prefix}': {e}"
+                )
+                continue
+
+        # Finally, add local tools, which always take precedence
+        all_tools.update(self._tools)
+        return all_tools
+
+    async def has_tool(self, key: str) -> bool:
         """Check if a tool exists."""
-        return key in self._tools
+        tools = await self.get_tools()
+        return key in tools
 
-    def get_tool(self, key: str) -> Tool:
+    async def get_tool(self, key: str) -> Tool:
         """Get tool by key."""
-        if key in self._tools:
-            return self._tools[key]
-        raise NotFoundError(f"Unknown tool: {key}")
+        tools = await self.get_tools()
+        if key in tools:
+            return tools[key]
+        raise NotFoundError(f"Tool {key!r} not found")
 
-    def get_tools(self) -> dict[str, Tool]:
-        """Get all registered tools, indexed by registered key."""
-        return self._tools
+    async def get_tools(self) -> dict[str, Tool]:
+        """
+        Gets the complete, unfiltered inventory of all tools.
+        """
+        return await self._load_tools(mode="inventory")
 
-    def list_tools(self) -> list[Tool]:
-        """List all registered tools."""
-        return list(self.get_tools().values())
+    async def list_tools(self) -> list[Tool]:
+        """
+        Lists all tools, applying protocol filtering.
+        """
+        tools_dict = await self._load_tools(mode="protocol")
+        return list(tools_dict.values())
 
     def add_tool_from_fn(
         self,
@@ -119,28 +175,44 @@ class ToolManager:
         if key in self._tools:
             del self._tools[key]
         else:
-            raise NotFoundError(f"Unknown tool: {key}")
+            raise NotFoundError(f"Tool {key!r} not found")
 
     async def call_tool(self, key: str, arguments: dict[str, Any]) -> list[MCPContent]:
-        """Call a tool by name with arguments."""
-        tool = self.get_tool(key)
-        if not tool:
-            raise NotFoundError(f"Unknown tool: {key}")
-
-        try:
-            return await tool.run(arguments)
-
-        # raise ToolErrors as-is
-        except ToolError as e:
-            logger.exception(f"Error calling tool {key!r}: {e}")
-            raise e
-
-        # Handle other exceptions
-        except Exception as e:
-            logger.exception(f"Error calling tool {key!r}: {e}")
-            if self.mask_error_details:
-                # Mask internal details
-                raise ToolError(f"Error calling tool {key!r}") from e
-            else:
-                # Include original error details
-                raise ToolError(f"Error calling tool {key!r}: {e}") from e
+        """
+        Internal API for servers: Finds and calls a tool, respecting the
+        filtered protocol path.
+        """
+        # 1. Check local tools first. The server will have already applied its filter.
+        if key in self._tools:
+            tool = await self.get_tool(key)
+            if not tool:
+                raise NotFoundError(f"Tool {key!r} not found")
+
+            try:
+                return await tool.run(arguments)
+
+            # raise ToolErrors as-is
+            except ToolError as e:
+                logger.exception(f"Error calling tool {key!r}: {e}")
+                raise e
+
+            # Handle other exceptions
+            except Exception as e:
+                logger.exception(f"Error calling tool {key!r}: {e}")
+                if self.mask_error_details:
+                    # Mask internal details
+                    raise ToolError(f"Error calling tool {key!r}") from e
+                else:
+                    # Include original error details
+                    raise ToolError(f"Error calling tool {key!r}: {e}") from e
+
+        # 2. Check mounted servers using the filtered protocol path.
+        for mounted in reversed(self._mounted_sources):
+            if mounted.prefix and key.startswith(f"{mounted.prefix}_"):
+                key_on_child = key.removeprefix(f"{mounted.prefix}_")
+                try:
+                    return await mounted.server._call_tool(key_on_child, arguments)
+                except NotFoundError:
+                    continue
+
+        raise NotFoundError(f"Tool {key!r} not found.")

tests/tools/test_tool_manager.py

@@ -17,7 +17,7 @@ from fastmcp.utilities.types import Image
 
 
 class TestAddTools:
-    def test_basic_function(self):
+    async def test_basic_function(self):
         """Test registering and running a basic function."""
 
         def add(a: int, b: int) -> int:
@@ -28,7 +28,7 @@ class TestAddTools:
         tool = Tool.from_function(add)
         manager.add_tool(tool)
 
-        tool = manager.get_tool("add")
+        tool = await manager.get_tool("add")
         assert tool is not None
         assert tool.name == "add"
         assert tool.description == "Add two numbers."
@@ -46,13 +46,13 @@ class TestAddTools:
         tool = Tool.from_function(fetch_data)
         manager.add_tool(tool)
 
-        tool = manager.get_tool("fetch_data")
+        tool = await manager.get_tool("fetch_data")
         assert tool is not None
         assert tool.name == "fetch_data"
         assert tool.description == "Fetch data from URL."
         assert tool.parameters["properties"]["url"]["type"] == "string"
 
-    def test_pydantic_model_function(self):
+    async def test_pydantic_model_function(self):
         """Test registering a function that takes a Pydantic model."""
 
         class UserInput(BaseModel):
@@ -67,7 +67,7 @@ class TestAddTools:
         tool = Tool.from_function(create_user)
         manager.add_tool(tool)
 
-        tool = manager.get_tool("create_user")
+        tool = await manager.get_tool("create_user")
         assert tool is not None
         assert tool.name == "create_user"
         assert tool.description == "Create a new user."
@@ -75,7 +75,7 @@ class TestAddTools:
         assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
         assert "flag" in tool.parameters["properties"]
 
-    def test_callable_object(self):
+    async def test_callable_object(self):
         class Adder:
             """Adds two numbers."""
 
@@ -87,7 +87,7 @@ class TestAddTools:
         tool = Tool.from_function(Adder())
         manager.add_tool(tool)
 
-        tool = manager.get_tool("Adder")
+        tool = await manager.get_tool("Adder")
         assert tool is not None
         assert tool.name == "Adder"
         assert tool.description == "Adds two numbers."
@@ -95,7 +95,7 @@ class TestAddTools:
         assert tool.parameters["properties"]["x"]["type"] == "integer"
         assert tool.parameters["properties"]["y"]["type"] == "integer"
 
-    def test_async_callable_object(self):
+    async def test_async_callable_object(self):
         class Adder:
             """Adds two numbers."""
 
@@ -107,7 +107,7 @@ class TestAddTools:
         tool = Tool.from_function(Adder())
         manager.add_tool(tool)
 
-        tool = manager.get_tool("Adder")
+        tool = await manager.get_tool("Adder")
         assert tool is not None
         assert tool.name == "Adder"
         assert tool.description == "Adds two numbers."
@@ -123,7 +123,7 @@ class TestAddTools:
         tool = Tool.from_function(image_tool)
         manager.add_tool(tool)
 
-        tool = manager.get_tool("image_tool")
+        tool = await manager.get_tool("image_tool")
         result = await tool.run({"data": "test.png"})
         assert tool.parameters["properties"]["data"]["type"] == "string"
         assert isinstance(result[0], ImageContent)
@@ -148,7 +148,7 @@ class TestAddTools:
             tool = Tool.from_function(lambda x: x)
             manager.add_tool(tool)
 
-    def test_remove_tool_successfully(self):
+    async def test_remove_tool_successfully(self):
         """Test removing an added tool by key."""
         manager = ToolManager()
 
@@ -157,19 +157,19 @@ class TestAddTools:
 
         tool = Tool.from_function(add)
         manager.add_tool(tool)
-        assert manager.get_tool("add") is not None
+        assert await manager.get_tool("add") is not None
 
         manager.remove_tool("add")
         with pytest.raises(NotFoundError):
-            manager.get_tool("add")
+            await manager.get_tool("add")
 
     def test_remove_tool_missing_key(self):
         """Test removing a tool that does not exist raises NotFoundError."""
         manager = ToolManager()
-        with pytest.raises(NotFoundError, match=f"Unknown tool: {'missing'}"):
+        with pytest.raises(NotFoundError, match="Tool 'missing' not found"):
             manager.remove_tool("missing")
 
-    def test_warn_on_duplicate_tools(self, caplog):
+    async def test_warn_on_duplicate_tools(self, caplog):
         """Test warning on duplicate tools."""
         manager = ToolManager(duplicate_behavior="warn")
 
@@ -183,7 +183,7 @@ class TestAddTools:
 
         assert "Tool already exists: test_tool" in caplog.text
         # Should have the tool
-        assert manager.get_tool("test_tool") is not None
+        assert await manager.get_tool("test_tool") is not None
 
     def test_disable_warn_on_duplicate_tools(self, caplog):
         """Test disabling warning on duplicate tools."""
@@ -213,7 +213,7 @@ class TestAddTools:
             tool2 = Tool.from_function(test_fn, name="test_tool")
             manager.add_tool(tool2)
 
-    def test_replace_duplicate_tools(self):
+    async def test_replace_duplicate_tools(self):
         """Test replacing duplicate tools."""
         manager = ToolManager(duplicate_behavior="replace")
 
@@ -229,12 +229,12 @@ class TestAddTools:
         manager.add_tool(result)
 
         # Should have replaced with the new tool
-        tool = manager.get_tool("test_tool")
+        tool = await manager.get_tool("test_tool")
         assert tool is not None
         assert isinstance(tool, FunctionTool)
         assert tool.fn.__name__ == "replacement_fn"
 
-    def test_ignore_duplicate_tools(self):
+    async def test_ignore_duplicate_tools(self):
         """Test ignoring duplicate tools."""
         manager = ToolManager(duplicate_behavior="ignore")
 
@@ -250,7 +250,7 @@ class TestAddTools:
         manager.add_tool(result)
 
         # Should keep the original
-        tool = manager.get_tool("test_tool")
+        tool = await manager.get_tool("test_tool")
         assert tool is not None
         assert isinstance(tool, FunctionTool)
         assert tool.fn.__name__ == "original_fn"
@@ -262,7 +262,7 @@ class TestAddTools:
 class TestToolTags:
     """Test functionality related to tool tags."""
 
-    def test_add_tool_with_tags(self):
+    async def test_add_tool_with_tags(self):
         """Test adding tags to a tool."""
 
         def example_tool(x: int) -> int:
@@ -274,11 +274,11 @@ class TestToolTags:
         manager.add_tool(tool)
 
         assert tool.tags == {"math", "utility"}
-        tool = manager.get_tool("example_tool")
+        tool = await manager.get_tool("example_tool")
         assert tool is not None
         assert tool.tags == {"math", "utility"}
 
-    def test_add_tool_with_empty_tags(self):
+    async def test_add_tool_with_empty_tags(self):
         """Test adding a tool with empty tags set."""
 
         def example_tool(x: int) -> int:
@@ -291,7 +291,7 @@ class TestToolTags:
 
         assert tool.tags == set()
 
-    def test_add_tool_with_none_tags(self):
+    async def test_add_tool_with_none_tags(self):
         """Test adding a tool with None tags."""
 
         def example_tool(x: int) -> int:
@@ -304,7 +304,7 @@ class TestToolTags:
 
         assert tool.tags == set()
 
-    def test_list_tools_with_tags(self):
+    async def test_list_tools_with_tags(self):
         """Test listing tools with specific tags."""
 
         def math_tool(x: int) -> int:
@@ -328,12 +328,16 @@ class TestToolTags:
         manager.add_tool(tool3)
 
         # Check if we can filter by tags when listing tools
-        math_tools = [tool for tool in manager.list_tools() if "math" in tool.tags]
+        math_tools = [
+            tool for tool in (await manager.get_tools()).values() if "math" in tool.tags
+        ]
         assert len(math_tools) == 2
         assert {tool.name for tool in math_tools} == {"math_tool", "mixed_tool"}
 
         utility_tools = [
-            tool for tool in manager.list_tools() if "utility" in tool.tags
+            tool
+            for tool in (await manager.get_tools()).values()
+            if "utility" in tool.tags
         ]
         assert len(utility_tools) == 2
         assert {tool.name for tool in utility_tools} == {"string_tool", "mixed_tool"}
@@ -416,7 +420,7 @@ class TestCallTools:
 
     async def test_call_unknown_tool(self):
         manager = ToolManager()
-        with pytest.raises(NotFoundError, match="Unknown tool: unknown"):
+        with pytest.raises(NotFoundError, match="Tool 'unknown' not found"):
             await manager.call_tool("unknown", {"a": 1})
 
     async def test_call_tool_with_list_int_input(self):
@@ -728,7 +732,7 @@ class TestContextHandling:
 class TestCustomToolNames:
     """Test adding tools with custom names that differ from their function names."""
 
-    def test_add_tool_with_custom_name(self):
+    async def test_add_tool_with_custom_name(self):
         """Test adding a tool with a custom name parameter using add_tool_from_fn."""
 
         def original_fn(x: int) -> int:
@@ -739,15 +743,15 @@ class TestCustomToolNames:
         manager.add_tool(tool)
 
         # The tool is stored under the custom name and its .name is also set to custom_name
-        assert manager.get_tool("custom_name") is not None
+        assert await manager.get_tool("custom_name") is not None
         assert tool.name == "custom_name"
         assert isinstance(tool, FunctionTool)
         assert tool.fn.__name__ == "original_fn"
         # The tool should not be accessible via its original function name
-        with pytest.raises(NotFoundError, match="Unknown tool: original_fn"):
-            manager.get_tool("original_fn")
+        with pytest.raises(NotFoundError, match="Tool 'original_fn' not found"):
+            await manager.get_tool("original_fn")
 
-    def test_add_tool_object_with_custom_key(self):
+    async def test_add_tool_object_with_custom_key(self):
         """Test adding a Tool object with a custom key using add_tool()."""
 
         def fn(x: int) -> int:
@@ -759,13 +763,13 @@ class TestCustomToolNames:
         # Store it under a different name
         manager.add_tool(tool, key="proxy_tool")
         # The tool is accessible under the key
-        stored = manager.get_tool("proxy_tool")
+        stored = await manager.get_tool("proxy_tool")
         assert stored is not None
         # But the tool's .name is unchanged
         assert stored.name == "my_tool"
         # The tool is not accessible under its original name
-        with pytest.raises(NotFoundError, match="Unknown tool: my_tool"):
-            manager.get_tool("my_tool")
+        with pytest.raises(NotFoundError, match="Tool 'my_tool' not found"):
+            await manager.get_tool("my_tool")
 
     async def test_call_tool_with_custom_name(self):
         """Test calling a tool added with a custom name."""
@@ -783,10 +787,10 @@ class TestCustomToolNames:
         assert result[0].text == "15"  # type: ignore[attr-defined]
 
         # Original name should not be registered
-        with pytest.raises(NotFoundError, match="Unknown tool: multiply"):
+        with pytest.raises(NotFoundError, match="Tool 'multiply' not found"):
             await manager.call_tool("multiply", {"a": 5, "b": 3})
 
-    def test_replace_tool_keeps_original_name(self):
+    async def test_replace_tool_keeps_original_name(self):
         """Test that replacing a tool with "replace" keeps the original name."""
 
         def original_fn(x: int) -> int:
@@ -808,7 +812,7 @@ class TestCustomToolNames:
         manager.add_tool(replacement_tool)
 
         # The tool object should have been replaced
-        stored_tool = manager.get_tool("test_tool")
+        stored_tool = await manager.get_tool("test_tool")
         assert stored_tool is not None
         assert stored_tool == replacement_tool