mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(runtime): backfill run.terminal after confirmed recovery takeover (#7500)
* fix(runtime): backfill run.terminal on confirmed recovery takeover * fix(runtime): isolate recovery terminal publication failures * fix(runtime): fence recovery terminal publication lifecycle * fix(runtime): unify recovery receipt publication * docs(runtime): document lifecycle guard lock constraints * fix(runtime): reuse event transaction for recovery guards Read Run lifecycle through the active event session when SQL stores share an engine, preserving deletion fencing without a second pool checkout. Keep custom lookups and different engines on their ordinary read path. Scope memory recovery guards and deletions to per-thread locks. Add SQLite/PostgreSQL pool pressure, stale identity, fallback, and independent-thread regressions, and document the contracts. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
f6da662c68
commit
33f7b43400
@ -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
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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,
|
||||
*,
|
||||
|
||||
@ -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"}:
|
||||
|
||||
@ -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"}:
|
||||
|
||||
@ -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`).
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
142
backend/tests/test_recovery_terminal_pool.py
Normal file
142
backend/tests/test_recovery_terminal_pool.py
Normal file
@ -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
|
||||
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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",
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user