github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-27 13:40:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: make MCP notification recovery rollback-safe (#5908)

Browse Source 
Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com>
This commit is contained in:
RongJie G 2026-09-27 20:46:56 +08:00 committed by GitHub
parent 457ac8cabd
commit 81c5c83384
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 637 additions and 121 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

216
backend/app/mcp_tasks/service.py
View File

@ -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,

8
backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py
View File

@ -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",

204
backend/tests/test_mcp_task_repository.py
View File

@ -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",

325
backend/tests/test_mcp_task_service.py
View File

@ -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()

3
frontend/src/core/background-tasks/types.ts
View File

@ -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<BackgroundTaskStatus>
new Set(["submitted", "working", "input_required"]);
export const ACTIVE_BACKGROUND_TASK_NOTIFICATION_STATUSES: ReadonlySet<BackgroundTaskNotificationStatus> =
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);

2
frontend/tests/unit/core/background-tasks/types.test.ts
View File

@ -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(