diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index 0cf37403d..476300942 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/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 diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index d0ec24fbd..6cb6ac3a3 100644 --- a/src/fastmcp/resources/template.py +++ b/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 diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 6f67e20a6..9d0058b6e 100644 --- a/src/fastmcp/tools/tool.py +++ b/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 diff --git a/src/fastmcp/utilities/types.py b/src/fastmcp/utilities/types.py index 98197a922..3f0f272a8 100644 --- a/src/fastmcp/utilities/types.py +++ b/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: diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index 4b6063186..25e3a99e1 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/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" diff --git a/tests/resources/test_resource_template.py b/tests/resources/test_resource_template.py index 86fd24df4..ea086119a 100644 --- a/tests/resources/test_resource_template.py +++ b/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" diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index c58406c6d..08c4dd446 100644 --- a/tests/tools/test_tool_manager.py +++ b/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.""" diff --git a/tests/utilities/test_types.py b/tests/utilities/test_types.py new file mode 100644 index 000000000..e317e50fa --- /dev/null +++ b/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")