diff --git a/CLAUDE.md b/CLAUDE.md index a85214784..0fd0608ef 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e948f0c0d..dcbd6a60c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -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 diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index 956fce711..4d2c2b1c4 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -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, diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index b36296f38..034f133dd 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/src/fastmcp/server/tasks/handlers.py @@ -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) diff --git a/tests/server/tasks/test_task_mount.py b/tests/server/tasks/test_task_mount.py index 875612ada..1b1558579 100644 --- a/tests/server/tasks/test_task_mount.py +++ b/tests/server/tasks/test_task_mount.py @@ -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: