fix(mcp): compensate cancelled task submissions (#4933)

* fix(mcp): compensate cancelled task submissions

* fix(mcp): shield submission compensation

* docs(mcp): preserve notification lifecycle contract

* fix(mcp): bound submission compensation wait

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

backend/app/mcp_tasks/service.py

@@ -33,6 +33,7 @@ logger = logging.getLogger(__name__)
 _MAX_PERSISTED_ERROR_CHARS = 4_000
 _MAX_INPUT_REQUIRED_BYTES = 65_536
 _MAX_NOTIFICATION_ATTEMPTS = 5
+_UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS = 5.0
 
 
 def _bound_error(error: str | None) -> str | None:
@@ -74,6 +75,7 @@ class McpTaskService:
         self._get_run = get_run
         self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}"
         self._task: asyncio.Task[None] | None = None
+        self._compensation_tasks: set[asyncio.Task[Any]] = set()
         self._stop = asyncio.Event()
 
     @property
@@ -138,17 +140,79 @@ class McpTaskService:
             # This handle already has a durable owner. Cancelling it as
             # compensation would terminate the pre-existing tracked task.
             raise
-        except Exception:
-            try:
-                await driver.cancel(task_reference)
-            except Exception:  # noqa: BLE001 - preserve the original persistence failure
-                logger.exception(
-                    "Failed to cancel untracked MCP task after persistence failure (task_id=%s, driver=%s, remote_task_id=%s)",
-                    local_task_id,
-                    driver_name,
-                    submission.remote_task_id,
-                )
+        except asyncio.CancelledError:
+            # Cancellation can race with a successful database commit. If it
+            # did, the durable row will converge to cancelled on its next poll;
+            # compensating is safer than leaving a live remote task untracked.
+            await self._cancel_untracked_task(
+                driver=driver,
+                task_reference=task_reference,
+                driver_name=driver_name,
+                reason="caller cancellation during local persistence",
+            )
             raise
+        except Exception:
+            await self._cancel_untracked_task(
+                driver=driver,
+                task_reference=task_reference,
+                driver_name=driver_name,
+                reason="local submission finalization failure",
+            )
+            raise
+
+    async def _cancel_untracked_task(
+        self,
+        *,
+        driver,
+        task_reference: TaskReference,
+        driver_name: str,
+        reason: str,
+    ) -> None:
+        compensation = asyncio.create_task(
+            driver.cancel(task_reference),
+            name=f"mcp-submit-compensation-{task_reference.local_task_id}",
+        )
+        self._compensation_tasks.add(compensation)
+
+        def finalize(task: asyncio.Task[Any]) -> None:
+            self._compensation_tasks.discard(task)
+            try:
+                error = task.exception()
+            except asyncio.CancelledError as exc:
+                error = exc
+            if error is None:
+                return
+            logger.error(
+                "Failed to cancel untracked MCP task after %s (task_id=%s, driver=%s, remote_task_id=%s)",
+                reason,
+                task_reference.local_task_id,
+                driver_name,
+                task_reference.remote_task_id,
+                exc_info=(type(error), error, error.__traceback__),
+            )
+
+        compensation.add_done_callback(finalize)
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + _UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS
+        while not compensation.done():
+            remaining = deadline - loop.time()
+            if remaining <= 0:
+                logger.warning(
+                    "Timed out after %.1f seconds waiting for untracked MCP task compensation after %s; cancellation continues in the background (task_id=%s, driver=%s, remote_task_id=%s)",
+                    _UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS,
+                    reason,
+                    task_reference.local_task_id,
+                    driver_name,
+                    task_reference.remote_task_id,
+                )
+                return
+            try:
+                await asyncio.wait({compensation}, timeout=remaining)
+            except asyncio.CancelledError:
+                # Repeated caller cancellation does not propagate through
+                # asyncio.wait() to the compensation task. Keep waiting only
+                # until the original deadline.
+                continue
 
     async def run_once(self, *, now: datetime) -> None:
         await self._run_cancellations(now=now)

backend/packages/harness/deerflow/mcp/AGENTS.md

@@ -1,7 +1,7 @@
 ### MCP System (`packages/harness/deerflow/mcp/`)
 
 - Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management
-- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails, the service best-effort cancels the remote task and preserves the original persistence error if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own.
+- **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own.
 - **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts.
 - **Long-running ordinary task driver**: `extensions_config.json -> mcpServers.<server>.task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same `(server_name, user_id:thread_id)` stdio session scope; HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task.
 - **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior.

backend/tests/test_mcp_task_service.py

@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock
 
 import pytest
 
+import app.mcp_tasks.service as service_module
 from app.mcp_tasks.errors import PermanentNotificationError
 from app.mcp_tasks.service import McpTaskService
 from deerflow.mcp.tasks import (
@@ -61,6 +62,17 @@ class FailingCreateRepository(FakeRepository):
         raise RuntimeError("database unavailable")
 
 
+class BlockingCreateRepository(FakeRepository):
+    def __init__(self):
+        super().__init__()
+        self.create_started = asyncio.Event()
+
+    async def create(self, **kwargs):
+        self.created.append(kwargs)
+        self.create_started.set()
+        await asyncio.Event().wait()
+
+
 class DuplicateCreateRepository(FakeRepository):
     async def create(self, **kwargs):
         self.created.append(kwargs)
@@ -119,6 +131,29 @@ class HangingDriver(FakeDriver):
             raise
 
 
+class BlockingCancelDriver(FakeDriver):
+    def __init__(self, *, submission):
+        super().__init__(submission=submission)
+        self.cancel_started = asyncio.Event()
+        self.finish_cancel = asyncio.Event()
+        self.cancel_finished = asyncio.Event()
+        self.cancel_completed = False
+        self.cancel_interrupted = False
+
+    async def cancel(self, task):
+        self.cancel_calls.append(task)
+        self.cancel_started.set()
+        try:
+            await self.finish_cancel.wait()
+        except asyncio.CancelledError:
+            self.cancel_interrupted = True
+            self.cancel_finished.set()
+            raise
+        self.cancel_completed = True
+        self.cancel_finished.set()
+        return TaskSnapshot(status=TaskStatus.CANCELLED)
+
+
 def _claimed_row(*, driver_name="fake"):
     return {
         "id": "task-1",
@@ -221,6 +256,188 @@ async def test_submit_cancels_remote_task_when_persistence_fails():
     }
 
 
+@pytest.mark.asyncio
+async def test_submit_cancellation_during_persistence_cancels_remote_task():
+    repo = BlockingCreateRepository()
+    driver = FakeDriver(
+        submission=TaskSubmission(
+            remote_task_id="remote-1",
+            snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
+            driver_data={"cancel_tool": "cancel"},
+        )
+    )
+    registry = McpTaskDriverRegistry()
+    registry.register("fake", driver)
+    service = McpTaskService(
+        repository=repo,
+        drivers=registry,
+        poll_interval_seconds=5,
+        lease_seconds=120,
+        max_concurrent_polls=3,
+    )
+    request = TaskSubmitRequest(
+        user_id="user-1",
+        thread_id="thread-1",
+        run_id="run-1",
+        tool_call_id="call-1",
+        server_name="reports",
+        task_name="Generate report",
+        arguments={"topic": "MCP"},
+        local_task_id="task-1",
+    )
+
+    submit_task = asyncio.create_task(service.submit(driver_name="fake", request=request))
+    await repo.create_started.wait()
+    submit_task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await submit_task
+
+    assert len(driver.cancel_calls) == 1
+    cancelled = driver.cancel_calls[0]
+    assert cancelled.local_task_id == "task-1"
+    assert cancelled.remote_task_id == "remote-1"
+
+
+@pytest.mark.asyncio
+async def test_submit_repeated_cancellation_does_not_interrupt_compensation():
+    repo = BlockingCreateRepository()
+    driver = BlockingCancelDriver(
+        submission=TaskSubmission(
+            remote_task_id="remote-1",
+            snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
+            driver_data={"cancel_tool": "cancel"},
+        )
+    )
+    registry = McpTaskDriverRegistry()
+    registry.register("fake", driver)
+    service = McpTaskService(
+        repository=repo,
+        drivers=registry,
+        poll_interval_seconds=5,
+        lease_seconds=120,
+        max_concurrent_polls=3,
+    )
+    request = TaskSubmitRequest(
+        user_id="user-1",
+        thread_id="thread-1",
+        run_id="run-1",
+        tool_call_id="call-1",
+        server_name="reports",
+        task_name="Generate report",
+        arguments={"topic": "MCP"},
+        local_task_id="task-1",
+    )
+
+    submit_task = asyncio.create_task(service.submit(driver_name="fake", request=request))
+    await repo.create_started.wait()
+    submit_task.cancel()
+    await driver.cancel_started.wait()
+
+    submit_task.cancel()
+    driver.finish_cancel.set()
+
+    with pytest.raises(asyncio.CancelledError):
+        await submit_task
+
+    assert len(driver.cancel_calls) == 1
+    assert driver.cancel_completed
+    assert not driver.cancel_interrupted
+
+
+@pytest.mark.asyncio
+async def test_submit_stops_waiting_for_hung_compensation_without_cancelling_it(monkeypatch, caplog):
+    monkeypatch.setattr(service_module, "_UNTRACKED_TASK_COMPENSATION_WAIT_SECONDS", 0)
+    repo = BlockingCreateRepository()
+    driver = BlockingCancelDriver(
+        submission=TaskSubmission(
+            remote_task_id="remote-1",
+            snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
+            driver_data={"cancel_tool": "cancel"},
+        )
+    )
+    registry = McpTaskDriverRegistry()
+    registry.register("fake", driver)
+    service = McpTaskService(
+        repository=repo,
+        drivers=registry,
+        poll_interval_seconds=5,
+        lease_seconds=120,
+        max_concurrent_polls=3,
+    )
+    request = TaskSubmitRequest(
+        user_id="user-1",
+        thread_id="thread-1",
+        run_id="run-1",
+        tool_call_id="call-1",
+        server_name="reports",
+        task_name="Generate report",
+        arguments={"topic": "MCP"},
+        local_task_id="task-1",
+    )
+
+    submit_task = asyncio.create_task(service.submit(driver_name="fake", request=request))
+    await repo.create_started.wait()
+    submit_task.cancel()
+
+    with caplog.at_level(logging.WARNING), pytest.raises(asyncio.CancelledError):
+        await submit_task
+
+    assert "cancellation continues in the background" in caplog.text
+    await driver.cancel_started.wait()
+    assert not driver.cancel_interrupted
+    assert not driver.cancel_completed
+
+    driver.finish_cancel.set()
+    await driver.cancel_finished.wait()
+
+    assert len(driver.cancel_calls) == 1
+    assert driver.cancel_completed
+    assert not driver.cancel_interrupted
+
+
+@pytest.mark.asyncio
+async def test_submit_cancellation_preserves_cancelled_error_when_compensation_fails(caplog):
+    repo = BlockingCreateRepository()
+    driver = FakeDriver(
+        submission=TaskSubmission(
+            remote_task_id="remote-1",
+            snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED),
+        ),
+        cancel_error=RuntimeError("cancel unavailable"),
+    )
+    registry = McpTaskDriverRegistry()
+    registry.register("fake", driver)
+    service = McpTaskService(
+        repository=repo,
+        drivers=registry,
+        poll_interval_seconds=5,
+        lease_seconds=120,
+        max_concurrent_polls=3,
+    )
+    request = TaskSubmitRequest(
+        user_id="user-1",
+        thread_id="thread-1",
+        run_id="run-1",
+        tool_call_id="call-1",
+        server_name="reports",
+        task_name="Generate report",
+        arguments={},
+        local_task_id="task-1",
+    )
+
+    submit_task = asyncio.create_task(service.submit(driver_name="fake", request=request))
+    await repo.create_started.wait()
+    submit_task.cancel()
+
+    with caplog.at_level(logging.ERROR), pytest.raises(asyncio.CancelledError):
+        await submit_task
+
+    assert len(driver.cancel_calls) == 1
+    assert "Failed to cancel untracked MCP task" in caplog.text
+    assert "cancel unavailable" in caplog.text
+
+
 @pytest.mark.asyncio
 async def test_submit_cancels_remote_task_when_its_id_exceeds_storage_limit():
     repo = FakeRepository()