From d642dfdd2f86cfd14daff7b42296520d6f413abc Mon Sep 17 00:00:00 2001 From: RongJie G <111257566+CorgiBoyG@users.noreply.github.com> Date: Wed, 30 Sep 2026 09:12:11 +0800 Subject: [PATCH] fix(mcp): reserve notification launches atomically (#5965) * fix(mcp): reserve notification launches atomically Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> * fix(mcp): classify busy notification conflicts by type Signed-off-by: Willem Jiang --------- Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Signed-off-by: Willem Jiang Co-authored-by: Willem Jiang --- backend/app/gateway/AGENTS.md | 2 +- backend/app/gateway/services.py | 13 +- backend/app/mcp_tasks/service.py | 41 ++++ .../packages/harness/deerflow/mcp/AGENTS.md | 4 +- .../deerflow/persistence/mcp_tasks/sql.py | 41 +++- backend/tests/test_gateway_services.py | 127 +++++++++- backend/tests/test_mcp_task_repository.py | 139 ++++++++++- backend/tests/test_mcp_task_service.py | 218 +++++++++++++++++- 8 files changed, 556 insertions(+), 29 deletions(-) diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index 69cab1d358..6f165b4d9d 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -53,7 +53,7 @@ hot-reloaded top-level `recursion_limit` setting. A valid request-level value takes precedence; invalid values fall back to that default, and `max_recursion_limit` caps both sources. -Durable MCP notifications use internal Agent runs. Keep their trusted delivery instruction outside the user-input boundary, and frame serialized remote events as untrusted before model invocation. Strict thread existence/ownership admission dead-letters events whose task outlives its deleted chat instead of recreating the thread. +Durable MCP notifications use internal Agent runs. Keep their trusted delivery instruction outside the user-input boundary, and frame serialized remote events as untrusted before model invocation. Strict thread existence/ownership admission dead-letters events whose task outlives its deleted chat instead of recreating the thread. `start_run` marks run-manager admission conflicts with `BusyThreadConflict`; the notification launcher retries only that typed 409 and treats other 409 responses as permanent. CORS is same-origin by default when requests enter through nginx on port 2026. Split-origin or port-forwarded browser clients must opt in with `GATEWAY_CORS_ORIGINS` (exact origins); Gateway `CORSMiddleware` and `CSRFMiddleware` both read that variable so browser CORS and auth-origin checks stay aligned. Those clients also need `CORS_EXPOSED_HEADERS` (`csrf_middleware.py`): run-creating routes return the run's id in `Content-Location`, which is not CORS-safelisted, so JS cannot read it unless it is exposed — and the LangGraph SDK resolves run metadata from that header alone, so withholding it breaks `useStream`'s `onCreated` and thread-gated actions. diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 3f8c94e3cf..05e69043c6 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -104,6 +104,13 @@ from deerflow.utils.thread_id import validate_thread_id logger = logging.getLogger(__name__) +class BusyThreadConflict(HTTPException): + """A retryable run-manager admission conflict exposed as HTTP 409.""" + + def __init__(self, detail: str) -> None: + super().__init__(status_code=409, detail=detail) + + @asynccontextmanager async def reserve_checkpoint_write( request: Request, @@ -2125,7 +2132,7 @@ async def start_run( ) raise except ConflictError as exc: - raise HTTPException(status_code=409, detail=str(exc)) from exc + raise BusyThreadConflict(str(exc)) from exc except UnsupportedStrategyError as exc: raise HTTPException(status_code=501, detail=str(exc)) from exc @@ -2286,9 +2293,9 @@ async def launch_mcp_task_notification_run( require_existing_thread=True, ) except HTTPException as exc: - if exc.status_code == 409: + if isinstance(exc, BusyThreadConflict): raise ConflictError(str(exc.detail)) from exc - if exc.status_code == 404: + if exc.status_code in {400, 401, 403, 404, 409, 422, 501}: raise PermanentNotificationError(str(exc.detail)) from exc raise return {"run_id": record.run_id, "thread_id": record.thread_id} diff --git a/backend/app/mcp_tasks/service.py b/backend/app/mcp_tasks/service.py index 46546370aa..a44b384d52 100644 --- a/backend/app/mcp_tasks/service.py +++ b/backend/app/mcp_tasks/service.py @@ -936,6 +936,47 @@ class McpTaskService: return source_run = await self._get_run(record.get("run_id"), user_id=record["user_id"]) if record.get("run_id") else None + launch_started_at = _notification_completion_time(not_before=now) + recovering_launch = record.get("notification_status") == "launching" + # Reservation commit outcome is ambiguous under cancellation. Prefer a + # phase-preserving lease release until the await returns definitively. + record["notification_status"] = "launching" + launch_reserved = await self._repository.begin_notification_launch( + task_id, + lease_owner=self._lease_owner, + notification_lease_token=record["notification_lease_token"], + dispatch_version=dispatch_version, + lease_seconds=self._lease_seconds, + now=launch_started_at, + ) + if not launch_reserved: + if recovering_launch: + 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=launch_started_at, + error=record.get("notification_error"), + count_failure=False, + ), + action="release expired notification launch", + ) + return + 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=launch_started_at, + error=record.get("notification_error"), + replace_with_latest=True, + ), + action="release stale notification launch", + ) + return try: result = await self._launch_notification( thread_id=record["thread_id"], diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index da280bbc51..bd31851ae5 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -62,9 +62,9 @@ **This is defense in depth, not a trust boundary.** `npx`/`uvx` exist to fetch and execute remote packages, so an admin can still point one at a package they published; the boundary is admin authentication plus network reachability. Do not add a check here on the assumption that it makes MCP registration safe for untrusted admins — it does not, and the fix for that is not a bigger denylist. **Durable MCP task runtime** (`mcp_tasks`, `McpTaskService`; summarized in -[backend/AGENTS.md](../../../../AGENTS.md)): Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry or after a cancel request must be discarded even when the owner token still matches. The first cancel request fences an in-flight poll lease, while repeats preserve an active cancellation lease so they cannot issue concurrent remote cancels; cancellation backoff starts when the remote attempt finishes, so a slow timeout cannot consume the retry delay. Cancellation, polling, and notification batches isolate per-task exceptions; an unexpected cancellation/poll failure leaves that record's lease to expire, while notification failures release only the affected lease for retry. Input-required and terminal event snapshots are delivered by idempotent Agent runs and marked delivered only after run success; the trusted notification instruction stays outside the input boundary while the serialized remote event is framed as untrusted data. A busy-thread conflict is normalized back to the service boundary so the queued snapshot coalesces to the latest task event. A missing dispatched run becomes a failed delivery attempt, while transient run-store hydration errors stay distinguishable and retry the same lookup. The database is the source of truth; `ThreadState` receives only a bounded current-thread projection, and display names are neutralized at that model-state boundary. The installed process-local submitter is the source of truth for management-tool exposure; hot `mcp_tasks` edits take effect only after restart, and active skills must explicitly declare the list/cancel business tools. +[backend/AGENTS.md](../../../../AGENTS.md)): Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry or after a cancel request must be discarded even when the owner token still matches. The first cancel request fences an in-flight poll lease, while repeats preserve an active cancellation lease so they cannot issue concurrent remote cancels; cancellation backoff starts when the remote attempt finishes, so a slow timeout cannot consume the retry delay. Cancellation, polling, and notification batches isolate per-task exceptions; an unexpected cancellation/poll failure leaves that record's lease to expire, while notification failures release only the affected lease for retry. Input-required and terminal event snapshots are delivered by idempotent Agent runs and marked delivered only after run success; the trusted notification instruction stays outside the input boundary while the serialized remote event is framed as untrusted data. Before creating an Agent notification run, one owner/token/version/expiry-fenced update reserves `notification_status=launching`; from that point, cancellation, ambiguous launch outcomes, and busy-thread conflicts preserve the reserved snapshot and idempotency key for recovery instead of coalescing to a newer event. A missing dispatched run becomes a failed delivery attempt, while transient run-store hydration errors stay distinguishable and retry the same lookup. The database is the source of truth; `ThreadState` receives only a bounded current-thread projection, and display names are neutralized at that model-state boundary. The installed process-local submitter is the source of truth for management-tool exposure; hot `mcp_tasks` edits take effect only after restart, and active skills must explicitly declare the list/cancel business tools. -**Task notification failure handling**: MCP notification failures use a consecutive counter separate from the idempotency-key `dispatch_attempt`, capped exponential backoff, latest-event rebuilding before a run launches, and a five-attempt budget before `dead_letter`. A permanently missing/mismatched target thread is dead-lettered immediately instead of being recreated or reclaimed. HTTP and Agent cancellation requests return after the durable cancel fence; the background loop alone owns the potentially slow remote call and retry schedule. The HTTP cancel endpoint rejects requests with 503 when the loop is not running (`mcp_tasks_available` false, e.g. `mcp_tasks.enabled=false` with SQL persistence), so a cancellation is never acknowledged without a worker to perform it. The bounded notification error/count/status join poll and cancellation diagnostics in the task detail API and expanded card. +**Task notification failure handling**: MCP notification failures use a consecutive counter separate from the idempotency-key `dispatch_attempt`, capped exponential backoff, latest-event rebuilding before a run launches, and a five-attempt budget for pre-reservation and dispatched-run failures before `dead_letter`. A permanently rejected launch, including a missing/mismatched target thread, is dead-lettered immediately instead of being recreated or reclaimed. An ambiguous launch result or busy-thread conflict preserves `launching`, the snapshot, and the same idempotency key without consuming that budget so a retry can reconcile a run that may already exist. HTTP and Agent cancellation requests return after the durable cancel fence; the background loop alone owns the potentially slow remote call and retry schedule. The HTTP cancel endpoint rejects requests with 503 when the loop is not running (`mcp_tasks_available` false, e.g. `mcp_tasks.enabled=false` with SQL persistence), so a cancellation is never acknowledged without a worker to perform it. The bounded notification error/count/status join poll and cancellation diagnostics in the task detail API and expanded card. - Personal MCP configs live under the authenticated user root (`user_config.py`), use literal values, and never enter the deployment tool cache. `user_tools.py` guards owner and current revision before every call. Runtime server names bind owner/name/revision; durable task `driver_data.connection_scope` records personal or deployment ownership at submission, and background calls resolve a personal key only in the persisted task owner's file. Older rows without that field are deployment-owned. When a deployment name matches a personal runtime name, Agent assembly omits the personal tool rather than exposing two tools under one name; an already submitted personal task still uses its recorded personal scope. Keep catalog IDs stable across edits; never fall back by display name. - Personal definitions without `personal_public_network: true` require current administrator authority at discovery and every tool/durable call. `personal_access.py` uses a host-installed async lookup, never a captured role or cached decision; absent/failed lookup denies privileged connections. The Gateway binds the lookup to fresh account records on its database-owning loop for its lifespan. Deployment connections are unaffected. This is admission control, not termination of operations already admitted before a role change. diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index 93295c6d35..32ace7f401 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -6,7 +6,7 @@ import uuid from datetime import UTC, datetime, timedelta from typing import Any -from sqlalchemy import case, or_, select, text, update +from sqlalchemy import and_, case, or_, select, text, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -675,6 +675,45 @@ class McpTaskRepository: await session.commit() return [self._row_to_dict(row) for row in rows] + async def begin_notification_launch( + self, + task_id: str, + *, + lease_owner: str, + notification_lease_token: str, + dispatch_version: int, + lease_seconds: int, + now: datetime, + ) -> bool: + """Reserve one idempotent Agent launch before starting the side effect.""" + launchable = or_( + McpTaskRow.notification_status == "launching", + and_( + McpTaskRow.notification_status.in_(("claimed", "retry")), + McpTaskRow.event_version == dispatch_version, + ), + ) + stmt = ( + update(McpTaskRow) + .where( + McpTaskRow.id == task_id, + McpTaskRow.notification_lease_owner == lease_owner, + McpTaskRow.notification_lease_token == notification_lease_token, + McpTaskRow.notification_lease_expires_at >= now, + McpTaskRow.dispatch_version == dispatch_version, + launchable, + ) + .values( + notification_status="launching", + notification_lease_expires_at=now + timedelta(seconds=lease_seconds), + updated_at=now, + ) + ) + async with self._sf() as session: + result = await session.execute(stmt) + await session.commit() + return bool(result.rowcount) + async def mark_notification_dispatched( self, task_id: str, diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index 45e22f3e7a..7d0f00d9e2 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -3683,23 +3683,78 @@ def test_launch_mcp_task_notification_run_hides_internal_prompt(_stub_app_config assert result == {"run_id": "run-notification", "thread_id": "thread-notification"} +def test_start_run_marks_run_manager_conflict_as_busy(_stub_app_config): + import asyncio + from unittest.mock import AsyncMock, patch + + from app.gateway.services import BusyThreadConflict, start_run + from deerflow.runtime.runs.manager import ConflictError + + async def _scenario(): + request, _run_store, _thread_store = _make_start_run_persistence_context() + request.app.state.run_manager.create_or_reject = AsyncMock(side_effect=ConflictError("Thread already has an active run")) + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=object()), + patch("app.gateway.services.ensure_checkpoint_history_seeded", new_callable=AsyncMock), + pytest.raises(BusyThreadConflict) as exc_info, + ): + await start_run(_run_create_request(), "thread-busy", request) + assert exc_info.value.status_code == 409 + assert exc_info.value.detail == "Thread already has an active run" + + asyncio.run(_scenario()) + + def test_launch_mcp_task_notification_run_restores_busy_thread_conflict(_stub_app_config): import asyncio from types import SimpleNamespace from unittest.mock import patch + from app.gateway.services import BusyThreadConflict, launch_mcp_task_notification_run + from deerflow.runtime.runs.manager import ConflictError + + async def _scenario(): + with ( + patch( + "app.gateway.services.start_run", + side_effect=BusyThreadConflict("Thread already has an active run"), + ), + pytest.raises(ConflictError, match="Thread already has an active run"), + ): + await launch_mcp_task_notification_run( + app=SimpleNamespace(state=SimpleNamespace()), + thread_id="thread-notification", + assistant_id="lead_agent", + owner_user_id="user-1", + task_id="task-1", + dispatch_version=2, + dispatch_attempt=3, + event={"status": "completed", "result": "done"}, + ) + + asyncio.run(_scenario()) + + +def test_launch_mcp_task_notification_run_dead_letters_idempotency_conflict(_stub_app_config): + import asyncio + from types import SimpleNamespace + from unittest.mock import patch + from fastapi import HTTPException from app.gateway.services import launch_mcp_task_notification_run - from deerflow.runtime.runs.manager import ConflictError + from app.mcp_tasks.errors import PermanentNotificationError async def _scenario(): with ( patch( "app.gateway.services.start_run", - side_effect=HTTPException(status_code=409, detail="Thread already has an active run"), + side_effect=HTTPException( + status_code=409, + detail="Idempotency-Key already used with a different request", + ), ), - pytest.raises(ConflictError, match="Thread already has an active run"), + pytest.raises(PermanentNotificationError, match="Idempotency-Key"), ): await launch_mcp_task_notification_run( app=SimpleNamespace(state=SimpleNamespace()), @@ -3747,6 +3802,72 @@ def test_launch_mcp_task_notification_run_dead_letters_missing_thread(_stub_app_ asyncio.run(_scenario()) +@pytest.mark.parametrize("status_code", [400, 401, 403, 422, 501]) +def test_launch_mcp_task_notification_run_dead_letters_deterministic_rejection(_stub_app_config, status_code): + import asyncio + from types import SimpleNamespace + from unittest.mock import patch + + from fastapi import HTTPException + + from app.gateway.services import launch_mcp_task_notification_run + from app.mcp_tasks.errors import PermanentNotificationError + + async def _scenario(): + with ( + patch( + "app.gateway.services.start_run", + side_effect=HTTPException(status_code=status_code, detail="request rejected"), + ), + pytest.raises(PermanentNotificationError, match="request rejected"), + ): + await launch_mcp_task_notification_run( + app=SimpleNamespace(state=SimpleNamespace()), + thread_id="thread-notification", + assistant_id="lead_agent", + owner_user_id="user-1", + task_id="task-1", + dispatch_version=2, + dispatch_attempt=3, + event={"status": "completed", "result": "done"}, + ) + + asyncio.run(_scenario()) + + +@pytest.mark.parametrize("status_code", [408, 429, 500]) +def test_launch_mcp_task_notification_run_preserves_ambiguous_http_failure(_stub_app_config, status_code): + import asyncio + from types import SimpleNamespace + from unittest.mock import patch + + from fastapi import HTTPException + + from app.gateway.services import launch_mcp_task_notification_run + + async def _scenario(): + with ( + patch( + "app.gateway.services.start_run", + side_effect=HTTPException(status_code=status_code, detail="launch outcome unknown"), + ), + pytest.raises(HTTPException) as exc_info, + ): + await launch_mcp_task_notification_run( + app=SimpleNamespace(state=SimpleNamespace()), + thread_id="thread-notification", + assistant_id="lead_agent", + owner_user_id="user-1", + task_id="task-1", + dispatch_version=2, + dispatch_attempt=3, + event={"status": "completed", "result": "done"}, + ) + assert exc_info.value.status_code == status_code + + asyncio.run(_scenario()) + + def test_start_run_strict_mode_rejects_missing_thread(_stub_app_config): import asyncio from types import SimpleNamespace diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index ad523ac95b..8b805a52b7 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -1024,14 +1024,14 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_ ) 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() + assert await repo.begin_notification_launch( + "task-launching-compat", + lease_owner="notifier-a", + notification_lease_token=first[0]["notification_lease_token"], + dispatch_version=first_version, + lease_seconds=60, + now=now, + ) second_poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) await repo.apply_snapshot( @@ -1049,8 +1049,9 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_ polled_at=now, ) + reclaimed_at = now + timedelta(seconds=61) reclaimed = await repo.claim_notification_work( - now=now, + now=reclaimed_at, lease_owner="notifier-b", lease_seconds=60, limit=1, @@ -1072,9 +1073,9 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_ get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), ) service._lease_owner = "notifier-b" - await service._notify_one_claimed(reclaimed[0], now=now) + await service._notify_one_claimed(reclaimed[0], now=reclaimed_at) - retry_at = now + timedelta(seconds=10) + retry_at = reclaimed_at + timedelta(seconds=10) after_conflict = await repo.claim_notification_work( now=retry_at, lease_owner="notifier-c", @@ -1097,6 +1098,122 @@ async def test_reclaimed_launching_notification_preserves_reserved_snapshot(tmp_ ) +@pytest.mark.asyncio +async def test_begin_notification_launch_rejects_unstarted_stale_snapshot(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-notify-stale", now=now) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-notify-stale", + 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", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + second_poll = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-notify-stale", + 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, + ) + + reserved = await repo.begin_notification_launch( + "task-notify-stale", + lease_owner="notifier", + notification_lease_token=first[0]["notification_lease_token"], + dispatch_version=first[0]["dispatch_version"], + lease_seconds=120, + now=now, + ) + + assert reserved is False + stored = await repo.get("task-notify-stale", user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert stored is not None + assert stored["notification_status"] == "claimed" + assert stored["event_version"] > stored["dispatch_version"] + + +@pytest.mark.asyncio +async def test_begin_notification_launch_is_fenced_by_reclaimed_token(tmp_path, monkeypatch): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="task-launch-token-fence", now=now) + poll_claim = await repo.claim_due_tasks(now=now, lease_owner="poller", lease_seconds=60, limit=1) + await repo.apply_snapshot( + "task-launch-token-fence", + 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, + ) + first = await repo.claim_notification_work( + now=now, + lease_owner="same-worker", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + operation = repo.begin_notification_launch( + "task-launch-token-fence", + lease_owner="same-worker", + notification_lease_token=first[0]["notification_lease_token"], + dispatch_version=first[0]["dispatch_version"], + lease_seconds=120, + now=now, + ) + + async with _pause_claim_mutation(monkeypatch, operation) as (stale_launch, resume): + reclaimed = await repo.claim_notification_work( + now=now + timedelta(seconds=61), + lease_owner="same-worker", + lease_seconds=60, + limit=1, + tracking_degraded_after_errors=3, + ) + assert len(reclaimed) == 1 + new_token = reclaimed[0]["notification_lease_token"] + assert new_token != first[0]["notification_lease_token"] + resume.set() + assert await stale_launch is False + + stored = await repo.get("task-launch-token-fence", user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert stored is not None + assert stored["notification_status"] == "claimed" + assert stored["notification_lease_owner"] == "same-worker" + assert stored["notification_lease_token"] == new_token + + @pytest.mark.asyncio async def test_recovered_launching_at_retry_budget_reconciles_successful_run(tmp_path): repo = await _make_repo(tmp_path) diff --git a/backend/tests/test_mcp_task_service.py b/backend/tests/test_mcp_task_service.py index 95e1dc447f..d719d16f7c 100644 --- a/backend/tests/test_mcp_task_service.py +++ b/backend/tests/test_mcp_task_service.py @@ -65,6 +65,9 @@ class FakeRepository: self.released.append((task_id, kwargs)) return True + async def begin_notification_launch(self, *_args, **_kwargs): + return True + class FailingApplyRepository(FakeRepository): async def apply_snapshot(self, task_id, **kwargs): @@ -695,6 +698,7 @@ async def test_cancel_recovery_failures_are_isolated_and_later_phases_continue(c @pytest.mark.asyncio async def test_notification_delivery_waits_for_successful_agent_run(): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), mark_notification_dispatched=AsyncMock(return_value=True), finish_notification_run=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), @@ -893,9 +897,11 @@ async def test_notification_failure_releases_at_its_own_completion_without_waiti @pytest.mark.asyncio -async def test_notification_busy_thread_replaces_claim_with_latest_event(monkeypatch): +async def test_reserved_notification_busy_thread_preserves_snapshot(monkeypatch): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), ) service = McpTaskService( repository=repo, @@ -920,14 +926,16 @@ async def test_notification_busy_thread_replaces_claim_with_latest_event(monkeyp now=now, ) - released = repo.release_notification_claim.await_args.kwargs - assert released["replace_with_latest"] is True + 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_reclaimed_launching_busy_thread_preserves_reserved_snapshot(monkeypatch): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), release_notification_lease=AsyncMock(return_value=True), ) @@ -962,9 +970,11 @@ async def test_reclaimed_launching_busy_thread_preserves_reserved_snapshot(monke @pytest.mark.asyncio -async def test_notification_launch_failure_backs_off_and_replaces_with_latest_event(monkeypatch): +async def test_notification_launch_failure_backs_off_without_replacing_reserved_snapshot(monkeypatch): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), + release_notification_lease=AsyncMock(return_value=True), ) service = McpTaskService( repository=repo, @@ -991,15 +1001,16 @@ async def test_notification_launch_failure_backs_off_and_replaces_with_latest_ev now=now, ) - released = repo.release_notification_claim.await_args.kwargs - assert released["replace_with_latest"] is True - assert released["count_failure"] is True + 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=40) @pytest.mark.asyncio async def test_permanently_rejected_notification_is_dead_lettered(): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), dead_letter_notification=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), ) @@ -1037,6 +1048,7 @@ async def test_permanently_rejected_notification_is_dead_lettered(): @pytest.mark.asyncio async def test_permanently_rejected_reclaimed_launching_notification_is_dead_lettered(): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), dead_letter_notification=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), release_notification_lease=AsyncMock(return_value=True), @@ -1114,7 +1126,10 @@ async def test_notification_dispatch_uses_launch_completion_time(monkeypatch): return {"run_id": "notify-run-1"} monkeypatch.setattr(service_module, "datetime", _MutableDateTime) - repo = SimpleNamespace(mark_notification_dispatched=AsyncMock(return_value=False)) + repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), + mark_notification_dispatched=AsyncMock(return_value=False), + ) service = McpTaskService( repository=repo, drivers=McpTaskDriverRegistry(), @@ -1139,9 +1154,80 @@ async def test_notification_dispatch_uses_launch_completion_time(monkeypatch): assert repo.mark_notification_dispatched.await_args.kwargs["now"] == launch_completed_at +@pytest.mark.asyncio +async def test_notification_does_not_launch_when_preflight_fence_is_rejected(): + launch_notification = AsyncMock(return_value={"run_id": "stale-notify-run"}) + repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=False), + mark_notification_dispatched=AsyncMock(return_value=False), + release_notification_claim=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=datetime.now(UTC), + ) + + repo.begin_notification_launch.assert_awaited_once() + repo.release_notification_claim.assert_awaited_once() + launch_notification.assert_not_awaited() + repo.mark_notification_dispatched.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_reclaimed_launching_notification_preserves_snapshot_when_preflight_lease_expires(): + launch_notification = AsyncMock(return_value={"run_id": "stale-notify-run"}) + repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=False), + release_notification_claim=AsyncMock(return_value=False), + 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")), + ) + + await service._notify_one_claimed( + { + **_claimed_row(), + "notification_status": "launching", + "dispatch_version": 2, + "dispatch_attempt": 0, + "dispatch_event": {"status": "completed"}, + }, + now=datetime.now(UTC), + ) + + repo.begin_notification_launch.assert_awaited_once() + repo.release_notification_lease.assert_awaited_once() + repo.release_notification_claim.assert_not_awaited() + launch_notification.assert_not_awaited() + + @pytest.mark.asyncio async def test_reclaimed_launching_notification_preserves_snapshot_on_uncertain_failure_at_retry_budget(): repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), release_notification_lease=AsyncMock(return_value=True), dead_letter_notification=AsyncMock(return_value=True), @@ -2386,6 +2472,9 @@ class NotificationPersistenceRepo: } ] + async def begin_notification_launch(self, *_args, **_kwargs): + return True + async def mark_notification_dispatched(self, *_args, **_kwargs): self.mark_started.set() await asyncio.Event().wait() @@ -3167,6 +3256,10 @@ class SelfCancellingNotificationRepository(FakeRepository): self.notification_releases.append((task_id, kwargs)) return True + async def release_notification_lease(self, task_id, **kwargs): + self.notification_releases.append((task_id, kwargs)) + return True + @pytest.mark.asyncio async def test_notification_child_self_cancellation_releases_once(): @@ -3201,6 +3294,10 @@ class BatchNotificationRepository(FakeRepository): self.notification_releases.append((task_id, kwargs)) return True + async def release_notification_lease(self, task_id, **kwargs): + self.notification_releases.append((task_id, kwargs)) + return True + @pytest.mark.asyncio async def test_notification_outer_cancellation_releases_started_and_never_started_once(): @@ -3537,6 +3634,9 @@ class OrdinaryReleaseBatchRepository: async def claim_notification_work(self, **_kwargs): return self._records("notification") + async def begin_notification_launch(self, *_args, **_kwargs): + return True + async def _release(self, task_id, **_kwargs): self.release_calls.append(task_id) self.release_started.set() @@ -3570,6 +3670,9 @@ class OrdinaryReleaseBatchRepository: async def release_notification_claim(self, task_id, **kwargs): return await self._release(task_id, **kwargs) + async def release_notification_lease(self, task_id, **kwargs): + return await self._release(task_id, **kwargs) + async def release_poll_claim_after_cancellation(self, task_id, **kwargs): return await self._release(task_id, **kwargs) @@ -4382,6 +4485,104 @@ async def test_notification_cancellation_during_source_run_lookup_releases_claim repo.release_notification_claim.assert_awaited_once() +@pytest.mark.asyncio +async def test_notification_cancellation_during_launch_reservation_preserves_phase(): + reservation_started = asyncio.Event() + + async def begin_notification_launch(*_args, **_kwargs): + reservation_started.set() + await asyncio.Event().wait() + + repo = SimpleNamespace( + begin_notification_launch=begin_notification_launch, + 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(), + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + record = { + **_claimed_row(), + "notification_status": "claimed", + "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 item: service._notify_one_claimed(item, now=now), + release=service._release_notification_after_cancellation, + action="notification", + ) + ) + await reservation_started.wait() + assert record["notification_status"] == "launching" + 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_notification_cancellation_after_launch_reservation_preserves_launching_phase(): + launch_started = asyncio.Event() + + async def launch_notification(**_kwargs): + launch_started.set() + await asyncio.Event().wait() + + repo = SimpleNamespace( + begin_notification_launch=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=launch_notification, + get_run=AsyncMock(return_value=SimpleNamespace(assistant_id="lead_agent")), + ) + record = { + **_claimed_row(), + "notification_status": "claimed", + "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 item: service._notify_one_claimed(item, now=now), + release=service._release_notification_after_cancellation, + action="notification", + ) + ) + await launch_started.wait() + assert record["notification_status"] == "launching" + 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_reclaimed_launching_notification_cancellation_preserves_phase(): launch_started = asyncio.Event() @@ -4391,6 +4592,7 @@ async def test_reclaimed_launching_notification_cancellation_preserves_phase(): await asyncio.Event().wait() repo = SimpleNamespace( + begin_notification_launch=AsyncMock(return_value=True), release_notification_claim=AsyncMock(return_value=True), release_notification_lease=AsyncMock(return_value=True), )