Dry out retrieving context kwarg

src/fastmcp/prompts/prompt.py

@@ -12,9 +12,11 @@ from mcp.types import Prompt as MCPPrompt
 from mcp.types import PromptArgument as MCPPromptArgument
 from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call
 
+from fastmcp.utilities.json_schema import prune_params
 from fastmcp.utilities.types import (
     _convert_set_defaults,
-    is_class_member_of_type,
+    find_kwarg_by_type,
+    get_cached_typeadapter,
 )
 
 if TYPE_CHECKING:
@@ -110,34 +112,24 @@ class Prompt(BaseModel):
         if func_name == "<lambda>":
             raise ValueError("You must provide a name for lambda functions")
 
+        type_adapter = get_cached_typeadapter(fn)
+        parameters = type_adapter.json_schema()
+
         # Auto-detect context parameter if not provided
         if context_kwarg is None:
-            if inspect.ismethod(fn) and hasattr(fn, "__func__"):
-                sig = inspect.signature(fn.__func__)
-            else:
-                sig = inspect.signature(fn)
-            for param_name, param in sig.parameters.items():
-                if is_class_member_of_type(param.annotation, Context):
-                    context_kwarg = param_name
-                    break
-
-        # Get schema from TypeAdapter - will fail if function isn't properly typed
-        parameters = TypeAdapter(fn).json_schema()
+            context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context)
+        if context_kwarg:
+            parameters = prune_params(parameters, params=[context_kwarg])
 
         # Convert parameters to PromptArguments
         arguments: list[PromptArgument] = []
         if "properties" in parameters:
             for param_name, param in parameters["properties"].items():
-                # Skip context parameter
-                if param_name == context_kwarg:
-                    continue
-
-                required = param_name in parameters.get("required", [])
                 arguments.append(
                     PromptArgument(
                         name=param_name,
                         description=param.get("description"),
-                        required=required,
+                        required=param_name in parameters.get("required", []),
                     )
                 )
 

src/fastmcp/resources/template.py

@@ -22,7 +22,7 @@ from pydantic import (
 from fastmcp.resources.types import FunctionResource, Resource
 from fastmcp.utilities.types import (
     _convert_set_defaults,
-    is_class_member_of_type,
+    find_kwarg_by_type,
 )
 
 if TYPE_CHECKING:
@@ -111,14 +111,7 @@ class ResourceTemplate(BaseModel):
 
         # Auto-detect context parameter if not provided
         if context_kwarg is None:
-            if inspect.ismethod(fn) and hasattr(fn, "__func__"):
-                sig = inspect.signature(fn.__func__)
-            else:
-                sig = inspect.signature(fn)
-            for param_name, param in sig.parameters.items():
-                if is_class_member_of_type(param.annotation, Context):
-                    context_kwarg = param_name
-                    break
+            context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context)
 
         # Validate that URI params match function params
         uri_params = set(re.findall(r"{(\w+)(?:\*)?}", uri_template))

src/fastmcp/tools/tool.py

@@ -16,8 +16,8 @@ from fastmcp.utilities.logging import get_logger
 from fastmcp.utilities.types import (
     Image,
     _convert_set_defaults,
+    find_kwarg_by_type,
     get_cached_typeadapter,
-    is_class_member_of_type,
 )
 
 if TYPE_CHECKING:
@@ -74,18 +74,11 @@ class Tool(BaseModel):
 
         func_doc = description or fn.__doc__ or ""
 
-        if inspect.ismethod(fn) and hasattr(fn, "__func__"):
-            sig = inspect.signature(fn.__func__)
-        else:
-            sig = inspect.signature(fn)
-        if context_kwarg is None:
-            for param_name, param in sig.parameters.items():
-                if is_class_member_of_type(param.annotation, Context):
-                    context_kwarg = param_name
-                    break
-
         type_adapter = get_cached_typeadapter(fn)
         schema = type_adapter.json_schema()
+
+        if context_kwarg is None:
+            context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context)
         if context_kwarg:
             schema = prune_params(schema, params=[context_kwarg])
 

src/fastmcp/utilities/types.py

@@ -1,6 +1,8 @@
 """Common types used across FastMCP."""
 
 import base64
+import inspect
+from collections.abc import Callable
 from functools import lru_cache
 from pathlib import Path
 from types import UnionType
@@ -34,7 +36,12 @@ def issubclass_safe(cls: type, base: type) -> bool:
 
 
 def is_class_member_of_type(cls: type, base: type) -> bool:
-    """Check if cls is a member of base, even if cls is a type variable."""
+    """
+    Check if cls is a member of base, even if cls is a type variable.
+
+    Base can be a type, a UnionType, or an Annotated type. Generic types are not
+    considered members (e.g. T is not a member of list[T]).
+    """
     origin = get_origin(cls)
     # Handle both types of unions: UnionType (from types module, used with | syntax)
     # and typing.Union (used with Union[] syntax)
@@ -50,6 +57,23 @@ def is_class_member_of_type(cls: type, base: type) -> bool:
         return issubclass_safe(cls, base)
 
 
+def find_kwarg_by_type(fn: Callable, kwarg_type: type) -> str | None:
+    """
+    Find the name of the kwarg that is of type kwarg_type.
+
+    Includes union types that contain the kwarg_type, as well as Annotated types.
+    """
+    if inspect.ismethod(fn) and hasattr(fn, "__func__"):
+        sig = inspect.signature(fn.__func__)
+    else:
+        sig = inspect.signature(fn)
+
+    for name, param in sig.parameters.items():
+        if is_class_member_of_type(param.annotation, kwarg_type):
+            return name
+    return None
+
+
 def _convert_set_defaults(maybe_set: set[T] | list[T] | None) -> set[T]:
     """Convert a set or list to a set, defaulting to an empty set if None."""
     if maybe_set is None: