Fix task execution for tools with custom names (#2645)

Uses pydocket 0.16.0's `names=` parameter for explicit registration keys:
- Tools/prompts: registered by `.key`
- Resources/templates: registered by `.name`

Removes `_create_named_fn_wrapper` function.

Closes #2642
This commit is contained in:
Jeremiah Lowin 2025-12-23 18:05:29 -05:00
commit bf60ac4e44
4 changed files with 44 additions and 35 deletions

View file

@ -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",

View file

@ -105,23 +105,6 @@ 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"]
@ -437,7 +420,7 @@ class FastMCP(Generic[LifespanResultT]):
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 (
@ -445,21 +428,24 @@ class FastMCP(Generic[LifespanResultT]):
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])
# Also register functions from mounted servers so tasks can
# execute in the parent's Docket context
@ -535,8 +521,7 @@ class FastMCP(Generic[LifespanResultT]):
fn_name = f"{prefix}_{tool.key}"
else:
fn_name = tool.key
named_fn = _create_named_fn_wrapper(tool.fn, fn_name)
docket.register(named_fn)
docket.register(tool.fn, names=[fn_name])
# Register prompts with prefixed names
for prompt in server._prompt_manager._prompts.values():
@ -545,10 +530,10 @@ class FastMCP(Generic[LifespanResultT]):
and prompt.task_config.mode != "forbidden"
):
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(
cast(Callable[..., Awaitable[Any]], prompt.fn),
names=[fn_name],
)
docket.register(named_fn)
# Register resources with prefixed names (use name, not key/URI)
for resource in server._resource_manager._resources.values():
@ -557,8 +542,7 @@ class FastMCP(Generic[LifespanResultT]):
and resource.task_config.mode != "forbidden"
):
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)
docket.register(resource.fn, names=[fn_name])
# Register resource templates with prefixed names (use name, not key/URI)
for template in server._resource_manager._templates.values():
@ -567,8 +551,7 @@ class FastMCP(Generic[LifespanResultT]):
and template.task_config.mode != "forbidden"
):
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)
docket.register(template.fn, names=[fn_name])
# Recursively register from nested mounted servers with accumulated prefix
for nested in server._mounted_servers:

View file

@ -308,3 +308,29 @@ async def test_multiple_components_same_name_different_tasks():
# 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"

8
uv.lock generated
View file

@ -752,7 +752,7 @@ requires-dist = [
{ name = "platformdirs", specifier = ">=4.0.0" },
{ name = "py-key-value-aio", extras = ["disk", "keyring", "memory"], specifier = ">=0.3.0,<0.4.0" },
{ name = "pydantic", extras = ["email"], specifier = ">=2.11.7" },
{ name = "pydocket", specifier = ">=0.15.5" },
{ name = "pydocket", specifier = ">=0.16.0" },
{ name = "pyperclip", specifier = ">=1.9.0" },
{ name = "python-dotenv", specifier = ">=1.1.0" },
{ name = "rich", specifier = ">=13.9.4" },
@ -1775,7 +1775,7 @@ wheels = [
[[package]]
name = "pydocket"
version = "0.15.5"
version = "0.16.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "cloudpickle" },
@ -1792,9 +1792,9 @@ dependencies = [
{ name = "typer" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/86/77/842e41be3cf3592b971cf42b24cae76e282294f474dc2dbf7cd6808d1b09/pydocket-0.15.5.tar.gz", hash = "sha256:b3af47702a293dd1da2e5e0f8f73f27fd3b3c95e36de72a2f71026d16908d5ba", size = 277245, upload-time = "2025-12-12T22:28:47.32Z" }
sdist = { url = "https://files.pythonhosted.org/packages/f4/b6/b9b5af725be7ffb195750e1cbe0ad4aa413858f16af258a49ad84cff82ac/pydocket-0.16.2.tar.gz", hash = "sha256:b48e59ad247cbf23fd7064cd765b75e5b870451a7778afa3b208b3aae6474c1f", size = 295623, upload-time = "2025-12-23T19:24:46.751Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e5/c0/fdbc6e04e3369b90c6bf6567bc62871cf59e88550b94529821500dc807c1/pydocket-0.15.5-py3-none-any.whl", hash = "sha256:ad0d86c9a1bea394e875bcf8c793be2d0a7ebd1891bfe99e2e9eaf99ef0cb42e", size = 58517, upload-time = "2025-12-12T22:28:45.598Z" },
{ url = "https://files.pythonhosted.org/packages/ef/33/81132e8ca41e11000807cc1f457a391d5a364dd1eeb7de40f2ff971301bc/pydocket-0.16.2-py3-none-any.whl", hash = "sha256:044a74deec2062c305bf1d452a1348a4e3739b3a2cf0a411ca1e11368cef5baf", size = 66812, upload-time = "2025-12-23T19:24:45.643Z" },
]
[[package]]