Merge pull request #291 from jlowin/typing

Improve context injection type checks

src/fastmcp/prompts/prompt.py

@@ -13,7 +13,10 @@ 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.types import _convert_set_defaults
+from fastmcp.utilities.types import (
+    _convert_set_defaults,
+    is_class_member_of_type,
+)
 
 if TYPE_CHECKING:
     from mcp.server.session import ServerSessionT
@@ -115,7 +118,7 @@ class Prompt(BaseModel):
             else:
                 sig = inspect.signature(fn)
             for param_name, param in sig.parameters.items():
-                if param.annotation is Context:
+                if is_class_member_of_type(param.annotation, Context):
                     context_kwarg = param_name
                     break
 

src/fastmcp/resources/template.py

@@ -20,7 +20,10 @@ from pydantic import (
 )
 
 from fastmcp.resources.types import FunctionResource, Resource
-from fastmcp.utilities.types import _convert_set_defaults
+from fastmcp.utilities.types import (
+    _convert_set_defaults,
+    is_class_member_of_type,
+)
 
 if TYPE_CHECKING:
     from mcp.server.session import ServerSessionT
@@ -113,7 +116,7 @@ class ResourceTemplate(BaseModel):
             else:
                 sig = inspect.signature(fn)
             for param_name, param in sig.parameters.items():
-                if param.annotation is Context:
+                if is_class_member_of_type(param.annotation, Context):
                     context_kwarg = param_name
                     break
 

src/fastmcp/tools/tool.py

@@ -12,7 +12,11 @@ from pydantic import BaseModel, BeforeValidator, Field
 
 from fastmcp.exceptions import ToolError
 from fastmcp.utilities.func_metadata import FuncMetadata, func_metadata
-from fastmcp.utilities.types import Image, _convert_set_defaults
+from fastmcp.utilities.types import (
+    Image,
+    _convert_set_defaults,
+    is_class_member_of_type,
+)
 
 if TYPE_CHECKING:
     from mcp.server.session import ServerSessionT
@@ -66,7 +70,7 @@ class Tool(BaseModel):
             else:
                 sig = inspect.signature(fn)
             for param_name, param in sig.parameters.items():
-                if param.annotation is Context:
+                if is_class_member_of_type(param.annotation, Context):
                     context_kwarg = param_name
                     break
 

src/fastmcp/utilities/types.py

@@ -2,13 +2,41 @@
 
 import base64
 from pathlib import Path
-from typing import TypeVar
+from types import UnionType
+from typing import Annotated, TypeVar, Union, get_args, get_origin
 
 from mcp.types import ImageContent
 
 T = TypeVar("T")
 
 
+def issubclass_safe(cls: type, base: type) -> bool:
+    """Check if cls is a subclass of base, even if cls is a type variable."""
+    try:
+        if origin := get_origin(cls):
+            return issubclass_safe(origin, base)
+        return issubclass(cls, base)
+    except TypeError:
+        return False
+
+
+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."""
+    origin = get_origin(cls)
+    # Handle both types of unions: UnionType (from types module, used with | syntax)
+    # and typing.Union (used with Union[] syntax)
+    if origin is UnionType or origin == Union:
+        return any(is_class_member_of_type(arg, base) for arg in get_args(cls))
+    elif origin is Annotated:
+        # For Annotated[T, ...], check if T is a member of base
+        args = get_args(cls)
+        if args:
+            return is_class_member_of_type(args[0], base)
+        return False
+    else:
+        return issubclass_safe(cls, base)
+
+
 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/prompts/test_prompt_manager.py

@@ -1,5 +1,10 @@
-import pytest
+from typing import Annotated
 
+import pytest
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT
+
+from fastmcp import Context
 from fastmcp.exceptions import NotFoundError
 from fastmcp.prompts import Prompt
 from fastmcp.prompts.prompt import TextContent, UserMessage
@@ -259,3 +264,97 @@ class TestPromptTags:
         nlp_prompts = [p for p in manager.get_prompts().values() if "nlp" in p.tags]
         assert len(nlp_prompts) == 1
         assert nlp_prompts[0].name == "summary"
+
+
+class TestContextHandling:
+    """Test context handling in prompts."""
+
+    def test_context_parameter_detection(self):
+        """Test that context parameters are properly detected in
+        Prompt.from_function()."""
+
+        def prompt_with_context(x: int, ctx: Context) -> str:
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_with_context)
+        assert prompt.context_kwarg == "ctx"
+
+        def prompt_without_context(x: int) -> str:
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_without_context)
+        assert prompt.context_kwarg is None
+
+    def test_parameterized_context_parameter_detection(self):
+        """Test that parameterized context parameters are properly detected in
+        Prompt.from_function()."""
+
+        def prompt_with_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT]
+        ) -> str:
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_with_context)
+        assert prompt.context_kwarg == "ctx"
+
+    def test_parameterized_union_context_parameter_detection(self):
+        """Test that context parameters in a union are properly detected in
+        Prompt.from_function()."""
+
+        def prompt_with_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT] | None
+        ) -> str:
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_with_context)
+        assert prompt.context_kwarg == "ctx"
+
+    async def test_context_injection(self):
+        """Test that context is properly injected during prompt rendering."""
+
+        def prompt_with_context(x: int, ctx: Context) -> str:
+            assert isinstance(ctx, Context)
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_with_context)
+        assert prompt.context_kwarg == "ctx"
+
+        from fastmcp import FastMCP
+
+        mcp = FastMCP()
+        ctx = mcp.get_context()
+
+        messages = await prompt.render(
+            arguments={"x": 42},
+            context=ctx,
+        )
+        assert len(messages) == 1
+        assert isinstance(messages[0].content, TextContent)
+        assert messages[0].content.text == "42"
+
+    async def test_context_optional(self):
+        """Test that context is optional when rendering prompts."""
+
+        def prompt_with_context(x: int, ctx: Context | None = None) -> str:
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_with_context)
+        assert prompt.context_kwarg == "ctx"
+
+        # Should not raise an error when context is not provided
+        messages = await prompt.render(
+            arguments={"x": 42},
+        )
+        assert len(messages) == 1
+        assert isinstance(messages[0].content, TextContent)
+        assert messages[0].content.text == "42"
+
+    async def test_annotated_context_parameter_detection(self):
+        """Test that annotated context parameters are properly detected in
+        Prompt.from_function()."""
+
+        def prompt_with_context(x: int, ctx: Annotated[Context, "ctx"]) -> str:
+            return str(x)
+
+        prompt = Prompt.from_function(prompt_with_context)
+        assert prompt.context_kwarg == "ctx"

tests/resources/test_resource_template.py

@@ -2,8 +2,11 @@ import json
 from urllib.parse import quote
 
 import pytest
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT
 from pydantic import BaseModel
 
+from fastmcp import Context
 from fastmcp.resources import FunctionResource, ResourceTemplate
 from fastmcp.resources.template import match_uri_template
 
@@ -520,3 +523,113 @@ class TestMatchUriTemplate:
         uri_template = "file://abc/{path*}.py"
         result = match_uri_template(uri=uri, uri_template=uri_template)
         assert result == expected_params
+
+
+class TestContextHandling:
+    """Test context handling in resource templates."""
+
+    def test_context_parameter_detection(self):
+        """Test that context parameters are properly detected in
+        ResourceTemplate.from_function()."""
+
+        def template_with_context(x: int, ctx: Context) -> str:
+            return str(x)
+
+        template = ResourceTemplate.from_function(
+            fn=template_with_context,
+            uri_template="test://{x}",
+            name="test",
+        )
+        assert template.context_kwarg == "ctx"
+
+        def template_without_context(x: int) -> str:
+            return str(x)
+
+        template = ResourceTemplate.from_function(
+            fn=template_without_context,
+            uri_template="test://{x}",
+            name="test",
+        )
+        assert template.context_kwarg is None
+
+    def test_parameterized_context_parameter_detection(self):
+        """Test that parameterized context parameters are properly detected in
+        ResourceTemplate.from_function()."""
+
+        def template_with_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT]
+        ) -> str:
+            return str(x)
+
+        template = ResourceTemplate.from_function(
+            fn=template_with_context,
+            uri_template="test://{x}",
+            name="test",
+        )
+        assert template.context_kwarg == "ctx"
+
+    def test_parameterized_union_context_parameter_detection(self):
+        """Test that context parameters in a union are properly detected in
+        ResourceTemplate.from_function()."""
+
+        def template_with_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT] | None
+        ) -> str:
+            return str(x)
+
+        template = ResourceTemplate.from_function(
+            fn=template_with_context,
+            uri_template="test://{x}",
+            name="test",
+        )
+        assert template.context_kwarg == "ctx"
+
+    async def test_context_injection(self):
+        """Test that context is properly injected during resource creation."""
+
+        def resource_with_context(x: int, ctx: Context) -> str:
+            assert isinstance(ctx, Context)
+            return str(x)
+
+        template = ResourceTemplate.from_function(
+            fn=resource_with_context,
+            uri_template="test://{x}",
+            name="test",
+        )
+        assert template.context_kwarg == "ctx"
+
+        from fastmcp import FastMCP
+
+        mcp = FastMCP()
+        ctx = mcp.get_context()
+
+        resource = await template.create_resource(
+            "test://42",
+            {"x": 42},
+            context=ctx,
+        )
+        assert isinstance(resource, FunctionResource)
+        content = await resource.read()
+        assert content == "42"
+
+    async def test_context_optional(self):
+        """Test that context is optional when creating resources."""
+
+        def resource_with_context(x: int, ctx: Context | None = None) -> str:
+            return str(x)
+
+        template = ResourceTemplate.from_function(
+            fn=resource_with_context,
+            uri_template="test://{x}",
+            name="test",
+        )
+        assert template.context_kwarg == "ctx"
+
+        # Should not raise an error when context is not provided
+        resource = await template.create_resource(
+            "test://42",
+            {"x": 42},
+        )
+        assert isinstance(resource, FunctionResource)
+        content = await resource.read()
+        assert content == "42"

tests/tools/test_tool_manager.py

@@ -1,7 +1,10 @@
 import json
 import logging
+from typing import Annotated
 
 import pytest
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT
 from mcp.types import ImageContent, TextContent
 from pydantic import BaseModel
 
@@ -441,7 +444,8 @@ class TestContextHandling:
             return str(x)
 
         manager = ToolManager()
-        manager.add_tool_from_fn(tool_with_context)
+        tool = manager.add_tool_from_fn(tool_with_context)
+        assert tool.context_kwarg == "ctx"
 
         mcp = FastMCP()
         ctx = mcp.get_context()
@@ -459,7 +463,8 @@ class TestContextHandling:
             return str(x)
 
         manager = ToolManager()
-        manager.add_tool_from_fn(async_tool)
+        tool = manager.add_tool_from_fn(async_tool)
+        assert tool.context_kwarg == "ctx"
 
         mcp = FastMCP()
         ctx = mcp.get_context()
@@ -473,11 +478,12 @@ class TestContextHandling:
         """Test that context is optional when calling tools."""
         from mcp.types import TextContent
 
-        def tool_with_context(x: int, ctx: Context | None = None) -> str:
-            return str(x)
+        def tool_with_context(x: int, ctx: Context | None) -> int:
+            return x
 
         manager = ToolManager()
-        manager.add_tool_from_fn(tool_with_context)
+        tool = manager.add_tool_from_fn(tool_with_context)
+        assert tool.context_kwarg == "ctx"
         # Should not raise an error when context is not provided
         result = await manager.call_tool("tool_with_context", {"x": 42})
         assert isinstance(result, list)
@@ -485,6 +491,40 @@ class TestContextHandling:
         assert isinstance(result[0], TextContent)
         assert result[0].text == "42"
 
+    def test_parameterized_context_parameter_detection(self):
+        """Test that context parameters are properly detected in
+        Tool.from_function()."""
+
+        def tool_with_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT]
+        ) -> str:
+            return str(x)
+
+        manager = ToolManager()
+        tool = manager.add_tool_from_fn(tool_with_context)
+        assert tool.context_kwarg == "ctx"
+
+    def test_annotated_context_parameter_detection(self):
+        def tool_with_context(x: int, ctx: Annotated[Context, "ctx"]) -> str:
+            return str(x)
+
+        manager = ToolManager()
+        tool = manager.add_tool_from_fn(tool_with_context)
+        assert tool.context_kwarg == "ctx"
+
+    def test_parameterized_union_context_parameter_detection(self):
+        """Test that context parameters are properly detected in
+        Tool.from_function()."""
+
+        def tool_with_context(
+            x: int, ctx: Context[ServerSessionT, LifespanContextT] | None
+        ) -> str:
+            return str(x)
+
+        manager = ToolManager()
+        tool = manager.add_tool_from_fn(tool_with_context)
+        assert tool.context_kwarg == "ctx"
+
     async def test_context_error_handling(self):
         """Test error handling when context injection fails."""
 

tests/utilities/test_types.py

@@ -0,0 +1,145 @@
+from typing import Annotated, Any
+
+import pytest
+
+from fastmcp.utilities.types import Image, is_class_member_of_type, issubclass_safe
+
+
+class BaseClass:
+    pass
+
+
+class ChildClass(BaseClass):
+    pass
+
+
+class OtherClass:
+    pass
+
+
+class TestIsClassMemberOfType:
+    def test_basic_subclass_check(self):
+        """Test that a subclass is recognized as a member of the base class."""
+        assert is_class_member_of_type(ChildClass, BaseClass)
+
+    def test_self_is_member(self):
+        """Test that a class is a member of itself."""
+        assert is_class_member_of_type(BaseClass, BaseClass)
+
+    def test_unrelated_class_is_not_member(self):
+        """Test that an unrelated class is not a member of the base class."""
+        assert not is_class_member_of_type(OtherClass, BaseClass)
+
+    def test_typing_union_with_member_is_member(self):
+        """Test that Union type with a member class is detected as a member."""
+        union_type1: Any = ChildClass | OtherClass
+        union_type2: Any = OtherClass | ChildClass
+
+        assert is_class_member_of_type(union_type1, BaseClass)
+        assert is_class_member_of_type(union_type2, BaseClass)
+
+    def test_typing_union_without_member_is_not_member(self):
+        """Test that Union type without any member class is not a member."""
+        union_type: Any = OtherClass | str
+        assert not is_class_member_of_type(union_type, BaseClass)
+
+    def test_pipe_union_with_member_is_member(self):
+        """Test that pipe syntax union with a member class is detected as a member."""
+        union_pipe1: Any = ChildClass | OtherClass
+        union_pipe2: Any = OtherClass | ChildClass
+
+        assert is_class_member_of_type(union_pipe1, BaseClass)
+        assert is_class_member_of_type(union_pipe2, BaseClass)
+
+    def test_pipe_union_without_member_is_not_member(self):
+        """Test that pipe syntax union without any member class is not a member."""
+        union_pipe: Any = OtherClass | str
+        assert not is_class_member_of_type(union_pipe, BaseClass)
+
+    def test_annotated_member_is_member(self):
+        """Test that Annotated with a member class is detected as a member."""
+        annotated1: Any = Annotated[ChildClass, "metadata"]
+        annotated2: Any = Annotated[BaseClass, "metadata"]
+
+        assert is_class_member_of_type(annotated1, BaseClass)
+        assert is_class_member_of_type(annotated2, BaseClass)
+
+    def test_annotated_non_member_is_not_member(self):
+        """Test that Annotated with a non-member class is not a member."""
+        annotated: Any = Annotated[OtherClass, "metadata"]
+        assert not is_class_member_of_type(annotated, BaseClass)
+
+    def test_annotated_with_union_member_is_member(self):
+        """Test that Annotated with a Union containing a member class is a member."""
+        # Test with both Union styles
+        annotated1: Any = Annotated[ChildClass | OtherClass, "metadata"]
+        annotated2: Any = Annotated[ChildClass | OtherClass, "metadata"]
+
+        assert is_class_member_of_type(annotated1, BaseClass)
+        assert is_class_member_of_type(annotated2, BaseClass)
+
+    def test_nested_annotated_with_member_is_member(self):
+        """Test that nested Annotated with a member class is a member."""
+        annotated: Any = Annotated[Annotated[ChildClass, "inner"], "outer"]
+        assert is_class_member_of_type(annotated, BaseClass)
+
+    def test_none_is_not_member(self):
+        """Test that None is not a member of any class."""
+        assert not is_class_member_of_type(None, BaseClass)  # type: ignore
+
+    def test_generic_type_is_not_member(self):
+        """Test that generic types are not members based on their parameter types."""
+        list_type: Any = list[ChildClass]
+        assert not is_class_member_of_type(list_type, BaseClass)
+
+
+class TestIsSubclassSafe:
+    def test_child_is_subclass_of_parent(self):
+        """Test that a child class is recognized as a subclass of its parent."""
+        assert issubclass_safe(ChildClass, BaseClass)
+
+    def test_class_is_subclass_of_itself(self):
+        """Test that a class is a subclass of itself."""
+        assert issubclass_safe(BaseClass, BaseClass)
+
+    def test_unrelated_class_is_not_subclass(self):
+        """Test that an unrelated class is not a subclass."""
+        assert not issubclass_safe(OtherClass, BaseClass)
+
+    def test_none_type_handled_safely(self):
+        """Test that None type is handled safely without raising TypeError."""
+        assert not issubclass_safe(None, BaseClass)  # type: ignore
+
+
+class TestImage:
+    def test_image_initialization_with_path(self):
+        """Test image initialization with a path."""
+        # Mock test - we're not actually going to read a file
+        image = Image(path="test.png")
+        assert image.path is not None
+        assert image.data is None
+        assert image._mime_type == "image/png"
+
+    def test_image_initialization_with_data(self):
+        """Test image initialization with data."""
+        image = Image(data=b"test")
+        assert image.path is None
+        assert image.data == b"test"
+        assert image._mime_type == "image/png"  # Default for raw data
+
+    def test_image_initialization_with_format(self):
+        """Test image initialization with a specific format."""
+        image = Image(data=b"test", format="jpeg")
+        assert image._mime_type == "image/jpeg"
+
+    def test_missing_data_and_path_raises_error(self):
+        """Test that error is raised when neither path nor data is provided."""
+        with pytest.raises(ValueError, match="Either path or data must be provided"):
+            Image()
+
+    def test_both_data_and_path_raises_error(self):
+        """Test that error is raised when both path and data are provided."""
+        with pytest.raises(
+            ValueError, match="Only one of path or data can be provided"
+        ):
+            Image(path="test.png", data=b"test")