Merge pull request #2575 from jlowin/multi-mount-docket-isolation

Prefix Docket function names to avoid collisions in multi-mount setups

Author: chrisguidry

src/fastmcp/server/server.py

@@ -102,6 +102,24 @@
 
 logger = get_logger(__name__)
 
+
+def _create_named_fn_wrapper(fn: Callable[..., Any], name: str) -> Callable[..., Any]:
+    """Create a wrapper function with a custom __name__ for Docket registration.
+
+    Docket uses fn.__name__ as the key for function registration and lookup.
+    When mounting servers, we need unique names to avoid collisions between
+    mounted servers that have identically-named functions.
+    """
+    import functools
+
+    @functools.wraps(fn)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        return await fn(*args, **kwargs)
+
+    wrapper.__name__ = name
+    return wrapper
+
+
 DuplicateBehavior = Literal["warn", "error", "replace", "ignore"]
 Transport = Literal["stdio", "http", "sse", "streamable-http"]
 
@@ -449,7 +467,7 @@ async def _docket_lifespan(self) -> AsyncIterator[None]:
                 # execute in the parent's Docket context
                 for mounted in self._mounted_servers:
                     await self._register_mounted_server_functions(
-                        mounted.server, docket
+                        mounted.server, docket, mounted.prefix
                     )
 
                 # Set Docket in ContextVar so CurrentDocket can access it
@@ -491,45 +509,69 @@ async def _docket_lifespan(self) -> AsyncIterator[None]:
             _current_server.reset(server_token)
 
     async def _register_mounted_server_functions(
-        self, server: FastMCP, docket: Docket
+        self, server: FastMCP, docket: Docket, prefix: str | None
     ) -> None:
         """Register task-enabled functions from a mounted server with Docket.
 
         This enables background task execution for mounted server components
         through the parent server's Docket context.
+
+        Args:
+            server: The mounted server whose functions to register
+            docket: The Docket instance to register with
+            prefix: The mount prefix to prepend to function names (matches
+                    client-facing tool/prompt names)
         """
-        # Register tools
+        # Register tools with prefixed names to avoid collisions
         for tool in server._tool_manager._tools.values():
             if isinstance(tool, FunctionTool) and tool.task_config.mode != "forbidden":
-                docket.register(tool.fn)
+                # Use same naming as client-facing tool keys
+                fn_name = f"{prefix}_{tool.key}" if prefix else tool.key
+                named_fn = _create_named_fn_wrapper(tool.fn, fn_name)
+                docket.register(named_fn)
 
-        # Register prompts
+        # Register prompts with prefixed names
         for prompt in server._prompt_manager._prompts.values():
             if (
                 isinstance(prompt, FunctionPrompt)
                 and prompt.task_config.mode != "forbidden"
             ):
-                docket.register(cast(Callable[..., Awaitable[Any]], prompt.fn))
+                fn_name = f"{prefix}_{prompt.key}" if prefix else prompt.key
+                named_fn = _create_named_fn_wrapper(
+                    cast(Callable[..., Awaitable[Any]], prompt.fn), fn_name
+                )
+                docket.register(named_fn)
 
-        # Register resources
+        # Register resources with prefixed names (use name, not key/URI)
         for resource in server._resource_manager._resources.values():
             if (
                 isinstance(resource, FunctionResource)
                 and resource.task_config.mode != "forbidden"
             ):
-                docket.register(resource.fn)
+                fn_name = f"{prefix}_{resource.name}" if prefix else resource.name
+                named_fn = _create_named_fn_wrapper(resource.fn, fn_name)
+                docket.register(named_fn)
 
-        # Register resource templates
+        # Register resource templates with prefixed names (use name, not key/URI)
         for template in server._resource_manager._templates.values():
             if (
                 isinstance(template, FunctionResourceTemplate)
                 and template.task_config.mode != "forbidden"
             ):
-                docket.register(template.fn)
+                fn_name = f"{prefix}_{template.name}" if prefix else template.name
+                named_fn = _create_named_fn_wrapper(template.fn, fn_name)
+                docket.register(named_fn)
 
-        # Recursively register from nested mounted servers
+        # Recursively register from nested mounted servers with accumulated prefix
         for nested in server._mounted_servers:
-            await self._register_mounted_server_functions(nested.server, docket)
+            nested_prefix = (
+                f"{prefix}_{nested.prefix}"
+                if prefix and nested.prefix
+                else (prefix or nested.prefix)
+            )
+            await self._register_mounted_server_functions(
+                nested.server, docket, nested_prefix
+            )
 
     @asynccontextmanager
     async def _lifespan_manager(self) -> AsyncIterator[None]:

src/fastmcp/server/tasks/handlers.py

@@ -99,9 +99,10 @@ async def handle_tool_as_task(
         # Don't let notification failures break task creation
         await ctx.session.send_notification(notification)  # type: ignore[arg-type]
 
-    # Queue function to Docket (result storage via execution_ttl)
+    # Queue function to Docket by name (result storage via execution_ttl)
+    # Use tool.key which matches what was registered - prefixed for mounted tools
     await docket.add(
-        tool.fn,  # type: ignore[attr-defined]
+        tool.key,
         key=task_key,
     )(**arguments)
 
@@ -204,9 +205,10 @@ async def handle_prompt_as_task(
     with suppress(Exception):
         await ctx.session.send_notification(notification)  # type: ignore[arg-type]
 
-    # Queue function to Docket (result storage via execution_ttl)
+    # Queue function to Docket by name (result storage via execution_ttl)
+    # Use prompt.key which matches what was registered - prefixed for mounted prompts
     await docket.add(
-        prompt.fn,  # type: ignore[attr-defined]
+        prompt.key,
         key=task_key,
     )(**(arguments or {}))
 
@@ -307,19 +309,20 @@ async def handle_resource_as_task(
     with suppress(Exception):
         await ctx.session.send_notification(notification)  # type: ignore[arg-type]
 
-    # Queue function to Docket (result storage via execution_ttl)
+    # Queue function to Docket by name (result storage via execution_ttl)
+    # Use resource.name which matches what was registered - prefixed for mounted resources
     # For templates, extract URI params and pass them to the function
     from fastmcp.resources.template import FunctionResourceTemplate, match_uri_template
 
     if isinstance(resource, FunctionResourceTemplate):
         params = match_uri_template(uri, resource.uri_template) or {}
         await docket.add(
-            resource.fn,  # type: ignore[attr-defined]
+            resource.name,
             key=task_key,
         )(**params)
     else:
         await docket.add(
-            resource.fn,  # type: ignore[attr-defined]
+            resource.name,
             key=task_key,
         )()
 

tests/server/tasks/test_task_mount.py

@@ -368,6 +368,80 @@ async def subtract(a: int, b: int) -> int:
             assert result2.data == 5
 
 
+class TestMountedFunctionNameCollisions:
+    """Test task execution when mounted servers have identically-named functions."""
+
+    async def test_multiple_mounts_with_same_function_names(self):
+        """Two mounted servers with identically-named functions don't collide."""
+        child1 = FastMCP("child1")
+        child2 = FastMCP("child2")
+
+        @child1.tool(task=True)
+        async def process(value: int) -> int:
+            return value * 2  # Double
+
+        @child2.tool(task=True)
+        async def process(value: int) -> int:  # noqa: F811
+            return value * 3  # Triple
+
+        parent = FastMCP("parent")
+        parent.mount(child1, prefix="c1")
+        parent.mount(child2, prefix="c2")
+
+        async with Client(parent) as client:
+            # Both should execute their own implementation
+            task1 = await client.call_tool("c1_process", {"value": 10}, task=True)
+            task2 = await client.call_tool("c2_process", {"value": 10}, task=True)
+
+            result1 = await task1.result()
+            result2 = await task2.result()
+
+            assert result1.data == 20  # child1's process (doubles)
+            assert result2.data == 30  # child2's process (triples)
+
+    async def test_no_prefix_mount_collision(self):
+        """No-prefix mounts with same tool name - last mount wins."""
+        child1 = FastMCP("child1")
+        child2 = FastMCP("child2")
+
+        @child1.tool(task=True)
+        async def process(value: int) -> int:
+            return value * 2
+
+        @child2.tool(task=True)
+        async def process(value: int) -> int:  # noqa: F811
+            return value * 3
+
+        parent = FastMCP("parent")
+        parent.mount(child1)  # No prefix
+        parent.mount(child2)  # No prefix - overwrites child1's "process"
+
+        async with Client(parent) as client:
+            # Last mount wins - child2's process should execute
+            task = await client.call_tool("process", {"value": 10}, task=True)
+            result = await task.result()
+            assert result.data == 30  # child2's process (triples)
+
+    async def test_nested_mount_prefix_accumulation(self):
+        """Nested mounts accumulate prefixes correctly for tasks."""
+        grandchild = FastMCP("gc")
+        child = FastMCP("child")
+        parent = FastMCP("parent")
+
+        @grandchild.tool(task=True)
+        async def deep_tool() -> str:
+            return "deep"
+
+        child.mount(grandchild, prefix="gc")
+        parent.mount(child, prefix="child")
+
+        async with Client(parent) as client:
+            # Tool should be accessible and execute correctly
+            task = await client.call_tool("child_gc_deep_tool", {}, task=True)
+            result = await task.result()
+            assert result.data == "deep"
+
+
 class TestMountedTaskList:
     """Test task listing with mounted servers."""