diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 63af58bf4..2d001738d 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -102,6 +102,24 @@ if TYPE_CHECKING: 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 @@ class FastMCP(Generic[LifespanResultT]): # 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 @@ class FastMCP(Generic[LifespanResultT]): _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]: diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index d3dcd9e3e..569330449 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/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, )() diff --git a/tests/server/tasks/test_task_mount.py b/tests/server/tasks/test_task_mount.py index 39c125333..ffff3cfc0 100644 --- a/tests/server/tasks/test_task_mount.py +++ b/tests/server/tasks/test_task_mount.py @@ -368,6 +368,80 @@ class TestMultipleMounts: 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."""