feat: implement context injection for resources and prompts in FastMCP

This adds automatic context injection support for both resources and prompts
in the FastMCP framework, allowing these handlers to receive the request
context without explicitly passing it as an argument.

Changes:
- Add context parameter detection in ResourceTemplate and Prompt classes
- Use func_metadata utility to generate schemas while excluding Context params
- Pass context through resource_manager.get_resource() and prompt.render()
- Update server validation to exclude Context params from URI parameter checks
- Add comprehensive test coverage for context injection scenarios

The implementation follows the same pattern already used for tools,
ensuring consistency across all FastMCP handler types.

Author: dsp-ant

src/mcp/server/fastmcp/prompts/base.py

@@ -1,14 +1,22 @@
 """Base classes for FastMCP prompts."""
 
+from __future__ import annotations
+
 import inspect
 from collections.abc import Awaitable, Callable, Sequence
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, get_origin
 
 import pydantic_core
 from pydantic import BaseModel, Field, TypeAdapter, validate_call
 
+from mcp.server.fastmcp.utilities.func_metadata import func_metadata
 from mcp.types import ContentBlock, TextContent
 
+if TYPE_CHECKING:
+    from mcp.server.fastmcp.server import Context
+    from mcp.server.session import ServerSessionT
+    from mcp.shared.context import LifespanContextT, RequestT
+
 
 class Message(BaseModel):
     """Base class for all prompt messages."""
@@ -62,6 +70,7 @@ class Prompt(BaseModel):
     description: str | None = Field(None, description="Description of what the prompt does")
     arguments: list[PromptArgument] | None = Field(None, description="Arguments that can be passed to the prompt")
     fn: Callable[..., PromptResult | Awaitable[PromptResult]] = Field(exclude=True)
+    context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context", exclude=True)
 
     @classmethod
     def from_function(
@@ -70,7 +79,8 @@ def from_function(
         name: str | None = None,
         title: str | None = None,
         description: str | None = None,
-    ) -> "Prompt":
+        context_kwarg: str | None = None,
+    ) -> Prompt:
         """Create a Prompt from a function.
 
         The function can return:
@@ -84,8 +94,29 @@ def from_function(
         if func_name == "<lambda>":
             raise ValueError("You must provide a name for lambda functions")
 
-        # Get schema from TypeAdapter - will fail if function isn't properly typed
-        parameters = TypeAdapter(fn).json_schema()
+        # Find context parameter if it exists
+        if context_kwarg is None:
+            from mcp.server.fastmcp.server import Context
+
+            sig = inspect.signature(fn)
+            for param_name, param in sig.parameters.items():
+                if get_origin(param.annotation) is not None:
+                    continue
+                if param.annotation is not inspect.Parameter.empty:
+                    try:
+                        if issubclass(param.annotation, Context):
+                            context_kwarg = param_name
+                            break
+                    except TypeError:
+                        # issubclass raises TypeError for non-class types
+                        pass
+
+        # Get schema from func_metadata, excluding context parameter
+        func_arg_metadata = func_metadata(
+            fn,
+            skip_names=[context_kwarg] if context_kwarg is not None else [],
+        )
+        parameters = func_arg_metadata.arg_model.model_json_schema()
 
         # Convert parameters to PromptArguments
         arguments: list[PromptArgument] = []
@@ -109,9 +140,14 @@ def from_function(
             description=description or fn.__doc__ or "",
             arguments=arguments,
             fn=fn,
+            context_kwarg=context_kwarg,
         )
 
-    async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]:
+    async def render(
+        self,
+        arguments: dict[str, Any] | None = None,
+        context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
+    ) -> list[Message]:
         """Render the prompt with arguments."""
         # Validate required arguments
         if self.arguments:
@@ -122,8 +158,13 @@ async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]
                 raise ValueError(f"Missing required arguments: {missing}")
 
         try:
+            # Add context to arguments if needed
+            call_args = arguments or {}
+            if self.context_kwarg is not None and context is not None:
+                call_args = {**call_args, self.context_kwarg: context}
+
             # Call function and check if result is a coroutine
-            result = self.fn(**(arguments or {}))
+            result = self.fn(**call_args)
             if inspect.iscoroutine(result):
                 result = await result
 

src/mcp/server/fastmcp/prompts/manager.py

@@ -1,10 +1,17 @@
 """Prompt management functionality."""
 
-from typing import Any
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from mcp.server.fastmcp.prompts.base import Message, Prompt
 from mcp.server.fastmcp.utilities.logging import get_logger
 
+if TYPE_CHECKING:
+    from mcp.server.fastmcp.server import Context
+    from mcp.server.session import ServerSessionT
+    from mcp.shared.context import LifespanContextT, RequestT
+
 logger = get_logger(__name__)
 
 
@@ -39,10 +46,15 @@ def add_prompt(
         self._prompts[prompt.name] = prompt
         return prompt
 
-    async def render_prompt(self, name: str, arguments: dict[str, Any] | None = None) -> list[Message]:
+    async def render_prompt(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
+    ) -> list[Message]:
         """Render a prompt by name with arguments."""
         prompt = self.get_prompt(name)
         if not prompt:
             raise ValueError(f"Unknown prompt: {name}")
 
-        return await prompt.render(arguments)
+        return await prompt.render(arguments, context=context)

src/mcp/server/fastmcp/resources/resource_manager.py

@@ -1,14 +1,21 @@
 """Resource manager functionality."""
 
+from __future__ import annotations
+
 from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pydantic import AnyUrl
 
 from mcp.server.fastmcp.resources.base import Resource
 from mcp.server.fastmcp.resources.templates import ResourceTemplate
 from mcp.server.fastmcp.utilities.logging import get_logger
 
+if TYPE_CHECKING:
+    from mcp.server.fastmcp.server import Context
+    from mcp.server.session import ServerSessionT
+    from mcp.shared.context import LifespanContextT, RequestT
+
 logger = get_logger(__name__)
 
 
@@ -67,7 +74,11 @@ def add_template(
         self._templates[template.uri_template] = template
         return template
 
-    async def get_resource(self, uri: AnyUrl | str) -> Resource | None:
+    async def get_resource(
+        self,
+        uri: AnyUrl | str,
+        context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
+    ) -> Resource | None:
         """Get resource by URI, checking concrete resources first, then templates."""
         uri_str = str(uri)
         logger.debug("Getting resource", extra={"uri": uri_str})
@@ -80,7 +91,7 @@ async def get_resource(self, uri: AnyUrl | str) -> Resource | None:
         for template in self._templates.values():
             if params := template.matches(uri_str):
                 try:
-                    return await template.create_resource(uri_str, params)
+                    return await template.create_resource(uri_str, params, context=context)
                 except Exception as e:
                     raise ValueError(f"Error creating resource from template: {e}")
 

src/mcp/server/fastmcp/resources/templates.py

@@ -5,11 +5,17 @@
 import inspect
 import re
 from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any, get_origin
 
-from pydantic import BaseModel, Field, TypeAdapter, validate_call
+from pydantic import BaseModel, Field, validate_call
 
 from mcp.server.fastmcp.resources.types import FunctionResource, Resource
+from mcp.server.fastmcp.utilities.func_metadata import func_metadata
+
+if TYPE_CHECKING:
+    from mcp.server.fastmcp.server import Context
+    from mcp.server.session import ServerSessionT
+    from mcp.shared.context import LifespanContextT, RequestT
 
 
 class ResourceTemplate(BaseModel):
@@ -22,6 +28,7 @@ class ResourceTemplate(BaseModel):
     mime_type: str = Field(default="text/plain", description="MIME type of the resource content")
     fn: Callable[..., Any] = Field(exclude=True)
     parameters: dict[str, Any] = Field(description="JSON schema for function parameters")
+    context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context")
 
     @classmethod
     def from_function(
@@ -32,14 +39,36 @@ def from_function(
         title: str | None = None,
         description: str | None = None,
         mime_type: str | None = None,
+        context_kwarg: str | None = None,
     ) -> ResourceTemplate:
         """Create a template from a function."""
         func_name = name or fn.__name__
         if func_name == "<lambda>":
             raise ValueError("You must provide a name for lambda functions")
 
-        # Get schema from TypeAdapter - will fail if function isn't properly typed
-        parameters = TypeAdapter(fn).json_schema()
+        # Find context parameter if it exists
+        if context_kwarg is None:
+            from mcp.server.fastmcp.server import Context
+
+            sig = inspect.signature(fn)
+            for param_name, param in sig.parameters.items():
+                if get_origin(param.annotation) is not None:
+                    continue
+                if param.annotation is not inspect.Parameter.empty:
+                    try:
+                        if issubclass(param.annotation, Context):
+                            context_kwarg = param_name
+                            break
+                    except TypeError:
+                        # issubclass raises TypeError for non-class types
+                        pass
+
+        # Get schema from func_metadata, excluding context parameter
+        func_arg_metadata = func_metadata(
+            fn,
+            skip_names=[context_kwarg] if context_kwarg is not None else [],
+        )
+        parameters = func_arg_metadata.arg_model.model_json_schema()
 
         # ensure the arguments are properly cast
         fn = validate_call(fn)
@@ -52,6 +81,7 @@ def from_function(
             mime_type=mime_type or "text/plain",
             fn=fn,
             parameters=parameters,
+            context_kwarg=context_kwarg,
         )
 
     def matches(self, uri: str) -> dict[str, Any] | None:
@@ -63,9 +93,18 @@ def matches(self, uri: str) -> dict[str, Any] | None:
             return match.groupdict()
         return None
 
-    async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource:
+    async def create_resource(
+        self,
+        uri: str,
+        params: dict[str, Any],
+        context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
+    ) -> Resource:
         """Create a resource from the template with the given parameters."""
         try:
+            # Add context to params if needed
+            if self.context_kwarg is not None and context is not None:
+                params = {**params, self.context_kwarg: context}
+
             # Call function and check if result is a coroutine
             result = self.fn(**params)
             if inspect.iscoroutine(result):

src/mcp/server/fastmcp/server.py

@@ -327,7 +327,8 @@ async def list_resource_templates(self) -> list[MCPResourceTemplate]:
     async def read_resource(self, uri: AnyUrl | str) -> Iterable[ReadResourceContents]:
         """Read a resource by URI."""
 
-        resource = await self._resource_manager.get_resource(uri)
+        context = self.get_context()
+        resource = await self._resource_manager.get_resource(uri, context=context)
         if not resource:
             raise ResourceError(f"Unknown resource: {uri}")
 
@@ -515,9 +516,27 @@ def decorator(fn: AnyFunction) -> AnyFunction:
             has_func_params = bool(inspect.signature(fn).parameters)
 
             if has_uri_params or has_func_params:
-                # Validate that URI params match function params
+                # Check for Context parameter to exclude from validation
+                from typing import get_origin
+
+                sig = inspect.signature(fn)
+                context_param = None
+                for param_name, param in sig.parameters.items():
+                    if param.annotation is not inspect.Parameter.empty:
+                        # Check if it's a Context type
+                        if get_origin(param.annotation) is None:
+                            # Try to check if it's a Context without importing
+                            # (to avoid circular imports)
+                            annotation_str = str(param.annotation)
+                            if "Context" in annotation_str:
+                                context_param = param_name
+                                break
+
+                # Validate that URI params match function params (excluding context)
                 uri_params = set(re.findall(r"{(\w+)}", uri))
-                func_params = set(inspect.signature(fn).parameters.keys())
+                func_params = set(sig.parameters.keys())
+                if context_param:
+                    func_params.discard(context_param)
 
                 if uri_params != func_params:
                     raise ValueError(
@@ -983,7 +1002,8 @@ async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -
             if not prompt:
                 raise ValueError(f"Unknown prompt: {name}")
 
-            messages = await prompt.render(arguments)
+            context = self.get_context()
+            messages = await prompt.render(arguments, context=context)
 
             return GetPromptResult(
                 description=prompt.description,

tests/server/fastmcp/test_server.py

@@ -963,7 +963,7 @@ def resource_with_context(name: str, ctx: Context[ServerSession, None]) -> str:
             try:
                 request_id = ctx.request_id
                 return f"Resource {name} - request_id: {request_id}"
-            except:
+            except (AttributeError, ValueError):
                 # Context was injected but request context not available
                 return f"Resource {name} - context injected"
 
@@ -976,7 +976,7 @@ def resource_with_context(name: str, ctx: Context[ServerSession, None]) -> str:
 
         # Test via client
         async with client_session(mcp._mcp_server) as client:
-            result = await client.read_resource("resource://context/test")
+            result = await client.read_resource(AnyUrl("resource://context/test"))
             assert len(result.contents) == 1
             content = result.contents[0]
             assert isinstance(content, TextResourceContents)
@@ -1003,7 +1003,7 @@ def resource_no_context(name: str) -> str:
 
         # Test via client
         async with client_session(mcp._mcp_server) as client:
-            result = await client.read_resource("resource://nocontext/test")
+            result = await client.read_resource(AnyUrl("resource://nocontext/test"))
             assert len(result.contents) == 1
             content = result.contents[0]
             assert isinstance(content, TextResourceContents)
@@ -1029,7 +1029,7 @@ def resource_custom_ctx(id: str, my_ctx: Context[ServerSession, None]) -> str:
 
         # Test via client
         async with client_session(mcp._mcp_server) as client:
-            result = await client.read_resource("resource://custom/123")
+            result = await client.read_resource(AnyUrl("resource://custom/123"))
             assert len(result.contents) == 1
             content = result.contents[0]
             assert isinstance(content, TextResourceContents)
@@ -1043,9 +1043,12 @@ async def test_prompt_with_context(self):
         @mcp.prompt("prompt_with_ctx")
         def prompt_with_context(text: str, ctx: Context[ServerSession, None]) -> str:
             """Prompt that expects context."""
-            if ctx and hasattr(ctx, "request_id"):
-                return f"Prompt '{text}' with context: {ctx.request_id}"
-            return f"Prompt '{text}' - no context"
+            assert ctx is not None
+            try:
+                request_id = ctx.request_id
+                return f"Prompt '{text}' with context: {request_id}"
+            except (AttributeError, ValueError):
+                return f"Prompt '{text}' - context injected"
 
         # Check if prompt has context parameter detection
         prompts = mcp._prompt_manager.list_prompts()
@@ -1065,7 +1068,7 @@ def prompt_with_context(text: str, ctx: Context[ServerSession, None]) -> str:
                 message = result.messages[0]
                 content = message.content
                 assert isinstance(content, TextContent)
-                if "with context:" in content.text:
+                if "context injected" in content.text or "with context:" in content.text:
                     # Context injection is working for prompts
                     assert has_context_kwarg, "Prompt should have context_kwarg attribute"
                 else: