mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-20 12:34:17 +02:00
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:
commit
659ec38d01
3 changed files with 138 additions and 19 deletions
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)()
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue