diff --git a/src/fastmcp/server/providers/base.py b/src/fastmcp/server/providers/base.py index c65c266a0..8074a8ce5 100644 --- a/src/fastmcp/server/providers/base.py +++ b/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 diff --git a/tests/server/providers/test_base_provider.py b/tests/server/providers/test_base_provider.py new file mode 100644 index 000000000..211aa50e1 --- /dev/null +++ b/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"