Merge pull request #316 from jlowin/context-kwarg

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:

tests/utilities/test_types.py

@@ -2,7 +2,12 @@ from typing import Annotated, Any
 
 import pytest
 
-from fastmcp.utilities.types import Image, is_class_member_of_type, issubclass_safe
+from fastmcp.utilities.types import (
+    Image,
+    find_kwarg_by_type,
+    is_class_member_of_type,
+    issubclass_safe,
+)
 
 
 class BaseClass:
@@ -143,3 +148,119 @@ class TestImage:
             ValueError, match="Only one of path or data can be provided"
         ):
             Image(path="test.png", data=b"test")
+
+
+class TestFindKwargByType:
+    def test_exact_type_match(self):
+        """Test finding parameter with exact type match."""
+
+        def func(a: int, b: str, c: BaseClass):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) == "c"
+
+    def test_no_matching_parameter(self):
+        """Test finding parameter when no match exists."""
+
+        def func(a: int, b: str, c: OtherClass):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) is None
+
+    def test_parameter_with_no_annotation(self):
+        """Test with a parameter that has no type annotation."""
+
+        def func(a: int, b, c: BaseClass):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) == "c"
+
+    def test_union_type_match_pipe_syntax(self):
+        """Test finding parameter with union type using pipe syntax."""
+
+        def func(a: int, b: str | BaseClass, c: str):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) == "b"
+
+    def test_union_type_match_typing_union(self):
+        """Test finding parameter with union type using Union."""
+
+        def func(a: int, b: str | BaseClass, c: str):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) == "b"
+
+    def test_annotated_type_match(self):
+        """Test finding parameter with Annotated type."""
+
+        def func(a: int, b: Annotated[BaseClass, "metadata"], c: str):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) == "b"
+
+    def test_method_parameter(self):
+        """Test finding parameter in a class method."""
+
+        class TestClass:
+            def method(self, a: int, b: BaseClass):
+                pass
+
+        instance = TestClass()
+        assert find_kwarg_by_type(instance.method, BaseClass) == "b"
+
+    def test_static_method_parameter(self):
+        """Test finding parameter in a static method."""
+
+        class TestClass:
+            @staticmethod
+            def static_method(a: int, b: BaseClass, c: str):
+                pass
+
+        assert find_kwarg_by_type(TestClass.static_method, BaseClass) == "b"
+
+    def test_class_method_parameter(self):
+        """Test finding parameter in a class method."""
+
+        class TestClass:
+            @classmethod
+            def class_method(cls, a: int, b: BaseClass, c: str):
+                pass
+
+        assert find_kwarg_by_type(TestClass.class_method, BaseClass) == "b"
+
+    def test_multiple_matching_parameters(self):
+        """Test finding first parameter when multiple matches exist."""
+
+        def func(a: BaseClass, b: str, c: BaseClass):
+            pass
+
+        # Should return the first match
+        assert find_kwarg_by_type(func, BaseClass) == "a"
+
+    def test_subclass_match(self):
+        """Test finding parameter with a subclass of the target type."""
+
+        def func(a: int, b: ChildClass, c: str):
+            pass
+
+        assert find_kwarg_by_type(func, BaseClass) == "b"
+
+    def test_nonstandard_annotation(self):
+        """Test finding parameter with a nonstandard annotation like an
+        instance. This is irregular."""
+
+        SENTINEL = object()
+
+        def func(a: int, b: SENTINEL, c: str):  # type: ignore
+            pass
+
+        assert find_kwarg_by_type(func, SENTINEL) is None  # type: ignore
+
+    def test_missing_type_annotation(self):
+        """Test finding parameter with a missing type annotation."""
+
+        def func(a: int, b, c: str):
+            pass
+
+        assert find_kwarg_by_type(func, str) == "c"