From bb7958886d1efcac001c8056c2f3245ab0cfdf99 Mon Sep 17 00:00:00 2001 From: Rei <2064470640@qq.com> Date: Wed, 4 Nov 2026 22:19:05 -0800 Subject: [PATCH] fix(mcp): fence notification claims on SQLite (#7859) * fix(mcp): fence notification claims on SQLite claim_notification_work() selects due generations with FOR UPDATE SKIP LOCKED, which SQLite ignores. Two workers could read the same due generation before either writes, both returned it with different lease tokens, and a stale flush could overwrite the winning owner's token or phase. On SQLite, claim each candidate with a conditional UPDATE fenced on the event version, notification status, dispatch version and the absent-or-expired notification lease the read saw, so a peer claim or a newer recorded event wins and the losing claim has no side effects. Pin the dispatch version on the empty-snapshot consumption fence too. PostgreSQL keeps its existing row-locking path. Coverage: a paused claim that loses a race to a second worker returns no row and leaves the winner's owner, token and phase stored. * fix(mcp): honor deferred notification backoff in SQLite claims claim_notification_work() fenced the observed event version, phase, dispatch version and lease, but not next_notification_at. A peer can claim a due row, then release the claim (release_notification_lease) or defer a dispatched one (defer_dispatched_notification) without changing phase, moving the next attempt into the future and clearing the lease. A paused claim still matched and returned the notification immediately, bypassing retry/defer backoff. Fence the claim on the observed next_notification_at value too, and cover the interleaving with a regression test. --------- Co-authored-by: REI-330 <173389910+REI-330@users.noreply.github.com> Co-authored-by: Willem Jiang --- .../harness/deerflow/persistence/AGENTS.md | 2 +- .../deerflow/persistence/mcp_tasks/sql.py | 72 +++++++++--- backend/tests/test_mcp_task_repository.py | 107 ++++++++++++++++++ 3 files changed, 164 insertions(+), 17 deletions(-) diff --git a/backend/packages/harness/deerflow/persistence/AGENTS.md b/backend/packages/harness/deerflow/persistence/AGENTS.md index a2c5abe98e..26a34dd53b 100644 --- a/backend/packages/harness/deerflow/persistence/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/AGENTS.md @@ -60,7 +60,7 @@ missing or strictly older leases remain recoverable. ## MCP task notifications -`mcp_tasks/sql.py::claim_notification_work` never returns a row without a dispatch event. A pending snapshot whose condition resolved before delivery (tracking recovered, or an answered `input_required` prompt returned the task to `working`) is consumed at claim time: `notified_version` advances, the content fingerprint clears so an identical later event is recorded again, stale error/attempt/dispatch state clears, status returns to `none`, and the claim lease is released instead of launching an empty event. A legacy row left without a dispatch event is consumed only when nothing newer exists; a newer recorded event is rebuilt and claimed instead. Consumption is a conditional UPDATE fenced on the event version and phase the claim read, so a concurrent newer event (SQLite has no row locks) is never acknowledged and stays claimable. Coverage: `tests/test_mcp_task_repository.py` (resolved snapshot, legacy retry, expired-claim reclaim, identical-event recurrence, interleaved newer completion, legacy launching/dispatched with and without a newer completion). +`mcp_tasks/sql.py::claim_notification_work` never returns a row without a dispatch event. A pending snapshot whose condition resolved before delivery (tracking recovered, or an answered `input_required` prompt returned the task to `working`) is consumed at claim time: `notified_version` advances, the content fingerprint clears so an identical later event is recorded again, stale error/attempt/dispatch state clears, status returns to `none`, and the claim lease is released instead of launching an empty event. A legacy row left without a dispatch event is consumed only when nothing newer exists; a newer recorded event is rebuilt and claimed instead. Consumption is a conditional UPDATE fenced on the event version, notification phase and dispatch version the claim read, and on SQLite the claim itself is a conditional UPDATE fenced on those columns, the observed next attempt time and an absent or expired notification lease, so a concurrent newer event, a peer claim or a release/defer moving the next attempt (SQLite has no row locks) is never acknowledged, overwritten, returned twice or delivered ahead of its backoff. Coverage: `tests/test_mcp_task_repository.py` (resolved snapshot, legacy retry, expired-claim reclaim, identical-event recurrence, interleaved newer completion, two concurrent claims, a deferred backoff, legacy launching/dispatched with and without a newer completion). ## MCP control leases diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index aa2ce12f0c..ff2ad188bd 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -719,6 +719,13 @@ class McpTaskRepository: async with self._sf() as session: rows = list((await session.execute(stmt)).scalars()) expires_at = now + timedelta(seconds=lease_seconds) + # SQLite ignores ``FOR UPDATE``, so a claim is also fenced on the + # state this read saw. A peer claim (future lease, changed phase), a + # release/defer moving the next attempt, or a newer recorded event + # changes one of those columns, the conditional UPDATE matches no + # row, and the paused claim has no side effects instead of + # overwriting the winning worker or bypassing its backoff. + sqlite = session.get_bind().dialect.name == "sqlite" claimed: list[McpTaskRow] = [] for row in rows: rebuild_snapshot = row.notification_status in ("pending", "claimed") or (row.notification_status == "retry" and row.dispatch_version != row.event_version) @@ -736,6 +743,8 @@ class McpTaskRepository: if rebuild_snapshot else row.dispatch_event ) + dispatch_version_fence = McpTaskRow.dispatch_version.is_(None) if row.dispatch_version is None else McpTaskRow.dispatch_version == row.dispatch_version + next_notification_fence = McpTaskRow.next_notification_at.is_(None) if row.next_notification_at is None else McpTaskRow.next_notification_at == row.next_notification_at if event is None: # The condition behind this snapshot resolved before delivery # (for example tracking recovered, or an input_required prompt @@ -743,12 +752,13 @@ class McpTaskRepository: # a pre-fix claim. Consume only what was proven empty — the # current event version for a rebuild, the empty dispatch # generation otherwise — through a conditional UPDATE fenced on - # the event version and phase this read saw. SQLite has no - # ``SELECT FOR UPDATE``, so a newer event recorded after the - # read must win: the fence then matches no row, this consume is - # a no-op, and the newer event stays claimable. Clearing the - # content fingerprint keeps an identical later event from being - # suppressed by a snapshot that was never delivered. + # the event version, phase and dispatch version this read saw. + # SQLite has no ``SELECT FOR UPDATE``, so a newer event recorded + # after the read must win: the fence then matches no row, this + # consume is a no-op, and the newer event stays claimable. + # Clearing the content fingerprint keeps an identical later + # event from being suppressed by a snapshot that was never + # delivered. consumed_version = row.event_version if rebuild_snapshot else (row.dispatch_version if row.dispatch_version is not None else row.event_version) status_at_read = row.notification_status session.expunge(row) @@ -758,6 +768,7 @@ class McpTaskRepository: McpTaskRow.id == row.id, McpTaskRow.event_version == consumed_version, McpTaskRow.notification_status == status_at_read, + dispatch_version_fence, ) .values( notified_version=consumed_version, @@ -776,18 +787,47 @@ class McpTaskRepository: ) ) continue - row.notification_lease_owner = lease_owner - row.notification_lease_expires_at = expires_at - row.notification_lease_token = _new_claim_token() + claim_token = _new_claim_token() + claim_values: dict[str, Any] = { + "notification_lease_owner": lease_owner, + "notification_lease_expires_at": expires_at, + "notification_lease_token": claim_token, + "updated_at": now, + } if rebuild_snapshot: + claim_values.update( + dispatch_version=row.event_version, + dispatch_event=event, + notification_run_id=None, + notification_status="claimed", + ) if row.dispatch_version != row.event_version: - row.dispatch_attempt = 0 - row.notification_attempt_count = 0 - row.dispatch_version = row.event_version - row.dispatch_event = event - row.notification_run_id = None - row.notification_status = "claimed" - row.updated_at = now + claim_values.update(dispatch_attempt=0, notification_attempt_count=0) + if sqlite: + claim_result = await session.execute( + update(McpTaskRow) + .where( + McpTaskRow.id == row.id, + McpTaskRow.event_version == row.event_version, + McpTaskRow.notification_status == row.notification_status, + dispatch_version_fence, + next_notification_fence, + or_( + McpTaskRow.notification_lease_expires_at.is_(None), + McpTaskRow.notification_lease_expires_at < now, + ), + ) + .values(**claim_values) + .execution_options(synchronize_session=False) + ) + if not claim_result.rowcount: + # A peer claimed the row or a newer event was recorded + # after this read; leave it to the next claim cycle. + session.expunge(row) + continue + session.expunge(row) + for field, value in claim_values.items(): + setattr(row, field, value) claimed.append(row) await session.commit() return [self._row_to_dict(row) for row in claimed] diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index da0ed659b5..b31bdc9e7c 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -295,6 +295,113 @@ async def test_interleaved_reclaim_fences_inflight_notification_completion(tmp_p assert applied is False +@pytest.mark.asyncio +async def test_concurrent_notification_claims_return_one_owner(tmp_path, monkeypatch): + repo = await _make_repo(tmp_path) + now = datetime(2026, 1, 1, tzinfo=UTC) + task_id = "concurrent-notification-claim" + await _create_working_task(repo, task_id=task_id, now=now) + poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + assert await repo.apply_snapshot( + task_id, + lease_owner="poller", + lease_token=poll[0]["lease_token"], + status="completed", + result={"done": True}, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + polled_at=now, + ) + claim_kwargs = {"now": now, "lease_seconds": 60, "limit": 1, "tracking_degraded_after_errors": 3} + + async with _pause_claim_mutation( + monkeypatch, + repo.claim_notification_work(lease_owner="worker-a", **claim_kwargs), + ) as (paused_claim, resume): + winner = await asyncio.wait_for(repo.claim_notification_work(lease_owner="worker-b", **claim_kwargs), timeout=5) + assert len(winner) == 1 + before_resume = await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert before_resume is not None + resume.set() + stale = await asyncio.wait_for(paused_claim, timeout=5) + + # The paused claim read the row before worker-b won it; its write must not + # overwrite the winning owner, token or dispatch phase. + assert stale == [] + stored = await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert stored is not None + assert stored["notification_lease_owner"] == "worker-b" + assert stored["notification_lease_token"] == winner[0]["notification_lease_token"] + assert stored["notification_status"] == winner[0]["notification_status"] + assert stored["dispatch_version"] == winner[0]["dispatch_version"] + + +@pytest.mark.asyncio +async def test_stale_claim_does_not_bypass_deferred_next_notification(tmp_path, monkeypatch): + repo = await _make_repo(tmp_path) + now = datetime(2026, 1, 1, tzinfo=UTC) + task_id = "deferred-notification-claim" + await _create_working_task(repo, task_id=task_id, now=now) + poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + assert await repo.apply_snapshot( + task_id, + lease_owner="poller", + lease_token=poll[0]["lease_token"], + status="completed", + result={"done": True}, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + polled_at=now, + ) + claim_kwargs = {"now": now, "lease_seconds": 60, "limit": 1, "tracking_degraded_after_errors": 3} + seeded = await repo.claim_notification_work(lease_owner="seed", **claim_kwargs) + assert len(seeded) == 1 + # Release the phase-preserving claim back to a due state: the status stays + # ``claimed`` while the row becomes claimable again. + assert await repo.release_notification_lease( + task_id, + lease_owner="seed", + notification_lease_token=seeded[0]["notification_lease_token"], + next_notification_at=now - timedelta(seconds=1), + error=None, + ) + deferred_until = now + timedelta(seconds=30) + + async with _pause_claim_mutation( + monkeypatch, + repo.claim_notification_work(lease_owner="worker-a", **claim_kwargs), + ) as (paused_claim, resume): + winner = await asyncio.wait_for(repo.claim_notification_work(lease_owner="worker-b", **claim_kwargs), timeout=5) + assert len(winner) == 1 + # Worker-b releases the claim without changing its phase, moving the + # next attempt into the future (retry/defer backoff). + assert await repo.release_notification_lease( + task_id, + lease_owner="worker-b", + notification_lease_token=winner[0]["notification_lease_token"], + next_notification_at=deferred_until, + error="backoff", + ) + resume.set() + stale = await asyncio.wait_for(paused_claim, timeout=5) + + # The paused reader saw a due claim; the backoff belongs to the winning + # worker, so the stale claim must not deliver the notification early. + assert stale == [] + stored = await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert stored is not None + assert stored["notification_lease_owner"] is None + assert datetime.fromisoformat(stored["next_notification_at"]) == deferred_until + + @pytest.mark.asyncio async def test_legacy_task_writer_leaves_thread_incarnation_null(tmp_path): repo = await _make_repo(tmp_path)