refactor: extract common context injection pattern

Eliminated code duplication across tools, resources, and prompts by creating
a centralized context injection utility module.

Changes:
- Add context_injection.py utility with find_context_parameter() and inject_context()
- Refactor tool, resource, and prompt implementations to use the shared utility
- Remove duplicated context detection logic from all three implementations
- Maintain backward compatibility with all existing tests passing

This refactoring reduces code duplication and makes the context injection
pattern consistent and maintainable across all FastMCP handler types.

Author: dsp-ant

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

@@ -4,11 +4,12 @@
 
 import inspect
 from collections.abc import Awaitable, Callable, Sequence
-from typing import TYPE_CHECKING, Any, Literal, get_origin
+from typing import TYPE_CHECKING, Any, Literal
 
 import pydantic_core
 from pydantic import BaseModel, Field, TypeAdapter, validate_call
 
+from mcp.server.fastmcp.utilities.context_injection import find_context_parameter, inject_context
 from mcp.server.fastmcp.utilities.func_metadata import func_metadata
 from mcp.types import ContentBlock, TextContent
 
@@ -96,20 +97,7 @@ def from_function(
 
         # 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
+            context_kwarg = find_context_parameter(fn)
 
         # Get schema from func_metadata, excluding context parameter
         func_arg_metadata = func_metadata(
@@ -159,9 +147,7 @@ async def render(
 
         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_args = inject_context(self.fn, arguments or {}, context, self.context_kwarg)
 
             # Call function and check if result is a coroutine
             result = self.fn(**call_args)

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

@@ -5,11 +5,12 @@
 import inspect
 import re
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, get_origin
+from typing import TYPE_CHECKING, Any
 
 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:
@@ -48,20 +49,7 @@ def from_function(
 
         # 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
+            context_kwarg = find_context_parameter(fn)
 
         # Get schema from func_metadata, excluding context parameter
         func_arg_metadata = func_metadata(
@@ -102,8 +90,7 @@ async def create_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}
+            params = inject_context(self.fn, params, context, self.context_kwarg)
 
             # Call function and check if result is a coroutine
             result = self.fn(**params)

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

@@ -4,11 +4,12 @@
 import inspect
 from collections.abc import Callable
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, get_origin
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel, Field
 
 from mcp.server.fastmcp.exceptions import ToolError
+from mcp.server.fastmcp.utilities.context_injection import find_context_parameter
 from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata, func_metadata
 from mcp.types import ToolAnnotations
 
@@ -49,8 +50,6 @@ def from_function(
         structured_output: bool | None = None,
     ) -> Tool:
         """Create a Tool from a function."""
-        from mcp.server.fastmcp.server import Context
-
         func_name = name or fn.__name__
 
         if func_name == "<lambda>":
@@ -60,13 +59,7 @@ def from_function(
         is_async = _is_async_callable(fn)
 
         if context_kwarg is None:
-            sig = inspect.signature(fn)
-            for param_name, param in sig.parameters.items():
-                if get_origin(param.annotation) is not None:
-                    continue
-                if issubclass(param.annotation, Context):
-                    context_kwarg = param_name
-                    break
+            context_kwarg = find_context_parameter(fn)
 
         func_arg_metadata = func_metadata(
             fn,

src/mcp/server/fastmcp/utilities/context_injection.py

@@ -0,0 +1,65 @@
+"""Context injection utilities for FastMCP."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, get_origin
+
+if TYPE_CHECKING:
+    pass
+
+
+def find_context_parameter(fn: Callable[..., Any]) -> str | None:
+    """Find the parameter that should receive the Context object.
+
+    Searches through the function's signature to find a parameter
+    with a Context type annotation.
+
+    Args:
+        fn: The function to inspect
+
+    Returns:
+        The name of the context parameter, or None if not found
+    """
+    from mcp.server.fastmcp.server import Context
+
+    sig = inspect.signature(fn)
+    for param_name, param in sig.parameters.items():
+        # Skip generic types
+        if get_origin(param.annotation) is not None:
+            continue
+
+        # Check if parameter has annotation
+        if param.annotation is not inspect.Parameter.empty:
+            try:
+                # Check if it's a Context subclass
+                if issubclass(param.annotation, Context):
+                    return param_name
+            except TypeError:
+                # issubclass raises TypeError for non-class types
+                pass
+
+    return None
+
+
+def inject_context(
+    fn: Callable[..., Any],
+    kwargs: dict[str, Any],
+    context: Any | None,
+    context_kwarg: str | None,
+) -> dict[str, Any]:
+    """Inject context into function kwargs if needed.
+
+    Args:
+        fn: The function that will be called
+        kwargs: The current keyword arguments
+        context: The context object to inject (if any)
+        context_kwarg: The name of the parameter to inject into
+
+    Returns:
+        Updated kwargs with context injected if applicable
+    """
+    if context_kwarg is not None and context is not None:
+        return {**kwargs, context_kwarg: context}
+    return kwargs