Add task_meta to prompts and centralize fn_key enrichment

src/fastmcp/prompts/prompt.py

@@ -5,7 +5,7 @@
 import inspect
 import json
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, overload
 
 import pydantic
 import pydantic_core
@@ -27,7 +27,7 @@
 
 from fastmcp.exceptions import PromptError
 from fastmcp.server.dependencies import without_injected_parameters
-from fastmcp.server.tasks.config import TaskConfig
+from fastmcp.server.tasks.config import TaskConfig, TaskMeta
 from fastmcp.utilities.components import FastMCPComponent
 from fastmcp.utilities.json_schema import compress_schema
 from fastmcp.utilities.logging import get_logger
@@ -302,26 +302,52 @@ def convert_result(self, raw_value: Any) -> PromptResult:
             f"got {type(raw_value).__name__}"
         )
 
+    @overload
     async def _render(
         self,
         arguments: dict[str, Any] | None = None,
+        task_meta: None = None,
+    ) -> PromptResult: ...
+
+    @overload
+    async def _render(
+        self,
+        arguments: dict[str, Any] | None,
+        task_meta: TaskMeta,
+    ) -> mcp.types.CreateTaskResult: ...
+
+    async def _render(
+        self,
+        arguments: dict[str, Any] | None = None,
+        task_meta: TaskMeta | None = None,
     ) -> PromptResult | mcp.types.CreateTaskResult:
         """Server entry point that handles task routing.
 
         This allows ANY Prompt subclass to support background execution by setting
         task_config.mode to "supported" or "required". The server calls this
         method instead of render() directly.
 
+        Args:
+            arguments: Prompt arguments
+            task_meta: If provided, execute as background task and return
+                CreateTaskResult. If None (default), execute synchronously and
+                return PromptResult.
+
+        Returns:
+            PromptResult when task_meta is None.
+            CreateTaskResult when task_meta is provided.
+
         Subclasses can override this to customize task routing behavior.
         For example, FastMCPProviderPrompt overrides to delegate to child
         middleware without submitting to Docket.
         """
-        from fastmcp.server.dependencies import _docket_fn_key
         from fastmcp.server.tasks.routing import check_background_task
 
-        key = _docket_fn_key.get() or self.key
         task_result = await check_background_task(
-            component=self, task_type="prompt", key=key, arguments=arguments
+            component=self,
+            task_type="prompt",
+            arguments=arguments,
+            task_meta=task_meta,
         )
         if task_result:
             return task_result

src/fastmcp/resources/resource.py

@@ -5,7 +5,6 @@
 import base64
 import inspect
 from collections.abc import Callable
-from dataclasses import replace
 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, overload
 
 import mcp.types
@@ -331,10 +330,6 @@ async def _read(
         """
         from fastmcp.server.tasks.routing import check_background_task
 
-        # Enrich task_meta with fn_key if not already set (fallback for programmatic API)
-        if task_meta is not None and task_meta.fn_key is None:
-            task_meta = replace(task_meta, fn_key=self.key)
-
         task_result = await check_background_task(
             component=self, task_type="resource", arguments=None, task_meta=task_meta
         )

src/fastmcp/resources/template.py

@@ -5,7 +5,6 @@
 import inspect
 import re
 from collections.abc import Callable
-from dataclasses import replace
 from typing import TYPE_CHECKING, Any, ClassVar, overload
 from urllib.parse import parse_qs, unquote
 
@@ -214,10 +213,6 @@ async def _read(
         """
         from fastmcp.server.tasks.routing import check_background_task
 
-        # Enrich task_meta with fn_key if not already set (fallback for programmatic API)
-        if task_meta is not None and task_meta.fn_key is None:
-            task_meta = replace(task_meta, fn_key=self.key)
-
         task_result = await check_background_task(
             component=self, task_type="template", arguments=params, task_meta=task_meta
         )
@@ -346,10 +341,6 @@ async def _read(
         """
         from fastmcp.server.tasks.routing import check_background_task
 
-        # Enrich task_meta with fn_key if not already set (fallback for programmatic API)
-        if task_meta is not None and task_meta.fn_key is None:
-            task_meta = replace(task_meta, fn_key=self.key)
-
         task_result = await check_background_task(
             component=self, task_type="template", arguments=params, task_meta=task_meta
         )

src/fastmcp/server/dependencies.py

@@ -40,16 +40,6 @@
     "server", default=None
 )
 
-# ContextVar for propagating task metadata through the async call stack
-# When set, indicates this call should be executed as a background task
-_task_metadata: ContextVar[dict[str, Any] | None] = ContextVar(
-    "task_metadata", default=None
-)
-
-# ContextVar for the component's Docket function lookup key (with namespace prefix)
-# Used by Tool._run(), Resource._read(), Prompt._render() to find the registered function
-_docket_fn_key: ContextVar[str | None] = ContextVar("docket_fn_key", default=None)
-
 __all__ = [
     "AccessToken",
     "CurrentContext",
@@ -62,7 +52,6 @@
     "get_http_headers",
     "get_http_request",
     "get_server",
-    "get_task_metadata",
     "resolve_dependencies",
     "without_injected_parameters",
 ]
@@ -277,16 +266,6 @@ def get_context() -> Context:
     return context
 
 
-def get_task_metadata() -> dict[str, Any] | None:
-    """Get the current task metadata from the context.
-
-    Returns:
-        The task metadata dict if this is a background task request,
-        or None if this is a normal execution.
-    """
-    return _task_metadata.get()
-
-
 class _CurrentContext(Dependency):
     """Internal dependency class for CurrentContext."""
 

src/fastmcp/server/providers/fastmcp_provider.py

@@ -13,7 +13,6 @@
 import re
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
-from dataclasses import replace
 from typing import TYPE_CHECKING, Any, overload
 
 import mcp.types
@@ -108,16 +107,9 @@ async def _run(
         """Delegate to child server's call_tool() with task_meta.
 
         Passes task_meta through to the child server so it can handle
-        backgrounding appropriately.
-
-        Enriches fn_key with self.key (the parent's namespaced tool name) so that
-        when the child tool delegates to Docket, it uses the correct lookup key
-        that was registered via get_tasks().
+        backgrounding appropriately. fn_key is already set by the parent
+        server before calling this method.
         """
-        # Enrich fn_key with parent's key before delegating to child
-        if task_meta is not None and task_meta.fn_key is None:
-            task_meta = replace(task_meta, fn_key=self.key)
-
         return await self._server.call_tool(
             self._original_name, arguments, task_meta=task_meta
         )
@@ -183,16 +175,9 @@ async def _read(
         """Delegate to child server's read_resource() with task_meta.
 
         Passes task_meta through to the child server so it can handle
-        backgrounding appropriately.
-
-        Enriches fn_key with self.key (the parent's namespaced URI) so that
-        when the child resource delegates to Docket, it uses the correct
-        lookup key that was registered via get_tasks().
+        backgrounding appropriately. fn_key is already set by the parent
+        server before calling this method.
         """
-        # Enrich fn_key with parent's URI before delegating to child
-        if task_meta is not None and task_meta.fn_key is None:
-            task_meta = TaskMeta(ttl=task_meta.ttl, fn_key=self.key)
-
         return await self._server.read_resource(self._original_uri, task_meta=task_meta)
 
 
@@ -229,28 +214,47 @@ def wrap(cls, server: Any, prompt: Prompt) -> FastMCPProviderPrompt:
             task_config=prompt.task_config,
         )
 
+    @overload
     async def _render(
-        self, arguments: dict[str, Any] | None = None
+        self,
+        arguments: dict[str, Any] | None = None,
+        task_meta: None = None,
+    ) -> PromptResult: ...
+
+    @overload
+    async def _render(
+        self,
+        arguments: dict[str, Any] | None,
+        task_meta: TaskMeta,
+    ) -> mcp.types.CreateTaskResult: ...
+
+    async def _render(
+        self,
+        arguments: dict[str, Any] | None = None,
+        task_meta: TaskMeta | None = None,
     ) -> PromptResult | mcp.types.CreateTaskResult:
-        """Skip task routing - delegate to render() which calls child middleware.
+        """Delegate to child server's render_prompt() with task_meta.
 
-        The actual underlying prompt will check _task_metadata contextvar and
-        submit to Docket if appropriate. This wrapper just passes through.
+        Passes task_meta through to the child server so it can handle
+        backgrounding appropriately. fn_key is already set by the parent
+        server before calling this method.
         """
-        return await self.render(arguments)
+        return await self._server.render_prompt(
+            self._original_name, arguments, task_meta=task_meta
+        )
 
     async def render(
         self, arguments: dict[str, Any] | None = None
     ) -> PromptResult | mcp.types.CreateTaskResult:  # type: ignore[override]
-        """Delegate to child server's render_prompt().
+        """Not implemented - use _render() which delegates to child server.
 
-        Note: The _docket_fn_key contextvar is intentionally NOT updated here.
-        The parent set it to the full namespaced name (e.g., c_gc_greet) which
-        is what the function is registered under in Docket. All provider layers
-        pass this through unchanged so the eventual prompt._render() uses the
-        correct Docket lookup key.
+        FastMCPProviderPrompt._render() handles all execution by delegating
+        to the child server's render_prompt() with task_meta.
         """
-        return await self._server.render_prompt(self._original_name, arguments)
+        raise NotImplementedError(
+            "FastMCPProviderPrompt.render() should not be called directly. "
+            "Use _render() which delegates to the child server's render_prompt()."
+        )
 
 
 class FastMCPProviderResourceTemplate(ResourceTemplate):
@@ -326,19 +330,12 @@ async def _read(
         """Delegate to child server's read_resource() with task_meta.
 
         Passes task_meta through to the child server so it can handle
-        backgrounding appropriately.
-
-        Enriches fn_key with self.key (the parent's namespaced template pattern)
-        so that when the child template delegates to Docket, it uses the correct
-        lookup key that was registered via get_tasks().
+        backgrounding appropriately. fn_key is already set by the parent
+        server before calling this method.
         """
         # Expand the original template with params to get internal URI
         original_uri = _expand_uri_template(self._original_uri_template or "", params)
 
-        # Enrich fn_key with parent's namespaced template pattern before delegating
-        if task_meta is not None and task_meta.fn_key is None:
-            task_meta = replace(task_meta, fn_key=self.key)
-
         return await self._server.read_resource(original_uri, task_meta=task_meta)
 
     async def read(self, arguments: dict[str, Any]) -> str | bytes | ResourceResult: