mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 07:09:11 +02:00
fix: resolve CurrentFastMCP/ctx.fastmcp to child server in mounted background tasks (#3651)
This commit is contained in:
parent
426572670e
commit
b9ea53618d
5 changed files with 139 additions and 11 deletions
|
|
@ -101,6 +101,14 @@ The handwritten notes are prepended above the auto-generated changelog and are t
|
|||
- Minor fixes: keep body short and concise
|
||||
- No "test plan" sections or testing summaries
|
||||
|
||||
### Code Review Guidelines
|
||||
|
||||
- **Fix causes, not symptoms.** When a PR works around a problem instead of addressing why it occurs, that's a red flag. A side-channel that compensates for a missing step adds permanent complexity. If the fix doesn't change the code path where the bug actually happens, ask why not.
|
||||
- Focus on API design and naming clarity
|
||||
- Identify confusing patterns (e.g., parameter values that contradict defaults) or non-idiomatic code (mutable defaults, etc.). Contributed code will need to be maintained indefinitely, and by someone other than the author (unless the author is a maintainer).
|
||||
- Suggest specific improvements, not generic "add more tests" comments
|
||||
- Think about API ergonomics from a user perspective
|
||||
|
||||
### Code Standards
|
||||
|
||||
- Python ≥ 3.10 with full type annotations
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ If you do open a PR:
|
|||
- **Keep it focused.** One logical change per PR. Don't bundle unrelated fixes or refactors.
|
||||
- **Match existing patterns.** Follow the code style, type annotation conventions, and test patterns you see in the codebase. Run `uv run prek run --all-files` before submitting.
|
||||
- **Write tests.** Bug fixes should include a test that fails without the fix. Enhancements should include tests for the new behavior.
|
||||
- **Fix the cause, not the symptom.** If the bug is that a code path skips a step, the fix should make it stop skipping that step — not add compensation elsewhere. Workaround-style fixes will be sent back for revision.
|
||||
- **Don't submit generated boilerplate.** We review every line. PRs that read like unedited LLM output — verbose descriptions, speculative changes, shotgun-style fixes — will be closed.
|
||||
|
||||
## What we'll close without review
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import contextlib
|
|||
import inspect
|
||||
import logging
|
||||
import weakref
|
||||
from collections import OrderedDict
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from contextvars import ContextVar, Token
|
||||
|
|
@ -72,6 +73,7 @@ __all__ = [
|
|||
"get_task_context",
|
||||
"get_task_session",
|
||||
"is_docket_available",
|
||||
"register_task_server",
|
||||
"register_task_session",
|
||||
"require_docket",
|
||||
"resolve_dependencies",
|
||||
|
|
@ -174,6 +176,32 @@ def get_task_session(session_id: str) -> ServerSession | None:
|
|||
_current_server: ContextVar[weakref.ref[FastMCP] | None] = ContextVar(
|
||||
"server", default=None
|
||||
)
|
||||
|
||||
# --- Background task server map ---
|
||||
# Maps task_id → server weakref so background workers can resolve the correct
|
||||
# server for mounted-child tasks. Follows the same pattern as _task_sessions.
|
||||
# Populated in submit_to_docket() where the child server is in context;
|
||||
# consulted in get_server() when running inside a Docket worker.
|
||||
|
||||
_task_server_map: OrderedDict[str, weakref.ref[FastMCP]] = OrderedDict()
|
||||
_TASK_SERVER_MAP_MAX_SIZE = 10_000
|
||||
|
||||
|
||||
def register_task_server(task_id: str, server: FastMCP) -> None:
|
||||
"""Register the server for a background task.
|
||||
|
||||
Called at task-submission time (inside the child server's call_tool
|
||||
context) so that background workers can resolve CurrentFastMCP() and
|
||||
ctx.fastmcp to the child server for mounted tasks.
|
||||
|
||||
The map is bounded to avoid unbounded growth in long-lived servers.
|
||||
Evicted entries fall back to the ContextVar (parent server).
|
||||
"""
|
||||
_task_server_map[task_id] = weakref.ref(server)
|
||||
while len(_task_server_map) > _TASK_SERVER_MAP_MAX_SIZE:
|
||||
_task_server_map.popitem(last=False)
|
||||
|
||||
|
||||
_current_docket: ContextVar[Docket | None] = ContextVar("docket", default=None)
|
||||
_current_worker: ContextVar[Worker | None] = ContextVar("worker", default=None)
|
||||
_task_access_token: ContextVar[AccessToken | None] = ContextVar(
|
||||
|
|
@ -378,12 +406,28 @@ def get_context() -> Context:
|
|||
def get_server() -> FastMCP:
|
||||
"""Get the current FastMCP server instance directly.
|
||||
|
||||
In a background-task worker, checks the task-server map first so that
|
||||
mounted-child tasks resolve to the child server (not the parent that
|
||||
started the worker).
|
||||
|
||||
Returns:
|
||||
The active FastMCP server
|
||||
|
||||
Raises:
|
||||
RuntimeError: If no server in context
|
||||
"""
|
||||
# In a task context, prefer the task-specific server mapping.
|
||||
# This handles mounted-child tasks where _current_server is the parent.
|
||||
task_info = get_task_context()
|
||||
if task_info is not None:
|
||||
ref = _task_server_map.get(task_info.task_id)
|
||||
if ref is not None:
|
||||
server = ref()
|
||||
if server is not None:
|
||||
return server
|
||||
# Server was garbage collected, clean up
|
||||
_task_server_map.pop(task_info.task_id, None)
|
||||
|
||||
server_ref = _current_server.get()
|
||||
if server_ref is None:
|
||||
raise RuntimeError("No FastMCP server instance in context")
|
||||
|
|
@ -1037,13 +1081,7 @@ class _CurrentFastMCP(Dependency["FastMCP"]):
|
|||
"""Async context manager for FastMCP server dependency."""
|
||||
|
||||
async def __aenter__(self) -> FastMCP:
|
||||
server_ref = _current_server.get()
|
||||
if server_ref is None:
|
||||
raise RuntimeError("No FastMCP server instance in context")
|
||||
server = server_ref()
|
||||
if server is None:
|
||||
raise RuntimeError("FastMCP server instance is no longer available")
|
||||
return server
|
||||
return get_server()
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -14,7 +14,12 @@ import mcp.types
|
|||
from mcp.shared.exceptions import McpError
|
||||
from mcp.types import INTERNAL_ERROR, ErrorData
|
||||
|
||||
from fastmcp.server.dependencies import _current_docket, get_access_token, get_context
|
||||
from fastmcp.server.dependencies import (
|
||||
_current_docket,
|
||||
get_access_token,
|
||||
get_context,
|
||||
register_task_server,
|
||||
)
|
||||
from fastmcp.server.tasks.config import TaskMeta
|
||||
from fastmcp.server.tasks.keys import build_task_key
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
|
@ -80,6 +85,12 @@ async def submit_to_docket(
|
|||
)
|
||||
)
|
||||
|
||||
# Register the current server so background workers resolve
|
||||
# CurrentFastMCP() / ctx.fastmcp to the correct (child) server
|
||||
# for mounted tasks. At this point ctx.fastmcp is the child because
|
||||
# we're inside the child's call_tool dispatch.
|
||||
register_task_server(server_task_id, ctx.fastmcp)
|
||||
|
||||
# Build full task key with embedded metadata
|
||||
task_key = build_task_key(session_id, server_task_id, task_type, key)
|
||||
|
||||
|
|
|
|||
|
|
@ -336,10 +336,80 @@ class TestMountedTaskDependencies:
|
|||
task = await client.call_tool("child_tool_with_server", {}, task=True)
|
||||
await task.result()
|
||||
|
||||
# The server should be the child server since that's where the tool is defined
|
||||
assert len(received_server) == 1
|
||||
# Note: It might be parent or child depending on implementation
|
||||
assert received_server[0] is not None
|
||||
assert received_server[0].name == "server-dep-child"
|
||||
|
||||
|
||||
class TestMountedTaskServerContext:
|
||||
"""Test that background tasks on mounted servers resolve to the child server (#3571)."""
|
||||
|
||||
async def test_current_fastmcp_resolves_to_child_server(self):
|
||||
"""CurrentFastMCP() inside a mounted background task returns the child server."""
|
||||
child = FastMCP("child")
|
||||
received_server: list[FastMCP] = []
|
||||
|
||||
@child.tool(task=True)
|
||||
async def whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form]
|
||||
received_server.append(server)
|
||||
return f"server name: {server.name}"
|
||||
|
||||
parent = FastMCP("parent")
|
||||
parent.mount(child, namespace="child")
|
||||
|
||||
async with Client(parent) as client:
|
||||
task = await client.call_tool("child_whoami", {}, task=True)
|
||||
result = await task.result()
|
||||
|
||||
assert len(received_server) == 1
|
||||
assert received_server[0].name == "child"
|
||||
assert "server name: child" in str(result)
|
||||
|
||||
async def test_context_fastmcp_resolves_to_child_server(self):
|
||||
"""ctx.fastmcp inside a mounted background task returns the child server."""
|
||||
from fastmcp import Context
|
||||
|
||||
child = FastMCP("child")
|
||||
received_server: list[FastMCP] = []
|
||||
|
||||
@child.tool(task=True)
|
||||
async def whoami_ctx(ctx: Context) -> str:
|
||||
received_server.append(ctx.fastmcp)
|
||||
return f"context server: {ctx.fastmcp.name}"
|
||||
|
||||
parent = FastMCP("parent")
|
||||
parent.mount(child, namespace="child")
|
||||
|
||||
async with Client(parent) as client:
|
||||
task = await client.call_tool("child_whoami_ctx", {}, task=True)
|
||||
result = await task.result()
|
||||
|
||||
assert len(received_server) == 1
|
||||
assert received_server[0].name == "child"
|
||||
assert "context server: child" in str(result)
|
||||
|
||||
async def test_nested_mount_resolves_to_innermost_server(self):
|
||||
"""Doubly-nested mounts resolve to the innermost child server."""
|
||||
grandchild = FastMCP("grandchild")
|
||||
received_server: list[FastMCP] = []
|
||||
|
||||
@grandchild.tool(task=True)
|
||||
async def deep_whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form]
|
||||
received_server.append(server)
|
||||
return f"server name: {server.name}"
|
||||
|
||||
child = FastMCP("child")
|
||||
child.mount(grandchild, namespace="gc")
|
||||
|
||||
parent = FastMCP("parent")
|
||||
parent.mount(child, namespace="child")
|
||||
|
||||
async with Client(parent) as client:
|
||||
task = await client.call_tool("child_gc_deep_whoami", {}, task=True)
|
||||
result = await task.result()
|
||||
|
||||
assert len(received_server) == 1
|
||||
assert received_server[0].name == "grandchild"
|
||||
assert "server name: grandchild" in str(result)
|
||||
|
||||
|
||||
class TestMultipleMounts:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue