Fix context detection by checking for subclass relationship

Previously, the code only checked for exact matches with `Context`. This update ensures that subclasses of `Context` are also correctly identified, improving flexibility and reliability.

Author: jkawamoto

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

@@ -2,7 +2,7 @@
 
 import inspect
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, get_origin
 
 from pydantic import BaseModel, Field
 
@@ -53,7 +53,9 @@ def from_function(
         if context_kwarg is None:
             sig = inspect.signature(fn)
             for param_name, param in sig.parameters.items():
-                if param.annotation is Context:
+                if get_origin(param.annotation) is not None:
+                    continue
+                if issubclass(param.annotation, Context):
                     context_kwarg = param_name
                     break
 

tests/server/fastmcp/test_tool_manager.py

@@ -242,6 +242,7 @@ def test_context_parameter_detection(self):
         """Test that context parameters are properly detected in
         Tool.from_function()."""
         from mcp.server.fastmcp import Context
+        from mcp.server.session import ServerSessionT
 
         def tool_with_context(x: int, ctx: Context) -> str:
             return str(x)
@@ -256,6 +257,16 @@ def tool_without_context(x: int) -> str:
         tool = manager.add_tool(tool_without_context)
         assert tool.context_kwarg is None
 
+        class LifespanParam: ...
+
+        def tool_with_specialized_context(
+            x: int, ctx: Context[ServerSessionT, LifespanParam]
+        ) -> str:
+            return str(x)
+
+        tool = manager.add_tool(tool_with_specialized_context)
+        assert tool.context_kwarg == "ctx"
+
     @pytest.mark.anyio
     async def test_context_injection(self):
         """Test that context is properly injected during tool execution."""