diff --git a/fastmcp_tasks/fastmcp_tasks/components.py b/fastmcp_tasks/fastmcp_tasks/components.py index 69b301e28..3fb7e5f01 100644 --- a/fastmcp_tasks/fastmcp_tasks/components.py +++ b/fastmcp_tasks/fastmcp_tasks/components.py @@ -59,7 +59,9 @@ def register_component_with_docket(component: FastMCPComponent, docket: Docket) # InputRequiredResult drives the reentrant in-task input cycle. The # wrapper is signature-preserving, so Docket's dependency injection is # unchanged for a body that never asks for input. - docket.register(reentrant_task_fn(component.fn), names=[component.key]) + docket.register( + reentrant_task_fn(component.fn, component.name), names=[component.key] + ) elif isinstance(component, Tool): docket.register(component.run, names=[component.key]) elif isinstance(component, FunctionResource): diff --git a/fastmcp_tasks/fastmcp_tasks/extension.py b/fastmcp_tasks/fastmcp_tasks/extension.py index 45ea7f6ec..aff882736 100644 --- a/fastmcp_tasks/fastmcp_tasks/extension.py +++ b/fastmcp_tasks/fastmcp_tasks/extension.py @@ -35,7 +35,11 @@ from mcp.shared.exceptions import MCPError from mcp_types.version import MODERN_PROTOCOL_VERSIONS from fastmcp.exceptions import NotFoundError -from fastmcp.server.extensions import MethodBinding, ServerExtension +from fastmcp.server.extensions import ( + MethodBinding, + ServerExtension, + read_client_extension_settings, +) from fastmcp.utilities.logging import get_logger from fastmcp.utilities.tasks import TASKS_EXTENSION_ID from fastmcp_tasks.creation import create_task @@ -130,19 +134,42 @@ class TasksExtension(ServerExtension): ), ] + def _require_tasks_capability(self, ctx: ServerRequestContext[Any, Any]) -> None: + """Reject a task method from a client that did not declare the extension. + + SEP-2663: a client issuing `tasks/get`/`tasks/update`/`tasks/cancel` + without the tasks capability in the request's `_meta` gets -32003. A + client normally only holds a taskId because it declared the capability + on the creating `tools/call`, but the method-level check is an explicit + MUST, so enforce it here rather than assume. + """ + if read_client_extension_settings(ctx, TASKS_EXTENSION_ID) is None: + raise MCPError( + code=MISSING_REQUIRED_CLIENT_CAPABILITY, + message=( + "This request targets the tasks extension " + f"({TASKS_EXTENSION_ID}); the client did not declare it for " + "this request." + ), + data=missing_capability_error_data(), + ) + async def _handle_get( self, ctx: ServerRequestContext[Any, Any], params: GetTaskParams ) -> GetTaskResult: + self._require_tasks_capability(ctx) return await tasks_get(self.server, params.task_id) async def _handle_update( self, ctx: ServerRequestContext[Any, Any], params: UpdateTaskParams ) -> UpdateTaskResult: + self._require_tasks_capability(ctx) return await tasks_update(self.server, params.task_id, params.input_responses) async def _handle_cancel( self, ctx: ServerRequestContext[Any, Any], params: CancelTaskParams ) -> CancelTaskResult: + self._require_tasks_capability(ctx) return await tasks_cancel(self.server, params.task_id) async def intercept_tool_call( diff --git a/fastmcp_tasks/fastmcp_tasks/handlers.py b/fastmcp_tasks/fastmcp_tasks/handlers.py index c65929013..c2f1bd8f4 100644 --- a/fastmcp_tasks/fastmcp_tasks/handlers.py +++ b/fastmcp_tasks/fastmcp_tasks/handlers.py @@ -29,16 +29,18 @@ from mcp.shared.exceptions import MCPError from mcp_types import INVALID_PARAMS from fastmcp.exceptions import NotFoundError -from fastmcp.tools.base import InputRequiredToolResult, Tool +from fastmcp.tools.base import InputRequiredToolResult, Tool, ToolResult from fastmcp.utilities.tasks import DEFAULT_POLL_INTERVAL_MS from fastmcp.utilities.versions import VersionSpec from fastmcp_tasks.context import get_task_scope from fastmcp_tasks.creation import enqueue_task_leg, registered_component_for_key from fastmcp_tasks.input_store import ( + acquire_update_lock, clear_outstanding, load_current_leg, load_task_args, read_outstanding_inputs, + release_update_lock, save_current_leg, store_input_responses, translate_responses, @@ -190,7 +192,13 @@ def _inline_result(tool: Tool, raw_value: Any) -> dict[str, Any]: "Guard-pattern tasks are supported for function tools." ), ) - mcp_result = tool.convert_result(raw_value).to_mcp_result() + # A raised tool error arrives as an is_error ToolResult the wrapper built + # (end-and-reenter G2); use it directly so isError round-trips. A normal + # return is converted through the tool's own result coercion. + if isinstance(raw_value, ToolResult): + mcp_result = raw_value.to_mcp_result() + else: + mcp_result = tool.convert_result(raw_value).to_mcp_result() if isinstance(mcp_result, mcp_types.CallToolResult): call_tool_result = mcp_result elif isinstance(mcp_result, tuple): @@ -299,36 +307,44 @@ async def tasks_update( docket, task_scope, task_id ) - translated = await translate_responses( - docket, task_scope, task_id, leg_number, input_responses - ) - if translated is None: - # Nothing matched the current leg's outstanding requests: the leg was - # already answered, or the keys are unknown. Idempotent no-op. + # Serialize concurrent updates for this task so two racing answers cannot + # each enqueue a next leg (double execution). A loser is an idempotent no-op. + if not await acquire_update_lock(docket, task_scope, task_id): return UpdateTaskResult() + try: + translated = await translate_responses( + docket, task_scope, task_id, leg_number, input_responses + ) + if translated is None: + # Nothing matched the current leg's outstanding requests: the leg was + # already answered, or the keys are unknown. Idempotent no-op. + return UpdateTaskResult() - # Store the answers for the next leg to read, then enqueue that leg. Ordering - # matters: the answers must be in Redis before the next leg's worker context - # loads them, and current_leg must not advance to an execution that is not - # yet durable — so enqueue (with its durable wait) precedes the pointer swap. - await store_input_responses(docket, task_scope, task_id, translated) + # Store the answers for the next leg to read, then enqueue that leg. + # Ordering matters: the answers must be in Redis before the next leg's + # worker context loads them, and current_leg must not advance to an + # execution that is not yet durable — so enqueue (with its durable wait) + # precedes the pointer swap. + await store_input_responses(docket, task_scope, task_id, translated) - component = await registered_component_for_key( - server, parse_task_key(base_task_key)["component_identifier"] - ) - raw_arguments = await load_task_args(docket, task_scope, task_id) - next_leg = leg_number + 1 - next_leg_key = leg_execution_key(base_task_key, next_leg) + component = await registered_component_for_key( + server, parse_task_key(base_task_key)["component_identifier"] + ) + raw_arguments = await load_task_args(docket, task_scope, task_id) + next_leg = leg_number + 1 + next_leg_key = leg_execution_key(base_task_key, next_leg) - await enqueue_task_leg(server, docket, component, raw_arguments, next_leg_key) - ttl_seconds = int(docket.execution_ttl.total_seconds()) - await save_current_leg( - docket, task_scope, task_id, next_leg_key, next_leg, ttl_seconds - ) - # The answered leg's surfaced keys are now superseded; drop them so they are - # never reused (SEP-2663 L350). - await clear_outstanding(docket, task_scope, task_id, leg_number) - return UpdateTaskResult() + await enqueue_task_leg(server, docket, component, raw_arguments, next_leg_key) + ttl_seconds = int(docket.execution_ttl.total_seconds()) + await save_current_leg( + docket, task_scope, task_id, next_leg_key, next_leg, ttl_seconds + ) + # The answered leg's surfaced keys are now superseded; drop them so they + # are never reused (SEP-2663 L350). + await clear_outstanding(docket, task_scope, task_id, leg_number) + return UpdateTaskResult() + finally: + await release_update_lock(docket, task_scope, task_id) async def tasks_cancel(server: FastMCP, task_id: str) -> CancelTaskResult: diff --git a/fastmcp_tasks/fastmcp_tasks/input_loop.py b/fastmcp_tasks/fastmcp_tasks/input_loop.py index ed36def55..ff08a9439 100644 --- a/fastmcp_tasks/fastmcp_tasks/input_loop.py +++ b/fastmcp_tasks/fastmcp_tasks/input_loop.py @@ -32,7 +32,8 @@ from typing import TYPE_CHECKING, Any import mcp_types -from fastmcp.tools.base import InputRequiredToolResult +from fastmcp.exceptions import FastMCPError +from fastmcp.tools.base import InputRequiredToolResult, ToolResult from fastmcp_tasks.context import get_task_context, get_task_leg_number from fastmcp_tasks.input_store import store_outstanding @@ -82,8 +83,40 @@ def _resolve_docket() -> Docket | None: return docket +def _mask_error_details() -> bool: + """The worker server's error-masking policy, mirroring the sync call path.""" + import fastmcp + from fastmcp.server.dependencies import get_context + + try: + return get_context().fastmcp._mask_error_details + except RuntimeError: + return fastmcp.settings.mask_error_details + + +def _error_result(tool_name: str, exc: Exception) -> ToolResult: + """An ``is_error`` result for a task tool that raised, mirroring foreground. + + A raised tool error is a *completed* task carrying an error result, never a + ``failed`` task (SEP-2663 reserves ``failed`` for protocol faults, and a live + ``tools/call`` returns the same `isError` result). A `FastMCPError` (e.g. + ``ToolError``) reaches the client verbatim, as the synchronous path re-raises + it unmasked; any other exception is masked per the server's policy. + """ + if isinstance(exc, FastMCPError): + message = str(exc) + elif _mask_error_details(): + message = f"Error calling tool {tool_name!r}" + else: + message = f"Error calling tool {tool_name!r}: {exc}" + return ToolResult( + content=[mcp_types.TextContent(type="text", text=message)], is_error=True + ) + + def reentrant_task_fn( fn: Callable[..., Awaitable[Any]], + tool_name: str, ) -> Callable[..., Awaitable[Any]]: """Wrap a task tool's callable to capture a guard leg's ask (end-and-reenter). @@ -91,12 +124,20 @@ def reentrant_task_fn( unchanged. The body runs exactly once: a real return is the leg's result; an `InputRequiredResult` is captured to Redis (outstanding requests + carried state) and the wrapper returns, ending the leg without blocking. The next - leg is enqueued by ``tasks/update`` when the client answers. + leg is enqueued by ``tasks/update`` when the client answers. A raised tool + error becomes a completed `is_error` result (not a failed task), matching the + synchronous `tools/call` path. """ @functools.wraps(fn) async def wrapper(*args: Any, **kwargs: Any) -> Any: - result = await fn(*args, **kwargs) + try: + result = await fn(*args, **kwargs) + except FastMCPError as exc: + return _error_result(tool_name, exc) + except Exception as exc: + logger.exception("background task tool %r raised", tool_name) + return _error_result(tool_name, exc) input_required = _as_input_required(result) if input_required is None: return result diff --git a/fastmcp_tasks/fastmcp_tasks/input_store.py b/fastmcp_tasks/fastmcp_tasks/input_store.py index 48af941ea..761a3806d 100644 --- a/fastmcp_tasks/fastmcp_tasks/input_store.py +++ b/fastmcp_tasks/fastmcp_tasks/input_store.py @@ -359,6 +359,44 @@ async def clear_outstanding( await redis.delete(_map_key(docket, task_scope, task_id, leg)) +# How long the per-task update lock lives if its holder dies mid-update. A +# generous ceiling: a single tasks/update is fast, so the lock is normally held +# for milliseconds; the TTL only guards against a crashed holder. +_UPDATE_LOCK_TTL_SECONDS = 30 + + +def _update_lock_key(docket: Docket, task_scope: str | None, task_id: str) -> str: + return docket.key(f"{_prefix(docket, task_scope, task_id)}:update_lock") + + +async def acquire_update_lock( + docket: Docket, task_scope: str | None, task_id: str +) -> bool: + """Take the per-task update lock, or return False if one is already held. + + Serializes concurrent ``tasks/update`` calls for a task so two racing + answers cannot each enqueue a next leg (double execution). A well-behaved + client polls sequentially and never contends; a loser is an idempotent + no-op, matching SEP-2663's "ignore already-satisfied" rule. + """ + async with docket.redis() as redis: + got = await redis.set( + _update_lock_key(docket, task_scope, task_id), + b"1", + nx=True, + ex=_UPDATE_LOCK_TTL_SECONDS, + ) + return bool(got) + + +async def release_update_lock( + docket: Docket, task_scope: str | None, task_id: str +) -> None: + """Release the per-task update lock.""" + async with docket.redis() as redis: + await redis.delete(_update_lock_key(docket, task_scope, task_id)) + + async def load_pending_input( docket: Docket, task_scope: str | None, task_id: str ) -> tuple[str | None, mcp_types.InputResponses | None]: diff --git a/tests/tasks/client/test_transparent_tasks.py b/tests/tasks/client/test_transparent_tasks.py index a0b95301e..43c3eb6eb 100644 --- a/tests/tasks/client/test_transparent_tasks.py +++ b/tests/tasks/client/test_transparent_tasks.py @@ -142,9 +142,7 @@ async def test_in_task_input_answered_transparently(guard_server: FastMCP): async def handle_elicitation(message, response_type, params, context): return DinnerPrefs(cuisine="Thai", vegetarian=True) - client = Client( - guard_server, mode="auto", elicitation_handler=handle_elicitation - ) + client = Client(guard_server, mode="auto", elicitation_handler=handle_elicitation) async with client: result = await client.call_tool("plan_dinner", {}) diff --git a/tests/tasks/server/test_context_background_task.py b/tests/tasks/server/test_context_background_task.py index a4ab4ed61..bdc7b939d 100644 --- a/tests/tasks/server/test_context_background_task.py +++ b/tests/tasks/server/test_context_background_task.py @@ -393,8 +393,11 @@ class TestBackgroundTaskIntegration: } async def test_imperative_elicit_fails_with_guard_guidance(self): - """A task=True tool that calls ctx.elicit() fails with the guard-pattern - error rather than parking a worker on a client round-trip.""" + """A task=True tool that calls ctx.elicit() errors with guard guidance. + + The ToolError it raises surfaces as a completed is_error result (like any + raised tool error, SEP-2663), never parking a worker on a round-trip. + """ mcp = FastMCP("elicit-forbidden") mcp.add_extension(TasksExtension()) @@ -407,9 +410,10 @@ class TestBackgroundTaskIntegration: created = await submit_task(mcp, "ask_name", {}) final = await wait_for_task(mcp, created.task_id) - assert final.status == "failed" - assert final.error is not None - assert "InputRequiredResult" in final.error["message"] + assert final.status == "completed" + assert final.result is not None + assert final.result["isError"] is True + assert "InputRequiredResult" in final.result["content"][0]["text"] class TestAccessTokenInBackgroundTasks: diff --git a/tests/tasks/server/test_extension.py b/tests/tasks/server/test_extension.py index 9d9d49029..eda1b557b 100644 --- a/tests/tasks/server/test_extension.py +++ b/tests/tasks/server/test_extension.py @@ -14,16 +14,18 @@ from contextlib import AsyncExitStack from types import SimpleNamespace from typing import cast +import mcp_types import pytest from fastmcp_tasks.models import ( MISSING_REQUIRED_CLIENT_CAPABILITY, CreateTaskResult, + GetTaskParams, ) from mcp.server.context import ServerRequestContext from mcp.server.session import ServerSession from mcp.shared.exceptions import MCPError -from fastmcp import FastMCP +from fastmcp import Context, FastMCP from fastmcp.client import Client from fastmcp.exceptions import ToolError from fastmcp.server.dependencies import bind_request_context @@ -40,6 +42,7 @@ from tests.tasks.task_helpers import ( run_task, running_task_server, submit_task, + update_task, wait_for_task, ) @@ -178,15 +181,22 @@ async def test_get_unknown_task_raises_not_found(): await get_task(mcp, "does-not-exist") -async def test_failed_task_surfaces_error_not_completed(): +async def test_raised_tool_error_completes_with_is_error(): + """A tool that RAISES is a completed task with an is_error result, not failed. + + SEP-2663 reserves `failed` for protocol faults; a raised tool error is the + same `isError` CallToolResult a live tools/call returns (the task path must + return exactly what the underlying request would). + """ mcp = _tasks_server() async with running_task_server(mcp): created = await submit_task(mcp, "boom", {}) final = await wait_for_task(mcp, created.task_id) - assert final.status == "failed" - assert final.error is not None - assert "kaboom" in final.error["message"] - assert final.result is None + assert final.status == "completed" + assert final.error is None + assert final.result is not None + assert final.result["isError"] is True + assert "kaboom" in final.result["content"][0]["text"] # --------------------------------------------------------------------------- @@ -369,3 +379,74 @@ async def test_worker_hooks_survive_sibling_server_shutdown(): assert core_dependencies._background_context_factory is not None # The last extension exited; hooks are cleared. assert core_dependencies._background_context_factory is None + + +# --------------------------------------------------------------------------- +# Compliance: -32003 on task methods for non-declaring clients (SEP-2663) +# --------------------------------------------------------------------------- + + +async def test_task_method_without_capability_raises_missing_capability(): + """tasks/get from a client that did not declare the extension gets -32003.""" + mcp = _tasks_server() + extension = cast(TasksExtension, mcp._extensions[TASKS_EXTENSION_ID]) + # A request context with no tasks capability in its _meta. + srctx = ServerRequestContext( + session=cast(ServerSession, SimpleNamespace()), + lifespan_context={}, + protocol_version="2026-07-28", + method="tasks/get", + params={"taskId": "whatever"}, + ) + params = GetTaskParams.model_validate({"taskId": "whatever"}) + async with running_task_server(mcp): + with pytest.raises(MCPError) as exc_info: + await extension._handle_get(srctx, params) + assert exc_info.value.error.code == MISSING_REQUIRED_CLIENT_CAPABILITY + + +# --------------------------------------------------------------------------- +# Compliance: concurrent tasks/update must not enqueue two next legs +# --------------------------------------------------------------------------- + + +async def test_concurrent_update_enqueues_a_single_next_leg(): + """Two racing tasks/update answers re-enter the task exactly once.""" + calls: list[int] = [] + mcp = FastMCP("race") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def guard(ctx: Context) -> str | mcp_types.InputRequiredResult: + calls.append(1) + if ctx.input_responses is None: + req = mcp_types.ElicitRequest( + params=mcp_types.ElicitRequestFormParams( + message="?", requested_schema={"type": "object"} + ) + ) + return mcp_types.InputRequiredResult( + result_type="input_required", input_requests={"k": req} + ) + return "done" + + async with running_task_server(mcp): + created = await submit_task(mcp, "guard", {}) + parked = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"input_required"}) + ) + assert parked.input_requests is not None + key = next(iter(parked.input_requests)) + answer = {key: {"action": "accept", "content": {}}} + # Fire two identical updates concurrently. + await asyncio.gather( + update_task(mcp, created.task_id, answer), + update_task(mcp, created.task_id, answer), + return_exceptions=True, + ) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + # Leg 1 (park) + exactly one re-entered leg 2 — never a third from a double + # enqueue. + assert calls == [1, 1] diff --git a/tests/tasks/server/test_guard_reentrant.py b/tests/tasks/server/test_guard_reentrant.py index 99e53c481..da2422f32 100644 --- a/tests/tasks/server/test_guard_reentrant.py +++ b/tests/tasks/server/test_guard_reentrant.py @@ -170,5 +170,5 @@ def test_reentrant_wrapper_preserves_signature(): async def fn(n: int, ctx: Any) -> int: return n - wrapped = reentrant_task_fn(fn) + wrapped = reentrant_task_fn(fn, "fn") assert list(inspect.signature(wrapped).parameters) == ["n", "ctx"] diff --git a/tests/tasks/server/test_task_methods.py b/tests/tasks/server/test_task_methods.py index d7bc39731..5a463cc28 100644 --- a/tests/tasks/server/test_task_methods.py +++ b/tests/tasks/server/test_task_methods.py @@ -68,15 +68,16 @@ async def test_tasks_get_includes_poll_interval(): assert got.poll_interval_ms == 5000 -async def test_tasks_get_returns_error_for_failed_task(): - """`tasks/get` surfaces the error for a failed task rather than a result.""" +async def test_tasks_get_returns_is_error_result_for_raised_tool(): + """A raised tool error is a completed task with an is_error result (SEP-2663).""" mcp = _methods_server() async with running_task_server(mcp): final = await run_task(mcp, "error_tool", {}) - assert final.status == "failed" - assert final.error is not None - assert "Task failed!" in final.error["message"] - assert final.result is None + assert final.status == "completed" + assert final.error is None + assert final.result is not None + assert final.result["isError"] is True + assert "Task failed!" in final.result["content"][0]["text"] async def test_tasks_get_unknown_id_raises_not_found(): diff --git a/tests/tasks/server/test_task_protocol.py b/tests/tasks/server/test_task_protocol.py index 784e0a51f..62892dcfd 100644 --- a/tests/tasks/server/test_task_protocol.py +++ b/tests/tasks/server/test_task_protocol.py @@ -42,12 +42,17 @@ async def test_task_metadata_includes_task_id_and_ttl(): assert created.ttl_ms is not None and created.ttl_ms > 0 -async def test_failed_task_stores_error(): - """A task whose tool raises reaches `failed` and stores the error.""" +async def test_raised_tool_error_completes_with_is_error(): + """A task whose tool raises completes with an is_error result (SEP-2663). + + `failed` is reserved for protocol faults; a raised tool error is the same + `isError` result a live tools/call returns. + """ mcp = _task_server() async with running_task_server(mcp): final = await run_task(mcp, "failing_tool", {}) - assert final.status == "failed" - assert final.error is not None - assert "This tool always fails" in final.error["message"] - assert final.result is None + assert final.status == "completed" + assert final.error is None + assert final.result is not None + assert final.result["isError"] is True + assert "This tool always fails" in final.result["content"][0]["text"]