fix(runtime): release cancellation compensations before re-raising CancelledError

[jamespud]

Related: #4933 (sibling merged fix that compensates cancelled task submission); no dedicated issue for the broader class.

Why

asyncio.CancelledError is a BaseException, so any compensation/cleanup placed inside an except Exception block is silently skipped when a task is cancelled.

DeerFlow already encodes the correct pattern in runtime/runs/manager.py ("Also covers cancellation, which bypasses except Exception"), and #4933 fixed one instance in McpTaskService.submit(). But the same gap remains in several other compensation paths: a run's ownership, MCP cancel/notification claims, and the run journal's event buffer are only released on normal errors — so a cancellation leaves them held until lease expiry / orphan reconciliation, and a cancelled journal flush can drop events. This PR closes those remaining leaks so recovery is immediate and no events are lost on cancellation.

What changed

From a caller/operator perspective: when a background run/task in flight is cancelled, its pending compensation now runs before the CancelledError propagates, so ownership/claims are released immediately instead of waiting on a lease lapse, and a cancelled journal flush keeps its batch.

• RunManager finalize path: on cancellation, _mark_ownership_lost(...) is invoked before re-raising.
• McpTaskService task-poll paths: _cancel_one, _notify_one, and the status-poll each release the cancel claim / notification claim / release-after-error before re-raising on cancellation.
• RunJournal _flush_async and the batched flush loop: on cancellation, the unsent batch is re-queued to the buffer before re-raising instead of being dropped.

No default behavior changes for non-cancelled runs; these paths only add the cancellation branch.

Surface area

• Frontend UI
• Backend API (backend/app: mcp_tasks/service.py)
• Agents / LangGraph (harness runtime: runs/manager.py, journal.py)
• Sandbox
• Skills
• Dependencies
• Default behavior change
• Docs / tests / CI only

Screenshots / Recording

N/A (no frontend change).

Bug fix verification

• Test path that reproduces the bug: backend/tests/test_mcp_task_service.py::test_cancel_one_releases_claim_when_cancelled
• Red on main, green on this branch? Yes — I temporarily removed the except asyncio.CancelledError handler from _cancel_one; the test then failed (release_cancel_claim was never called), and it passes again with the fix restored. The journal/manager paths share the same pattern and guarantee.
• Regression test: cancels _cancel_one mid-apply_cancel_snapshot and asserts release_cancel_claim is invoked.

Validation

cd backend
.venv/bin/python -m pytest tests/test_mcp_task_service.py tests/test_run_manager.py tests/test_run_journal.py -q # 195 passed
.venv/bin/ruff check app/mcp_tasks/service.py packages/harness/deerflow/runtime/runs/manager.py packages/harness/deerflow/runtime/journal.py tests/test_mcp_task_service.py # All checks passed
.venv/bin/ruff format --check # already formatted

Pre-commit hooks (ruff lint + format) passed on commit 624f456.

AI assistance

Tool(s) used: Codex

How you used it: Codex analysed the codebase for the except Exception (misses asyncio.CancelledError) compensation pattern via AST scanning plus manual reads, implemented the fixes and regression test following the existing manager.py pattern and the #4933 precedent, ran red→green verification, and ran backend tests and lint. The human (jamespud) directed the scope and reviewed the approach and result.

• I've read and understand every line of this change and take responsibility for it — it's not unreviewed AI output.

[AnnaSuSu]

Thanks for tackling this cancellation-safety gap. I found three correctness blockers that need to be addressed before this is safe to merge:

1. Preserve the original cancellation when compensation fails. The new CancelledError handlers directly await release_cancel_claim, release_notification_claim, or release_claim. If that repository call raises, the cleanup exception replaces the original CancelledError, despite the stated guarantee that cancellation is preserved. In a focused probe that cancels during apply_cancel_snapshot and makes release_cancel_claim fail, current main propagates CancelledError, while this PR propagates RuntimeError. Please protect/drain the cleanup, treat cleanup failure as best-effort, and always re-raise the original cancellation; repeated cancellation should not abandon the cleanup either.

2. Cover the full lifetime of each claimed record, not only the remote call. In _notify_one, the new handler only surrounds _launch_notification; cancellation during the preceding _get_run, the dispatched-status branch, or the following mark_notification_dispatched still leaves the notification lease held. In _poll_one, the handler only surrounds driver.get_status; cancellation during _apply_snapshot still leaves the poll lease held. This is observable through the real service.start() / service.stop() path: cancelling during either persistence stage produces zero release calls because cancellation of the outer gather prevents the post-gather fallback from running. Please make cancellation compensation span every await after the record is claimed and add shutdown-path regressions for both notification and polling.

3. Do not blindly requeue a journal batch after an ambiguous write. JsonlRunEventStore.put_batch uses asyncio.to_thread, so cancelling the await does not stop the file append. If the append completes before the task observes cancellation, the new handler requeues an already-persisted batch and a later flush() writes it again. A deterministic probe stores one event on current main but two identical events with this PR. The SQL commit boundary has the same general ambiguity. Please drain/shield the in-flight write or introduce an idempotent write boundary before requeueing, and cover the JSONL case.

For validation, I merged the PR head into current origin/main in an isolated worktree. The existing focused suite passes (195 passed) and ruff check/format pass, but four targeted cancellation probes fail on the merge candidate: two show the advertised lease-release class is still incomplete, and two isolate regressions introduced by the new compensation behavior.

[AnnaSuSu]

Thanks for addressing the previous cancellation-safety findings. The original blockers are now covered, but the new shield-and-drain paths introduce a shutdown liveness blocker.

Both _drain_cancellation_operation and _put_batch_cancellation_safe keep awaiting the shielded operation until it finishes, swallowing repeated CancelledError without any deadline. If a repository claim/release or journal put_batch call stalls indefinitely, McpTaskService.stop(), run cancellation, and worker finalization cannot complete. In addition, _release_claimed_records releases sequentially, so one hung release prevents every later claimed record in that batch from being released.

This is deterministic in focused probes against c250f24: the four probes for the previous review now pass, while cancellation of service.stop() during a hung release and cancellation of journal.flush() during a hung write both remain pending until the mocked operation is manually unblocked. The exact PR base passes those two liveness probes. The existing service already uses a bounded 5-second deadline for untracked-task compensation, which provides a suitable precedent.

Please bound these drains and keep any ambiguous in-flight operation tracked in the background without blindly requeueing it. Batch release should also ensure that one stalled record cannot starve the remaining records. The focused repository suite passes (220 tests) and ruff passes, but neither currently covers this unbounded-stall case.

[willem-bd]

Reviewed at head c250f24. The shield-and-drain compensation pattern is consistent across the three surfaces, the SQL releases are all lease-owner-guarded so the new compensation calls can't double-apply, and the journal drain correctly distinguishes a committed write from a failed one before requeueing. Two lower-severity observations below (the unbounded-drain liveness concern and the sequential batch release were already raised in earlier reviews, so I won't repeat them).

[willem-bd]

Suggestion: this routes a cancellation through _release_after_error(error="cancelled"), i.e. the same path as a genuine driver/poll failure. release_claim (persistence/mcp_tasks/sql.py) then increments consecutive_poll_error_count, writes last_error="cancelled", applies exponential backoff to next_poll_at, and can cross the tracking_degraded threshold in _record_event_if_changed. So a routine Gateway shutdown while a poll is in flight is recorded as a poll failure on that task (visible via last_error in the task detail API, and repeated occurrences accumulate toward the degraded-tracking event). The other two compensations added in this PR deliberately avoid failure counting (release_notification_lease(..., count_failure=False), release_cancel_claim without any counter). Consider a non-counting release variant for the cancelled-poll path too, so shutdown bookkeeping doesn't masquerade as task errors.

[willem-bd]

Informational: when cancellation lands on the asyncio.gather(...) itself, gather cancels the child tasks, and each child (_poll_one/_cancel_one/_notify_one) already releases its own claim in the CancelledError handler added in this PR - so this parent-level loop then re-releases every claimed record (including records whose child already completed successfully). With the SQL repository this is benign because every release is guarded by lease_owner/notification_lease_owner and the second call is a no-op, but the parent loop's usefulness is limited to children gather cancelled before their first step, and the redundancy silently relies on the owner guard. Fine to keep as-is (the guard makes it safe); just flagging so the double release traffic isn't mistaken for a required invariant later.