Fix Provider.get_tasks() to include custom component subclasses (#2729)

src/fastmcp/server/providers/base.py

@@ -245,35 +245,32 @@ class Provider:
         """Return components that should be registered as background tasks.
 
         Override to customize which components are task-eligible.
-        Default calls list_* methods and filters for function-based components
+        Default calls list_* methods and filters for components
         with task_config.mode != 'forbidden'.
 
         Used by the server during startup to register functions with Docket.
         """
-        from fastmcp.prompts.prompt import FunctionPrompt
-        from fastmcp.resources.resource import FunctionResource
-        from fastmcp.resources.template import FunctionResourceTemplate
-        from fastmcp.tools.tool import FunctionTool
+        from fastmcp.prompts.prompt import Prompt
+        from fastmcp.resources.resource import Resource
+        from fastmcp.resources.template import ResourceTemplate
+        from fastmcp.tools.tool import Tool
 
         components: list[FastMCPComponent] = []
 
         for t in await self.list_tools():
-            if isinstance(t, FunctionTool) and t.task_config.supports_tasks():
+            if isinstance(t, Tool) and t.task_config.supports_tasks():
                 components.append(t)
 
         for r in await self.list_resources():
-            if isinstance(r, FunctionResource) and r.task_config.supports_tasks():
+            if isinstance(r, Resource) and r.task_config.supports_tasks():
                 components.append(r)
 
         for t in await self.list_resource_templates():
-            if (
-                isinstance(t, FunctionResourceTemplate)
-                and t.task_config.supports_tasks()
-            ):
+            if isinstance(t, ResourceTemplate) and t.task_config.supports_tasks():
                 components.append(t)
 
         for p in await self.list_prompts():
-            if isinstance(p, FunctionPrompt) and p.task_config.supports_tasks():
+            if isinstance(p, Prompt) and p.task_config.supports_tasks():
                 components.append(p)
 
         return components

tests/server/providers/test_base_provider.py

@@ -0,0 +1,78 @@
+"""Tests for base Provider class behavior."""
+
+from typing import Any
+
+from fastmcp.server.providers.base import Provider
+from fastmcp.server.tasks.config import TaskConfig
+from fastmcp.tools.tool import Tool, ToolResult
+
+
+class CustomTool(Tool):
+    """A custom Tool subclass (not FunctionTool) with task support."""
+
+    task_config: TaskConfig = TaskConfig(mode="optional")
+    parameters: dict[str, Any] = {"type": "object", "properties": {}}
+
+    async def run(self, arguments: dict[str, Any]) -> ToolResult:
+        return ToolResult(content="custom result")
+
+
+class SimpleProvider(Provider):
+    """Minimal provider that returns custom components from list methods."""
+
+    def __init__(self, tools: list[Tool] | None = None):
+        self._tools = tools or []
+
+    async def list_tools(self) -> list[Tool]:
+        return self._tools
+
+
+class TestBaseProviderGetTasks:
+    """Tests for Provider.get_tasks() base implementation."""
+
+    async def test_get_tasks_includes_custom_tool_subclasses(self):
+        """Base Provider.get_tasks() should include custom Tool subclasses."""
+        custom_tool = CustomTool(name="custom", description="A custom tool")
+        provider = SimpleProvider(tools=[custom_tool])
+
+        tasks = await provider.get_tasks()
+
+        assert len(tasks) == 1
+        assert tasks[0].name == "custom"
+        assert tasks[0] is custom_tool
+
+    async def test_get_tasks_filters_forbidden_custom_tools(self):
+        """Base Provider.get_tasks() should exclude tools with forbidden task mode."""
+
+        class ForbiddenTool(Tool):
+            task_config: TaskConfig = TaskConfig(mode="forbidden")
+            parameters: dict[str, Any] = {"type": "object", "properties": {}}
+
+            async def run(self, arguments: dict[str, Any]) -> ToolResult:
+                return ToolResult(content="forbidden")
+
+        forbidden_tool = ForbiddenTool(name="forbidden", description="Forbidden tool")
+        provider = SimpleProvider(tools=[forbidden_tool])
+
+        tasks = await provider.get_tasks()
+
+        assert len(tasks) == 0
+
+    async def test_get_tasks_mixed_custom_and_forbidden(self):
+        """Base Provider.get_tasks() filters correctly with mixed task modes."""
+
+        class ForbiddenTool(Tool):
+            task_config: TaskConfig = TaskConfig(mode="forbidden")
+            parameters: dict[str, Any] = {"type": "object", "properties": {}}
+
+            async def run(self, arguments: dict[str, Any]) -> ToolResult:
+                return ToolResult(content="forbidden")
+
+        enabled_tool = CustomTool(name="enabled", description="Task enabled")
+        forbidden_tool = ForbiddenTool(name="forbidden", description="Task forbidden")
+        provider = SimpleProvider(tools=[enabled_tool, forbidden_tool])
+
+        tasks = await provider.get_tasks()
+
+        assert len(tasks) == 1
+        assert tasks[0].name == "enabled"