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

Prefix Docket function names to avoid collisions in multi-mount setups
This commit is contained in:
Chris Guidry 2025-12-08 16:10:28 -05:00 committed by GitHub
commit 659ec38d01
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 138 additions and 19 deletions

View file

@ -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]:

View file

@ -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,
)()

View file

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