diff --git a/fastmcp_tasks/fastmcp_tasks/handlers.py b/fastmcp_tasks/fastmcp_tasks/handlers.py index c2f1bd8f4..61006937c 100644 --- a/fastmcp_tasks/fastmcp_tasks/handlers.py +++ b/fastmcp_tasks/fastmcp_tasks/handlers.py @@ -37,8 +37,10 @@ from fastmcp_tasks.creation import enqueue_task_leg, registered_component_for_ke from fastmcp_tasks.input_store import ( acquire_update_lock, clear_outstanding, + is_cancelled, load_current_leg, load_task_args, + mark_cancelled, read_outstanding_inputs, release_update_lock, save_current_leg, @@ -247,6 +249,12 @@ async def tasks_get(server: FastMCP, task_id: str) -> GetTaskResult: **payload, ) + # A logical cancellation wins over the underlying execution state: a task + # parked on input has a COMPLETED execution, so without this the branches + # below would report input_required (or completed) for a cancelled task. + if await is_cancelled(docket, task_scope, task_id): + return build("cancelled") + if execution.state == ExecutionState.COMPLETED: # A guard leg ends its Docket execution and records outstanding input # requests to Redis: a completed leg with outstanding requests is the @@ -312,6 +320,12 @@ async def tasks_update( if not await acquire_update_lock(docket, task_scope, task_id): return UpdateTaskResult() try: + # A cancelled task never re-enters: clearing outstanding on cancel makes + # translate return None already, but check explicitly so a cancel that + # races between this update's lookup and lock acquisition still wins. + if await is_cancelled(docket, task_scope, task_id): + return UpdateTaskResult() + translated = await translate_responses( docket, task_scope, task_id, leg_number, input_responses ) @@ -348,14 +362,25 @@ async def tasks_update( async def tasks_cancel(server: FastMCP, task_id: str) -> CancelTaskResult: - """Handle ``tasks/cancel``: cooperatively cancel the current leg, empty ack.""" + """Handle ``tasks/cancel``: cooperatively cancel the task, empty ack. + + A durable cancellation marker is recorded so the logical task reports + ``cancelled`` and refuses re-entry even when it is parked on input — whose + current Docket execution is already ``COMPLETED``, making ``docket.cancel`` + on it a no-op. The current leg's outstanding requests are cleared so a + racing ``tasks/update`` naming them finds nothing, and the running + execution is still cancelled cooperatively for the ``working`` case. + """ docket = server._docket if docket is None: raise _task_not_found(task_id) task_scope = get_task_scope() - execution, _base_task_key, _leg, _created_at, _poll = await _lookup_task( + execution, _base_task_key, leg_number, _created_at, _poll = await _lookup_task( docket, task_scope, task_id ) + ttl_seconds = int(docket.execution_ttl.total_seconds()) + await mark_cancelled(docket, task_scope, task_id, ttl_seconds) + await clear_outstanding(docket, task_scope, task_id, leg_number) await docket.cancel(execution.key) return CancelTaskResult() diff --git a/fastmcp_tasks/fastmcp_tasks/input_store.py b/fastmcp_tasks/fastmcp_tasks/input_store.py index 761a3806d..92b56d9bf 100644 --- a/fastmcp_tasks/fastmcp_tasks/input_store.py +++ b/fastmcp_tasks/fastmcp_tasks/input_store.py @@ -359,6 +359,33 @@ async def clear_outstanding( await redis.delete(_map_key(docket, task_scope, task_id, leg)) +def _cancelled_key(docket: Docket, task_scope: str | None, task_id: str) -> str: + return docket.key(f"{_prefix(docket, task_scope, task_id)}:cancelled") + + +async def mark_cancelled( + docket: Docket, task_scope: str | None, task_id: str, ttl_seconds: int +) -> None: + """Record that a task was cancelled at the logical (not per-leg) level. + + An ``input_required`` task's current Docket execution is already + ``COMPLETED`` — the outstanding-input record is what keeps it parked — so + ``docket.cancel`` on that execution is a no-op. This durable marker lets + ``tasks/get`` report ``cancelled`` and ``tasks/update`` refuse to resume, + regardless of the underlying execution state. Expires with the task's TTL. + """ + async with docket.redis() as redis: + await redis.set( + _cancelled_key(docket, task_scope, task_id), b"1", ex=max(1, ttl_seconds) + ) + + +async def is_cancelled(docket: Docket, task_scope: str | None, task_id: str) -> bool: + """Whether the task was logically cancelled (see ``mark_cancelled``).""" + async with docket.redis() as redis: + return bool(await redis.exists(_cancelled_key(docket, task_scope, task_id))) + + # 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. diff --git a/tests/tasks/server/test_guard_reentrant.py b/tests/tasks/server/test_guard_reentrant.py index da2422f32..463de404b 100644 --- a/tests/tasks/server/test_guard_reentrant.py +++ b/tests/tasks/server/test_guard_reentrant.py @@ -18,6 +18,8 @@ import mcp_types from fastmcp import Context, FastMCP from fastmcp_tasks import TasksExtension from tests.tasks.task_helpers import ( + cancel_task, + get_task, running_task_server, submit_task, update_task, @@ -65,6 +67,49 @@ async def _park_key(mcp: FastMCP, task_id: str) -> str: return next(iter(parked.input_requests)) +async def test_cancel_parked_task_reports_cancelled_and_refuses_resume(): + """Cancelling an `input_required` task actually cancels it. + + A parked guard leg's Docket execution is already COMPLETED, so cancelling + only that execution would leave `tasks/get` reporting `input_required` + forever and let a later `tasks/update` resume the task. The logical + cancellation marker must make `tasks/get` report `cancelled` and turn a + subsequent answer into a no-op that never re-enters the tool. + """ + mcp = FastMCP("guard-cancel") + mcp.add_extension(TasksExtension()) + + ran_after_cancel = False + + @mcp.tool(task=True) + async def greet(ctx: Context) -> str | mcp_types.InputRequiredResult: + nonlocal ran_after_cancel + responses = ctx.input_responses + if responses is None: + return _input_required({"name": _elicit_request("Your name?")}) + ran_after_cancel = True + return f"Hello, {_answer(responses, 'name')}!" + + async with running_task_server(mcp): + created = await submit_task(mcp, "greet", {}) + key = await _park_key(mcp, created.task_id) + + await cancel_task(mcp, created.task_id) + cancelled = await get_task(mcp, created.task_id) + assert cancelled.status == "cancelled" + + # Answering a cancelled task is an idempotent no-op: it must not resume. + await update_task( + mcp, + created.task_id, + {key: {"action": "accept", "content": {"value": "Ada"}}}, + ) + still_cancelled = await get_task(mcp, created.task_id) + assert still_cancelled.status == "cancelled" + + assert ran_after_cancel is False + + async def test_guard_return_single_round_completes(): """A tool that returns InputRequiredResult once is driven to completion.""" mcp = FastMCP("guard")