From d6654a2379c74ffc3172c756dd83834745708aec Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 21 Dec 2025 15:41:46 -0500 Subject: [PATCH] Add supports_tasks() method to replace string mode checks (#2664) * Add supports_tasks() method to replace string mode checks Consolidates task config mode checks into a readable method on TaskConfig. Instead of `task_config.mode == "forbidden"` or `task_config.mode != "forbidden"`, code now uses `task_config.supports_tasks()` for clearer intent. Updated 20 instances across the codebase and added type assertions in tests to resolve type checker warnings. * Update test to match new error message --- .cursor/worktrees.json | 6 ------ .github/workflows/marvin.yml | 2 +- src/fastmcp/prompts/prompt.py | 4 ++-- src/fastmcp/resources/resource.py | 4 ++-- src/fastmcp/resources/template.py | 4 ++-- src/fastmcp/server/providers/base.py | 8 ++++---- src/fastmcp/server/providers/fastmcp_provider.py | 8 ++++---- src/fastmcp/server/tasks/config.py | 10 +++++++++- src/fastmcp/server/tasks/routing.py | 2 +- src/fastmcp/tools/tool.py | 8 ++++---- src/fastmcp/utilities/components.py | 4 ++-- tests/prompts/test_prompt_manager.py | 14 +++++++++++++- tests/resources/test_resource_manager.py | 6 ++++++ tests/resources/test_resource_template.py | 8 ++++++++ tests/server/tasks/test_custom_subclass_tasks.py | 2 +- 15 files changed, 59 insertions(+), 31 deletions(-) delete mode 100644 .cursor/worktrees.json diff --git a/.cursor/worktrees.json b/.cursor/worktrees.json deleted file mode 100644 index 3321da678..000000000 --- a/.cursor/worktrees.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "setup-worktree": [ - "uv sync", - "uv run pre-commit install" - ] -} diff --git a/.github/workflows/marvin.yml b/.github/workflows/marvin.yml index 2693e42f4..d63eb27fb 100644 --- a/.github/workflows/marvin.yml +++ b/.github/workflows/marvin.yml @@ -65,7 +65,7 @@ jobs: trigger_phrase: "/marvin" allowed_bots: "*" claude_args: | - --allowedTools WebSearch,WebFetch,Bash(uv:*),Bash(pre-commit:*),Bash(pytest:*),Bash(ruff:*),Bash(ty:*),Bash(git:*),Bash(gh:*),mcp__github__add_issue_comment,mcp__github__create_issue,mcp__github__get_issue,mcp__github__list_issues,mcp__github__search_issues,mcp__github__update_issue,mcp__github__update_issue_comment,mcp__github__create_pull_request,mcp__github__get_pull_request,mcp__github__get_pull_request_comments,mcp__github__get_pull_request_files,mcp__github__get_pull_request_reviews,mcp__github__get_pull_request_status,mcp__github__list_pull_requests,mcp__github__update_pull_request,mcp__github__update_pull_request_branch,mcp__github__update_pull_request_comment,mcp__github__merge_pull_request + --allowedTools WebSearch,WebFetch,Bash(uv:*),Bash(pre-commit:*),Bash(prek:*),Bash(pytest:*),Bash(ruff:*),Bash(ty:*),Bash(git:*),Bash(gh:*),mcp__github__add_issue_comment,mcp__github__create_issue,mcp__github__get_issue,mcp__github__list_issues,mcp__github__search_issues,mcp__github__update_issue,mcp__github__update_issue_comment,mcp__github__create_pull_request,mcp__github__get_pull_request,mcp__github__get_pull_request_comments,mcp__github__get_pull_request_files,mcp__github__get_pull_request_reviews,mcp__github__get_pull_request_status,mcp__github__list_pull_requests,mcp__github__update_pull_request,mcp__github__update_pull_request_branch,mcp__github__update_pull_request_comment,mcp__github__merge_pull_request additional_permissions: | actions: read settings: | diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index 02f88df0a..cef1d23e2 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/src/fastmcp/prompts/prompt.py @@ -294,7 +294,7 @@ class Prompt(FastMCPComponent): def register_with_docket(self, docket: Docket) -> None: """Register this prompt with docket for background execution.""" - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.render, names=[self.key]) @@ -519,7 +519,7 @@ class FunctionPrompt(Prompt): FunctionPrompt registers the underlying function, which has the user's Depends parameters for docket to resolve. """ - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.fn, names=[self.key]) # type: ignore[arg-type] diff --git a/src/fastmcp/resources/resource.py b/src/fastmcp/resources/resource.py index d6a98b1b6..a0ed8f192 100644 --- a/src/fastmcp/resources/resource.py +++ b/src/fastmcp/resources/resource.py @@ -305,7 +305,7 @@ class Resource(FastMCPComponent): def register_with_docket(self, docket: Docket) -> None: """Register this resource with docket for background execution.""" - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.read, names=[self.key]) @@ -421,6 +421,6 @@ class FunctionResource(Resource): FunctionResource registers the underlying function, which has the user's Depends parameters for docket to resolve. """ - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.fn, names=[self.key]) diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index c03258e7b..9d3bbe89e 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -270,7 +270,7 @@ class ResourceTemplate(FastMCPComponent): def register_with_docket(self, docket: Docket) -> None: """Register this template with docket for background execution.""" - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.read, names=[self.key]) @@ -384,7 +384,7 @@ class FunctionResourceTemplate(ResourceTemplate): FunctionResourceTemplate registers the underlying function, which has the user's Depends parameters for docket to resolve. """ - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.fn, names=[self.key]) diff --git a/src/fastmcp/server/providers/base.py b/src/fastmcp/server/providers/base.py index c1d1d3dba..f7eff819f 100644 --- a/src/fastmcp/server/providers/base.py +++ b/src/fastmcp/server/providers/base.py @@ -255,23 +255,23 @@ class Provider: tools=[ t for t in all_tools - if isinstance(t, FunctionTool) and t.task_config.mode != "forbidden" + if isinstance(t, FunctionTool) and t.task_config.supports_tasks() ], resources=[ r for r in all_resources - if isinstance(r, FunctionResource) and r.task_config.mode != "forbidden" + if isinstance(r, FunctionResource) and r.task_config.supports_tasks() ], templates=[ t for t in all_templates if isinstance(t, FunctionResourceTemplate) - and t.task_config.mode != "forbidden" + and t.task_config.supports_tasks() ], prompts=[ p for p in all_prompts - if isinstance(p, FunctionPrompt) and p.task_config.mode != "forbidden" + if isinstance(p, FunctionPrompt) and p.task_config.supports_tasks() ], ) diff --git a/src/fastmcp/server/providers/fastmcp_provider.py b/src/fastmcp/server/providers/fastmcp_provider.py index 06f3a415b..1fa827a16 100644 --- a/src/fastmcp/server/providers/fastmcp_provider.py +++ b/src/fastmcp/server/providers/fastmcp_provider.py @@ -570,22 +570,22 @@ class FastMCPProvider(Provider): tools: list[Tool] = [ t for t in self.server._tool_manager._tools.values() - if t.task_config.mode != "forbidden" + if t.task_config.supports_tasks() ] resources: list[Resource] = [ r for r in self.server._resource_manager._resources.values() - if r.task_config.mode != "forbidden" + if r.task_config.supports_tasks() ] templates: list[ResourceTemplate] = [ t for t in self.server._resource_manager._templates.values() - if t.task_config.mode != "forbidden" + if t.task_config.supports_tasks() ] prompts: list[Prompt] = [ p for p in self.server._prompt_manager._prompts.values() - if p.task_config.mode != "forbidden" + if p.task_config.supports_tasks() ] # Recursively get tasks from nested providers diff --git a/src/fastmcp/server/tasks/config.py b/src/fastmcp/server/tasks/config.py index ac2670da6..4a939af50 100644 --- a/src/fastmcp/server/tasks/config.py +++ b/src/fastmcp/server/tasks/config.py @@ -59,6 +59,14 @@ class TaskConfig: """ return cls(mode="optional" if value else "forbidden") + def supports_tasks(self) -> bool: + """Check if this component supports task execution. + + Returns: + True if mode is "optional" or "required", False if "forbidden". + """ + return self.mode != "forbidden" + def validate_function(self, fn: Callable[..., Any], name: str) -> None: """Validate that function is compatible with this task config. @@ -72,7 +80,7 @@ class TaskConfig: Raises: ValueError: If task execution is enabled but function is sync. """ - if self.mode == "forbidden": + if not self.supports_tasks(): return # Unwrap callable classes and staticmethods diff --git a/src/fastmcp/server/tasks/routing.py b/src/fastmcp/server/tasks/routing.py index f0e0d804c..39f16b52a 100644 --- a/src/fastmcp/server/tasks/routing.py +++ b/src/fastmcp/server/tasks/routing.py @@ -60,7 +60,7 @@ async def check_background_task( ) # Enforce mode="forbidden" - cannot be called with task metadata - if task_config.mode == "forbidden" and task_meta: + if not task_config.supports_tasks() and task_meta: raise McpError( ErrorData( code=METHOD_NOT_FOUND, diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 3a8baed6e..4ca32c765 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -307,7 +307,7 @@ class Tool(FastMCPComponent): def register_with_docket(self, docket: Docket) -> None: """Register this tool with docket for background execution.""" - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.run, names=[self.key]) @@ -388,8 +388,8 @@ class FunctionTool(Tool): ) # Add task execution mode per SEP-1686 - # Only set execution if not overridden and mode is not "forbidden" - if self.task_config.mode != "forbidden" and "execution" not in overrides: + # Only set execution if not overridden and task execution is supported + if self.task_config.supports_tasks() and "execution" not in overrides: mcp_tool.execution = ToolExecution(taskSupport=self.task_config.mode) return mcp_tool @@ -483,7 +483,7 @@ class FunctionTool(Tool): FunctionTool registers the underlying function, which has the user's Depends parameters for docket to resolve. """ - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): return docket.register(self.fn, names=[self.key]) diff --git a/src/fastmcp/utilities/components.py b/src/fastmcp/utilities/components.py index 679b85266..2e5059f62 100644 --- a/src/fastmcp/utilities/components.py +++ b/src/fastmcp/utilities/components.py @@ -143,10 +143,10 @@ class FastMCPComponent(FastMCPBaseModel): The **kwargs are passed through to docket.add() (e.g., key=task_key). """ - if self.task_config.mode == "forbidden": + if not self.task_config.supports_tasks(): raise RuntimeError( f"Cannot add {self.__class__.__name__} '{self.name}' to docket: " - f"task_config.mode is 'forbidden'" + f"task execution not supported" ) raise NotImplementedError( f"{self.__class__.__name__} does not implement add_to_docket()" diff --git a/tests/prompts/test_prompt_manager.py b/tests/prompts/test_prompt_manager.py index b9e666926..b53c26ae4 100644 --- a/tests/prompts/test_prompt_manager.py +++ b/tests/prompts/test_prompt_manager.py @@ -6,7 +6,12 @@ import pytest from fastmcp import Context, FastMCP from fastmcp.exceptions import NotFoundError, PromptError from fastmcp.prompts import Prompt -from fastmcp.prompts.prompt import FunctionPrompt, PromptMessage, TextContent +from fastmcp.prompts.prompt import ( + FunctionPrompt, + PromptMessage, + PromptResult, + TextContent, +) from fastmcp.prompts.prompt_manager import PromptManager from fastmcp.utilities.tests import caplog_for_fastmcp from tests.conftest import get_fn_name @@ -165,6 +170,7 @@ class TestRenderPrompt: prompt = Prompt.from_function(fn) manager.add_prompt(prompt) result = await manager.render_prompt("fn") + assert isinstance(result, PromptResult) assert result.description == "An example prompt." assert result.messages == [ PromptMessage( @@ -183,6 +189,7 @@ class TestRenderPrompt: prompt = Prompt.from_function(fn) manager.add_prompt(prompt) result = await manager.render_prompt("fn", arguments={"name": "World"}) + assert isinstance(result, PromptResult) assert result.description == "An example prompt." assert result.messages == [ PromptMessage( @@ -204,6 +211,7 @@ class TestRenderPrompt: prompt = Prompt.from_function(MyPrompt()) manager.add_prompt(prompt) result = await manager.render_prompt("MyPrompt", arguments={"name": "World"}) + assert isinstance(result, PromptResult) assert result.description == "A callable object that can be used as a prompt." assert result.messages == [ PromptMessage( @@ -225,6 +233,7 @@ class TestRenderPrompt: prompt = Prompt.from_function(MyPrompt()) manager.add_prompt(prompt) result = await manager.render_prompt("MyPrompt", arguments={"name": "World"}) + assert isinstance(result, PromptResult) assert result.description == "A callable object that can be used as a prompt." assert result.messages == [ PromptMessage( @@ -402,6 +411,7 @@ class TestContextHandling: async with context: result = await prompt.render(arguments={"x": 42}) + assert isinstance(result, PromptResult) assert len(result.messages) == 1 assert result.messages[0].content.text == "42" # type: ignore[attr-defined] @@ -424,6 +434,7 @@ class TestContextHandling: arguments={"x": 42}, ) + assert isinstance(result, PromptResult) assert len(result.messages) == 1 assert result.messages[0].content.text == "42" # type: ignore[attr-defined] @@ -461,4 +472,5 @@ class TestContextHandling: async with context: result = await prompt.render(arguments={"topic": "cats"}) + assert isinstance(result, PromptResult) assert result.messages[0].content.text == "Write about cats" # type: ignore[attr-defined] diff --git a/tests/resources/test_resource_manager.py b/tests/resources/test_resource_manager.py index 50d72018d..8b9e48d39 100644 --- a/tests/resources/test_resource_manager.py +++ b/tests/resources/test_resource_manager.py @@ -304,6 +304,7 @@ class TestResourceManager: resource = await manager.get_resource(AnyUrl("greet://world")) assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == "Hello, world!" async def test_get_unknown_resource(self): @@ -481,6 +482,7 @@ class TestQueryOnlyTemplates: # Should also work via read_resource result = await manager.read_resource("data://config") + assert isinstance(result, ResourceContent) assert result.content == "Config in json format" async def test_template_with_only_query_params_with_query_string(self): @@ -506,6 +508,7 @@ class TestQueryOnlyTemplates: # Should also work via read_resource result = await manager.read_resource("data://config?format=xml") + assert isinstance(result, ResourceContent) assert result.content == "Config in xml format" async def test_template_with_only_multiple_query_params(self): @@ -524,14 +527,17 @@ class TestQueryOnlyTemplates: # No query params - use all defaults result = await manager.read_resource("data://items") + assert isinstance(result, ResourceContent) assert result.content == "Data in json (limit: 10)" # Partial query params result = await manager.read_resource("data://items?format=xml") + assert isinstance(result, ResourceContent) assert result.content == "Data in xml (limit: 10)" # All query params result = await manager.read_resource("data://items?format=xml&limit=20") + assert isinstance(result, ResourceContent) assert result.content == "Data in xml (limit: 20)" async def test_has_resource_with_query_only_template(self): diff --git a/tests/resources/test_resource_template.py b/tests/resources/test_resource_template.py index aff30ef7c..21979c01d 100644 --- a/tests/resources/test_resource_template.py +++ b/tests/resources/test_resource_template.py @@ -184,6 +184,7 @@ class TestResourceTemplate: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert isinstance(result.content, str) data = json.loads(result.content) assert data == {"key": "foo", "value": 123} @@ -207,6 +208,7 @@ class TestResourceTemplate: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == "Hello, world!" async def test_async_binary_resource(self): @@ -228,6 +230,7 @@ class TestResourceTemplate: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == b"test" async def test_basemodel_conversion(self): @@ -253,6 +256,7 @@ class TestResourceTemplate: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert isinstance(result.content, str) data = json.loads(result.content) assert data == {"key": "foo", "value": 123} @@ -283,6 +287,7 @@ class TestResourceTemplate: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == '"hello"' async def test_wildcard_param_can_create_resource(self): @@ -393,6 +398,7 @@ class TestResourceTemplate: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == "X was foo" @@ -679,6 +685,7 @@ class TestContextHandling: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == "42" async def test_context_optional(self): @@ -705,6 +712,7 @@ class TestContextHandling: assert isinstance(resource, FunctionResource) result = await resource.read() + assert isinstance(result, ResourceContent) assert result.content == "42" async def test_context_with_functools_wraps_decorator(self): diff --git a/tests/server/tasks/test_custom_subclass_tasks.py b/tests/server/tasks/test_custom_subclass_tasks.py index 39538836d..3bb6f834a 100644 --- a/tests/server/tasks/test_custom_subclass_tasks.py +++ b/tests/server/tasks/test_custom_subclass_tasks.py @@ -168,7 +168,7 @@ class TestFastMCPComponentDocketMethods: component = FastMCPComponent(name="test") mock_docket = MagicMock() - with pytest.raises(RuntimeError, match="task_config.mode is 'forbidden'"): + with pytest.raises(RuntimeError, match="task execution not supported"): await component.add_to_docket(mock_docket) async def test_add_to_docket_raises_not_implemented_when_allowed(self):