From 81c5c833843ac165e634899f7e7671e98e305bbe Mon Sep 17 00:00:00 2001 From: RongJie G <111257566+CorgiBoyG@users.noreply.github.com> Date: Sun, 27 Sep 2026 20:46:56 +0800 Subject: [PATCH] fix: make MCP notification recovery rollback-safe (#5908) Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> --- backend/app/mcp_tasks/service.py | 216 ++++++------ .../deerflow/persistence/mcp_tasks/sql.py | 8 +- backend/tests/test_mcp_task_repository.py | 204 ++++++++++- backend/tests/test_mcp_task_service.py | 325 +++++++++++++++++- frontend/src/core/background-tasks/types.ts | 3 +- .../unit/core/background-tasks/types.test.ts | 2 +- 6 files changed, 637 insertions(+), 121 deletions(-) diff --git a/backend/app/mcp_tasks/service.py b/backend/app/mcp_tasks/service.py index 24a7fbd9a4..46546370aa 100644 --- a/backend/app/mcp_tasks/service.py +++ b/backend/app/mcp_tasks/service.py @@ -70,6 +70,11 @@ def _bound_error(error: str | None) -> str | None: return error[:_MAX_PERSISTED_ERROR_CHARS] +def _notification_completion_time(*, not_before: datetime) -> datetime: + """Return notification completion time without moving before the claim.""" + return max(not_before, datetime.now(UTC)) + + def _consume_task_error(task: asyncio.Future[Any]) -> BaseException | None: try: return task.exception() @@ -843,43 +848,33 @@ class McpTaskService: release=self._release_notification_after_cancellation, ) if records: - states, results = await self._run_claimed_batch( + await self._run_claimed_batch( records, - operation=lambda record: self._notify_one_claimed(record, now=now), + operation=lambda record: self._notify_one_claimed_safely(record, now=now), release=self._release_notification_after_cancellation, action="notification", ) - for state, result in zip(states, results, strict=True): - if not isinstance(result, BaseException) or isinstance(result, asyncio.CancelledError): - continue - record = state.record - error = _bound_error(str(result) or type(result).__name__) or type(result).__name__ - logger.error( - "Unexpected MCP task notification failure (task_id=%s)", - record.get("id"), - exc_info=(type(result), result, result.__traceback__), - ) - await self._release_notification_failure(record, now=now, error=error) + + async def _notify_one_claimed_safely(self, record: dict[str, Any], *, now: datetime) -> None: + try: + await self._notify_one_claimed(record, now=now) + except Exception as exc: # noqa: BLE001 - isolate one notification from its claimed siblings + error = _bound_error(str(exc) or type(exc).__name__) or type(exc).__name__ + logger.error( + "Unexpected MCP task notification failure (task_id=%s)", + record.get("id"), + exc_info=(type(exc), exc, exc.__traceback__), + ) + await self._release_notification_failure(record, now=now, error=error) async def _notify_one_claimed(self, record: dict[str, Any], *, now: datetime) -> None: task_id = record["id"] dispatch_version = int(record.get("dispatch_version") or 0) notification_attempts = max(0, int(record.get("notification_attempt_count") or 0)) - if notification_attempts >= _MAX_NOTIFICATION_ATTEMPTS: - previous_error = record.get("notification_error") or "delivery failed" - await self._repository.dead_letter_notification( - task_id, - lease_owner=self._lease_owner, - notification_lease_token=record["notification_lease_token"], - dispatch_version=dispatch_version, - error=_bound_error(f"Notification delivery stopped after {notification_attempts} failed attempts: {previous_error}"), - count_failure=False, - now=now, - ) - return if record.get("notification_status") == "dispatched": run = await self._get_run(record.get("notification_run_id"), user_id=record["user_id"]) + completed_at = _notification_completion_time(not_before=now) status = getattr(run, "status", None) if run is None: run_id = record.get("notification_run_id") @@ -889,9 +884,9 @@ class McpTaskService: notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, delivered=False, - next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), + next_notification_at=completed_at + timedelta(seconds=self._notification_retry_seconds(record)), error=_bound_error(f"Notification run {run_id!r} was not found"), - now=now, + now=completed_at, ) elif status == RunStatus.success: await self._repository.finish_notification_run( @@ -902,7 +897,7 @@ class McpTaskService: delivered=True, next_notification_at=None, error=None, - now=now, + now=completed_at, ) elif status in {RunStatus.error, RunStatus.timeout, RunStatus.interrupted}: await self._repository.finish_notification_run( @@ -911,9 +906,9 @@ class McpTaskService: notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, delivered=False, - next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), + next_notification_at=completed_at + timedelta(seconds=self._notification_retry_seconds(record)), error=_bound_error(getattr(run, "error", None) or f"Notification run ended with {status}"), - now=now, + now=completed_at, ) else: await self._repository.defer_dispatched_notification( @@ -921,11 +916,25 @@ class McpTaskService: lease_owner=self._lease_owner, notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, - next_notification_at=now + timedelta(seconds=self._poll_interval_seconds), - now=now, + next_notification_at=completed_at + timedelta(seconds=self._poll_interval_seconds), + now=completed_at, ) return + if record.get("notification_status") != "launching" and notification_attempts >= _MAX_NOTIFICATION_ATTEMPTS: + previous_error = record.get("notification_error") or "delivery failed" + completed_at = _notification_completion_time(not_before=now) + await self._repository.dead_letter_notification( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + dispatch_version=dispatch_version, + error=_bound_error(f"Notification delivery stopped after {notification_attempts} failed attempts: {previous_error}"), + count_failure=False, + now=completed_at, + ) + return + source_run = await self._get_run(record.get("run_id"), user_id=record["user_id"]) if record.get("run_id") else None try: result = await self._launch_notification( @@ -938,6 +947,7 @@ class McpTaskService: event=dict(record.get("dispatch_event") or {}), ) except PermanentNotificationError as exc: + completed_at = _notification_completion_time(not_before=now) await self._repository.dead_letter_notification( task_id, lease_owner=self._lease_owner, @@ -945,47 +955,78 @@ class McpTaskService: dispatch_version=dispatch_version, error=_bound_error(str(exc) or type(exc).__name__), count_failure=True, - now=now, + now=completed_at, ) return except ConflictError as exc: + completed_at = _notification_completion_time(not_before=now) retry_error = _bound_error(str(exc)) - await self._release_ordinary_batch_record( - record, - release=lambda: self._repository.release_notification_claim( - task_id, - lease_owner=self._lease_owner, - notification_lease_token=record["notification_lease_token"], - next_notification_at=now + timedelta(seconds=self._poll_interval_seconds), - error=retry_error, - replace_with_latest=True, - ), - action="release notification conflict retry", - ) + if record.get("notification_status") == "launching": + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_lease( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + next_notification_at=completed_at + timedelta(seconds=self._poll_interval_seconds), + error=retry_error, + count_failure=False, + ), + action="release recovered notification conflict", + ) + else: + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_claim( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + next_notification_at=completed_at + timedelta(seconds=self._poll_interval_seconds), + error=retry_error, + replace_with_latest=True, + ), + action="release notification conflict retry", + ) return except Exception as exc: # noqa: BLE001 - retry the same idempotency key + completed_at = _notification_completion_time(not_before=now) retry_error = _bound_error(str(exc) or type(exc).__name__) - await self._release_ordinary_batch_record( - record, - release=lambda: self._repository.release_notification_claim( - task_id, - lease_owner=self._lease_owner, - notification_lease_token=record["notification_lease_token"], - next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), - error=retry_error, - replace_with_latest=True, - count_failure=True, - ), - action="release notification retry", - ) + if record.get("notification_status") == "launching": + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_lease( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + next_notification_at=completed_at + timedelta(seconds=self._notification_retry_seconds(record)), + error=retry_error, + count_failure=False, + ), + action="release uncertain notification launch", + ) + else: + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_claim( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + next_notification_at=completed_at + timedelta(seconds=self._notification_retry_seconds(record)), + error=retry_error, + replace_with_latest=True, + count_failure=True, + ), + action="release notification retry", + ) return + completed_at = _notification_completion_time(not_before=now) await self._repository.mark_notification_dispatched( task_id, lease_owner=self._lease_owner, notification_lease_token=record["notification_lease_token"], dispatch_version=dispatch_version, run_id=result["run_id"], - now=now, + now=completed_at, ) def _notification_retry_seconds(self, record: dict[str, Any]) -> int: @@ -1176,52 +1217,21 @@ class McpTaskService: now: datetime, error: str, ) -> None: - task = asyncio.create_task( - self._repository.release_notification_lease( - record["id"], - lease_owner=self._lease_owner, - notification_lease_token=record["notification_lease_token"], - next_notification_at=now + timedelta(seconds=self._notification_retry_seconds(record)), - error=error, - count_failure=True, - ), - name=f"mcp-release-notification-failure-{record.get('id', 'unknown')}", - ) + completed_at = _notification_completion_time(not_before=now) try: - await asyncio.wait_for( - asyncio.shield(task), - timeout=_CANCELLATION_DRAIN_TIMEOUT_SECONDS, - ) - except TimeoutError: - self._track_compensation_task( - task, + await self._release_ordinary_batch_record( + record, + release=lambda: self._repository.release_notification_lease( + record["id"], + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + next_notification_at=completed_at + timedelta(seconds=self._notification_retry_seconds(record)), + error=error, + count_failure=record.get("notification_status") != "launching", + ), action="release notification failure", - task_id=record["id"], ) - logger.warning( - "Timed out after %.1f seconds waiting for MCP task notification release; it continues in the background (task_id=%s)", - _CANCELLATION_DRAIN_TIMEOUT_SECONDS, - record["id"], - ) - return except asyncio.CancelledError: - caller_cancelling = asyncio.current_task().cancelling() - release_cancelled = _task_has_cancelled_terminal_state(task) - if release_cancelled and not caller_cancelling: - error = _consume_task_error(task) - if error is not None: - self._log_batch_release_error( - error, - action="release notification failure", - task_id=record["id"], - ) - return - await self._drain_cancellation_task( - task, - action="release notification failure", - task_id=record["id"], - deadline=asyncio.get_running_loop().time() + _CANCELLATION_DRAIN_TIMEOUT_SECONDS, - ) raise except Exception: # noqa: BLE001 - retain the task-scoped failure logger.exception( @@ -1247,7 +1257,7 @@ class McpTaskService: record: dict[str, Any], ) -> None: task_id = record["id"] - if record.get("notification_status") == "dispatched": + if record.get("notification_status") in {"launching", "dispatched"}: compensation = self._repository.release_notification_lease( task_id, lease_owner=self._lease_owner, @@ -1256,7 +1266,7 @@ class McpTaskService: error=record.get("notification_error"), count_failure=False, ) - action = "release dispatched notification lease" + action = "release in-flight notification lease" else: compensation = self._repository.release_notification_claim( task_id, diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index 5145a833ca..93295c6d35 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -31,7 +31,7 @@ _TIMESTAMP_FIELDS = ( "updated_at", ) -_INFLIGHT_NOTIFICATION_STATUSES = frozenset({"claimed", "dispatched", "retry"}) +_INFLIGHT_NOTIFICATION_STATUSES = frozenset({"claimed", "launching", "dispatched", "retry"}) def _new_claim_token() -> str: @@ -639,7 +639,7 @@ class McpTaskRepository: limit: int, tracking_degraded_after_errors: int, ) -> list[dict[str, Any]]: - statuses = ("pending", "claimed", "retry", "dispatched") + statuses = ("pending", "claimed", "launching", "retry", "dispatched") stmt = ( select(McpTaskRow) .where( @@ -693,7 +693,7 @@ class McpTaskRepository: McpTaskRow.notification_lease_token == notification_lease_token, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, - McpTaskRow.notification_status.in_(("claimed", "retry")), + McpTaskRow.notification_status.in_(("claimed", "launching", "retry")), ) .values( notification_status="dispatched", @@ -863,7 +863,7 @@ class McpTaskRepository: McpTaskRow.notification_lease_token == notification_lease_token, McpTaskRow.notification_lease_expires_at >= now, McpTaskRow.dispatch_version == dispatch_version, - McpTaskRow.notification_status.in_(("claimed", "retry", "dispatched")), + McpTaskRow.notification_status.in_(("claimed", "launching", "retry", "dispatched")), ) dead_letter_values: dict[str, Any] = { "notification_status": "dead_letter", diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index dc7f891488..ad523ac95b 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -2,6 +2,8 @@ import asyncio import contextlib import sqlite3 from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest import pytest_asyncio @@ -9,7 +11,9 @@ from sqlalchemy import event from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession +from app.mcp_tasks.service import McpTaskService from deerflow.config.database_config import DatabaseConfig +from deerflow.mcp.tasks import McpTaskDriverRegistry from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config from deerflow.persistence.mcp_tasks import ( DuplicateMcpRemoteTaskError, @@ -18,6 +22,8 @@ from deerflow.persistence.mcp_tasks import ( ) from deerflow.persistence.mcp_tasks.model import McpTaskRow from deerflow.persistence.thread_meta.model import ThreadMetaRow +from deerflow.runtime.runs.manager import ConflictError +from deerflow.runtime.runs.schemas import RunStatus @pytest_asyncio.fixture(autouse=True) @@ -989,6 +995,195 @@ async def test_finish_notification_run_after_reclaim_cannot_clear_new_claim(tmp_ assert stored["notification_lease_token"] == new_notify_token +@pytest.mark.asyncio +async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-launching-compat", now=now) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-launching-compat", + lease_owner="poller", + lease_token=poll_claim[0]["lease_token"], + status="input_required", + result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required={"prompt": "Approve?"}, + next_poll_at=now, + polled_at=now, + ) + first = await repo.claim_notification_work( + now=now, + lease_owner="notifier-a", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + first_version = first[0]["dispatch_version"] + first_event = first[0]["dispatch_event"] + async with repo._sf() as session: + row = await session.get(McpTaskRow, "task-launching-compat") + assert row is not None + row.notification_status = "launching" + row.notification_lease_owner = None + row.notification_lease_expires_at = None + row.notification_lease_token = None + await session.commit() + + second_poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-launching-compat", + lease_owner="poller", + lease_token=second_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, + ) + + reclaimed = await repo.claim_notification_work( + now=now, + lease_owner="notifier-b", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + + assert len(reclaimed) == 1 + assert reclaimed[0]["notification_status"] == "launching" + assert reclaimed[0]["dispatch_version"] == first_version + assert reclaimed[0]["dispatch_event"] == first_event + assert reclaimed[0]["event_version"] > first_version + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=60, + max_concurrent_polls=1, + launch_notification=AsyncMock(side_effect=ConflictError("thread busy")), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + service._lease_owner = "notifier-b" + await service._notify_one_claimed(reclaimed[0], now=now) + + retry_at = now + timedelta(seconds=10) + after_conflict = await repo.claim_notification_work( + now=retry_at, + lease_owner="notifier-c", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + assert len(after_conflict) == 1 + assert after_conflict[0]["notification_status"] == "launching" + assert after_conflict[0]["dispatch_version"] == first_version + assert after_conflict[0]["dispatch_event"] == first_event + assert after_conflict[0]["notification_attempt_count"] == 0 + assert await repo.mark_notification_dispatched( + "task-launching-compat", + lease_owner="notifier-c", + notification_lease_token=after_conflict[0]["notification_lease_token"], + dispatch_version=first_version, + run_id="notify-run-1", + now=retry_at, + ) + + +@pytest.mark.asyncio +async def test_recovered_launching_at_retry_budget_reconciles_successful_run(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-launching-success", now=now) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-launching-success", + lease_owner="poller", + lease_token=poll_claim[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, + ) + claimed = await repo.claim_notification_work( + now=now, + lease_owner="notifier-a", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + dispatch_version = claimed[0]["dispatch_version"] + async with repo._sf() as session: + row = await session.get(McpTaskRow, "task-launching-success") + assert row is not None + row.notification_status = "launching" + row.notification_attempt_count = 5 + await session.commit() + + launch = AsyncMock(return_value={"run_id": "notify-run-success"}) + launching_service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=60, + max_concurrent_polls=1, + launch_notification=launch, + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + launching_service._lease_owner = "notifier-a" + await launching_service._notify_one_claimed( + { + **claimed[0], + "notification_status": "launching", + "notification_attempt_count": 5, + }, + now=now, + ) + + reconcile_at = now + timedelta(seconds=1) + dispatched = await repo.claim_notification_work( + now=reconcile_at, + lease_owner="notifier-b", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + assert len(dispatched) == 1 + assert dispatched[0]["notification_status"] == "dispatched" + assert dispatched[0]["notification_attempt_count"] == 5 + get_run = AsyncMock(return_value=SimpleNamespace(status=RunStatus.success)) + dispatched_service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=60, + max_concurrent_polls=1, + launch_notification=AsyncMock(), + get_run=get_run, + ) + dispatched_service._lease_owner = "notifier-b" + await dispatched_service._notify_one_claimed(dispatched[0], now=reconcile_at) + + get_run.assert_awaited_once_with("notify-run-success", user_id="user-1") + stored = await repo.get("task-launching-success", user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert stored is not None + assert stored["notification_status"] == "delivered" + assert stored["notified_version"] == dispatch_version + assert stored["notification_attempt_count"] == 0 + + @pytest.mark.asyncio async def test_release_poll_claim_after_cancellation_preserves_poll_failure_state(tmp_path): repo = await _make_repo(tmp_path) @@ -1363,7 +1558,8 @@ async def test_notification_launch_failure_counts_and_reclaims_latest_snapshot(t @pytest.mark.asyncio -async def test_permanent_notification_failure_is_not_reclaimed(tmp_path): +@pytest.mark.parametrize("notification_status", ["claimed", "launching"]) +async def test_permanent_notification_failure_is_not_reclaimed(tmp_path, notification_status): repo = await _make_repo(tmp_path) now = datetime.now(UTC) await _create_working_task(repo, task_id="task-dead-letter", now=now) @@ -1389,6 +1585,12 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path): limit=1, tracking_degraded_after_errors=3, ) + if notification_status == "launching": + async with repo._sf() as session: + row = await session.get(McpTaskRow, "task-dead-letter") + assert row is not None + row.notification_status = "launching" + await session.commit() assert await repo.dead_letter_notification( "task-dead-letter", diff --git a/backend/tests/test_mcp_task_service.py b/backend/tests/test_mcp_task_service.py index bb53a3422d..cef7b4fd11 100644 --- a/backend/tests/test_mcp_task_service.py +++ b/backend/tests/test_mcp_task_service.py @@ -27,6 +27,14 @@ from deerflow.runtime.runs.manager import ConflictError from deerflow.runtime.runs.schemas import RunStatus +class _MutableDateTime(datetime): + current = datetime(2000, 1, 1, tzinfo=UTC) + + @classmethod + def now(cls, tz=None): + return cls.current + + class FakeRepository: def __init__(self, rows=None): self.rows = list(rows or []) @@ -731,7 +739,7 @@ async def test_notification_delivery_waits_for_successful_agent_run(): @pytest.mark.asyncio -async def test_missing_dispatched_notification_run_retries_delivery(): +async def test_missing_dispatched_notification_run_retries_delivery(monkeypatch): repo = SimpleNamespace( finish_notification_run=AsyncMock(return_value=True), defer_dispatched_notification=AsyncMock(return_value=True), @@ -746,6 +754,7 @@ async def test_missing_dispatched_notification_run_retries_delivery(): get_run=AsyncMock(return_value=None), ) now = datetime.now(UTC) + monkeypatch.setattr(service_module, "_notification_completion_time", lambda *, not_before: not_before) await service._notify_one_claimed( { @@ -767,7 +776,7 @@ async def test_missing_dispatched_notification_run_retries_delivery(): @pytest.mark.asyncio -async def test_notification_failures_are_isolated_and_release_their_lease(caplog): +async def test_notification_failures_are_isolated_and_release_their_lease(caplog, monkeypatch): records = [ { **_claimed_row(), @@ -806,6 +815,7 @@ async def test_notification_failures_are_isolated_and_release_their_lease(caplog get_run=get_run, ) now = datetime.now(UTC) + monkeypatch.setattr(service_module, "_notification_completion_time", lambda *, not_before: not_before) with caplog.at_level(logging.ERROR): await service._run_notifications(now=now) @@ -821,7 +831,69 @@ async def test_notification_failures_are_isolated_and_release_their_lease(caplog @pytest.mark.asyncio -async def test_notification_busy_thread_replaces_claim_with_latest_event(): +async def test_notification_failure_releases_at_its_own_completion_without_waiting_for_sibling(monkeypatch): + sibling_gate = asyncio.Event() + failure_released = asyncio.Event() + scan_started_at = datetime(2000, 1, 1, tzinfo=UTC) + failure_completed_at = scan_started_at + timedelta(seconds=7) + records = [ + { + **_claimed_row(), + "id": "task-broken", + "notification_status": "dispatched", + "notification_run_id": "run-broken", + "dispatch_version": 2, + }, + { + **_claimed_row(), + "id": "task-blocked", + "notification_status": "dispatched", + "notification_run_id": "run-blocked", + "dispatch_version": 3, + }, + ] + + async def get_run(run_id, **_kwargs): + if run_id == "run-broken": + raise RuntimeError("run store unavailable") + await sibling_gate.wait() + return SimpleNamespace(status=RunStatus.success) + + async def release_notification_lease(*_args, **_kwargs): + failure_released.set() + return True + + repo = SimpleNamespace( + claim_notification_work=AsyncMock(return_value=records), + finish_notification_run=AsyncMock(return_value=True), + defer_dispatched_notification=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(side_effect=release_notification_lease), + ) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(), + get_run=get_run, + ) + monkeypatch.setattr(service_module, "_notification_completion_time", lambda *, not_before: failure_completed_at) + + task = asyncio.create_task(service._run_notifications(now=scan_started_at)) + try: + await asyncio.wait_for(failure_released.wait(), timeout=1) + assert not task.done() + released = repo.release_notification_lease.await_args + assert released.args[0] == "task-broken" + assert released.kwargs["next_notification_at"] == failure_completed_at + timedelta(seconds=5) + finally: + sibling_gate.set() + await asyncio.wait_for(task, timeout=1) + + +@pytest.mark.asyncio +async def test_notification_busy_thread_replaces_claim_with_latest_event(monkeypatch): repo = SimpleNamespace( release_notification_claim=AsyncMock(return_value=True), ) @@ -835,6 +907,7 @@ async def test_notification_busy_thread_replaces_claim_with_latest_event(): get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), ) now = datetime.now(UTC) + monkeypatch.setattr(service_module, "_notification_completion_time", lambda *, not_before: not_before) await service._notify_one_claimed( { @@ -853,7 +926,43 @@ async def test_notification_busy_thread_replaces_claim_with_latest_event(): @pytest.mark.asyncio -async def test_notification_launch_failure_backs_off_and_replaces_with_latest_event(): +async def test_reclaimed_launching_busy_thread_preserves_reserved_snapshot(monkeypatch): + repo = SimpleNamespace( + release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), + ) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(side_effect=ConflictError("thread busy")), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + now = datetime.now(UTC) + monkeypatch.setattr(service_module, "_notification_completion_time", lambda *, not_before: not_before) + + await service._notify_one_claimed( + { + **_claimed_row(), + "notification_status": "launching", + "dispatch_version": 2, + "dispatch_attempt": 0, + "notification_attempt_count": 5, + "dispatch_event": {"status": "input_required"}, + }, + now=now, + ) + + repo.release_notification_claim.assert_not_awaited() + released = repo.release_notification_lease.await_args.kwargs + assert released["count_failure"] is False + assert released["next_notification_at"] == now + timedelta(seconds=5) + + +@pytest.mark.asyncio +async def test_notification_launch_failure_backs_off_and_replaces_with_latest_event(monkeypatch): repo = SimpleNamespace( release_notification_claim=AsyncMock(return_value=True), ) @@ -868,6 +977,7 @@ async def test_notification_launch_failure_backs_off_and_replaces_with_latest_ev get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), ) now = datetime.now(UTC) + monkeypatch.setattr(service_module, "_notification_completion_time", lambda *, not_before: not_before) await service._notify_one_claimed( { @@ -924,6 +1034,149 @@ async def test_permanently_rejected_notification_is_dead_lettered(): repo.release_notification_claim.assert_not_awaited() +@pytest.mark.asyncio +async def test_permanently_rejected_reclaimed_launching_notification_is_dead_lettered(): + repo = SimpleNamespace( + dead_letter_notification=AsyncMock(return_value=True), + release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), + ) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(side_effect=PermanentNotificationError("Thread thread-1 not found")), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + + await service._notify_one_claimed( + { + **_claimed_row(), + "notification_status": "launching", + "dispatch_version": 2, + "dispatch_attempt": 0, + "notification_attempt_count": 5, + "dispatch_event": {"status": "completed"}, + }, + now=datetime.now(UTC), + ) + + repo.dead_letter_notification.assert_awaited_once() + assert repo.dead_letter_notification.await_args.kwargs["count_failure"] is True + repo.release_notification_claim.assert_not_awaited() + repo.release_notification_lease.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reclaimed_launching_source_lookup_failure_releases_without_counting_failure(): + record = { + **_claimed_row(), + "notification_status": "launching", + "dispatch_version": 2, + "dispatch_attempt": 0, + "notification_attempt_count": 5, + "dispatch_event": {"status": "completed"}, + } + repo = SimpleNamespace( + claim_notification_work=AsyncMock(return_value=[record]), + release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), + dead_letter_notification=AsyncMock(return_value=True), + ) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=AsyncMock(), + get_run=AsyncMock(side_effect=RuntimeError("run store unavailable")), + ) + + await service._run_notifications(now=datetime.now(UTC)) + + repo.dead_letter_notification.assert_not_awaited() + repo.release_notification_claim.assert_not_awaited() + repo.release_notification_lease.assert_awaited_once() + assert repo.release_notification_lease.await_args.kwargs["count_failure"] is False + + +@pytest.mark.asyncio +async def test_notification_dispatch_uses_launch_completion_time(monkeypatch): + scan_started_at = datetime(2000, 1, 1, tzinfo=UTC) + launch_completed_at = scan_started_at + timedelta(seconds=121) + _MutableDateTime.current = scan_started_at + + async def launch_notification(**_kwargs): + _MutableDateTime.current = launch_completed_at + return {"run_id": "notify-run-1"} + + monkeypatch.setattr(service_module, "datetime", _MutableDateTime) + repo = SimpleNamespace(mark_notification_dispatched=AsyncMock(return_value=False)) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=launch_notification, + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + + await service._notify_one_claimed( + { + **_claimed_row(), + "notification_status": "claimed", + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + }, + now=scan_started_at, + ) + + assert repo.mark_notification_dispatched.await_args.kwargs["now"] == launch_completed_at + + +@pytest.mark.asyncio +async def test_reclaimed_launching_notification_preserves_snapshot_on_uncertain_failure_at_retry_budget(): + repo = SimpleNamespace( + release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), + dead_letter_notification=AsyncMock(return_value=True), + ) + launch_notification = AsyncMock(side_effect=RuntimeError("run result unavailable")) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=launch_notification, + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + + await service._notify_one_claimed( + { + **_claimed_row(), + "notification_status": "launching", + "notification_error": "previous result unavailable", + "dispatch_version": 2, + "dispatch_attempt": 0, + "notification_attempt_count": 5, + "dispatch_event": {"status": "completed"}, + }, + now=datetime.now(UTC), + ) + + launch_notification.assert_awaited_once() + repo.dead_letter_notification.assert_not_awaited() + repo.release_notification_lease.assert_awaited_once() + assert repo.release_notification_lease.await_args.kwargs["count_failure"] is False + repo.release_notification_claim.assert_not_awaited() + + @pytest.mark.asyncio async def test_notification_retry_budget_dead_letters_before_creating_another_run(): repo = SimpleNamespace( @@ -964,11 +1217,12 @@ async def test_notification_retry_budget_dead_letters_before_creating_another_ru @pytest.mark.asyncio -async def test_dispatched_notification_retry_budget_dead_letters_before_hydrating_run(): +async def test_dispatched_notification_retry_budget_reconciles_run_before_dead_lettering(): repo = SimpleNamespace( dead_letter_notification=AsyncMock(return_value=True), + finish_notification_run=AsyncMock(return_value=True), ) - get_run = AsyncMock() + get_run = AsyncMock(return_value=SimpleNamespace(status=RunStatus.success)) service = McpTaskService( repository=repo, drivers=McpTaskDriverRegistry(), @@ -992,11 +1246,12 @@ async def test_dispatched_notification_retry_budget_dead_letters_before_hydratin now=now, ) - get_run.assert_not_awaited() - dead_lettered = repo.dead_letter_notification.await_args.kwargs - assert dead_lettered["dispatch_version"] == 2 - assert dead_lettered["count_failure"] is False - assert "5 failed attempts" in dead_lettered["error"] + get_run.assert_awaited_once_with("notify-run-1", user_id="user-1") + repo.dead_letter_notification.assert_not_awaited() + finished = repo.finish_notification_run.await_args.kwargs + assert finished["dispatch_version"] == 2 + assert finished["delivered"] is True + assert finished["next_notification_at"] is None @pytest.mark.asyncio @@ -4125,6 +4380,54 @@ async def test_notification_cancellation_during_source_run_lookup_releases_claim repo.release_notification_claim.assert_awaited_once() +@pytest.mark.asyncio +async def test_reclaimed_launching_notification_cancellation_preserves_phase(): + launch_started = asyncio.Event() + + async def launch_notification(**_kwargs): + launch_started.set() + await asyncio.Event().wait() + + repo = SimpleNamespace( + release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), + ) + service = McpTaskService( + repository=repo, + drivers=McpTaskDriverRegistry(), + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + launch_notification=launch_notification, + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + record = { + **_claimed_row(), + "notification_status": "launching", + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + } + now = datetime.now(UTC) + + task = asyncio.create_task( + service._run_claimed_batch( + [record], + operation=lambda r: service._notify_one_claimed(r, now=now), + release=service._release_notification_after_cancellation, + action="notification", + ) + ) + await launch_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + repo.release_notification_lease.assert_awaited_once() + repo.release_notification_claim.assert_not_awaited() + + @pytest.mark.asyncio async def test_dispatched_notification_cancellation_preserves_phase_when_releasing_lease(): lookup_started = asyncio.Event() diff --git a/frontend/src/core/background-tasks/types.ts b/frontend/src/core/background-tasks/types.ts index fb31589a77..61feaa0cf8 100644 --- a/frontend/src/core/background-tasks/types.ts +++ b/frontend/src/core/background-tasks/types.ts @@ -10,6 +10,7 @@ export type BackgroundTaskNotificationStatus = | "none" | "pending" | "claimed" + | "launching" | "retry" | "dispatched" | "delivered" @@ -45,7 +46,7 @@ export const ACTIVE_BACKGROUND_TASK_STATUSES: ReadonlySet new Set(["submitted", "working", "input_required"]); export const ACTIVE_BACKGROUND_TASK_NOTIFICATION_STATUSES: ReadonlySet = - new Set(["pending", "claimed", "retry", "dispatched"]); + new Set(["pending", "claimed", "launching", "retry", "dispatched"]); export function isActiveBackgroundTask(task: BackgroundTask): boolean { return ACTIVE_BACKGROUND_TASK_STATUSES.has(task.status); diff --git a/frontend/tests/unit/core/background-tasks/types.test.ts b/frontend/tests/unit/core/background-tasks/types.test.ts index b08218d5a7..5d3b6a5483 100644 --- a/frontend/tests/unit/core/background-tasks/types.test.ts +++ b/frontend/tests/unit/core/background-tasks/types.test.ts @@ -29,7 +29,7 @@ const TERMINAL_TASK: BackgroundTaskDetail = { }; describe("background task detail polling", () => { - it.each(["pending", "claimed", "retry", "dispatched"] as const)( + it.each(["pending", "claimed", "launching", "retry", "dispatched"] as const)( "keeps polling a terminal task while notification status is %s", (notificationStatus) => { expect(