Fix task execution for tools with custom names

pyproject.toml

@@ -10,7 +10,7 @@ dependencies = [
     "mcp>=1.24.0",
     "openapi-pydantic>=0.5.1",
     "platformdirs>=4.0.0",
-    "pydocket>=0.15.5",
+    "pydocket>=0.16.0",
     "rich>=13.9.4",
     "cyclopts>=4.0.0",
     "authlib>=1.6.5",

src/fastmcp/server/server.py

@@ -112,23 +112,6 @@
 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"]
 
@@ -441,50 +424,48 @@ async def _docket_lifespan(self) -> AsyncIterator[None]:
                         isinstance(tool, FunctionTool)
                         and tool.task_config.mode != "forbidden"
                     ):
-                        docket.register(tool.fn)
+                        docket.register(tool.fn, names=[tool.key])
 
                 for prompt in self._prompt_manager._prompts.values():
                     if (
                         isinstance(prompt, FunctionPrompt)
                         and prompt.task_config.mode != "forbidden"
                     ):
                         # task execution requires async fn (validated at creation time)
-                        docket.register(cast(Callable[..., Awaitable[Any]], prompt.fn))
+                        docket.register(
+                            cast(Callable[..., Awaitable[Any]], prompt.fn),
+                            names=[prompt.key],
+                        )
 
                 for resource in self._resource_manager._resources.values():
                     if (
                         isinstance(resource, FunctionResource)
                         and resource.task_config.mode != "forbidden"
                     ):
-                        docket.register(resource.fn)
+                        docket.register(resource.fn, names=[resource.name])
 
                 for template in self._resource_manager._templates.values():
                     if (
                         isinstance(template, FunctionResourceTemplate)
                         and template.task_config.mode != "forbidden"
                     ):
-                        docket.register(template.fn)
+                        docket.register(template.fn, names=[template.name])
 
                 # Register provider components
                 for provider in self._providers:
                     try:
                         tasks = await provider.get_tasks()
                         for tool in tasks.tools:
-                            named_fn = _create_named_fn_wrapper(tool.fn, tool.key)
-                            docket.register(named_fn)
+                            docket.register(tool.fn, names=[tool.key])
                         for resource in tasks.resources:
-                            named_fn = _create_named_fn_wrapper(
-                                resource.fn, resource.name
-                            )
-                            docket.register(named_fn)
+                            docket.register(resource.fn, names=[resource.name])
                         for template in tasks.templates:
-                            named_fn = _create_named_fn_wrapper(
-                                template.fn, template.name
-                            )
-                            docket.register(named_fn)
+                            docket.register(template.fn, names=[template.name])
                         for prompt in tasks.prompts:
-                            named_fn = _create_named_fn_wrapper(prompt.fn, prompt.key)
-                            docket.register(named_fn)
+                            docket.register(
+                                cast(Callable[..., Awaitable[Any]], prompt.fn),
+                                names=[prompt.key],
+                            )
                     except Exception as e:
                         provider_name = getattr(
                             provider, "server", provider

tests/server/tasks/test_server_tasks_parameter.py

@@ -308,3 +308,79 @@ async def shared_name_prompt() -> str:
         # Prompt inheriting False (mode="forbidden") raises McpError
         with pytest.raises(McpError):
             await client.get_prompt("shared_name_prompt", task=True)
+
+
+async def test_task_with_custom_tool_name():
+    """Tools with custom names work correctly as tasks (issue #2642).
+
+    When a tool is registered with a custom name different from the function
+    name, task execution should use the custom name for Docket lookup.
+    """
+    mcp = FastMCP("test", tasks=True)
+
+    async def my_function() -> str:
+        return "result from custom-named tool"
+
+    mcp.tool(my_function, name="custom-tool-name")
+
+    async with Client(mcp) as client:
+        # Verify the tool is registered with its custom name in Docket
+        docket = mcp.docket
+        assert docket is not None
+        assert "custom-tool-name" in docket.tasks
+
+        # Call the tool as a task using its custom name
+        task = await client.call_tool("custom-tool-name", task=True)
+        assert not task.returned_immediately
+        result = await task
+        assert result.data == "result from custom-named tool"
+
+
+async def test_task_with_custom_resource_name():
+    """Resources with custom names work correctly as tasks.
+
+    When a resource is registered with a custom name different from the function
+    name, task execution should use the custom name for Docket lookup.
+    """
+    mcp = FastMCP("test", tasks=True)
+
+    @mcp.resource("test://resource", name="custom-resource-name")
+    async def my_resource_func() -> str:
+        return "result from custom-named resource"
+
+    async with Client(mcp) as client:
+        # Verify the resource is registered with its custom name in Docket
+        docket = mcp.docket
+        assert docket is not None
+        assert "custom-resource-name" in docket.tasks
+
+        # Call the resource as a task
+        task = await client.read_resource("test://resource", task=True)
+        assert not task.returned_immediately
+        result = await task.result()
+        assert result[0].text == "result from custom-named resource"
+
+
+async def test_task_with_custom_template_name():
+    """Resource templates with custom names work correctly as tasks.
+
+    When a template is registered with a custom name different from the function
+    name, task execution should use the custom name for Docket lookup.
+    """
+    mcp = FastMCP("test", tasks=True)
+
+    @mcp.resource("test://{item_id}", name="custom-template-name")
+    async def my_template_func(item_id: str) -> str:
+        return f"result for {item_id}"
+
+    async with Client(mcp) as client:
+        # Verify the template is registered with its custom name in Docket
+        docket = mcp.docket
+        assert docket is not None
+        assert "custom-template-name" in docket.tasks
+
+        # Call the template as a task
+        task = await client.read_resource("test://123", task=True)
+        assert not task.returned_immediately
+        result = await task.result()
+        assert result[0].text == "result for 123"