Fix context injection for resources and prompts (#1336)

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

@@ -5,11 +5,18 @@ from __future__ import annotations
 import inspect
 import re
 from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-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.context_injection import find_context_parameter, inject_context
+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 +29,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 +40,23 @@ class ResourceTemplate(BaseModel):
         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:
+            context_kwarg = find_context_parameter(fn)
+
+        # 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 +69,7 @@ class ResourceTemplate(BaseModel):
             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 +81,17 @@ class ResourceTemplate(BaseModel):
             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
+            params = inject_context(self.fn, params, context, self.context_kwarg)
+
             # Call function and check if result is a coroutine
             result = self.fn(**params)
             if inspect.iscoroutine(result):