diff --git a/README.md b/README.md index 6234c78971..11696958ba 100644 --- a/README.md +++ b/README.md @@ -138,7 +138,9 @@ With `run_events.backend: db`, caller-supplied ISO 8601 event timestamps are nor Run event streams include `run.terminal` for the owning worker's confirmed final Run status, while `run.end` continues to describe each root graph completion. Terminal publication is best effort; event persistence failure does not undo the -Run outcome or prevent stream completion. See [Run Event Stream](backend/docs/RUN_EVENT_STREAM.md#authoritative-run-outcome). +Run outcome or prevent stream completion. Recovery with built-in SQL Run and Event stores +sharing an engine uses one pooled connection for each guarded event write, so it +can publish with only one pool slot available. See [Run Event Stream](backend/docs/RUN_EVENT_STREAM.md#authoritative-run-outcome). Operators can extend lead-agent, subagent, and DeerMem extraction prompts with literal prepend/append configuration without editing source templates. See diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index 0502fbf1e9..d4b580c575 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -898,6 +898,9 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen run_events_config = getattr(config, "run_events", None) app.state.run_events_config = run_events_config app.state.run_event_store = make_run_event_store(run_events_config) + from deerflow.runtime.events.terminal import RecoveryTerminalPublisher + + app.state.recovery_terminal_publisher = RecoveryTerminalPublisher(app.state.run_event_store, app.state.run_store) from deerflow.extensions.run_evidence import StoreRunEvidenceReader, StoreRunEvidenceReaderFactory @@ -974,10 +977,12 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen app.state.scheduled_task_repo = ScheduledTaskRepository( sf, run_repository=app.state.run_store, + recovery_terminal_publisher=app.state.recovery_terminal_publisher, ) app.state.scheduled_task_run_repo = ScheduledTaskRunRepository( sf, run_repository=app.state.run_store, + recovery_terminal_publisher=app.state.recovery_terminal_publisher, ) app.state.scheduled_task_event_repo = ScheduledTaskEventRepository(sf) # Per-user unread state and the activity feed's read clock. @@ -1020,6 +1025,7 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen store=app.state.run_store, run_ownership_config=run_ownership_config, event_store=app.state.run_event_store, + recovery_terminal_publisher=app.state.recovery_terminal_publisher, on_orphans_recovered=terminalize_recovered_runs, ) # Startup recovery: mark inflight runs whose lease has expired as error. diff --git a/backend/docs/RUN_EVENT_STREAM.md b/backend/docs/RUN_EVENT_STREAM.md index d834aef0de..20a2bc109f 100644 --- a/backend/docs/RUN_EVENT_STREAM.md +++ b/backend/docs/RUN_EVENT_STREAM.md @@ -75,8 +75,10 @@ outcome. It is a singleton identified by `(thread_id, run_id, event_type)` and confirmed through `put_if_absent()`. Content contains only `status`: `success`, `error`, `interrupted`, or `timeout`; graph outputs and error details are excluded. Metadata requires `protocol_version: 1` and a boolean `recovered`. -Current Workers emit `recovered: false`; the marker describes evidence origin -and grants no publication authority. Field definitions remain in the +Owner Workers emit `recovered: false`; a confirmed recovery takeover emits +`recovered: true` only when it creates a missing terminal. An existing matching +event retains its original marker. The marker describes evidence origin and +grants no publication authority. Field definitions remain in the [JSON contract](../../contracts/run_event_stream_contract.json). `run.start` and `run.end` retain their root-invocation producers, counts, @@ -97,6 +99,18 @@ Rollback's `interrupted` ACK is intermediate while restoration remains pending; publication waits for restoration and the final `error` ACK. Ownership loss, rejected or unconfirmed commits, and unknown predecessor writes fail closed. Reading a historical terminal row does not create local publication authority. +Recovery publication requires the boolean result of a successful atomic takeover +claim, with the target bound to the claimed Run identity and its committed +`error` outcome. The target is one-shot and cannot be reconstructed with another +Run identity. A later same-status RunRow read cannot create that authority. +Recovery receipts and terminal confirmation re-check the Run lifecycle inside +the Event Store's mutation fence; a completed thread deletion wins, +so a stale target cannot recreate an orphan event. +When SQL Run and Event stores share an engine, the built-in guard reads the Run +through the active event transaction, so recovery needs only one available pooled +connection. It retains the mutation fence through the check and event commit. +Different engines and custom Run lookups keep their normal read path. Memory +guard/deletion fences are per-thread; a slow guard does not block unrelated threads. ### Evidence Ordering @@ -123,6 +137,11 @@ appended later cannot be claimed as covered by an earlier Seal; an event alone does not prove complete evidence. Barrier release before END preserves immediate next-message admission, including when later cleanup is slow. +Recovery publishes its zero-delivery receipt before `run.terminal`, matching the +normal delivery-before-terminal order. Receipt failure remains non-fatal and does +not prevent the terminal confirmation attempt. Both recovery writes are fenced +by the current Run lifecycle. + ### Idempotence and Failure Semantics Confirmation validates the returned event's identity, category, status, version, @@ -132,7 +151,7 @@ required field types, positive sequence, and database user ownership: | --- | --- | | `CONFIRMED_NEW` | Insertion is confirmed and the returned event matches; Receipt capture may proceed. | | `CONFIRMED_EXISTING` | The existing event itself matches; accept idempotently and preserve its sequence and original source marker. | -| `CONFLICT` | Identity or protocol fields contradict expectations or are malformed; record the conflict, do not overwrite, and do not seal as complete. | +| `CONFLICT` | Identity or protocol fields contradict expectations or are malformed, or the guarded Run lifecycle no longer admits publication; record the conflict, do not overwrite, and do not seal as complete. | | `UNKNOWN` | The operation outcome cannot be established; preserve uncertainty and do not seal as complete. | `created=False` alone proves neither success nor conflict. Additional optional @@ -162,11 +181,16 @@ Confirmation is limited to the Store's declared failure and concurrency model: | SQLite | Database-backed confirmation with per-thread in-process serialization; this does not establish PostgreSQL or multi-Gateway concurrency semantics. | | PostgreSQL | Database-backed confirmation with per-thread transaction advisory locking across independent connections/processes; concurrency validation requires real sessions. | -Singleton insertion and RunRow terminal commit are separate transactions. A crash -between them can leave a terminal Run without an event. Recovery publication, -takeover backfill, historical migration, and durable outbox/reconciliation are -outside the current owner-side implementation; eventual evidence after every -terminal commit is not guaranteed. Missing events must not gate SSE recovery. +Singleton insertion and RunRow terminal commit are separate transactions. +`RecoveryTerminalPublisher` backfills an error terminal only after a successful +takeover claim. A crash between a normal terminal commit and event confirmation +can still leave a terminal Run without an event when no later takeover claim +exists; closing that remaining gap requires durable outbox/reconciliation and +is outside this implementation. Historical migration is also out of scope. +Later recovery appends do not upgrade an earlier Evidence Seal's coverage, and +terminal singleton confirmation does not deduplicate Redis END. Eventual +evidence after every terminal commit is not yet guaranteed. Missing events must +not gate SSE recovery. ## Producers diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index ce229dc7b1..229e093c5b 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -255,6 +255,26 @@ class RunRepository(RunStore): return None return self._row_to_dict(row) + async def get_in_session( + self, + session: AsyncSession, + run_id: str, + *, + user_id: str | None | _AutoSentinel = AUTO, + ): + """Read through a compatible caller transaction without another checkout. + + Keep custom ``get`` implementations and differently bound databases + on their normal read path. The caller owns this session's lifetime. + """ + if type(self).get is not RunRepository.get or session.bind is not self._sf.kw.get("bind"): + return await self.get(run_id, user_id=user_id) + resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.get") + row = await session.get(RunRow, run_id) + if row is None or (resolved_user_id is not None and row.user_id != resolved_user_id): + return None + return self._row_to_dict(row) + async def list_changed( self, *, diff --git a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py index 98287fc459..37e4bd6a44 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py @@ -21,6 +21,7 @@ from deerflow.persistence.scheduled_tasks.model import ( ScheduledTaskRow, ScheduledTaskRunStatus, ) +from deerflow.runtime.events.terminal import RecoveryTerminalPublisher, RecoveryTerminalTarget from deerflow.scheduler.host_notes import RUN_ERROR_END_REACHED from deerflow.scheduler.schedules import next_run_at as compute_next_run_at from deerflow.scheduler.schedules import should_advance_schedule_cursor @@ -153,9 +154,11 @@ class ScheduledTaskRunRepository: session_factory: async_sessionmaker[AsyncSession], *, run_repository: RunRepository | None = None, + recovery_terminal_publisher: RecoveryTerminalPublisher | None = None, ) -> None: self._sf = session_factory self._run_repository = run_repository or RunRepository(session_factory) + self._recovery_terminal_publisher = recovery_terminal_publisher self._finalization_observer: FinalizationObserver | None = None def set_finalization_observer(self, callback: FinalizationObserver | None) -> None: @@ -1029,6 +1032,27 @@ class ScheduledTaskRunRepository: underlying run, or a queued row whose parent task still has a dispatch lease, belongs to another process and must survive this startup. """ + recovery_targets: list[RecoveryTerminalTarget] = [] + try: + return await self._reconcile_active_runs( + error=error, + now=now, + lease_grace_seconds=lease_grace_seconds, + recovery_targets=recovery_targets, + ) + finally: + if self._recovery_terminal_publisher is not None: + for target in recovery_targets: + await self._recovery_terminal_publisher.publish(target) + + async def _reconcile_active_runs( + self, + *, + error: str, + now: datetime, + lease_grace_seconds: int, + recovery_targets: list[RecoveryTerminalTarget], + ) -> int: async with self._sf() as session: result = await session.execute( select( @@ -1091,6 +1115,16 @@ class ScheduledTaskRunRepository: error=error, stop_reason="scheduled_task_orphan_recovered", ) + if claimed and self._recovery_terminal_publisher is not None: + target = self._recovery_terminal_publisher.authorize_takeover( + claimed=claimed, + thread_id=candidate.thread_id, + run_id=candidate.run_id, + user_id=candidate.user_id, + operation_kind=candidate.operation_kind, + ) + if target is not None: + recovery_targets.append(target) if not claimed: refreshed = await self._run_repository.get(candidate.run_id, user_id=None) if refreshed is not None and refreshed.get("status") in {"pending", "running"}: diff --git a/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py b/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py index 9a0b39fbbc..5882cccce9 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py @@ -15,6 +15,7 @@ from deerflow.persistence.scheduled_task_runs.finalization import FinalizationOb from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow from deerflow.persistence.scheduled_task_runs.projection import account_launch, can_project from deerflow.persistence.scheduled_tasks.model import ACTIVE_RUN_STATUSES, LIVE_TASK_STATUSES, ONCE_TASK_STATUS_BY_RUN_STATUS, TERMINAL_RUN_STATUSES, TERMINAL_TASK_STATUSES, ScheduledTaskRow +from deerflow.runtime.events.terminal import RecoveryTerminalPublisher, RecoveryTerminalTarget from deerflow.scheduler.host_notes import RUN_ERROR_END_REACHED, RUN_ERROR_INTERRUPTED from deerflow.utils.goal_objective import normalize_goal_objective from deerflow.utils.time import coerce_iso @@ -91,9 +92,11 @@ class ScheduledTaskRepository: session_factory: async_sessionmaker[AsyncSession], *, run_repository: RunRepository | None = None, + recovery_terminal_publisher: RecoveryTerminalPublisher | None = None, ) -> None: self._sf = session_factory self._run_repository = run_repository or RunRepository(session_factory) + self._recovery_terminal_publisher = recovery_terminal_publisher self._finalization_observer: FinalizationObserver | None = None def set_finalization_observer(self, callback: FinalizationObserver | None) -> None: @@ -1017,6 +1020,27 @@ class ScheduledTaskRepository: lease_grace_seconds: int = 10, ) -> int: """Cancel once tasks only after their underlying run is no longer live.""" + recovery_targets: list[RecoveryTerminalTarget] = [] + try: + return await self._reconcile_stuck_once_tasks( + error=error, + now=now, + lease_grace_seconds=lease_grace_seconds, + recovery_targets=recovery_targets, + ) + finally: + if self._recovery_terminal_publisher is not None: + for target in recovery_targets: + await self._recovery_terminal_publisher.publish(target) + + async def _reconcile_stuck_once_tasks( + self, + *, + error: str, + now: datetime, + lease_grace_seconds: int, + recovery_targets: list[RecoveryTerminalTarget], + ) -> int: async with self._sf() as session: result = await session.execute( select(ScheduledTaskRow.id).where( @@ -1056,6 +1080,16 @@ class ScheduledTaskRepository: error=error, stop_reason="scheduled_task_orphan_recovered", ) + if claimed and self._recovery_terminal_publisher is not None: + target = self._recovery_terminal_publisher.authorize_takeover( + claimed=claimed, + thread_id=candidate.thread_id, + run_id=candidate.run_id, + user_id=candidate.user_id, + operation_kind=candidate.operation_kind, + ) + if target is not None: + recovery_targets.append(target) if not claimed: refreshed = await self._run_repository.get(candidate.run_id, user_id=None) if refreshed is not None and refreshed.get("status") in {"pending", "running"}: diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index dc751f8939..7e66fb5a58 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -435,3 +435,14 @@ freeze/drain → singleton confirmation → Receipt → Seal → barrier release Force close invalidates Receipt bounds. Preserve ownership checks, four-state confirmation, and partial evidence on failure without changing business status or SSE liveness. See [Run outcome protocol](../../../../docs/RUN_EVENT_STREAM.md#authoritative-run-outcome). +Recovery uses a non-replaceable one-shot target, writes the zero-delivery receipt +before `run.terminal`, and passes a Run-lifecycle guard into both singleton +writes. Stores evaluate that guard under the same store mutation fence as +event deletion; completed thread deletion wins over a stale recovery target. +DB recovery guards use `RunRepository.get_in_session` when both stores share +an engine, retaining that fence without a second pool checkout. Different engines +and overridden Run lookups retain their ordinary read path. Plain callable guards +remain supported; they must not acquire another connection from the writer's pool. +Memory guard/deletion locks are per-thread, so a slow guard cannot block another +thread. Coverage: `tests/test_recovery_terminal_pool.py` (PostgreSQL via +`TEST_POSTGRES_URI`). diff --git a/backend/packages/harness/deerflow/runtime/events/store/base.py b/backend/packages/harness/deerflow/runtime/events/store/base.py index 9f85d8a7e9..5340433616 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/base.py +++ b/backend/packages/harness/deerflow/runtime/events/store/base.py @@ -13,7 +13,11 @@ Implementations: from __future__ import annotations import abc -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession from deerflow.runtime.user_context import AUTO, _AutoSentinel @@ -37,6 +41,26 @@ class RunEventStoreNotCommittedError(RuntimeError): """ +class RunEventLifecycleRejected(RuntimeError): + """A guarded singleton write was rejected because its Run is no longer live.""" + + +class SessionAwareRunEventLifecycleGuard(abc.ABC): + """A lifecycle guard that can reuse a DB event writer's active session. + + Other stores call it normally. DB stores call ``check_in_session`` under + their mutation fence; implementations must fall back to their ordinary + read when that session does not belong to the guarded store's database. + Neither entry point may re-enter the event store or commit the session. + """ + + @abc.abstractmethod + async def __call__(self) -> bool: ... + + @abc.abstractmethod + async def check_in_session(self, session: AsyncSession) -> bool: ... + + def normalize_message_ids(message_ids: set[str]) -> set[str]: """Return the non-empty string IDs that can participate in a lookup.""" return {message_id for message_id in message_ids if isinstance(message_id, str) and message_id} @@ -104,6 +128,7 @@ class RunEventStore(abc.ABC): metadata: dict | None = None, created_at: str | None = None, user_id: str | None | _AutoSentinel = AUTO, + lifecycle_guard: Callable[[], Awaitable[bool]] | None = None, ) -> tuple[dict, bool]: """Write one event unless this run already has the same event type. @@ -115,6 +140,16 @@ class RunEventStore(abc.ABC): Explicit ``user_id`` stamps the owner for contextless worker writes. Stores without user isolation accept it for interface parity. Existing events retain their original owner and must be validated by the caller. + + When ``lifecycle_guard`` is provided, stores evaluate it while holding + the same store mutation fence used by event deletion. A false + result raises :class:`RunEventLifecycleRejected` without writing. + That fence is non-reentrant: the guard must not call back into this + event store or acquire locks in an order that can invert the fence's + lock ordering, as doing so may deadlock. The built-in recovery guard + only reads the RunStore. DB stores evaluate + :class:`SessionAwareRunEventLifecycleGuard` with their active session + to avoid a nested connection checkout from the same pool. """ @abc.abstractmethod diff --git a/backend/packages/harness/deerflow/runtime/events/store/db.py b/backend/packages/harness/deerflow/runtime/events/store/db.py index 22bbd39c7e..60b27961b3 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/db.py +++ b/backend/packages/harness/deerflow/runtime/events/store/db.py @@ -22,7 +22,7 @@ from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeq from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunRow from deerflow.persistence.run.sql import RunRepository from deerflow.runtime.events.message_identity import message_identity -from deerflow.runtime.events.store.base import RunEventStore +from deerflow.runtime.events.store.base import RunEventLifecycleRejected, RunEventStore, SessionAwareRunEventLifecycleGuard from deerflow.runtime.user_context import AUTO, _AutoSentinel, get_current_user, resolve_user_id from deerflow.utils.time import coerce_iso @@ -294,6 +294,7 @@ class DbRunEventStore(RunEventStore): metadata=None, created_at=None, user_id: str | None | _AutoSentinel = AUTO, + lifecycle_guard=None, ): """Idempotently insert a run-scoped singleton event. @@ -310,6 +311,10 @@ class DbRunEventStore(RunEventStore): async with self._sf() as session: async with session.begin(): max_seq = await self._max_seq_for_thread(session, thread_id) + if lifecycle_guard is not None: + admitted = await lifecycle_guard.check_in_session(session) if isinstance(lifecycle_guard, SessionAwareRunEventLifecycleGuard) else await lifecycle_guard() + if not admitted: + raise RunEventLifecycleRejected(f"Run lifecycle no longer admits terminal evidence for {thread_id}/{run_id}") stmt = ( select(RunEventRow) .where( diff --git a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py index b660da359a..470602f2d6 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py +++ b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py @@ -41,7 +41,7 @@ from pathlib import Path from typing import Any from deerflow.runtime.events.message_identity import message_identity -from deerflow.runtime.events.store.base import RunEventStore, match_ai_message_run_id, normalize_message_ids +from deerflow.runtime.events.store.base import RunEventLifecycleRejected, RunEventStore, match_ai_message_run_id, normalize_message_ids from deerflow.runtime.user_context import AUTO, _AutoSentinel from deerflow.utils.thread_id import validate_thread_id @@ -287,8 +287,11 @@ class JsonlRunEventStore(RunEventStore): metadata=None, created_at=None, user_id: str | None | _AutoSentinel = AUTO, + lifecycle_guard=None, ): async def mutate(): + if lifecycle_guard is not None and not await lifecycle_guard(): + raise RunEventLifecycleRejected(f"Run lifecycle no longer admits terminal evidence for {thread_id}/{run_id}") existing = await asyncio.to_thread(self._read_run_events, thread_id, run_id) for event in existing: if event.get("event_type") == event_type: diff --git a/backend/packages/harness/deerflow/runtime/events/store/memory.py b/backend/packages/harness/deerflow/runtime/events/store/memory.py index 64347b9efb..7e9dc89cf4 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/events/store/memory.py @@ -6,11 +6,13 @@ since all mutations happen within the same event loop). from __future__ import annotations +import asyncio import bisect +import weakref from datetime import UTC, datetime from deerflow.runtime.events.message_identity import message_identity -from deerflow.runtime.events.store.base import RunEventStore +from deerflow.runtime.events.store.base import RunEventLifecycleRejected, RunEventStore from deerflow.runtime.user_context import AUTO, _AutoSentinel @@ -30,6 +32,16 @@ class MemoryRunEventStore(RunEventStore): self._events_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted events self._messages_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted messages self._seq_counters: dict[str, int] = {} # thread_id -> last assigned seq + # Guarded singleton writes may await a lifecycle callback, so deletion + # needs an explicit fence rather than relying on no-await atomicity. + self._mutation_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + + def _get_mutation_lock(self, thread_id: str) -> asyncio.Lock: + lock = self._mutation_locks.get(thread_id) + if lock is None: + lock = asyncio.Lock() + self._mutation_locks[thread_id] = lock + return lock def _next_seq(self, thread_id: str) -> int: current = self._seq_counters.get(thread_id, 0) @@ -105,24 +117,26 @@ class MemoryRunEventStore(RunEventStore): metadata=None, created_at=None, user_id: str | None | _AutoSentinel = AUTO, + lifecycle_guard=None, ): - # No await occurs between the lookup and append, so this is atomic for - # the backend's documented single-event-loop concurrency model. - for event in self._events_by_run.get(thread_id, {}).get(run_id, []): - if event["event_type"] == event_type: - return event, False - return ( - self._put_one( - thread_id=thread_id, - run_id=run_id, - event_type=event_type, - category=category, - content=content, - metadata=metadata, - created_at=created_at, - ), - True, - ) + async with self._get_mutation_lock(thread_id): + if lifecycle_guard is not None and not await lifecycle_guard(): + raise RunEventLifecycleRejected(f"Run lifecycle no longer admits terminal evidence for {thread_id}/{run_id}") + for event in self._events_by_run.get(thread_id, {}).get(run_id, []): + if event["event_type"] == event_type: + return event, False + return ( + self._put_one( + thread_id=thread_id, + run_id=run_id, + event_type=event_type, + category=category, + content=content, + metadata=metadata, + created_at=created_at, + ), + True, + ) async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None, user_id: str | None | _AutoSentinel = AUTO): # ``messages`` is messages-only and seq-sorted, so the seq window is a @@ -221,24 +235,26 @@ class MemoryRunEventStore(RunEventStore): accepted for interface parity with the user-scoped backends and ignored — the same convention as this store's read methods. """ - events = self._events.pop(thread_id, []) - self._messages.pop(thread_id, None) - self._events_by_run.pop(thread_id, None) - self._messages_by_run.pop(thread_id, None) - self._seq_counters.pop(thread_id, None) - return len(events) + async with self._get_mutation_lock(thread_id): + events = self._events.pop(thread_id, []) + self._messages.pop(thread_id, None) + self._events_by_run.pop(thread_id, None) + self._messages_by_run.pop(thread_id, None) + self._seq_counters.pop(thread_id, None) + return len(events) async def delete_by_run(self, thread_id, run_id, *, user_id: str | None | _AutoSentinel = AUTO): """Delete one run's events; ``user_id`` is accepted for parity only.""" - all_events = self._events.get(thread_id, []) - if not all_events: - return 0 - remaining = [e for e in all_events if e["run_id"] != run_id] - removed = len(all_events) - len(remaining) - self._events[thread_id] = remaining - # Keep the message projection in lockstep (same surviving dict objects). - self._messages[thread_id] = [e for e in remaining if e["category"] == "message"] - # Drop the deleted run from the run-keyed projections. - self._events_by_run.get(thread_id, {}).pop(run_id, None) - self._messages_by_run.get(thread_id, {}).pop(run_id, None) - return removed + async with self._get_mutation_lock(thread_id): + all_events = self._events.get(thread_id, []) + if not all_events: + return 0 + remaining = [e for e in all_events if e["run_id"] != run_id] + removed = len(all_events) - len(remaining) + self._events[thread_id] = remaining + # Keep the message projection in lockstep (same surviving dict objects). + self._messages[thread_id] = [e for e in remaining if e["category"] == "message"] + # Drop the deleted run from the run-keyed projections. + self._events_by_run.get(thread_id, {}).pop(run_id, None) + self._messages_by_run.get(thread_id, {}).pop(run_id, None) + return removed diff --git a/backend/packages/harness/deerflow/runtime/events/terminal.py b/backend/packages/harness/deerflow/runtime/events/terminal.py index 36def224df..0cf170424c 100644 --- a/backend/packages/harness/deerflow/runtime/events/terminal.py +++ b/backend/packages/harness/deerflow/runtime/events/terminal.py @@ -1,17 +1,230 @@ -"""Confirmation of the owning worker's Run terminal evidence.""" +"""Confirmation and recovery publication of authoritative Run terminal evidence.""" from __future__ import annotations -from collections.abc import Callable +import logging +import threading +from collections.abc import Awaitable, Callable +from dataclasses import dataclass from enum import StrEnum from deerflow.runtime.events.catalog import RUN_TERMINAL_EVENT -from deerflow.runtime.events.store.base import RunEventStore -from deerflow.runtime.events.store.db import DbRunEventStore +from deerflow.runtime.events.store.base import RunEventLifecycleRejected, RunEventStore, SessionAwareRunEventLifecycleGuard +from deerflow.runtime.runs.store.base import RunStore TERMINAL_PROTOCOL_VERSION = 1 TERMINAL_STATUSES = frozenset({"success", "error", "interrupted", "timeout"}) +logger = logging.getLogger(__name__) + +_RECOVERY_TARGET_ISSUER = object() + + +class RecoveryTerminalTargetConsumed(RuntimeError): + """A one-shot recovery terminal target was published more than once.""" + + +class _RecoveryTargetClaimState: + __slots__ = ("_consumed", "_lock") + + def __init__(self) -> None: + object.__setattr__(self, "_consumed", False) + object.__setattr__(self, "_lock", threading.Lock()) + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError("Recovery target claim state is immutable") + + def consume(self) -> None: + with self._lock: + if self._consumed: + raise RecoveryTerminalTargetConsumed("Recovery terminal target was already consumed") + object.__setattr__(self, "_consumed", True) + + +@dataclass(frozen=True) +class _RecoveryTargetIdentity: + thread_id: str + run_id: str + user_id: str | None + operation_kind: str + committed_status: str + + +class RecoveryTerminalTarget: + """One-shot authority to publish an error terminal after a takeover claim.""" + + __slots__ = ("_claim_state", "_identity") + + def __init__( + self, + *, + thread_id: str, + run_id: str, + user_id: str | None, + operation_kind: str, + committed_status: str, + _issuer: object, + ) -> None: + if _issuer is not _RECOVERY_TARGET_ISSUER: + raise TypeError("RecoveryTerminalTarget must be issued by RecoveryTerminalPublisher") + if operation_kind != "run": + raise ValueError("Recovery terminal evidence is limited to ordinary runs") + if committed_status != "error": + raise ValueError("Takeover recovery terminal evidence must be error") + if not thread_id or not run_id: + raise ValueError("Recovery terminal target requires thread_id and run_id") + if user_id is not None and not isinstance(user_id, str): + raise TypeError("Recovery terminal target user_id must be a string or None") + self._identity = _RecoveryTargetIdentity(thread_id, run_id, user_id, operation_kind, committed_status) + self._claim_state = _RecoveryTargetClaimState() + + def __setattr__(self, name: str, value: object) -> None: + if name in self.__slots__ and hasattr(self, name): + raise AttributeError("RecoveryTerminalTarget is immutable") + object.__setattr__(self, name, value) + + @property + def thread_id(self) -> str: + return self._identity.thread_id + + @property + def run_id(self) -> str: + return self._identity.run_id + + @property + def user_id(self) -> str | None: + return self._identity.user_id + + @property + def operation_kind(self) -> str: + return self._identity.operation_kind + + @property + def committed_status(self) -> str: + return self._identity.committed_status + + def _consume(self) -> _RecoveryTargetIdentity: + try: + self._claim_state.consume() + except RecoveryTerminalTargetConsumed: + raise RecoveryTerminalTargetConsumed(f"Recovery terminal target for run {self.run_id} was already consumed") from None + return self._identity + + +@dataclass(frozen=True) +class _RecoveryRunLifecycleGuard(SessionAwareRunEventLifecycleGuard): + run_store: RunStore + thread_id: str + run_id: str + user_id: str | None + + def _matches(self, row: dict | None) -> bool: + return row is not None and row.get("run_id") == self.run_id and row.get("thread_id") == self.thread_id and row.get("operation_kind", "run") == "run" and row.get("status") == "error" + + async def __call__(self) -> bool: + return self._matches(await self.run_store.get(self.run_id, user_id=self.user_id)) + + async def check_in_session(self, session) -> bool: + # Import lazily so memory/JSONL recovery does not load SQL persistence. + from deerflow.persistence.run.sql import RunRepository + + if isinstance(self.run_store, RunRepository): + return self._matches(await self.run_store.get_in_session(session, self.run_id, user_id=self.user_id)) + return await self() + + +class RecoveryTerminalPublisher: + """Publish one confirmed error terminal for each successful takeover claim. + + ``RunStore.claim_for_takeover()`` atomically commits ``status=error``. The + target therefore pins that known outcome rather than deriving authority + from a later RunRow read. + """ + + def __init__(self, store: RunEventStore, run_store: RunStore) -> None: + self._store = store + self._run_store = run_store + + def lifecycle_guard(self, *, thread_id: str, run_id: str, user_id: str | None) -> Callable[[], Awaitable[bool]]: + """Return a guard that rejects publication after Run deletion or identity drift.""" + + return _RecoveryRunLifecycleGuard(self._run_store, thread_id, run_id, user_id) + + def authorize_takeover( + self, + *, + claimed: bool, + thread_id: str, + run_id: str, + user_id: str | None, + operation_kind: str, + ) -> RecoveryTerminalTarget | None: + """Issue one-shot publication authority only for a committed run takeover.""" + if claimed is not True or operation_kind != "run": + return None + return RecoveryTerminalTarget( + thread_id=thread_id, + run_id=run_id, + user_id=user_id, + operation_kind=operation_kind, + committed_status="error", + _issuer=_RECOVERY_TARGET_ISSUER, + ) + + async def publish(self, target: RecoveryTerminalTarget) -> TerminalEventConfirmation: + """Consume *target* once, backfill delivery, then confirm terminal evidence.""" + if not isinstance(target, RecoveryTerminalTarget): + raise TypeError("RecoveryTerminalPublisher.publish requires a RecoveryTerminalTarget") + identity = target._consume() + guard = self.lifecycle_guard(thread_id=identity.thread_id, run_id=identity.run_id, user_id=identity.user_id) + await self._ensure_delivery_receipt(identity, guard) + try: + confirmation = await confirm_terminal_event( + self._store, + thread_id=identity.thread_id, + run_id=identity.run_id, + status=identity.committed_status, + user_id=identity.user_id, + recovered=True, + authorized=lambda: True, + lifecycle_guard=guard, + ) + except Exception: + # Cancellation is BaseException and must retain its independent + # propagation semantics after the one-shot target was consumed. + logger.warning( + "Unexpected recovery terminal confirmation failure for run %s; preserving the committed RunRow", + identity.run_id, + exc_info=True, + ) + return TerminalEventConfirmation.UNKNOWN + if not confirmation.confirmed: + logger.warning( + "Recovery terminal publication for run %s returned %s; preserving the committed RunRow", + identity.run_id, + confirmation.value, + ) + return confirmation + + async def _ensure_delivery_receipt(self, identity: _RecoveryTargetIdentity, lifecycle_guard: Callable[[], Awaitable[bool]]) -> None: + """Best-effort zero-delivery receipt; terminal publication remains independent.""" + try: + await self._store.put_if_absent( + thread_id=identity.thread_id, + run_id=identity.run_id, + event_type="run.delivery", + category="outputs", + content={"presented": 0, "paths": [], "by_tool": {}}, + user_id=identity.user_id, + lifecycle_guard=lifecycle_guard, + ) + except Exception: + logger.warning( + "Failed to backfill delivery receipt for recovered run %s; continuing with terminal confirmation", + identity.run_id, + exc_info=True, + ) + class TerminalEventConfirmation(StrEnum): CONFIRMED_NEW = "CONFIRMED_NEW" @@ -53,28 +266,41 @@ async def confirm_terminal_event( status: str, user_id: str | None, authorized: Callable[[], bool], + recovered: bool = False, + lifecycle_guard: Callable[[], Awaitable[bool]] | None = None, ) -> TerminalEventConfirmation: """Confirm a singleton without replaying any Journal batch. One bounded retry uses the store's atomic check/insert to resolve a lost ACK. Both attempts require the same final commit authority. Cancellation propagates to the worker's owned finalizer; it cannot authorize a receipt. + ``recovered`` marks only a newly inserted recovery event; a matching + existing event retains its original source marker. """ if status not in TERMINAL_STATUSES: raise ValueError("Terminal evidence requires a final Run status") + if type(recovered) is not bool: + raise ValueError("Terminal evidence recovered marker must be a boolean") + from deerflow.runtime.events.store.db import DbRunEventStore + for attempt in range(2): if not authorized(): raise RuntimeError("Terminal event lacks an attributable final commit ACK") try: - result = await store.put_if_absent( - thread_id=thread_id, - run_id=run_id, - event_type=RUN_TERMINAL_EVENT.event_type, - category=RUN_TERMINAL_EVENT.category, - content={"status": status}, - metadata={"protocol_version": TERMINAL_PROTOCOL_VERSION, "recovered": False}, - user_id=user_id, - ) + put_kwargs = { + "thread_id": thread_id, + "run_id": run_id, + "event_type": RUN_TERMINAL_EVENT.event_type, + "category": RUN_TERMINAL_EVENT.category, + "content": {"status": status}, + "metadata": {"protocol_version": TERMINAL_PROTOCOL_VERSION, "recovered": recovered}, + "user_id": user_id, + } + if lifecycle_guard is not None: + put_kwargs["lifecycle_guard"] = lifecycle_guard + result = await store.put_if_absent(**put_kwargs) + except RunEventLifecycleRejected: + return TerminalEventConfirmation.CONFLICT except Exception: if attempt == 0: continue diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 1762da793d..434e5b0540 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any from sqlalchemy.exc import IntegrityError as SAIntegrityError +from deerflow.runtime.events.terminal import RecoveryTerminalPublisher from deerflow.runtime.user_context import AUTO, _AutoSentinel, get_current_user, resolve_user_id from deerflow.utils.time import is_lease_expired from deerflow.utils.time import now_iso as _now_iso @@ -342,6 +343,7 @@ class RunManager: worker_id: str | None = None, run_ownership_config: RunOwnershipConfig | None = None, event_store: RunEventStore | None = None, + recovery_terminal_publisher: RecoveryTerminalPublisher | None = None, on_orphans_recovered: OrphanRecoveryCallback | None = None, ) -> None: self._runs: dict[str, RunRecord] = {} @@ -356,6 +358,9 @@ class RunManager: self._worker_id = worker_id or _generate_worker_id() self._run_ownership_config = run_ownership_config self._event_store = event_store + self._recovery_terminal_publisher = recovery_terminal_publisher + if self._recovery_terminal_publisher is None and event_store is not None and store is not None: + self._recovery_terminal_publisher = RecoveryTerminalPublisher(event_store, store) self._on_orphans_recovered = on_orphans_recovered self._heartbeat_task: asyncio.Task | None = None self._heartbeat_stop: asyncio.Event | None = None @@ -1500,27 +1505,6 @@ class RunManager: await self._fence_unconfirmed_terminal(record) return None - async def _ensure_delivery_receipt(self, record: RunRecord) -> bool: - """Idempotently persist a zero-delivery receipt during recovery.""" - if self._event_store is None: - return True - try: - await self._event_store.put_if_absent( - thread_id=record.thread_id, - run_id=record.run_id, - event_type="run.delivery", - category="outputs", - content={"presented": 0, "paths": [], "by_tool": {}}, - ) - return True - except Exception: - logger.warning( - "Failed to backfill delivery receipt for recovered run %s; preserving its terminal status", - record.run_id, - exc_info=True, - ) - return False - async def set_finalizing(self, run_id: str, finalizing: bool) -> None: """Mark whether a run is performing post-cancel cleanup.""" async with self._lock: @@ -1894,6 +1878,16 @@ class RunManager: return CancelOutcome.unknown if taken: + if self._recovery_terminal_publisher is not None: + target = self._recovery_terminal_publisher.authorize_takeover( + claimed=taken, + thread_id=row["thread_id"], + run_id=run_id, + user_id=row.get("user_id"), + operation_kind=row.get("operation_kind", "run"), + ) + if target is not None: + await self._recovery_terminal_publisher.publish(target) logger.warning("Run %s taken over by worker %s (action=%s)", run_id, self._worker_id, action) return CancelOutcome.taken_over @@ -2459,12 +2453,17 @@ class RunManager: record.stop_reason = stop_reason record.updated_at = now if record.operation_kind == ThreadOperationKind.run: - # The atomic takeover above must win before writing a zero-delivery - # receipt; otherwise a stale scan could race a heartbeat renewal and - # permanently overwrite a live run's later detailed receipt. The - # receipt remains best-effort, matching normal terminal delivery - # when its event store is unavailable. - await self._ensure_delivery_receipt(record) + publisher = self._recovery_terminal_publisher + if publisher is not None: + target = publisher.authorize_takeover( + claimed=claimed, + thread_id=record.thread_id, + run_id=record.run_id, + user_id=record.user_id, + operation_kind=record.operation_kind.value, + ) + if target is not None: + await publisher.publish(target) recovered.append(record) if recovered: diff --git a/backend/tests/AGENTS.md b/backend/tests/AGENTS.md index 3b7d5721e9..598886cb32 100644 --- a/backend/tests/AGENTS.md +++ b/backend/tests/AGENTS.md @@ -29,11 +29,17 @@ unrelated parent config that stays untouched. Keep argument parsing at the CLI entry point so in-process `doctor.main()` callers do not consume pytest's argv. `test_run_terminal_protocol.py` exercises authoritative `run.terminal` -evidence through real Worker finalization. Keep final commit authority, -producer drain, singleton confirmation, and Seal boundaries observable. Release all storage/producer gates before bounded cleanup, -preserve healthy controls in refusal cases, and retain Journal/Seal compatibility -coverage. `test_terminal_event_confirmation.py` covers supported stores; use -`TEST_POSTGRES_URI` for independent PostgreSQL pools and real advisory contention. Protocol: [Run outcome](../docs/RUN_EVENT_STREAM.md#authoritative-run-outcome). +evidence through real Worker finalization and recovery takeover. Keep final +commit authority, producer drain, singleton confirmation, and Seal boundaries +observable. Recovery coverage must block after the committed claim, delete the +Run/events, then prove the lifecycle guard cannot recreate terminal history; +pin delivery-before-terminal order across all takeover entry points, and keep a +legacy `put_if_absent` signature on the owner path. `test_terminal_event_confirmation.py` +covers supported stores, non-replaceable one-shot recovery targets, and guarded +singleton rejection; use `TEST_POSTGRES_URI` for independent PostgreSQL pools +and real advisory contention. Release all storage/producer gates before bounded +cleanup, preserve healthy controls in refusal cases, and retain Journal/Seal +compatibility coverage. Protocol: [Run outcome](../docs/RUN_EVENT_STREAM.md#authoritative-run-outcome). `test_goal_loop_stop_probe.py` runs the standalone Goal-loop reproduction through the native graph, loop middleware, worker and checkpoints. Only model responses are scripted; do not replace the Goal decision or simulate the stop-reason write. InfoQuest fetch log assertions install the real `UrlRedactionFilter` on the diff --git a/backend/tests/test_gateway_run_recovery.py b/backend/tests/test_gateway_run_recovery.py index 66498c27f2..b2df08ff7e 100644 --- a/backend/tests/test_gateway_run_recovery.py +++ b/backend/tests/test_gateway_run_recovery.py @@ -40,11 +40,13 @@ class _FakeRunManager: store, run_ownership_config=None, event_store=None, + recovery_terminal_publisher=None, on_orphans_recovered=None, ): self.store = store self.run_ownership_config = run_ownership_config self.event_store = event_store + self.recovery_terminal_publisher = recovery_terminal_publisher self.on_orphans_recovered = on_orphans_recovered self.reconcile_calls: list[dict] = [] self.list_by_thread_calls: list[dict] = [] @@ -322,6 +324,7 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch): assert _FakeRunManager.instances[0].reconcile_calls assert _FakeRunManager.instances[0].reconcile_calls[0]["error"] assert _FakeRunManager.instances[0].reconcile_calls[0]["stop_reason"] == runtime_module.ORPHAN_RECOVERY_STOP_REASON + assert _FakeRunManager.instances[0].recovery_terminal_publisher is app.state.recovery_terminal_publisher assert _FakeRunManager.instances[0].list_by_thread_calls == [{"thread_id": "thread-1", "user_id": None, "limit": 1}] assert thread_store.status_updates == [("thread-1", "error", None)] assert stream_bridge.publish_end_calls == ["run-1"] @@ -359,6 +362,9 @@ async def test_sql_runtime_shares_run_repository_with_scheduler(monkeypatch): async with gateway_deps.langgraph_runtime(app, config): assert app.state.scheduled_task_repo._run_repository is app.state.run_store assert app.state.scheduled_task_run_repo._run_repository is app.state.run_store + assert app.state.scheduled_task_repo._recovery_terminal_publisher is app.state.recovery_terminal_publisher + assert app.state.scheduled_task_run_repo._recovery_terminal_publisher is app.state.recovery_terminal_publisher + assert _FakeRunManager.instances[-1].recovery_terminal_publisher is app.state.recovery_terminal_publisher @pytest.mark.anyio diff --git a/backend/tests/test_recovery_terminal_pool.py b/backend/tests/test_recovery_terminal_pool.py new file mode 100644 index 0000000000..73aca485ba --- /dev/null +++ b/backend/tests/test_recovery_terminal_pool.py @@ -0,0 +1,142 @@ +"""Recovery publication must work with only one free pooled connection.""" + +from __future__ import annotations + +import asyncio +import os +from contextlib import AsyncExitStack +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +import pytest_asyncio +from sqlalchemy import text +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from support.postgres import asyncpg_test_url + +from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow +from deerflow.persistence.run.model import RunChangeClockRow, RunRow +from deerflow.persistence.run.sql import RunRepository +from deerflow.runtime.events.store.base import RunEventLifecycleRejected +from deerflow.runtime.events.store.db import DbRunEventStore +from deerflow.runtime.events.store.memory import MemoryRunEventStore +from deerflow.runtime.events.terminal import RecoveryTerminalPublisher, TerminalEventConfirmation + + +@pytest_asyncio.fixture(params=["sqlite", "postgresql"]) +async def pooled_database(request, tmp_path): + schema = None + options = {} + if request.param == "postgresql": + uri = os.getenv("TEST_POSTGRES_URI") + if not uri: + pytest.skip("requires TEST_POSTGRES_URI for real pooled PostgreSQL recovery") + schema = f"recovery_pool_{uuid4().hex}" + url = asyncpg_test_url(uri) + options["connect_args"] = {"server_settings": {"search_path": schema}} + else: + url = f"sqlite+aiosqlite:///{tmp_path / 'recovery.db'}" + engine = create_async_engine(url, pool_size=5, max_overflow=10, pool_timeout=0.1, **options) + try: + async with engine.begin() as connection: + if schema: + await connection.execute(text(f'CREATE SCHEMA "{schema}"')) + await connection.run_sync(lambda c: RunRow.metadata.create_all(c, tables=[RunRow.__table__, RunChangeClockRow.__table__, RunEventRow.__table__, RunEventThreadSeqRow.__table__])) + # Separate factories sharing an engine are still the same connection pool. + run_store = RunRepository(async_sessionmaker(engine, expire_on_commit=False)) + events = DbRunEventStore(async_sessionmaker(engine, expire_on_commit=False)) + yield SimpleNamespace(engine=engine, runs=run_store, events=events) + finally: + if schema: + async with engine.begin() as connection: + await connection.execute(text(f'DROP SCHEMA "{schema}" CASCADE')) + await engine.dispose() + + +@pytest.mark.asyncio +async def test_recovery_uses_last_available_pool_connection(pooled_database): + db = pooled_database + await db.runs.put("run", thread_id="thread", user_id="alice", status="running") + claimed = await db.runs.claim_for_takeover("run", grace_seconds=0, error="owner expired") + assert claimed + publisher = RecoveryTerminalPublisher(db.events, db.runs) + target = publisher.authorize_takeover(claimed=claimed, thread_id="thread", run_id="run", user_id="alice", operation_kind="run") + async with AsyncExitStack() as occupied: + for _ in range(14): + await occupied.enter_async_context(db.engine.connect()) + assert await publisher.publish(target) == TerminalEventConfirmation.CONFIRMED_NEW + events = await db.events.list_events("thread", "run", user_id="alice") + assert [event["event_type"] for event in events] == ["run.delivery", "run.terminal"] + assert events[-1]["content"] == {"status": "error"} + assert events[-1]["metadata"]["recovered"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["deleted", "owner", "thread", "operation", "status"]) +async def test_pooled_guard_still_rejects_stale_run(pooled_database, change): + db = pooled_database + row = {"thread_id": "thread", "user_id": "alice", "operation_kind": "run", "status": "error"} + if change != "deleted": + key, value = {"owner": ("user_id", "bob"), "thread": ("thread_id", "other"), "operation": ("operation_kind", "checkpoint_write"), "status": ("status", "success")}[change] + row[key] = value + await db.runs.put("run", **row) + guard = RecoveryTerminalPublisher(db.events, db.runs).lifecycle_guard(thread_id="thread", run_id="run", user_id="alice") + async with AsyncExitStack() as occupied: + for _ in range(14): + await occupied.enter_async_context(db.engine.connect()) + with pytest.raises(RunEventLifecycleRejected): + await db.events.put_if_absent(thread_id="thread", run_id="run", event_type="run.terminal", category="trace", user_id="alice", lifecycle_guard=guard) + assert await db.events.list_events("thread", "run", user_id="alice") == [] + + +@pytest.mark.asyncio +async def test_transaction_guard_keeps_differently_bound_run_store(pooled_database): + db = pooled_database + await db.runs.put("run", thread_id="thread", user_id="alice", status="error") + guard = RecoveryTerminalPublisher(db.events, db.runs).lifecycle_guard(thread_id="thread", run_id="run", user_id="alice") + # No tables: using the caller's unrelated database would fail or read the + # wrong Run. Compatibility is engine identity, not session-factory identity. + unrelated_engine = create_async_engine("sqlite+aiosqlite:///:memory:") + try: + async with async_sessionmaker(unrelated_engine)() as session: + assert await guard.check_in_session(session) + finally: + await unrelated_engine.dispose() + + +@pytest.mark.asyncio +async def test_transaction_guard_preserves_custom_run_lookup(pooled_database): + db = pooled_database + await db.runs.put("run", thread_id="thread", user_id="alice", status="error") + + class RejectingRunRepository(RunRepository): + async def get(self, run_id, *, user_id): + assert (run_id, user_id) == ("run", "alice") + return None + + custom = RejectingRunRepository(async_sessionmaker(db.engine, expire_on_commit=False)) + guard = RecoveryTerminalPublisher(db.events, custom).lifecycle_guard(thread_id="thread", run_id="run", user_id="alice") + with pytest.raises(RunEventLifecycleRejected): + await db.events.put_if_absent(thread_id="thread", run_id="run", event_type="run.terminal", category="trace", lifecycle_guard=guard) + assert await db.events.list_events("thread", "run", user_id="alice") == [] + + +@pytest.mark.asyncio +async def test_memory_guard_does_not_block_another_thread(): + store = MemoryRunEventStore() + entered, release = asyncio.Event(), asyncio.Event() + + async def guard(): + entered.set() + await release.wait() + return True + + pending = asyncio.create_task(store.put_if_absent(thread_id="slow", run_id="run", event_type="run.terminal", category="trace", lifecycle_guard=guard)) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + await asyncio.wait_for(store.put_if_absent(thread_id="other", run_id="run", event_type="run.terminal", category="trace"), timeout=1) + assert await asyncio.wait_for(store.delete_by_run("other", "run"), timeout=1) == 1 + assert await asyncio.wait_for(store.delete_by_thread("other"), timeout=1) == 0 + finally: + release.set() + await pending diff --git a/backend/tests/test_run_terminal_protocol.py b/backend/tests/test_run_terminal_protocol.py index 4e9bda5ac9..2fa4e58823 100644 --- a/backend/tests/test_run_terminal_protocol.py +++ b/backend/tests/test_run_terminal_protocol.py @@ -1,7 +1,8 @@ """Behavioral regressions for authoritative Run terminal evidence. Pinned baseline: fee0cd0c180874325305e995197db99f03fa6a6f. These tests -exercise owner-side run.terminal publication through actual finalization. +exercise owner-side run.terminal publication through actual finalization and +confirmed recovery takeover publication. See backend/docs/RUN_EVENT_STREAM.md for the protocol. """ @@ -23,6 +24,7 @@ from langgraph.graph.message import add_messages from sqlalchemy import update from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +import deerflow.runtime.events.terminal as terminal_module import deerflow.runtime.runs.worker as worker_module from deerflow.config.paths import Paths from deerflow.config.run_ownership_config import RunOwnershipConfig @@ -74,17 +76,41 @@ async def terminal_events(runtime, record): return await runtime.events.list_events(record.thread_id, record.run_id, event_types=["run.terminal"], user_id=None) -def confirmed_terminal(events, record, status): +def confirmed_terminal(events, record, status, *, recovered=False): assert len(events) == 1, f"expected one authoritative run.terminal({status}), got {events!r}" terminal = events[0] assert (terminal["thread_id"], terminal["run_id"], terminal["event_type"]) == (record.thread_id, record.run_id, "run.terminal") assert terminal["category"] == "trace" assert terminal["content"] == {"status": status} assert terminal["metadata"]["protocol_version"] == 1 - assert terminal["metadata"]["recovered"] is False + assert terminal["metadata"]["recovered"] is recovered return terminal +async def put_expired_run(runtime, *, operation_kind="run", user_id="alice"): + run_id = f"recovery-run-{uuid4()}" + thread_id = f"recovery-thread-{uuid4()}" + await runtime.runs.put( + run_id, + thread_id=thread_id, + user_id=user_id, + status="running", + operation_kind=operation_kind, + owner_worker_id="dead-owner", + lease_expires_at=(datetime.now(UTC) - timedelta(seconds=60)).isoformat(), + ) + return SimpleNamespace(run_id=run_id, thread_id=thread_id, user_id=user_id) + + +def recovery_manager(runtime, *, worker_id=None): + return RunManager( + store=runtime.runs, + event_store=runtime.events, + worker_id=worker_id or f"recovery-{uuid4()}", + run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=300, grace_seconds=0), + ) + + @pytest_asyncio.fixture async def runtime(tmp_path, monkeypatch): """Real shared SQLite run/event stores, with bounded owned-task cleanup.""" @@ -604,3 +630,324 @@ async def test_process_local_terminal_proof_allows_later_staged_outcome(heartbea assert await manager.persist_current_status(record.run_id) assert record.status == status assert manager.terminal_event_status(record) == status.value + + +@pytest.mark.asyncio +@pytest.mark.no_auto_user +async def test_orphan_recovery_claim_backfills_error_terminal(runtime): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert [run.run_id for run in recovered] == [record.run_id] + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "error" + confirmed_terminal(await terminal_events(runtime, record), record, "error", recovered=True) + + +@pytest.mark.asyncio +async def test_losing_orphan_recovery_claim_does_not_publish_terminal(runtime, monkeypatch): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + original_scan = runtime.runs.list_inflight_with_expired_lease + + async def stale_scan_then_owner_renews(**kwargs): + rows = await original_scan(**kwargs) + async with runtime.sf() as session: + await session.execute(update(RunRow).where(RunRow.run_id == record.run_id).values(lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat())) + await session.commit() + return rows + + monkeypatch.setattr(runtime.runs, "list_inflight_with_expired_lease", stale_scan_then_owner_renews) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="must not win", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert recovered == [] + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "running" + assert await terminal_events(runtime, record) == [] + + +@pytest.mark.asyncio +async def test_non_local_cancel_takeover_backfills_error_terminal(runtime): + record = await put_expired_run(runtime, user_id=get_effective_user_id()) + manager = recovery_manager(runtime) + + outcome = await manager.cancel(record.run_id) + + assert outcome == CancelOutcome.taken_over + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "error" + delivery = await runtime.events.list_events(record.thread_id, record.run_id, event_types=["run.delivery"], user_id=None) + terminal = await terminal_events(runtime, record) + assert len(delivery) == 1 + confirmed_terminal(terminal, record, "error", recovered=True) + assert delivery[0]["seq"] < terminal[0]["seq"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["event_store", "unexpected"]) +async def test_non_local_cancel_returns_taken_over_when_terminal_publish_fails(runtime, monkeypatch, failure): + record = await put_expired_run(runtime, user_id=get_effective_user_id()) + manager = recovery_manager(runtime) + attempts = [] + if failure == "event_store": + original = runtime.events.put_if_absent + + async def fail_terminal(**kwargs): + if kwargs["event_type"] == "run.terminal": + attempts.append(kwargs) + raise OSError("terminal event unavailable") + return await original(**kwargs) + + monkeypatch.setattr(runtime.events, "put_if_absent", fail_terminal) + else: + + async def fail_confirmation(*args, **kwargs): + raise RuntimeError("unexpected confirmation failure") + + monkeypatch.setattr(terminal_module, "confirm_terminal_event", fail_confirmation) + + outcome = await manager.cancel(record.run_id) + + assert outcome == CancelOutcome.taken_over + assert len(attempts) == (2 if failure == "event_store" else 0) + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "error" + assert row["evidence_seal_state"] == "partial" + assert await terminal_events(runtime, record) == [] + + +@pytest.mark.asyncio +async def test_non_run_takeover_claim_does_not_publish_terminal(runtime): + record = await put_expired_run(runtime, operation_kind="checkpoint_write") + manager = recovery_manager(runtime) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="checkpoint write recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert recovered == [] + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "error" + assert await terminal_events(runtime, record) == [] + + +@pytest.mark.asyncio +async def test_concurrent_orphan_recovery_only_winner_publishes_terminal(runtime): + record = await put_expired_run(runtime) + managers = [recovery_manager(runtime), recovery_manager(runtime)] + + results = await asyncio.gather( + *( + manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + for manager in managers + ) + ) + + assert sum(len(recovered) for recovered in results) == 1 + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "error" + confirmed_terminal(await terminal_events(runtime, record), record, "error", recovered=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["unknown", "conflict", "lost_ack"]) +async def test_recovery_terminal_failure_never_changes_committed_takeover(runtime, monkeypatch, failure): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + attempts = [] + existing = None + if failure == "conflict": + existing = await runtime.events.put( + thread_id=record.thread_id, + run_id=record.run_id, + event_type="run.terminal", + category="trace", + content={"status": "success"}, + metadata={"protocol_version": 1, "recovered": False}, + ) + original = runtime.events.put_if_absent + + async def fail_terminal(**kwargs): + nonlocal existing + if kwargs["event_type"] != "run.terminal": + return await original(**kwargs) + if failure == "unknown": + attempts.append(kwargs) + raise OSError("terminal store unavailable") + if failure == "lost_ack": + result = await original(**kwargs) + attempts.append(result) + if len(attempts) == 1: + raise OSError("terminal ACK lost after commit") + return result + return await original(**kwargs) + + monkeypatch.setattr(runtime.events, "put_if_absent", fail_terminal) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert [run.run_id for run in recovered] == [record.run_id] + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None + assert row["status"] == "error" + assert row["evidence_seal_state"] == "partial" + terminal = await terminal_events(runtime, record) + if failure == "unknown": + assert len(attempts) == 2 + assert terminal == [] + elif failure == "conflict": + assert terminal == [existing] + else: + assert len(attempts) == 2 + confirmed_terminal(terminal, record, "error", recovered=True) + + +@pytest.mark.asyncio +async def test_orphan_recovery_unexpected_publisher_error_preserves_recovery_flow(runtime, monkeypatch): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + + async def fail_confirmation(*args, **kwargs): + raise RuntimeError("unexpected confirmation failure") + + monkeypatch.setattr(terminal_module, "confirm_terminal_event", fail_confirmation) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert [run.run_id for run in recovered] == [record.run_id] + row = await runtime.runs.get(record.run_id, user_id=record.user_id) + assert row is not None and row["status"] == "error" + assert row["evidence_seal_state"] == "partial" + assert await terminal_events(runtime, record) == [] + delivery = await runtime.events.list_events(record.thread_id, record.run_id, event_types=["run.delivery"], user_id=None) + assert len(delivery) == 1 + + +@pytest.mark.asyncio +async def test_recovery_publication_cannot_resurrect_deleted_thread(runtime, monkeypatch): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + publishing, release = asyncio.Event(), asyncio.Event() + original_confirm = terminal_module.confirm_terminal_event + + async def pause_before_event_write(*args, **kwargs): + publishing.set() + await release.wait() + return await original_confirm(*args, **kwargs) + + monkeypatch.setattr(terminal_module, "confirm_terminal_event", pause_before_event_write) + recovery = asyncio.create_task( + manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + ) + try: + await asyncio.wait_for(publishing.wait(), timeout=5) + await runtime.runs.delete_by_thread(record.thread_id, user_id=record.user_id) + await runtime.events.delete_by_thread(record.thread_id, user_id=record.user_id) + finally: + release.set() + recovered = await asyncio.wait_for(recovery, timeout=5) + + assert [run.run_id for run in recovered] == [record.run_id] + assert await runtime.runs.get(record.run_id, user_id=record.user_id) is None + assert await runtime.events.list_events(record.thread_id, record.run_id, user_id=None) == [] + assert await runtime.events.list_messages(record.thread_id, user_id=None) == [] + + +@pytest.mark.asyncio +async def test_recovery_delivery_receipt_precedes_terminal(runtime): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert [run.run_id for run in recovered] == [record.run_id] + delivery = await runtime.events.list_events(record.thread_id, record.run_id, event_types=["run.delivery"], user_id=None) + terminal = await terminal_events(runtime, record) + assert len(delivery) == 1 + assert delivery[0]["user_id"] == record.user_id + confirmed_terminal(terminal, record, "error", recovered=True) + assert delivery[0]["seq"] < terminal[0]["seq"] + + +@pytest.mark.asyncio +async def test_recovery_delivery_receipt_failure_does_not_block_terminal(runtime, monkeypatch): + record = await put_expired_run(runtime) + manager = recovery_manager(runtime) + original = runtime.events.put_if_absent + + async def fail_delivery(**kwargs): + if kwargs["event_type"] == "run.delivery": + raise OSError("delivery receipt unavailable") + return await original(**kwargs) + + monkeypatch.setattr(runtime.events, "put_if_absent", fail_delivery) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert [run.run_id for run in recovered] == [record.run_id] + assert await runtime.events.list_events(record.thread_id, record.run_id, event_types=["run.delivery"], user_id=None) == [] + confirmed_terminal(await terminal_events(runtime, record), record, "error", recovered=True) + + +@pytest.mark.asyncio +async def test_recovery_accepts_existing_owner_terminal_without_rewriting_source(runtime): + record = await put_expired_run(runtime) + owner_event, created = await runtime.events.put_if_absent( + thread_id=record.thread_id, + run_id=record.run_id, + event_type="run.terminal", + category="trace", + content={"status": "error"}, + metadata={"protocol_version": 1, "recovered": False}, + user_id=record.user_id, + ) + assert created is True + manager = recovery_manager(runtime) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="orphan recovered", + before=datetime.now(UTC).isoformat(), + stop_reason="orphan_recovered", + ) + + assert [run.run_id for run in recovered] == [record.run_id] + terminal = await terminal_events(runtime, record) + assert terminal == [owner_event] + assert terminal[0]["metadata"]["recovered"] is False diff --git a/backend/tests/test_scheduled_task_postgres.py b/backend/tests/test_scheduled_task_postgres.py index fc8c1eb63f..d7b3bdb947 100644 --- a/backend/tests/test_scheduled_task_postgres.py +++ b/backend/tests/test_scheduled_task_postgres.py @@ -16,13 +16,16 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from deerflow.config.database_config import DatabaseConfig from deerflow.persistence.base import Base from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow from deerflow.persistence.postgres_schema import build_asyncpg_connect_args from deerflow.persistence.run import RunRepository -from deerflow.persistence.run.model import RunRow +from deerflow.persistence.run.model import RunChangeClockRow, RunRow from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow +from deerflow.runtime.events.store.db import DbRunEventStore +from deerflow.runtime.events.terminal import RecoveryTerminalPublisher POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI") @@ -166,6 +169,70 @@ async def test_postgres_reconciliation_uses_metadata_and_atomically_claims_expir assert recovered["stop_reason"] == "scheduled_task_orphan_recovered" +@pytest.mark.asyncio +async def test_postgres_recovery_takeover_has_one_publisher(): + async with _contender_engines(2) as factories: + seed_repo = RunRepository(factories[0]) + event_store = DbRunEventStore(factories[0]) + now = datetime.now(UTC) + run_id, thread_id = "run-recovery-postgres", "thread-recovery-postgres" + await seed_repo.put( + run_id, + thread_id=thread_id, + user_id="user-1", + status="running", + owner_worker_id="dead-owner", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + + both_ready = asyncio.Event() + start_claiming = asyncio.Event() + ready_count = 0 + ready_lock = asyncio.Lock() + + async def recover_once(factory): + nonlocal ready_count + repo = RunRepository(factory) + publisher = RecoveryTerminalPublisher(DbRunEventStore(factory), repo) + async with ready_lock: + ready_count += 1 + if ready_count == 2: + both_ready.set() + await both_ready.wait() + await start_claiming.wait() + claimed = await repo.claim_for_takeover( + run_id, + grace_seconds=10, + error="orphan recovered", + stop_reason="orphan_recovered", + ) + target = publisher.authorize_takeover( + claimed=claimed, + thread_id=thread_id, + run_id=run_id, + user_id="user-1", + operation_kind="run", + ) + if target is not None: + await publisher.publish(target) + return target is not None + + tasks = [asyncio.create_task(recover_once(factory)) for factory in factories] + try: + await asyncio.wait_for(both_ready.wait(), timeout=5) + start_claiming.set() + outcomes = await asyncio.wait_for(asyncio.gather(*tasks), timeout=10) + finally: + start_claiming.set() + await asyncio.gather(*tasks, return_exceptions=True) + + assert sum(outcomes) == 1 + terminal = await event_store.list_events(thread_id, run_id, event_types=["run.terminal"], user_id=None) + assert len(terminal) == 1 + assert terminal[0]["content"] == {"status": "error"} + assert terminal[0]["metadata"]["recovered"] is True + + @pytest.mark.asyncio @pytest.mark.parametrize("recovery_method", ["cancel_stuck_once_tasks", "reconcile_stuck_once_tasks"]) @pytest.mark.parametrize("newer_status", ["skipped", "queued", "launching", "running"]) @@ -246,7 +313,19 @@ async def _contender_engines(count: int): try: async with engines[0].begin() as conn: await conn.execute(text(f'CREATE SCHEMA "{schema}"')) - await conn.run_sync(lambda sync: Base.metadata.create_all(sync, tables=[ScheduledTaskRow.__table__, ScheduledTaskRunRow.__table__, RunRow.__table__])) + await conn.run_sync( + lambda sync: Base.metadata.create_all( + sync, + tables=[ + ScheduledTaskRow.__table__, + ScheduledTaskRunRow.__table__, + RunRow.__table__, + RunChangeClockRow.__table__, + RunEventRow.__table__, + RunEventThreadSeqRow.__table__, + ], + ) + ) yield [async_sessionmaker(engine, expire_on_commit=False) for engine in engines] finally: async with engines[0].begin() as conn: diff --git a/backend/tests/test_scheduled_task_repository.py b/backend/tests/test_scheduled_task_repository.py index 733fce0fae..9cf11b6024 100644 --- a/backend/tests/test_scheduled_task_repository.py +++ b/backend/tests/test_scheduled_task_repository.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime, timedelta import pytest +import deerflow.runtime.events.terminal as terminal_module from deerflow.config.database_config import DatabaseConfig from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config from deerflow.persistence.run import RunRepository @@ -15,6 +16,8 @@ from deerflow.persistence.scheduled_task_runs.finalization import AGENT_STOP_LAS from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow from deerflow.persistence.scheduled_tasks import ActiveScheduledTaskMutationConflict, ScheduledTaskRepository from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow +from deerflow.runtime.events.store.db import DbRunEventStore +from deerflow.runtime.events.terminal import RecoveryTerminalPublisher @pytest.mark.asyncio @@ -1441,3 +1444,218 @@ async def test_pausing_a_finished_task_reports_finished(tmp_path, terminal): assert await tasks.get("task-1", user_id="user-1") == before finally: await close_engine() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("recovery_path", ["once", "active"]) +async def test_scheduler_recovery_claim_backfills_error_terminal(tmp_path, recovery_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + event_store = DbRunEventStore(sf) + run_repo = RunRepository(sf) + publisher = RecoveryTerminalPublisher(event_store, run_repo) + task_repo = ScheduledTaskRepository(sf, run_repository=run_repo, recovery_terminal_publisher=publisher) + task_run_repo = ScheduledTaskRunRepository(sf, run_repository=run_repo, recovery_terminal_publisher=publisher) + now = datetime.now(UTC) + run_id, thread_id = f"run-{recovery_path}", f"thread-{recovery_path}" + + if recovery_path == "once": + await task_repo.create( + task_id="task-once-terminal", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="once terminal", + prompt="p", + schedule_type="once", + schedule_spec={"run_at": now.isoformat()}, + timezone="UTC", + next_run_at=None, + ) + await task_repo.update("task-once-terminal", user_id="user-1", updates={"status": "running", "last_run_id": run_id}) + await run_repo.put( + run_id, + thread_id=thread_id, + user_id="user-1", + status="running", + owner_worker_id="dead-owner", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + reconciled = await task_repo.reconcile_stuck_once_tasks(error="restart", now=now) + else: + await task_repo.create( + task_id="task-active-terminal", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="active terminal", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await task_run_repo.create( + run_record_id="task-run-active-terminal", + task_id="task-active-terminal", + thread_id=thread_id, + scheduled_for=now, + trigger="scheduled", + status="running", + ) + await task_run_repo.update_status("task-run-active-terminal", status="running", run_id=run_id) + await run_repo.put( + run_id, + thread_id=thread_id, + user_id="user-1", + status="running", + metadata={"scheduled_task_id": "task-active-terminal", "scheduled_task_run_id": "task-run-active-terminal"}, + owner_worker_id="dead-owner", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + reconciled = await task_run_repo.reconcile_active_runs(error="restart", now=now) + + assert reconciled == 1 + row = await run_repo.get(run_id, user_id="user-1") + assert row is not None + assert row["status"] == "error" + assert row["evidence_seal_state"] == "partial" + delivery = await event_store.list_events(thread_id, run_id, event_types=["run.delivery"], user_id=None) + terminal = await event_store.list_events(thread_id, run_id, event_types=["run.terminal"], user_id=None) + assert len(delivery) == 1 + assert len(terminal) == 1 + assert delivery[0]["seq"] < terminal[0]["seq"] + assert terminal[0]["content"] == {"status": "error"} + assert terminal[0]["metadata"]["protocol_version"] == 1 + assert terminal[0]["metadata"]["recovered"] is True + finally: + await close_engine() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("publisher_failure", [False, True]) +async def test_scheduler_outer_transaction_failure_still_publishes_committed_claim(tmp_path, monkeypatch, publisher_failure): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + event_store = DbRunEventStore(sf) + run_repo = RunRepository(sf) + publisher = RecoveryTerminalPublisher(event_store, run_repo) + task_repo = ScheduledTaskRepository(sf, run_repository=run_repo, recovery_terminal_publisher=publisher) + now = datetime.now(UTC) + run_id, thread_id = "run-outer-rollback", "thread-outer-rollback" + await task_repo.create( + task_id="task-outer-rollback", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="outer rollback", + prompt="p", + schedule_type="once", + schedule_spec={"run_at": now.isoformat()}, + timezone="UTC", + next_run_at=None, + ) + await task_repo.update("task-outer-rollback", user_id="user-1", updates={"status": "running", "last_run_id": run_id}) + task = await task_repo.get("task-outer-rollback", user_id="user-1") + assert task is not None and task["status"] == "running" + await run_repo.put( + run_id, + thread_id=thread_id, + user_id="user-1", + status="running", + owner_worker_id="dead-owner", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + + def fail_after_claim(*args, **kwargs): + raise RuntimeError("scheduler outer transaction failed") + + if publisher_failure: + + async def fail_confirmation(*args, **kwargs): + raise RuntimeError("unexpected confirmation failure") + + monkeypatch.setattr(terminal_module, "confirm_terminal_event", fail_confirmation) + monkeypatch.setattr(ScheduledTaskRepository, "_finalise_once_task_from_run", fail_after_claim) + with pytest.raises(RuntimeError, match="scheduler outer transaction failed"): + await asyncio.wait_for(task_repo.reconcile_stuck_once_tasks(error="restart", now=now), timeout=5) + + row = await run_repo.get(run_id, user_id="user-1") + assert row is not None + assert row["status"] == "error" + task = await task_repo.get("task-outer-rollback", user_id="user-1") + assert task is not None and task["status"] == "running" + terminal = await event_store.list_events(thread_id, run_id, event_types=["run.terminal"], user_id=None) + if publisher_failure: + assert terminal == [] + else: + assert len(terminal) == 1 + assert terminal[0]["metadata"]["recovered"] is True + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_scheduler_recovery_event_failure_preserves_committed_state(tmp_path, monkeypatch): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + event_store = DbRunEventStore(sf) + run_repo = RunRepository(sf) + publisher = RecoveryTerminalPublisher(event_store, run_repo) + task_repo = ScheduledTaskRepository(sf, run_repository=run_repo, recovery_terminal_publisher=publisher) + now = datetime.now(UTC) + run_id, thread_id = "run-event-failure", "thread-event-failure" + await task_repo.create( + task_id="task-event-failure", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="event failure", + prompt="p", + schedule_type="once", + schedule_spec={"run_at": now.isoformat()}, + timezone="UTC", + next_run_at=None, + ) + await task_repo.update("task-event-failure", user_id="user-1", updates={"status": "running", "last_run_id": run_id}) + task = await task_repo.get("task-event-failure", user_id="user-1") + assert task is not None and task["status"] == "running" + await run_repo.put( + run_id, + thread_id=thread_id, + user_id="user-1", + status="running", + owner_worker_id="dead-owner", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + original = event_store.put_if_absent + attempts = [] + + async def fail_terminal(**kwargs): + if kwargs["event_type"] == "run.terminal": + attempts.append(kwargs) + raise OSError("terminal event unavailable") + return await original(**kwargs) + + monkeypatch.setattr(event_store, "put_if_absent", fail_terminal) + + assert await task_repo.reconcile_stuck_once_tasks(error="restart", now=now) == 1 + + assert len(attempts) == 2 + row = await run_repo.get(run_id, user_id="user-1") + assert row is not None and row["status"] == "error" + task = await task_repo.get("task-event-failure", user_id="user-1") + assert task is not None and task["status"] == "cancelled" + assert await event_store.list_events(thread_id, run_id, event_types=["run.terminal"], user_id=None) == [] + finally: + await close_engine() diff --git a/backend/tests/test_terminal_event_confirmation.py b/backend/tests/test_terminal_event_confirmation.py index a138de372e..a95051301d 100644 --- a/backend/tests/test_terminal_event_confirmation.py +++ b/backend/tests/test_terminal_event_confirmation.py @@ -4,7 +4,8 @@ from __future__ import annotations import asyncio import os -from copy import deepcopy +from copy import copy, deepcopy +from dataclasses import replace from types import SimpleNamespace from uuid import uuid4 @@ -14,11 +15,21 @@ from sqlalchemy import func, select, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from support.postgres import asyncpg_test_url +import deerflow.runtime.events.terminal as terminal_module from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow +from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunRow +from deerflow.runtime.events.store.base import RunEventLifecycleRejected from deerflow.runtime.events.store.db import DbRunEventStore from deerflow.runtime.events.store.jsonl import JsonlRunEventStore from deerflow.runtime.events.store.memory import MemoryRunEventStore -from deerflow.runtime.events.terminal import TerminalEventConfirmation, confirm_terminal_event +from deerflow.runtime.events.terminal import ( + RecoveryTerminalPublisher, + RecoveryTerminalTargetConsumed, + TerminalEventConfirmation, + confirm_terminal_event, +) +from deerflow.runtime.runs.store.memory import MemoryRunStore +from deerflow.runtime.user_context import AUTO @pytest_asyncio.fixture(params=["memory", "jsonl", "sqlite", "postgresql"]) @@ -45,7 +56,18 @@ async def stores(request, tmp_path): async with first_engine.begin() as conn: if schema: await conn.execute(text(f'CREATE SCHEMA "{schema}"')) - await conn.run_sync(lambda c: RunEventRow.metadata.create_all(c, tables=[RunEventRow.__table__, RunEventThreadSeqRow.__table__])) + await conn.run_sync( + lambda c: RunEventRow.metadata.create_all( + c, + tables=[ + RunEventRow.__table__, + RunEventThreadSeqRow.__table__, + RunRow.__table__, + RunChangeClockRow.__table__, + CompletedRunSnapshotRow.__table__, + ], + ) + ) first = DbRunEventStore(async_sessionmaker(first_engine, expire_on_commit=False)) second = DbRunEventStore(async_sessionmaker(second_engine, expire_on_commit=False)) yield SimpleNamespace(first=first, second=second, backend=backend) @@ -61,6 +83,12 @@ async def confirm(store, status="success", *, authorized=lambda: True): return await confirm_terminal_event(store, thread_id="thread", run_id="run", status=status, user_id="alice", authorized=authorized) +async def recovery_publisher(store): + run_store = MemoryRunStore() + await run_store.put("run", thread_id="thread", user_id="alice", status="error", operation_kind="run") + return RecoveryTerminalPublisher(store, run_store) + + @pytest.mark.asyncio @pytest.mark.parametrize("status", ["success", "error", "interrupted", "timeout"]) async def test_confirmed_existing_retains_singleton_and_original_source(stores, status): @@ -198,3 +226,259 @@ async def test_database_ack_must_retain_owner_identity(stores, monkeypatch): monkeypatch.setattr(stores.first, "put_if_absent", omit_owner) assert await confirm(stores.first) == TerminalEventConfirmation.CONFLICT + + +@pytest.mark.asyncio +async def test_guarded_singleton_rejects_false_lifecycle(stores): + async def reject(): + return False + + with pytest.raises(RunEventLifecycleRejected): + await stores.first.put_if_absent( + thread_id="thread", + run_id="run", + event_type="run.terminal", + category="trace", + content={"status": "error"}, + metadata={"protocol_version": 1, "recovered": True}, + user_id="alice", + lifecycle_guard=reject, + ) + assert await stores.first.list_events("thread", "run", user_id="alice") == [] + + +@pytest.mark.asyncio +async def test_lifecycle_guard_holds_store_fence_against_thread_delete(stores, monkeypatch): + entered, release, delete_started = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def guard(): + entered.set() + await release.wait() + return True + + original_delete = stores.first.delete_by_thread + + async def observed_delete(thread_id, **kwargs): + delete_started.set() + return await original_delete(thread_id, **kwargs) + + monkeypatch.setattr(stores.first, "delete_by_thread", observed_delete) + put_task = asyncio.create_task( + stores.first.put_if_absent( + thread_id="thread", + run_id="run", + event_type="run.terminal", + category="trace", + content={"status": "error"}, + metadata={"protocol_version": 1, "recovered": True}, + user_id="alice", + lifecycle_guard=guard, + ) + ) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + delete_task = asyncio.create_task(stores.first.delete_by_thread("thread", user_id="alice")) + await asyncio.wait_for(delete_started.wait(), timeout=5) + assert not delete_task.done() + finally: + release.set() + await asyncio.wait_for(put_task, timeout=5) + await asyncio.wait_for(delete_task, timeout=5) + assert await stores.first.list_events("thread", "run", user_id="alice") == [] + + +@pytest.mark.asyncio +async def test_owner_confirmation_supports_legacy_store_signature(): + class LegacyRunEventStore(MemoryRunEventStore): + async def put_if_absent( + self, + *, + thread_id, + run_id, + event_type, + category, + content="", + metadata=None, + created_at=None, + user_id: str | None | type(AUTO) = AUTO, + ): + return await super().put_if_absent( + thread_id=thread_id, + run_id=run_id, + event_type=event_type, + category=category, + content=content, + metadata=metadata, + created_at=created_at, + user_id=user_id, + ) + + store = LegacyRunEventStore() + confirmation = await confirm_terminal_event( + store, + thread_id="thread", + run_id="run", + status="error", + user_id="alice", + authorized=lambda: True, + ) + assert confirmation == TerminalEventConfirmation.CONFIRMED_NEW + assert len(await store.list_events("thread", "run", user_id="alice")) == 1 + + +@pytest.mark.asyncio +async def test_recovery_confirmation_fails_closed_for_legacy_store_signature(): + class LegacyRunEventStore(MemoryRunEventStore): + async def put_if_absent( + self, + *, + thread_id, + run_id, + event_type, + category, + content="", + metadata=None, + created_at=None, + user_id: str | None | type(AUTO) = AUTO, + ): + return await super().put_if_absent( + thread_id=thread_id, + run_id=run_id, + event_type=event_type, + category=category, + content=content, + metadata=metadata, + created_at=created_at, + user_id=user_id, + ) + + store = LegacyRunEventStore() + publisher = await recovery_publisher(store) + target = publisher.authorize_takeover(claimed=True, thread_id="thread", run_id="run", user_id="alice", operation_kind="run") + assert target is not None + + assert await publisher.publish(target) == TerminalEventConfirmation.UNKNOWN + assert await store.list_events("thread", "run", event_types=["run.terminal"], user_id="alice") == [] + + +@pytest.mark.asyncio +async def test_recovery_target_requires_successful_run_takeover(): + store = MemoryRunEventStore() + publisher = await recovery_publisher(store) + identity = dict(thread_id="thread", run_id="run", user_id="alice") + + assert publisher.authorize_takeover(claimed=False, operation_kind="run", **identity) is None + assert publisher.authorize_takeover(claimed=True, operation_kind="checkpoint_write", **identity) is None + target = publisher.authorize_takeover(claimed=True, operation_kind="run", **identity) + assert target is not None + assert (target.thread_id, target.run_id, target.user_id, target.operation_kind, target.committed_status) == ( + "thread", + "run", + "alice", + "run", + "error", + ) + + assert await publisher.publish(target) == TerminalEventConfirmation.CONFIRMED_NEW + delivery = await store.list_events("thread", "run", event_types=["run.delivery"], user_id="alice") + terminal = await store.list_events("thread", "run", event_types=["run.terminal"], user_id="alice") + assert len(delivery) == 1 + assert len(terminal) == 1 + assert delivery[0]["seq"] < terminal[0]["seq"] + assert terminal[0]["content"] == {"status": "error"} + assert terminal[0]["metadata"] == {"protocol_version": 1, "recovered": True} + + +@pytest.mark.asyncio +async def test_recovery_target_cannot_be_replaced_or_reauthorized(): + store = MemoryRunEventStore() + publisher = await recovery_publisher(store) + target = publisher.authorize_takeover(claimed=True, thread_id="thread", run_id="run", user_id="alice", operation_kind="run") + assert target is not None + + with pytest.raises(TypeError): + replace(target, run_id="another-run") + with pytest.raises(AttributeError): + target.run_id = "another-run" + with pytest.raises(AttributeError): + target._claim_state = type(target._claim_state)() + with pytest.raises(AttributeError): + target._claim_state._consumed = False + + copied = copy(target) + assert await publisher.publish(target) == TerminalEventConfirmation.CONFIRMED_NEW + with pytest.raises(RecoveryTerminalTargetConsumed): + await publisher.publish(copied) + + +@pytest.mark.asyncio +async def test_recovery_target_is_one_shot_across_lost_ack_retry(): + store = MemoryRunEventStore() + publisher = await recovery_publisher(store) + target = publisher.authorize_takeover(claimed=True, thread_id="thread", run_id="run", user_id="alice", operation_kind="run") + assert target is not None + entered, release = asyncio.Event(), asyncio.Event() + original = store.put_if_absent + attempts = [] + + async def commit_then_lose_first_terminal_ack(**kwargs): + if kwargs["event_type"] != "run.terminal": + return await original(**kwargs) + result = await original(**kwargs) + attempts.append(result) + if len(attempts) == 1: + entered.set() + await release.wait() + raise OSError("ACK lost after commit") + return result + + store.put_if_absent = commit_then_lose_first_terminal_ack + first = asyncio.create_task(publisher.publish(target)) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + with pytest.raises(RecoveryTerminalTargetConsumed): + await publisher.publish(target) + finally: + release.set() + assert await asyncio.wait_for(first, timeout=5) == TerminalEventConfirmation.CONFIRMED_EXISTING + assert [created for _event, created in attempts] == [True, False] + terminal = await store.list_events("thread", "run", event_types=["run.terminal"], user_id="alice") + assert len(terminal) == 1 + assert terminal[0]["metadata"]["recovered"] is True + + +@pytest.mark.asyncio +async def test_recovery_target_isolates_unexpected_confirmation_error(monkeypatch): + store = MemoryRunEventStore() + publisher = await recovery_publisher(store) + target = publisher.authorize_takeover(claimed=True, thread_id="thread", run_id="run", user_id="alice", operation_kind="run") + assert target is not None + + async def unexpected_failure(*args, **kwargs): + raise RuntimeError("unexpected confirmation failure") + + monkeypatch.setattr(terminal_module, "confirm_terminal_event", unexpected_failure) + + assert await publisher.publish(target) == TerminalEventConfirmation.UNKNOWN + with pytest.raises(RecoveryTerminalTargetConsumed): + await publisher.publish(target) + assert await store.list_events("thread", "run", event_types=["run.delivery"], user_id="alice") + assert await store.list_events("thread", "run", event_types=["run.terminal"], user_id="alice") == [] + + +@pytest.mark.asyncio +async def test_recovery_target_propagates_cancellation_after_consumption(monkeypatch): + store = MemoryRunEventStore() + publisher = await recovery_publisher(store) + target = publisher.authorize_takeover(claimed=True, thread_id="thread", run_id="run", user_id="alice", operation_kind="run") + assert target is not None + + async def cancelled(*args, **kwargs): + raise asyncio.CancelledError() + + monkeypatch.setattr(terminal_module, "confirm_terminal_event", cancelled) + + with pytest.raises(asyncio.CancelledError): + await publisher.publish(target) + with pytest.raises(RecoveryTerminalTargetConsumed): + await publisher.publish(target) diff --git a/contracts/run_event_stream_contract.json b/contracts/run_event_stream_contract.json index d393e0ddf9..e2a4cad1de 100644 --- a/contracts/run_event_stream_contract.json +++ b/contracts/run_event_stream_contract.json @@ -109,7 +109,7 @@ { "event_type": "run.terminal", "category": "trace", - "producer": "run_agent() finalization, after final commit ACK and Journal drain, before Evidence Receipt", + "producer": "run_agent() finalization, after final commit ACK and Journal drain, before Evidence Receipt; RecoveryTerminalPublisher after a successful atomic takeover claim", "identity": ["thread_id", "run_id", "event_type"], "content_schema": { "type": "object", @@ -126,7 +126,7 @@ }, "additionalProperties": true }, - "notes": "Owner-side terminal business outcome, confirmed as a run-scoped singleton. Current workers produce recovered=false. This does not redefine root-invocation run.end, prove complete producer coverage, or guarantee publication after a crash. Historical sealed records need not contain this additive event." + "notes": "Authoritative terminal business outcome, confirmed as a run-scoped singleton. Owner Workers emit recovered=false; confirmed recovery takeover emits recovered=true only when creating a missing terminal, guarded by the current Run lifecycle inside the event-store mutation fence. An existing matching event retains its original recovered marker. This does not redefine root-invocation run.end, prove complete producer coverage, or guarantee publication after every terminal commit. Historical sealed records need not contain this additive event." }, { "event_type": "run.error",