mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-12 22:20:33 +00:00
8850f8eb7e
* fix(runtime): keep assembly evidence from reviving a run a peer reclaimed In heartbeat mode every run-row write is conditional on this worker's live lease except assembly evidence, which went through RunStore.put(). That upsert rewrote status, owner, lease, error and stop_reason from the local record. When a peer had already reclaimed the run, it put the row back to running and erased the takeover error. If the owner's local lease was still live, its renewal then succeeded, the peer never reclaimed it again, and the run finished as success after the peer had published its terminal recovery event. If the write landed late, the row came back with an expired lease and was reclaimed a second time. Add RunStore.merge_owned_metadata(), which merges metadata keys only for the owner holding a live lease on a pending/running row and writes no other field. SQL follows the existing clock -> row lock -> conditional UPDATE order; MemoryRunStore mirrors it. With heartbeat enabled, record_assembly_evidence() uses it, refuses a locally fenced run before any write, fences the run when the store rejects the merge, and fails closed on stores without the method instead of falling back to put(). Single-worker mode keeps the full snapshot. * docs(changelog): reference #8046 for the assembly evidence takeover fix * test(runtime): scope the simulated lease expiry to the run's thread --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
2986 lines
116 KiB
Python
2986 lines
116 KiB
Python
"""Tests for multi-worker run ownership (work items 2–3).
|
||
|
||
Coverage:
|
||
- create_or_reject with reject strategy blocks duplicate active runs
|
||
- create_or_reject with interrupt strategy claims and cancels old runs
|
||
- create_thread_operation_atomic refuses to interrupt a run owned by another live worker
|
||
- reconcile_orphaned_inflight_runs uses lease-based detection
|
||
- periodic reconciliation notifies Gateway recovery orchestration
|
||
- Worker reconciliation skips runs with unexpired leases
|
||
- Lease heartbeat renews active run leases
|
||
- Lease heartbeat keeps renewing a staged terminal run until its commit
|
||
- GATEWAY_WORKERS=1 + heartbeat_enabled=false behaviour unchanged
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from datetime import UTC, datetime, timedelta
|
||
from types import SimpleNamespace
|
||
from unittest.mock import AsyncMock
|
||
|
||
import pytest
|
||
|
||
from deerflow.config.run_ownership_config import RunOwnershipConfig
|
||
from deerflow.runtime import ORPHAN_RECOVERY_STOP_REASON, RunManager, RunStatus, ThreadOperationKind
|
||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||
from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, _generate_worker_id
|
||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||
from deerflow.runtime.runs.worker import RunContext, run_agent
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _lease_config(**kwargs) -> RunOwnershipConfig:
|
||
return RunOwnershipConfig(
|
||
lease_seconds=kwargs.get("lease_seconds", 30),
|
||
grace_seconds=kwargs.get("grace_seconds", 10),
|
||
heartbeat_enabled=kwargs.get("heartbeat_enabled", False),
|
||
)
|
||
|
||
|
||
def _make_manager(store=None, **kwargs) -> RunManager:
|
||
return RunManager(
|
||
store=store or MemoryRunStore(),
|
||
run_ownership_config=kwargs.pop("run_ownership_config", _lease_config()),
|
||
**kwargs,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# create_or_reject — reject strategy
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reject_blocks_when_active_run_exists():
|
||
"""reject strategy must raise ConflictError when thread has an active run."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
await manager.create("thread-1")
|
||
await manager.set_status((await manager.list_by_thread("thread-1"))[0].run_id, RunStatus.running)
|
||
|
||
with pytest.raises(ConflictError, match="already has an active run"):
|
||
await manager.create_or_reject("thread-1", multitask_strategy="reject")
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reject_succeeds_when_no_active_run():
|
||
"""reject strategy must succeed when the thread has no active run."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
record = await manager.create_or_reject("thread-1", multitask_strategy="reject")
|
||
assert record is not None
|
||
assert record.status == RunStatus.pending
|
||
assert record.owner_worker_id is not None
|
||
assert record.lease_expires_at is not None
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_checkpoint_write_reservation_rejects_nonowning_worker_while_run_is_active():
|
||
"""A durable run owned by worker A must block worker B's checkpoint writer."""
|
||
store = MemoryRunStore()
|
||
owner = _make_manager(store=store, worker_id="worker-a")
|
||
non_owner = _make_manager(store=store, worker_id="worker-b")
|
||
active = await owner.create_or_reject("thread-1")
|
||
await owner.set_status(active.run_id, RunStatus.running)
|
||
|
||
with pytest.raises(ConflictError, match="already has an active run"):
|
||
async with non_owner.reserve_thread_operation("thread-1", kind=ThreadOperationKind.checkpoint_write):
|
||
pytest.fail("the checkpoint mutation guard must not be acquired")
|
||
|
||
stored = await store.get(active.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "running"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_checkpoint_write_reservation_blocks_new_runs_until_mutation_finishes():
|
||
"""The durable guard closes the check-then-write window in both directions."""
|
||
store = MemoryRunStore()
|
||
compaction_worker = _make_manager(store=store, worker_id="worker-a")
|
||
run_worker = _make_manager(store=store, worker_id="worker-b")
|
||
|
||
async with compaction_worker.reserve_thread_operation("thread-1", kind=ThreadOperationKind.checkpoint_write):
|
||
inflight = await store.list_inflight()
|
||
assert len(inflight) == 1
|
||
assert inflight[0]["operation_kind"] == ThreadOperationKind.checkpoint_write
|
||
assert inflight[0]["metadata"] == {}
|
||
|
||
with pytest.raises(ConflictError, match="checkpoint write"):
|
||
await run_worker.create_or_reject("thread-1", multitask_strategy="interrupt")
|
||
|
||
assert await store.list_inflight() == []
|
||
assert await store.list_by_thread("thread-1") == []
|
||
|
||
admitted = await run_worker.create_or_reject("thread-1")
|
||
assert admitted.status == RunStatus.pending
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_interrupt_reclaims_expired_checkpoint_write_reservation():
|
||
"""A dead checkpoint writer must not wait for periodic reconciliation."""
|
||
store = MemoryRunStore()
|
||
expired = (datetime.now(UTC) - timedelta(seconds=30)).isoformat()
|
||
await store.put(
|
||
"checkpoint-write-1",
|
||
thread_id="thread-1",
|
||
status="pending",
|
||
operation_kind=ThreadOperationKind.checkpoint_write,
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired,
|
||
created_at=expired,
|
||
)
|
||
manager = _make_manager(
|
||
store=store,
|
||
worker_id="worker-b",
|
||
run_ownership_config=_lease_config(grace_seconds=10),
|
||
)
|
||
|
||
admitted = await manager.create_or_reject("thread-1", multitask_strategy="interrupt")
|
||
|
||
assert admitted.status == RunStatus.pending
|
||
stale = await store.get("checkpoint-write-1")
|
||
assert stale is not None
|
||
assert stale["status"] == "interrupted"
|
||
assert stale["owner_worker_id"] == "worker-b"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reject_blocks_reentrant_same_thread_locally():
|
||
"""reject must also block when a local in-memory active run exists."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
await manager.create_or_reject("thread-1", multitask_strategy="reject")
|
||
|
||
with pytest.raises(ConflictError, match="already has an active run"):
|
||
await manager.create_or_reject("thread-1", multitask_strategy="reject")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# create_or_reject — cross-worker idempotent reuse
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_peer_idempotent_reuse_does_not_block_thread_after_owner_completes():
|
||
"""A peer's reuse handle must not stay behind as a local inflight record.
|
||
|
||
The peer never runs the task, so nothing on it would finalize or clean up
|
||
a registered copy: that copy keeps its admission-time status and rejects
|
||
every later admission for the thread until the worker restarts.
|
||
"""
|
||
store = MemoryRunStore()
|
||
owner = _make_manager(store=store, worker_id="worker-a")
|
||
peer = _make_manager(store=store, worker_id="worker-b")
|
||
first = await owner.create_or_reject("thread-1", idempotency_key="scheduled-task:occurrence-1")
|
||
await owner.set_status(first.run_id, RunStatus.running)
|
||
|
||
reused = await peer.create_or_reject("thread-1", idempotency_key="scheduled-task:occurrence-1")
|
||
|
||
assert reused.run_id == first.run_id
|
||
assert reused.store_only is True
|
||
assert reused.idempotency_reused is True
|
||
|
||
await owner.set_status(first.run_id, RunStatus.success)
|
||
await owner.cleanup(first.run_id, delay=0)
|
||
|
||
assert await peer.has_inflight("thread-1") is False
|
||
hydrated = await peer.get(first.run_id)
|
||
assert hydrated is not None
|
||
assert hydrated.status == RunStatus.success
|
||
retried = await peer.create_or_reject("thread-1", idempotency_key="scheduled-task:occurrence-1")
|
||
assert retried.run_id == first.run_id
|
||
assert retried.status == RunStatus.success
|
||
follow_up = await peer.create_or_reject("thread-1")
|
||
assert follow_up.run_id != first.run_id
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_peer_idempotent_reuse_does_not_shield_crashed_owner_from_reconciliation():
|
||
"""Reconciliation skips locally live records, so a reuse handle must not look like one."""
|
||
store = MemoryRunStore()
|
||
owner = _make_manager(store=store, worker_id="worker-a")
|
||
peer = _make_manager(store=store, worker_id="worker-b")
|
||
first = await owner.create_or_reject("thread-1", idempotency_key="mcp-task:task-1:1:0")
|
||
await owner.set_status(first.run_id, RunStatus.running)
|
||
await peer.create_or_reject("thread-1", idempotency_key="mcp-task:task-1:1:0")
|
||
|
||
# The owner crashes: its lease lapses past the grace window without renewal.
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
store._runs[first.run_id]["lease_expires_at"] = expired_lease
|
||
|
||
recovered = await peer.reconcile_orphaned_inflight_runs(error="owner expired")
|
||
|
||
assert [record.run_id for record in recovered] == [first.run_id]
|
||
stored = await store.get(first.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "error"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
@pytest.mark.parametrize(
|
||
("heartbeat_enabled", "expected_outcome", "expected_cancel_action"),
|
||
[(False, CancelOutcome.not_active_locally, None), (True, CancelOutcome.requested, "interrupt")],
|
||
)
|
||
async def test_peer_cancel_of_reused_run_leaves_live_owner_status_to_the_owner(heartbeat_enabled, expected_outcome, expected_cancel_action):
|
||
"""A reuse handle must not route the peer's cancel through the local-owner path.
|
||
|
||
That path would mark the owner's still-running row ``interrupted`` from a
|
||
worker that has no task to stop, releasing the thread while the owner runs.
|
||
"""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=heartbeat_enabled)
|
||
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
|
||
first = await owner.create_or_reject("thread-1", idempotency_key="http-run:retry")
|
||
await owner.set_status(first.run_id, RunStatus.running)
|
||
await peer.create_or_reject("thread-1", idempotency_key="http-run:retry")
|
||
|
||
outcome = await peer.cancel(first.run_id)
|
||
|
||
assert outcome == expected_outcome
|
||
stored = await store.get(first.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "running"
|
||
assert stored.get("cancel_action") == expected_cancel_action
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# create_or_reject — interrupt strategy
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_interrupt_cancels_old_run_and_creates_new():
|
||
"""interrupt must cancel the previous active run and create a new one."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
old = await manager.create_or_reject("thread-1", multitask_strategy="reject")
|
||
await manager.set_status(old.run_id, RunStatus.running)
|
||
|
||
new = await manager.create_or_reject("thread-1", multitask_strategy="interrupt")
|
||
|
||
assert new.run_id != old.run_id
|
||
assert new.status == RunStatus.pending
|
||
|
||
# Old run must be interrupted locally
|
||
assert old.status == RunStatus.interrupted
|
||
assert old.abort_event.is_set()
|
||
|
||
# Old run must be marked interrupted in-store (persist_status after local cancel)
|
||
old_after = await store.get(old.run_id)
|
||
assert old_after["status"] == "interrupted"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_interrupt_creates_new_when_old_completed():
|
||
"""interrupt must succeed when the previous run already reached a terminal status."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
old = await manager.create_or_reject("thread-1")
|
||
await manager.set_status(old.run_id, RunStatus.success)
|
||
|
||
new = await manager.create_or_reject("thread-1", multitask_strategy="interrupt")
|
||
assert new.run_id != old.run_id
|
||
assert new.status == RunStatus.pending
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_interrupt_exhausted_retries_surface_as_conflict_error():
|
||
"""When all retry attempts collide with a unique violation, the loop must
|
||
surface ConflictError (HTTP 409) — matching the reject branch — instead of
|
||
leaking the raw IntegrityError (HTTP 500).
|
||
|
||
Without the post-loop conversion, the last attempt's ``raise`` re-raises
|
||
the IntegrityError, giving callers an inconsistent signal depending on
|
||
which strategy they picked. The reject path already converts; this test
|
||
pins the symmetric behaviour for interrupt/rollback.
|
||
"""
|
||
import sqlite3
|
||
|
||
class _AlwaysUniqueViolationStore(MemoryRunStore):
|
||
"""MemoryRunStore whose ``create_thread_operation_atomic`` always raises a
|
||
real-flavoured unique-violation IntegrityError, simulating a worker
|
||
that keeps losing the cross-worker race for the same thread."""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.atomic_call_count = 0
|
||
|
||
async def create_thread_operation_atomic(self, *args, **kwargs):
|
||
self.atomic_call_count += 1
|
||
err = sqlite3.IntegrityError("UNIQUE constraint failed: runs.uq_runs_thread_active")
|
||
err.sqlite_errorcode = sqlite3.SQLITE_CONSTRAINT_UNIQUE
|
||
raise err
|
||
|
||
store = _AlwaysUniqueViolationStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
with pytest.raises(ConflictError, match="already has an active run"):
|
||
await manager.create_or_reject("thread-1", multitask_strategy="interrupt")
|
||
|
||
# Sanity: the loop actually retried 3 times before giving up.
|
||
assert store.atomic_call_count == 3
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# create_or_reject — run ownership metadata
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_run_record_stores_owner_and_lease():
|
||
"""Newly created runs must carry owner_worker_id and lease_expires_at (when heartbeat is on)."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
record = await manager.create_or_reject("thread-1")
|
||
|
||
assert record.owner_worker_id == manager.worker_id
|
||
assert isinstance(record.owner_worker_id, str) and len(record.owner_worker_id) > 0
|
||
assert record.lease_expires_at is not None
|
||
|
||
# Store row must also carry the fields
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["owner_worker_id"] == manager.worker_id
|
||
assert stored["lease_expires_at"] is not None
|
||
assert stored["operation_kind"] == ThreadOperationKind.run
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_store_row_roundtrips_ownership_fields():
|
||
"""Records hydrated from the store must surface ownership fields."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
record = await manager.create_or_reject("thread-1")
|
||
|
||
hydrated = await manager.get(record.run_id)
|
||
assert hydrated is not None
|
||
assert hydrated.owner_worker_id == manager.worker_id
|
||
assert hydrated.lease_expires_at is not None
|
||
assert hydrated.operation_kind == ThreadOperationKind.run
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_releases_expired_internal_operation_without_reporting_run():
|
||
"""Expired internal reservations release admission without becoming failed runs."""
|
||
store = MemoryRunStore()
|
||
expired = (datetime.now(UTC) - timedelta(seconds=30)).isoformat()
|
||
await store.put(
|
||
"checkpoint-write-1",
|
||
thread_id="thread-1",
|
||
status="pending",
|
||
operation_kind=ThreadOperationKind.checkpoint_write,
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired,
|
||
created_at=expired,
|
||
)
|
||
manager = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=10),
|
||
)
|
||
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(error="owner expired")
|
||
|
||
assert recovered == []
|
||
stored = await store.get("checkpoint-write-1")
|
||
assert stored is not None
|
||
assert stored["status"] == "error"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# reconcile_orphaned_inflight_runs — lease-based
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_claims_expired_lease_runs():
|
||
"""A run with an expired lease must be reclaimed as orphaned."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
# Insert a run with an already-expired lease
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"expired-run",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="worker-dead",
|
||
lease_expires_at=expired_lease,
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(),
|
||
)
|
||
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
assert len(recovered) == 1
|
||
assert recovered[0].run_id == "expired-run"
|
||
assert recovered[0].status == RunStatus.error
|
||
|
||
stored = await store.get("expired-run")
|
||
assert stored["status"] == "error"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_skips_active_lease_runs():
|
||
"""A run with a still-valid lease must NOT be reclaimed."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
# Insert a run with a still-valid lease
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"live-run",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="worker-alive",
|
||
lease_expires_at=valid_lease,
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=10)).isoformat(),
|
||
)
|
||
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
# Live run's lease is still valid — must not be reclaimed
|
||
assert all(r.run_id != "live-run" for r in recovered)
|
||
|
||
stored = await store.get("live-run")
|
||
assert stored["status"] == "running"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_claims_when_expired_owner_attempts_renewal_after_scan():
|
||
"""An expired owner cannot revive its lease between scan and takeover."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
|
||
await store.put(
|
||
"race-run",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="worker-alive",
|
||
lease_expires_at=expired_lease,
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(),
|
||
)
|
||
original_list = store.list_inflight_with_expired_lease
|
||
|
||
async def list_then_owner_renews(*, before=None, grace_seconds=10):
|
||
rows = [dict(row) for row in await original_list(before=before, grace_seconds=grace_seconds)]
|
||
renewed_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
updated = await store.update_lease(
|
||
"race-run",
|
||
owner_worker_id="worker-alive",
|
||
lease_expires_at=renewed_lease,
|
||
)
|
||
assert updated is False
|
||
return rows
|
||
|
||
store.list_inflight_with_expired_lease = list_then_owner_renews
|
||
manager = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace),
|
||
)
|
||
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
assert [r.run_id for r in recovered] == ["race-run"]
|
||
stored = await store.get("race-run")
|
||
assert stored["status"] == "error"
|
||
assert stored["lease_expires_at"] == expired_lease
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_concurrent_reconcilers_report_one_successful_claim():
|
||
"""Two reapers scanning the same candidate must report one recovery."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
run_id = "contended-orphan"
|
||
await store.put(
|
||
run_id,
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="worker-dead",
|
||
lease_expires_at=(datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat(),
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(),
|
||
)
|
||
|
||
original_scan = store.list_inflight_with_expired_lease
|
||
both_scanned = asyncio.Event()
|
||
scan_lock = asyncio.Lock()
|
||
scan_count = 0
|
||
|
||
async def synchronized_scan(*, before=None, grace_seconds=10):
|
||
nonlocal scan_count
|
||
rows = [dict(row) for row in await original_scan(before=before, grace_seconds=grace_seconds)]
|
||
async with scan_lock:
|
||
scan_count += 1
|
||
if scan_count == 2:
|
||
both_scanned.set()
|
||
await asyncio.wait_for(both_scanned.wait(), timeout=1)
|
||
return rows
|
||
|
||
store.list_inflight_with_expired_lease = synchronized_scan
|
||
managers = [
|
||
_make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace)),
|
||
_make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace)),
|
||
]
|
||
|
||
results = await asyncio.gather(*(manager.reconcile_orphaned_inflight_runs(error="orphaned") for manager in managers))
|
||
|
||
assert sorted(len(recovered) for recovered in results) == [0, 1]
|
||
assert [record.run_id for recovered in results for record in recovered] == [run_id]
|
||
row = await store.get(run_id)
|
||
assert row is not None
|
||
assert row["status"] == "error"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_claims_null_lease_runs():
|
||
"""Pre-ownership rows (NULL lease) must be reclaimed."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
await store.put(
|
||
"legacy-run",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(),
|
||
)
|
||
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
assert len(recovered) == 1
|
||
assert recovered[0].run_id == "legacy-run"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_heartbeat_disabled_crashed_run_reclaimed_immediately():
|
||
"""Single-worker regression: when heartbeat is off, a crashed run must be
|
||
reclaimed on the next restart without waiting for lease expiry.
|
||
|
||
The run is created with lease_expires_at=NULL (no heartbeat => no lease),
|
||
so reconciliation treats it as an orphan and reclaims it right away —
|
||
preserving the pre-ownership recovery latency.
|
||
"""
|
||
store = MemoryRunStore()
|
||
# Worker A: heartbeat disabled (single-worker default)
|
||
manager_a = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=False))
|
||
record = await manager_a.create("thread-1")
|
||
await manager_a.set_status(record.run_id, RunStatus.running)
|
||
|
||
# Verify the run was stored WITHOUT a lease (heartbeat off)
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["lease_expires_at"] is None
|
||
|
||
# Simulate crash: drop manager_a's local state, build a fresh manager
|
||
# (same store) as if Worker A restarted.
|
||
manager_b = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=False))
|
||
|
||
# Reconciliation must reclaim the run IMMEDIATELY — no lease to wait out.
|
||
recovered = await manager_b.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
assert len(recovered) == 1
|
||
assert recovered[0].run_id == record.run_id
|
||
assert recovered[0].status == RunStatus.error
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_skips_locally_active_runs():
|
||
"""An active local run (owned by this worker) must NOT be reclaimed even with an expired lease."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
# Create a live local run
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
|
||
# Its lease hasn't expired yet, so this is mostly testing the local-ownership guard
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
assert all(r.run_id != record.run_id for r in recovered)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_reconciliation_returns_empty_when_no_orphaned_runs():
|
||
"""Reconciliation must return empty when there are no orphaned runs."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
recovered = await manager.reconcile_orphaned_inflight_runs(
|
||
error="Gateway restarted before this run reached a durable final state.",
|
||
)
|
||
|
||
assert recovered == []
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_periodic_reconciliation_notifies_recovery_callback():
|
||
"""Periodic recovery must hand terminalized rows to Gateway orchestration."""
|
||
store = MemoryRunStore()
|
||
on_orphans_recovered = AsyncMock()
|
||
manager = _make_manager(
|
||
store=store,
|
||
on_orphans_recovered=on_orphans_recovered,
|
||
)
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"periodic-orphan",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(),
|
||
)
|
||
|
||
await manager._reconcile_orphans_periodic()
|
||
await asyncio.sleep(0)
|
||
|
||
on_orphans_recovered.assert_awaited_once()
|
||
recovered = on_orphans_recovered.await_args.args[0]
|
||
assert [record.run_id for record in recovered] == ["periodic-orphan"]
|
||
assert recovered[0].status == RunStatus.error
|
||
assert recovered[0].stop_reason == ORPHAN_RECOVERY_STOP_REASON
|
||
stored = await store.get("periodic-orphan")
|
||
assert stored is not None
|
||
assert stored["stop_reason"] == ORPHAN_RECOVERY_STOP_REASON
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_periodic_reconciliation_logs_recovered_run_ids_when_callback_fails(caplog):
|
||
"""Callback failures must identify every recovered run in the warning."""
|
||
store = MemoryRunStore()
|
||
on_orphans_recovered = AsyncMock(side_effect=RuntimeError("callback failed"))
|
||
manager = _make_manager(
|
||
store=store,
|
||
on_orphans_recovered=on_orphans_recovered,
|
||
)
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
created_at = (datetime.now(UTC) - timedelta(seconds=120)).isoformat()
|
||
for run_id in ("periodic-orphan-1", "periodic-orphan-2"):
|
||
await store.put(
|
||
run_id,
|
||
thread_id=f"thread-{run_id}",
|
||
status="running",
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
created_at=created_at,
|
||
)
|
||
|
||
with caplog.at_level("WARNING", logger="deerflow.runtime.runs.manager"):
|
||
await manager._reconcile_orphans_periodic()
|
||
await asyncio.sleep(0)
|
||
|
||
assert "Periodic orphan recovery callback failed for 2 run(s)" in caplog.text
|
||
assert "periodic-orphan-1" in caplog.text
|
||
assert "periodic-orphan-2" in caplog.text
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_periodic_terminalization_does_not_block_lease_renewal_or_shutdown():
|
||
"""The real heartbeat loop must keep renewing during a slow callback."""
|
||
store = MemoryRunStore()
|
||
callback_started = asyncio.Event()
|
||
callback_release = asyncio.Event()
|
||
callback_finished = asyncio.Event()
|
||
|
||
async def on_orphans_recovered(_recovered):
|
||
callback_started.set()
|
||
await callback_release.wait()
|
||
callback_finished.set()
|
||
|
||
manager = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=5),
|
||
on_orphans_recovered=on_orphans_recovered,
|
||
)
|
||
active = await manager.create_or_reject("active-thread")
|
||
await manager.set_status(active.run_id, RunStatus.running)
|
||
original_expiry = active.lease_expires_at
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"periodic-orphan",
|
||
thread_id="orphan-thread",
|
||
status="running",
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
created_at=expired_lease,
|
||
)
|
||
|
||
await manager.start_heartbeat()
|
||
await asyncio.wait_for(callback_started.wait(), timeout=4.5)
|
||
expiry_during_callback = active.lease_expires_at
|
||
await asyncio.sleep(1.2)
|
||
|
||
assert expiry_during_callback != original_expiry
|
||
assert active.lease_expires_at != expiry_during_callback
|
||
assert callback_finished.is_set() is False
|
||
|
||
shutdown_task = asyncio.create_task(manager.shutdown(timeout=1.0))
|
||
await asyncio.sleep(0)
|
||
assert shutdown_task.done() is False
|
||
callback_release.set()
|
||
await asyncio.wait_for(shutdown_task, timeout=1.0)
|
||
assert callback_finished.is_set() is True
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_periodic_store_scan_does_not_block_real_heartbeat_loop():
|
||
"""A slow orphan scan must not stall later lease-renewal cycles."""
|
||
|
||
class SlowScanStore(MemoryRunStore):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.scan_started = asyncio.Event()
|
||
self.scan_release = asyncio.Event()
|
||
|
||
async def list_inflight_with_expired_lease(
|
||
self,
|
||
*,
|
||
before=None,
|
||
grace_seconds=10,
|
||
):
|
||
self.scan_started.set()
|
||
await self.scan_release.wait()
|
||
return await super().list_inflight_with_expired_lease(
|
||
before=before,
|
||
grace_seconds=grace_seconds,
|
||
)
|
||
|
||
store = SlowScanStore()
|
||
manager = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=5),
|
||
)
|
||
active = await manager.create_or_reject("active-thread")
|
||
await manager.set_status(active.run_id, RunStatus.running)
|
||
|
||
await manager.start_heartbeat()
|
||
await asyncio.wait_for(store.scan_started.wait(), timeout=4.5)
|
||
expiry_during_scan = active.lease_expires_at
|
||
await asyncio.sleep(1.2)
|
||
|
||
assert active.lease_expires_at != expiry_during_scan
|
||
|
||
store.scan_release.set()
|
||
await manager.stop_heartbeat()
|
||
await manager._drain_orphan_recovery_task(timeout=0.5)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_periodic_recovery_is_single_flight():
|
||
"""A second trigger must not create another recovery pipeline."""
|
||
store = MemoryRunStore()
|
||
callback_started = asyncio.Event()
|
||
callback_release = asyncio.Event()
|
||
callback_calls = 0
|
||
|
||
async def on_orphans_recovered(_recovered):
|
||
nonlocal callback_calls
|
||
callback_calls += 1
|
||
callback_started.set()
|
||
await callback_release.wait()
|
||
|
||
manager = _make_manager(
|
||
store=store,
|
||
on_orphans_recovered=on_orphans_recovered,
|
||
)
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"periodic-orphan",
|
||
thread_id="orphan-thread",
|
||
status="running",
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
created_at=expired_lease,
|
||
)
|
||
|
||
manager._schedule_orphan_reconciliation()
|
||
await asyncio.wait_for(callback_started.wait(), timeout=0.5)
|
||
first_task = manager._orphan_recovery_task
|
||
manager._schedule_orphan_reconciliation()
|
||
|
||
assert manager._orphan_recovery_task is first_task
|
||
assert callback_calls == 1
|
||
|
||
callback_release.set()
|
||
await asyncio.wait_for(first_task, timeout=0.5)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_shutdown_cancels_recovery_that_exceeds_drain_budget():
|
||
"""The pending -> cancel -> gather branch must observe callback cancellation."""
|
||
store = MemoryRunStore()
|
||
callback_started = asyncio.Event()
|
||
callback_cancelled = asyncio.Event()
|
||
|
||
async def on_orphans_recovered(_recovered):
|
||
callback_started.set()
|
||
try:
|
||
await asyncio.Event().wait()
|
||
except asyncio.CancelledError:
|
||
callback_cancelled.set()
|
||
raise
|
||
|
||
manager = _make_manager(
|
||
store=store,
|
||
on_orphans_recovered=on_orphans_recovered,
|
||
)
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"periodic-orphan",
|
||
thread_id="orphan-thread",
|
||
status="running",
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
created_at=expired_lease,
|
||
)
|
||
manager._schedule_orphan_reconciliation()
|
||
await asyncio.wait_for(callback_started.wait(), timeout=0.5)
|
||
|
||
await manager.shutdown(timeout=0.01)
|
||
|
||
assert callback_cancelled.is_set()
|
||
assert manager._orphan_recovery_task is None
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_shutdown_applies_shared_deadline_to_heartbeat_stop():
|
||
"""A stuck heartbeat must not receive a separate five-second budget."""
|
||
manager = _make_manager(
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
heartbeat_cancelled = asyncio.Event()
|
||
|
||
async def stuck_heartbeat():
|
||
try:
|
||
await asyncio.Event().wait()
|
||
except asyncio.CancelledError:
|
||
heartbeat_cancelled.set()
|
||
raise
|
||
|
||
manager._heartbeat_stop = asyncio.Event()
|
||
manager._heartbeat_task = asyncio.create_task(stuck_heartbeat())
|
||
started = asyncio.get_running_loop().time()
|
||
|
||
await manager.shutdown(timeout=0.01)
|
||
|
||
elapsed = asyncio.get_running_loop().time() - started
|
||
assert heartbeat_cancelled.is_set()
|
||
assert elapsed < 0.5
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Lease heartbeat
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_heartbeat_renews_active_run_leases():
|
||
"""Heartbeat must extend the lease on active runs owned by this worker."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=config)
|
||
|
||
record = await manager.create_or_reject("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
|
||
original_lease = record.lease_expires_at
|
||
assert original_lease is not None
|
||
|
||
# Start heartbeat and let it tick once
|
||
await manager.start_heartbeat()
|
||
await asyncio.sleep(0.2) # heartbeat interval = 10s, too long; manually renew
|
||
|
||
await manager._renew_leases()
|
||
await manager.stop_heartbeat()
|
||
|
||
assert record.lease_expires_at is not None
|
||
# Lease should have been extended
|
||
assert record.lease_expires_at >= original_lease
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_heartbeat_renews_pending_run_before_task_is_spawned():
|
||
"""A run sitting in ``pending`` between ``create_thread_operation_atomic`` and task
|
||
spawn must still have its lease renewed.
|
||
|
||
Pre-fix the renewal filter required ``record.task is not None``, so a
|
||
pending run with no task yet (the brief window after
|
||
``create_thread_operation_atomic`` inserts the row before the worker layer spawns
|
||
the agent task) was silently skipped. If that window stretched past
|
||
``lease_seconds`` — e.g. event-loop saturation, slow checkpoint
|
||
hydrate — peer reconciliation reclaimed the run as an orphan and
|
||
marked it ``error`` even though this worker still intended to run it.
|
||
"""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=config)
|
||
|
||
record = await manager.create_or_reject("thread-1")
|
||
assert record.status == RunStatus.pending
|
||
# No task has been spawned — this is the regression sentinel.
|
||
assert record.task is None
|
||
|
||
original_lease = record.lease_expires_at
|
||
assert original_lease is not None
|
||
|
||
# Force a measurable gap so the renewed lease strictly post-dates the
|
||
# original — without this the two timestamps land in the same
|
||
# microsecond on fast hosts and the strict comparison fails trivially.
|
||
await asyncio.sleep(0.001)
|
||
|
||
store.update_lease = AsyncMock(wraps=store.update_lease)
|
||
|
||
await manager._renew_leases()
|
||
|
||
store.update_lease.assert_awaited_once()
|
||
assert record.lease_expires_at is not None
|
||
assert record.lease_expires_at > original_lease
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_transient_renewal_exception_before_deadline_keeps_run_alive():
|
||
"""A renewal error is retryable while the last confirmed lease is valid."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
original_update_lease = store.update_lease
|
||
attempts = 0
|
||
|
||
async def fail_once_then_renew(*args, **kwargs):
|
||
nonlocal attempts
|
||
attempts += 1
|
||
if attempts == 1:
|
||
raise OSError("temporary database outage")
|
||
return await original_update_lease(*args, **kwargs)
|
||
|
||
store.update_lease = fail_once_then_renew
|
||
# Keep the confirmed deadline inside the 30-second lease window. Renewal
|
||
# computes a fresh now+lease value; Windows' coarse clock can otherwise make
|
||
# two immediate samples identical and defeat the strict > below.
|
||
record.lease_expires_at = (datetime.now(UTC) + timedelta(seconds=10)).isoformat()
|
||
original_expiry = record.lease_expires_at
|
||
|
||
try:
|
||
await manager._renew_leases()
|
||
|
||
assert record.abort_event.is_set() is False
|
||
assert record.task.done() is False
|
||
assert record.lease_expires_at == original_expiry
|
||
|
||
await manager._renew_leases()
|
||
|
||
assert attempts == 2
|
||
assert record.abort_event.is_set() is False
|
||
assert record.task.done() is False
|
||
assert record.lease_expires_at is not None
|
||
assert original_expiry is not None
|
||
assert record.lease_expires_at > original_expiry
|
||
finally:
|
||
record.task.cancel()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_renewal_exception_through_confirmed_expiry_fail_stops_run():
|
||
"""The owner must stop once errors outlive its last confirmed lease."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
attempts = 0
|
||
|
||
async def fail_renewal(*_args, **_kwargs):
|
||
nonlocal attempts
|
||
attempts += 1
|
||
raise OSError("database unreachable")
|
||
|
||
store.update_lease = fail_renewal
|
||
|
||
await manager._renew_leases()
|
||
assert record.abort_event.is_set() is False
|
||
assert record.task.done() is False
|
||
|
||
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
|
||
record.lease_expires_at = expired
|
||
store._runs[record.run_id]["lease_expires_at"] = expired
|
||
|
||
await manager._renew_leases()
|
||
await asyncio.sleep(0)
|
||
|
||
assert attempts == 1
|
||
assert record.ownership_lost is True
|
||
assert record.abort_event.is_set() is True
|
||
assert record.task.cancelled()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_hung_renewal_is_bounded_by_confirmed_lease_deadline():
|
||
"""A blocked store call cannot keep execution alive beyond the lease."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
near_expiry = (datetime.now(UTC) + timedelta(milliseconds=50)).isoformat()
|
||
record.lease_expires_at = near_expiry
|
||
store._runs[record.run_id]["lease_expires_at"] = near_expiry
|
||
|
||
async def hang_renewal(*_args, **_kwargs):
|
||
await asyncio.Event().wait()
|
||
|
||
store.update_lease = hang_renewal
|
||
|
||
await asyncio.wait_for(manager._renew_leases(), timeout=1)
|
||
await asyncio.sleep(0)
|
||
|
||
assert record.ownership_lost is True
|
||
assert record.abort_event.is_set() is True
|
||
assert record.task.cancelled()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_late_successful_renewal_still_fences_local_run():
|
||
"""A renewal confirmed after the old deadline cannot revive local work."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
near_expiry = (datetime.now(UTC) + timedelta(milliseconds=50)).isoformat()
|
||
record.lease_expires_at = near_expiry
|
||
store._runs[record.run_id]["lease_expires_at"] = near_expiry
|
||
|
||
async def renew_after_timeout_cancellation(*args, **kwargs):
|
||
try:
|
||
await asyncio.sleep(3600)
|
||
except asyncio.CancelledError:
|
||
# Simulate a store operation that commits successfully despite the
|
||
# caller's deadline cancellation.
|
||
pass
|
||
# Deliberately bypass the built-in Store's expiry guard to exercise
|
||
# the Manager against an unsafe/legacy Store's late positive result.
|
||
store._runs[args[0]]["lease_expires_at"] = kwargs["lease_expires_at"]
|
||
return True
|
||
|
||
store.update_lease = renew_after_timeout_cancellation
|
||
|
||
await asyncio.wait_for(manager._renew_leases(), timeout=1)
|
||
await asyncio.sleep(0)
|
||
|
||
row = await store.get(record.run_id)
|
||
assert row is not None
|
||
assert row["lease_expires_at"] > near_expiry
|
||
assert record.lease_expires_at == near_expiry
|
||
assert record.ownership_lost is True
|
||
assert record.abort_event.is_set() is True
|
||
assert record.task.cancelled()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_heartbeat_skips_runs_not_owned_by_this_worker():
|
||
"""Heartbeat must only renew leases for runs owned by this worker."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=config)
|
||
|
||
# Create a run owned by a different worker
|
||
old_lease = (datetime.now(UTC) + timedelta(seconds=5)).isoformat()
|
||
await store.put(
|
||
"other-worker-run",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="other-worker",
|
||
lease_expires_at=old_lease,
|
||
created_at=(datetime.now(UTC) - timedelta(seconds=10)).isoformat(),
|
||
)
|
||
|
||
await manager._renew_leases()
|
||
|
||
stored = await store.get("other-worker-run")
|
||
# Lease should be unchanged (other worker's run)
|
||
assert stored["lease_expires_at"] == old_lease
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_staged_terminal_run_keeps_lease_through_worker_finalization():
|
||
"""A worker finalizing a staged success must keep its lease renewed.
|
||
|
||
With an event store the worker stages the terminal status locally and
|
||
commits it only after its receipt and duration writes. The durable row
|
||
stays active meanwhile, so an unrenewed lease would let a peer reclaim
|
||
the successful run as an orphan ``error``.
|
||
"""
|
||
config = _lease_config(lease_seconds=30, grace_seconds=0, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
|
||
record = await owner.create_or_reject("thread-finalize")
|
||
bridge = SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
|
||
finalizing = asyncio.Event()
|
||
release = asyncio.Event()
|
||
|
||
class SlowReceiptEventStore(MemoryRunEventStore):
|
||
async def put_if_absent(self, **event):
|
||
if event["event_type"] == "run.delivery":
|
||
finalizing.set()
|
||
await release.wait()
|
||
return await super().put_if_absent(**event)
|
||
|
||
class FinishingAgent:
|
||
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
||
del graph_input, config, stream_mode, subgraphs
|
||
yield {"messages": []}
|
||
|
||
task = asyncio.create_task(
|
||
run_agent(
|
||
bridge,
|
||
owner,
|
||
record,
|
||
ctx=RunContext(checkpointer=None, event_store=SlowReceiptEventStore()),
|
||
agent_factory=lambda **_kwargs: FinishingAgent(),
|
||
graph_input={},
|
||
config={},
|
||
)
|
||
)
|
||
record.task = task
|
||
|
||
try:
|
||
await asyncio.wait_for(finalizing.wait(), timeout=1)
|
||
assert record.status == RunStatus.success
|
||
assert (await store.get(record.run_id))["status"] == "running"
|
||
|
||
# Shrink the confirmed lease so it lapses during the test; only a
|
||
# heartbeat renewal can keep the peer from claiming the row.
|
||
near_expiry = (datetime.now(UTC) + timedelta(milliseconds=50)).isoformat()
|
||
record.lease_expires_at = near_expiry
|
||
store._runs[record.run_id]["lease_expires_at"] = near_expiry
|
||
await owner._renew_leases()
|
||
await asyncio.sleep(0.1)
|
||
|
||
assert await peer.reconcile_orphaned_inflight_runs(error="orphaned") == []
|
||
finally:
|
||
release.set()
|
||
await asyncio.wait_for(task, timeout=1)
|
||
|
||
assert record.ownership_lost is False
|
||
assert record.terminal_commit_pending is False
|
||
assert (await store.get(record.run_id))["status"] == "success"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_renewal_rejected_by_own_terminal_commit_releases_barrier():
|
||
"""A renewal racing this worker's terminal commit must not fence the run."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create_or_reject("thread-1")
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
record.terminal_commit_pending = True
|
||
await manager.set_status(record.run_id, RunStatus.success, persist=False)
|
||
# Only a Manager-owned operation supplies an attributable acknowledgment.
|
||
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
|
||
|
||
try:
|
||
await manager._renew_leases()
|
||
|
||
assert record.ownership_lost is False
|
||
assert record.terminal_commit_pending is False
|
||
assert record.status == RunStatus.success
|
||
finally:
|
||
record.task.cancel()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_renewal_rejected_by_peer_fences_staged_terminal_run():
|
||
"""A peer claim during finalization must fence the staged outcome."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create_or_reject("thread-1")
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
record.terminal_commit_pending = True
|
||
await manager.set_status(record.run_id, RunStatus.success, persist=False)
|
||
store._runs[record.run_id]["owner_worker_id"] = "worker-b"
|
||
|
||
await manager._renew_leases()
|
||
await asyncio.sleep(0)
|
||
|
||
assert record.ownership_lost is True
|
||
assert record.task.cancelled()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_fenced_staged_terminal_run_stops_renewing():
|
||
"""Fencing ends the terminal-commit barrier, so its lease can lapse."""
|
||
config = _lease_config(lease_seconds=30, heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
record = await manager.create_or_reject("thread-1")
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
record.terminal_commit_pending = True
|
||
await manager.set_status(record.run_id, RunStatus.success, persist=False)
|
||
await manager._mark_ownership_lost(record, reason="Test fence", require_active=False)
|
||
store.update_lease = AsyncMock(wraps=store.update_lease)
|
||
|
||
try:
|
||
await manager._renew_leases()
|
||
|
||
store.update_lease.assert_not_awaited()
|
||
assert record.terminal_commit_pending is False
|
||
finally:
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_heartbeat_not_started_when_disabled():
|
||
"""When heartbeat_enabled is False, start_heartbeat must be a no-op."""
|
||
config = _lease_config(heartbeat_enabled=False)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=config)
|
||
|
||
assert manager.heartbeat_enabled is False
|
||
await manager.start_heartbeat()
|
||
assert manager._heartbeat_task is None
|
||
assert manager._heartbeat_stop is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# cancel with cross-worker lease awareness
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_local_run_succeeds():
|
||
"""Cancel must succeed for a locally-owned active run."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
|
||
result = await manager.cancel(record.run_id)
|
||
assert result == CancelOutcome.cancelled
|
||
assert record.status == RunStatus.interrupted
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_unknown_run_returns_false():
|
||
"""Cancel must return not_active_locally for a run not known to this worker (heartbeat off)."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
|
||
result = await manager.cancel("nonexistent-run")
|
||
assert result == CancelOutcome.not_active_locally
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_idempotent():
|
||
"""Cancel must return cancelled when the run is already interrupted."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.interrupted)
|
||
|
||
result = await manager.cancel(record.run_id)
|
||
assert result == CancelOutcome.cancelled
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GATEWAY_WORKERS=1 backward compatibility
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_single_worker_default_config_behavior_unchanged():
|
||
"""With default config (heartbeat_enabled=False), behavior must match pre-ownership code."""
|
||
config = _lease_config(heartbeat_enabled=False)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=config)
|
||
|
||
# Create runs, cancel, create_or_reject — all must work
|
||
r1 = await manager.create("thread-1")
|
||
assert r1.owner_worker_id is not None
|
||
|
||
r2 = await manager.create_or_reject("thread-2", multitask_strategy="reject")
|
||
assert r2.owner_worker_id is not None
|
||
|
||
await manager.cancel(r2.run_id)
|
||
stored = await store.get(r2.run_id)
|
||
assert stored["status"] == "interrupted"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_manager_without_run_ownership_config():
|
||
"""Manager without run_ownership_config must still work (backward compat)."""
|
||
store = MemoryRunStore()
|
||
manager = RunManager(store=store) # no run_ownership_config
|
||
|
||
record = await manager.create_or_reject("thread-1")
|
||
assert record is not None
|
||
assert record.owner_worker_id is not None # always set, even without config
|
||
|
||
# Heartbeat must be a no-op without config
|
||
assert manager.heartbeat_enabled is False
|
||
await manager.start_heartbeat()
|
||
assert manager._heartbeat_task is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# worker_id uniqueness
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_worker_id_is_generated():
|
||
"""worker_id must be a non-empty string containing hostname."""
|
||
wid = _generate_worker_id()
|
||
assert isinstance(wid, str)
|
||
assert len(wid) > 0
|
||
assert ":" in wid
|
||
|
||
|
||
def test_two_managers_have_different_default_ids():
|
||
"""Two managers without explicit worker_id must get unique ids."""
|
||
m1 = RunManager()
|
||
m2 = RunManager()
|
||
assert m1.worker_id != m2.worker_id
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Store atomic methods
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_reject_prevents_duplicate():
|
||
"""Atomic thread-operation creation must reject a duplicate."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config()
|
||
|
||
store.create_thread_operation_atomic = AsyncMock(wraps=store.create_thread_operation_atomic)
|
||
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-1",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
with pytest.raises(ConflictError, match="already has an active run"):
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-2",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w2",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_interrupt_claims_and_creates():
|
||
"""Atomic thread-operation creation with interrupt must claim and replace."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config()
|
||
# Create an active run with an expired lease (simulating a crashed worker)
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-old",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=expired_lease,
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
new_row, claimed = await store.create_thread_operation_atomic(
|
||
run_id="run-new",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w2",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="interrupt",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
assert new_row["run_id"] == "run-new"
|
||
assert new_row["status"] == "pending"
|
||
assert len(claimed) == 1
|
||
assert claimed[0]["run_id"] == "run-old"
|
||
|
||
# Old run must be interrupted in-store
|
||
old_row = await store.get("run-old")
|
||
assert old_row["status"] == "interrupted"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_uses_one_change_position():
|
||
"""One atomic thread operation advances the change cursor once.
|
||
|
||
``runtime/AGENTS.md`` documents that an atomic thread operation uses one
|
||
position for its interrupted rows and its new row, with ``run_id``
|
||
ordering ties, and ``test_run_evidence_reader`` pins that for the SQL
|
||
store. The memory store must agree: it is the backend an install runs
|
||
whenever no durable database is configured. Advancing once per row splits
|
||
a single interrupt-and-replace into two positions in the
|
||
``(change_seq, run_id)`` cursor extensions page with.
|
||
"""
|
||
store = MemoryRunStore()
|
||
config = _lease_config()
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-old",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=expired_lease,
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
new_row, claimed = await store.create_thread_operation_atomic(
|
||
run_id="run-new",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w2",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="interrupt",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
assert [row["run_id"] for row in claimed] == ["run-old"]
|
||
assert claimed[0]["change_seq"] == new_row["change_seq"]
|
||
|
||
# The shared position is what a paging consumer observes: the interrupted
|
||
# row and its replacement surface under one cursor value, ordered by
|
||
# ``run_id``, exactly as the SQL store reports them.
|
||
changed = await store.list_changed(after_change_seq=-1, after_run_id="")
|
||
positions = {row["run_id"]: row["change_seq"] for row in changed}
|
||
assert positions["run-old"] == positions["run-new"]
|
||
assert [row["run_id"] for row in changed if row["change_seq"] == positions["run-new"]] == ["run-new", "run-old"]
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_rejection_consumes_no_change_position():
|
||
"""A ``ConflictError``-rejected operation advances no change position.
|
||
|
||
``create_thread_operation_atomic`` allocates its position only after the
|
||
raise-only candidate scan, so a rejected operation consumes nothing — the
|
||
memory-store counterpart of the SQL store's rollback leaving the clock
|
||
untouched. A consumed-but-unused position leaves no trace in the rows
|
||
themselves, so the assertion is on the position the next accepted
|
||
operation lands on. Hoisting the allocation above the scan keeps every
|
||
other test in this file green and shows up here as a gap.
|
||
"""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(grace_seconds=10)
|
||
|
||
accepted, _ = await store.create_thread_operation_atomic(
|
||
run_id="valid-lease-run",
|
||
thread_id="thread-1",
|
||
owner_worker_id="other-worker",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
with pytest.raises(ConflictError, match="another worker"):
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-new",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w2",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="interrupt",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
next_row, _ = await store.create_thread_operation_atomic(
|
||
run_id="run-after",
|
||
thread_id="thread-2",
|
||
owner_worker_id="w2",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
assert next_row["change_seq"] == accepted["change_seq"] + 1
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_interrupt_rejects_other_worker_valid_lease():
|
||
"""Interrupt must raise ConflictError when a valid-lease run is owned by another worker.
|
||
|
||
The partial unique index ``uq_runs_thread_active`` would reject the INSERT
|
||
anyway; surfacing ConflictError here gives the caller a clean signal
|
||
instead of a futile retry loop on IntegrityError.
|
||
"""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(grace_seconds=10)
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
|
||
|
||
await store.create_thread_operation_atomic(
|
||
run_id="valid-lease-run",
|
||
thread_id="thread-1",
|
||
owner_worker_id="other-worker",
|
||
lease_expires_at=valid_lease,
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
with pytest.raises(ConflictError, match="another worker"):
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-new",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w2",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="interrupt",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
# The valid-lease run must be untouched (transaction rolled back).
|
||
old_row = await store.get("valid-lease-run")
|
||
assert old_row["status"] == "pending"
|
||
assert old_row["owner_worker_id"] == "other-worker"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_interrupt_allows_self_owned_valid_lease():
|
||
"""Interrupt must succeed when the existing valid-lease run is owned by this worker."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(grace_seconds=10)
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
|
||
|
||
await store.create_thread_operation_atomic(
|
||
run_id="self-run",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=valid_lease,
|
||
multitask_strategy="reject",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
new_row, claimed = await store.create_thread_operation_atomic(
|
||
run_id="run-new",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w1", # same worker
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="interrupt",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
assert new_row["run_id"] == "run-new"
|
||
assert len(claimed) == 1
|
||
assert claimed[0]["run_id"] == "self-run"
|
||
assert claimed[0]["status"] == "interrupted"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_create_thread_operation_atomic_interrupt_rolls_back_earlier_mutations_on_conflict():
|
||
"""Interrupt must not leave earlier candidates interrupted when a later
|
||
candidate raises ConflictError.
|
||
|
||
Mirrors the SQL store's transactional semantics: the whole interrupt pass
|
||
is one transaction, so a raise on any candidate must roll back mutations
|
||
already applied to earlier candidates. Without this, the memory store
|
||
diverges from SQL (which the production path uses), and the
|
||
test_multi_worker_run_ownership.py suite gives false confidence by
|
||
passing against memory while SQL would behave differently.
|
||
|
||
Setup: expired-lease run (interruptible) inserted FIRST, then a
|
||
valid-lease run owned by another worker. Iteration order means the
|
||
expired run is mutated before the valid-lease run raises — so a naive
|
||
single-pass implementation would leave the expired run interrupted.
|
||
"""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(grace_seconds=10)
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
|
||
|
||
# Seed both active rows directly via ``put`` (bypassing atomic operation creation's
|
||
# reject check, which would refuse the second row). Insert the
|
||
# interruptible run first so dict iteration visits it first — that's the
|
||
# ordering that exposes the half-interrupted divergence in a naive
|
||
# single-pass implementation.
|
||
await store.put(
|
||
"expired-run",
|
||
thread_id="thread-1",
|
||
status="pending",
|
||
owner_worker_id="old-worker",
|
||
lease_expires_at=expired_lease,
|
||
)
|
||
await store.put(
|
||
"valid-lease-run",
|
||
thread_id="thread-1",
|
||
status="pending",
|
||
owner_worker_id="other-worker",
|
||
lease_expires_at=valid_lease,
|
||
)
|
||
|
||
with pytest.raises(ConflictError, match="another worker"):
|
||
await store.create_thread_operation_atomic(
|
||
run_id="run-new",
|
||
thread_id="thread-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(),
|
||
multitask_strategy="interrupt",
|
||
grace_seconds=config.grace_seconds,
|
||
)
|
||
|
||
# The expired run must be UNTOUCHED — the interrupt pass must roll back
|
||
# on ConflictError, not leave a half-interrupted store.
|
||
expired_row = await store.get("expired-run")
|
||
assert expired_row["status"] == "pending"
|
||
assert expired_row["owner_worker_id"] == "old-worker"
|
||
assert expired_row["error"] is None
|
||
|
||
# The valid-lease run that caused the conflict is also untouched.
|
||
valid_row = await store.get("valid-lease-run")
|
||
assert valid_row["status"] == "pending"
|
||
assert valid_row["owner_worker_id"] == "other-worker"
|
||
|
||
# The new run was never inserted.
|
||
assert await store.get("run-new") is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# update_lease
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_update_lease_renews_row():
|
||
"""update_lease must update the lease_expires_at on the stored row."""
|
||
store = MemoryRunStore()
|
||
old_lease = (datetime.now(UTC) + timedelta(seconds=5)).isoformat()
|
||
await store.put(
|
||
"run-1",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=old_lease,
|
||
)
|
||
|
||
new_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
|
||
updated = await store.update_lease(
|
||
"run-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=new_lease,
|
||
)
|
||
assert updated is True
|
||
|
||
stored = await store.get("run-1")
|
||
assert stored["lease_expires_at"] == new_lease
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_update_lease_returns_false_for_terminal_run():
|
||
"""update_lease must return False when the run is not pending/running."""
|
||
store = MemoryRunStore()
|
||
await store.put("run-1", thread_id="thread-1", status="success", owner_worker_id="w1")
|
||
|
||
new_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
|
||
updated = await store.update_lease(
|
||
"run-1",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=new_lease,
|
||
)
|
||
assert updated is False
|
||
|
||
stored = await store.get("run-1")
|
||
assert stored["status"] == "success"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_update_lease_returns_false_for_wrong_owner():
|
||
"""update_lease must reject renewal when owner_worker_id does not match."""
|
||
store = MemoryRunStore()
|
||
old_lease = (datetime.now(UTC) + timedelta(seconds=5)).isoformat()
|
||
await store.put(
|
||
"run-1",
|
||
thread_id="thread-1",
|
||
status="running",
|
||
owner_worker_id="w1",
|
||
lease_expires_at=old_lease,
|
||
)
|
||
|
||
new_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat()
|
||
updated = await store.update_lease(
|
||
"run-1",
|
||
owner_worker_id="w2", # different worker
|
||
lease_expires_at=new_lease,
|
||
)
|
||
assert updated is False
|
||
|
||
# The original lease must be untouched
|
||
stored = await store.get("run-1")
|
||
assert stored["owner_worker_id"] == "w1"
|
||
assert stored["lease_expires_at"] == old_lease
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# list_inflight_with_expired_lease
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_list_inflight_with_expired_lease_filters_correctly():
|
||
"""Only runs with expired or NULL leases must be returned."""
|
||
store = MemoryRunStore()
|
||
now = datetime.now(UTC)
|
||
grace = 10
|
||
|
||
# Expired lease
|
||
expired = (now - timedelta(seconds=60)).isoformat()
|
||
await store.put("expired-run", thread_id="t1", status="running", owner_worker_id="w1", lease_expires_at=expired, created_at=expired)
|
||
|
||
# Valid lease
|
||
valid = (now + timedelta(seconds=60)).isoformat()
|
||
await store.put("valid-run", thread_id="t2", status="running", owner_worker_id="w2", lease_expires_at=valid, created_at=valid)
|
||
|
||
# NULL lease (legacy)
|
||
await store.put("null-lease-run", thread_id="t3", status="running", created_at=(now - timedelta(seconds=30)).isoformat())
|
||
|
||
# Terminal status (should not appear)
|
||
await store.put("success-run", thread_id="t4", status="success", created_at=(now - timedelta(seconds=60)).isoformat())
|
||
|
||
results = await store.list_inflight_with_expired_lease(grace_seconds=grace)
|
||
|
||
result_ids = {r["run_id"] for r in results}
|
||
assert "expired-run" in result_ids
|
||
assert "null-lease-run" in result_ids
|
||
assert "valid-run" not in result_ids
|
||
assert "success-run" not in result_ids
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# MemoryRunStore — datetime comparison for created_at filtering
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_list_inflight_with_expired_lease_compares_created_at_as_datetime():
|
||
"""``before`` filter must use datetime comparison, not string lexical order.
|
||
|
||
ISO-8601 strings compare lexically only when every component is zero-padded
|
||
to the same width and the timezone suffix matches. Datetime parsing is
|
||
order-safe regardless of format.
|
||
"""
|
||
store = MemoryRunStore()
|
||
now = datetime.now(UTC)
|
||
grace = 10
|
||
|
||
# A run created "now" — should be included when before=None (defaults to now).
|
||
await store.put("recent-run", thread_id="t1", status="running", created_at=now.isoformat())
|
||
# A run created far in the future — should be excluded by the before filter
|
||
# even though the string "2300-01-01..." > "2025-..." lexically.
|
||
far_future = "2300-01-01T00:00:00+00:00"
|
||
await store.put("future-run", thread_id="t2", status="running", created_at=far_future)
|
||
|
||
results = await store.list_inflight_with_expired_lease(before=now.isoformat(), grace_seconds=grace)
|
||
result_ids = {r["run_id"] for r in results}
|
||
assert "recent-run" in result_ids
|
||
assert "future-run" not in result_ids
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_list_inflight_with_expired_lease_handles_malformed_created_at():
|
||
"""Malformed ``created_at`` values must not crash the listing."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
|
||
store._runs["bad-run"] = {
|
||
"run_id": "bad-run",
|
||
"thread_id": "t1",
|
||
"status": "running",
|
||
"created_at": "not-a-datetime",
|
||
}
|
||
store._runs["empty-run"] = {
|
||
"run_id": "empty-run",
|
||
"thread_id": "t2",
|
||
"status": "running",
|
||
"created_at": "",
|
||
}
|
||
|
||
results = await store.list_inflight_with_expired_lease(grace_seconds=grace)
|
||
# Both should be skipped because their created_at can't be parsed
|
||
result_ids = {r["run_id"] for r in results}
|
||
assert "bad-run" not in result_ids
|
||
assert "empty-run" not in result_ids
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_list_inflight_with_expired_lease_datetime_aware_naive_handling():
|
||
"""Lease comparison must handle aware and naive datetimes.
|
||
|
||
``lease_expires_at`` stored with a trailing ``+00:00`` (aware) and without
|
||
(naive) should both be comparable against the aware ``cutoff``. The MemoryRunStore
|
||
uses ``datetime.fromisoformat`` which preserves the offset, so both paths
|
||
must work.
|
||
"""
|
||
store = MemoryRunStore()
|
||
now = datetime.now(UTC)
|
||
grace = 10
|
||
|
||
# Naive datetime (no timezone suffix) — common on SQLite read-back
|
||
naive_expired = (now - timedelta(seconds=60)).isoformat() # "2025-01-01T00:00:00"
|
||
await store.put("naive-run", thread_id="t1", status="running", lease_expires_at=naive_expired, created_at=naive_expired)
|
||
|
||
# Aware datetime (with +00:00)
|
||
aware_expired = (now - timedelta(seconds=60)).replace(tzinfo=UTC).isoformat() # "2025-01-01T00:00:00+00:00"
|
||
await store.put("aware-run", thread_id="t2", status="running", lease_expires_at=aware_expired, created_at=aware_expired)
|
||
|
||
results = await store.list_inflight_with_expired_lease(grace_seconds=grace)
|
||
result_ids = {r["run_id"] for r in results}
|
||
# Both expired, both should be returned
|
||
assert "naive-run" in result_ids
|
||
assert "aware-run" in result_ids
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_list_inflight_with_expired_lease_null_lease_always_reclaimed():
|
||
"""NULL lease rows are always reclaimed regardless of created_at value."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
|
||
# NULL lease is the single-worker mode default — every inflight row
|
||
# must be returned so reconciliation can reclaim it.
|
||
await store.put("null-run", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat())
|
||
|
||
results = await store.list_inflight_with_expired_lease(grace_seconds=grace)
|
||
result_ids = {r["run_id"] for r in results}
|
||
assert "null-run" in result_ids
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# claim_for_takeover — store primitive
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_claim_for_takeover_succeeds_with_expired_lease():
|
||
"""claim_for_takeover must succeed when the lease has passed the grace window."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
|
||
await store.put("run-1", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=expired_lease)
|
||
|
||
ok = await store.claim_for_takeover(
|
||
"run-1",
|
||
grace_seconds=grace,
|
||
error="claimed",
|
||
stop_reason=ORPHAN_RECOVERY_STOP_REASON,
|
||
)
|
||
assert ok is True
|
||
|
||
row = await store.get("run-1")
|
||
assert row is not None
|
||
assert row["status"] == "error"
|
||
assert row["error"] == "claimed"
|
||
assert row["stop_reason"] == ORPHAN_RECOVERY_STOP_REASON
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_claim_for_takeover_fails_with_valid_lease():
|
||
"""claim_for_takeover must return False when the lease is still valid."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
await store.put("run-1", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=valid_lease)
|
||
|
||
ok = await store.claim_for_takeover("run-1", grace_seconds=grace, error="claimed")
|
||
assert ok is False
|
||
|
||
row = await store.get("run-1")
|
||
assert row is not None
|
||
assert row["status"] == "running"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_claim_for_takeover_succeeds_with_null_lease():
|
||
"""NULL-lease rows (pre-ownership data) must be claimable."""
|
||
store = MemoryRunStore()
|
||
await store.put("run-null", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat())
|
||
|
||
ok = await store.claim_for_takeover("run-null", grace_seconds=10, error="claimed")
|
||
assert ok is True
|
||
|
||
row = await store.get("run-null")
|
||
assert row["status"] == "error"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_claim_for_takeover_fails_on_terminal_status():
|
||
"""claim_for_takeover must return False for already-terminal runs."""
|
||
store = MemoryRunStore()
|
||
await store.put("run-done", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat())
|
||
|
||
ok = await store.claim_for_takeover("run-done", grace_seconds=10, error="claimed")
|
||
assert ok is False
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_claim_for_takeover_fails_for_nonexistent_run():
|
||
"""claim_for_takeover must return False when the run doesn't exist."""
|
||
store = MemoryRunStore()
|
||
ok = await store.claim_for_takeover("no-such-run", grace_seconds=10, error="claimed")
|
||
assert ok is False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# cancel() cross-worker takeover — work item 4
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_takeover_from_crashed_worker():
|
||
"""cancel must take over (mark error) when lease is expired and owner is another worker."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
|
||
await store.put("run-expired", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="dead-worker", lease_expires_at=expired_lease)
|
||
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
outcome = await manager.cancel("run-expired")
|
||
assert outcome == CancelOutcome.taken_over
|
||
|
||
row = await store.get("run-expired")
|
||
assert row is not None
|
||
assert row["status"] == "error"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_takeover_marks_orphan_recovery_for_waiting_consumers():
|
||
"""A crash takeover must carry the orphan-recovery stop reason.
|
||
|
||
The owner is presumed dead, so a cross-worker SSE/``wait`` consumer's only
|
||
liveness edge is the durable row: without the marker its heartbeat fallback
|
||
never fires and it waits forever for an END that no producer can publish.
|
||
"""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
|
||
await store.put("run-takeover-liveness", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="dead-worker", lease_expires_at=expired_lease)
|
||
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
outcome = await manager.cancel("run-takeover-liveness")
|
||
assert outcome == CancelOutcome.taken_over
|
||
|
||
row = await store.get("run-takeover-liveness")
|
||
assert row is not None
|
||
assert row["status"] == "error"
|
||
assert row["stop_reason"] == ORPHAN_RECOVERY_STOP_REASON
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_requests_active_lease_from_other_worker():
|
||
"""A cancel routed to a peer must durably notify the live owner."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
await store.put("run-alive", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="alive-worker", lease_expires_at=valid_lease)
|
||
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
outcome = await manager.cancel("run-alive")
|
||
assert outcome == CancelOutcome.requested
|
||
|
||
row = await store.get("run-alive")
|
||
assert row is not None
|
||
assert row["status"] == "running"
|
||
assert row["cancel_action"] == "interrupt"
|
||
assert row["cancel_requested_at"] is not None
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_non_owner_cancel_is_observed_by_owner_heartbeat():
|
||
"""Heartbeat signals the owner task without performing terminal writes."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=True, lease_seconds=30)
|
||
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
|
||
|
||
record = await owner.create_or_reject("thread-1")
|
||
await owner.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
|
||
try:
|
||
assert await peer.cancel(record.run_id, action="rollback") == CancelOutcome.requested
|
||
# Match local idempotency: once accepted, a later action cannot change
|
||
# whether the owner rolls back or merely interrupts.
|
||
assert await peer.cancel(record.run_id, action="interrupt") == CancelOutcome.requested
|
||
|
||
await owner._renew_leases()
|
||
await asyncio.sleep(0)
|
||
|
||
assert record.abort_event.is_set()
|
||
assert record.abort_action == "rollback"
|
||
assert record.status == RunStatus.running
|
||
assert record.task.cancelled()
|
||
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "running"
|
||
assert stored["cancel_action"] == "rollback"
|
||
|
||
# The worker's existing cancellation path, not heartbeat, owns the
|
||
# terminal status write and any rollback cleanup.
|
||
await owner.set_status(
|
||
record.run_id,
|
||
RunStatus.error,
|
||
error="Rolled back by user",
|
||
)
|
||
next_run = await peer.create_or_reject("thread-1")
|
||
assert next_run.status == RunStatus.pending
|
||
finally:
|
||
if not record.task.done():
|
||
record.task.cancel()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_first_cancel_action_wins_when_retry_lands_on_owner():
|
||
"""Routing a retry to the owner must not replace the durable first action."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=True, lease_seconds=30)
|
||
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
|
||
|
||
record = await owner.create_or_reject("thread-1")
|
||
await owner.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
|
||
try:
|
||
assert await peer.cancel(record.run_id, action="rollback") == CancelOutcome.requested
|
||
assert await owner.cancel(record.run_id, action="interrupt") == CancelOutcome.cancelled
|
||
await asyncio.sleep(0)
|
||
|
||
assert record.abort_action == "rollback"
|
||
assert record.task.cancelled()
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["cancel_action"] == "rollback"
|
||
finally:
|
||
if not record.task.done():
|
||
record.task.cancel()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_local_owner_cancel_falls_back_when_durable_request_fails():
|
||
"""A local owner can still abort its own task when durable cancel persistence fails."""
|
||
|
||
class FailingCancelStore(MemoryRunStore):
|
||
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
|
||
raise RuntimeError("store unavailable")
|
||
|
||
store = FailingCancelStore()
|
||
manager = _make_manager(
|
||
store=store,
|
||
worker_id="worker-a",
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=30),
|
||
)
|
||
|
||
record = await manager.create_or_reject("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
|
||
try:
|
||
assert await manager.cancel(record.run_id, action="rollback") == CancelOutcome.cancelled
|
||
assert record.abort_event.is_set()
|
||
assert record.abort_action == "rollback"
|
||
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "interrupted"
|
||
assert stored["cancel_action"] is None
|
||
finally:
|
||
if not record.task.done():
|
||
record.task.cancel()
|
||
with pytest.raises(asyncio.CancelledError):
|
||
await record.task
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_local_cancel_does_not_claim_unpersisted_terminal_status():
|
||
class FailingStatusStore(MemoryRunStore):
|
||
async def update_status(self, run_id, status, **kwargs):
|
||
if status == RunStatus.interrupted.value:
|
||
raise OSError("status store unavailable")
|
||
return await super().update_status(run_id, status, **kwargs)
|
||
|
||
store = FailingStatusStore()
|
||
manager = _make_manager(
|
||
store=store,
|
||
worker_id="worker-a",
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=5),
|
||
)
|
||
record = await manager.create_or_reject("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
record.task = asyncio.create_task(asyncio.sleep(3600))
|
||
|
||
try:
|
||
assert await manager.cancel(record.run_id) == CancelOutcome.cancelled
|
||
assert record.ownership_lost
|
||
assert record.status == RunStatus.error
|
||
assert record.terminal_commit_proof is None
|
||
assert record.terminal_status_persisted is False
|
||
assert (await store.get(record.run_id))["status"] == RunStatus.running.value
|
||
finally:
|
||
if not record.task.done():
|
||
record.task.cancel()
|
||
await asyncio.gather(record.task, return_exceptions=True)
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_owner_cancel_retry_on_peer_is_accepted():
|
||
"""A peer must treat the owner's interrupted status as an accepted cancel."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=True, lease_seconds=30)
|
||
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
|
||
|
||
record = await owner.create_or_reject("thread-1")
|
||
await owner.set_status(record.run_id, RunStatus.running)
|
||
|
||
assert await owner.cancel(record.run_id, action="rollback") == CancelOutcome.cancelled
|
||
assert await peer.cancel(record.run_id, action="interrupt") == CancelOutcome.requested
|
||
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "interrupted"
|
||
assert stored["cancel_action"] == "rollback"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_owner_cancel_uses_store_while_terminal_status_is_staged_locally():
|
||
"""A staged local terminal status must not bypass the durable cancel CAS."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
record = await manager.create_or_reject("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
|
||
# Event-store finalization stages success in memory before persisting it.
|
||
record.status = RunStatus.success
|
||
|
||
assert await manager.cancel(record.run_id, action="rollback") == CancelOutcome.cancelled
|
||
stored = await store.get(record.run_id)
|
||
assert stored is not None
|
||
assert stored["status"] == "running"
|
||
assert stored["cancel_action"] == "rollback"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_returns_unknown_when_no_store():
|
||
"""cancel must return unknown when there's no store and the run is not in memory."""
|
||
manager = _make_manager(run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
outcome = await manager.cancel("no-such-run")
|
||
assert outcome == CancelOutcome.unknown
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_returns_not_active_locally_when_heartbeat_disabled():
|
||
"""With heartbeat disabled, store-only runs must not be cancellable (old 409 path)."""
|
||
store = MemoryRunStore()
|
||
await store.put("store-only", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat())
|
||
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=False))
|
||
outcome = await manager.cancel("store-only")
|
||
assert outcome == CancelOutcome.not_active_locally
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_takeover_race_owner_renewed_lease():
|
||
"""When takeover loses to a renewal, cancellation must notify the live owner."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
|
||
await store.put("run-race", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=expired_lease)
|
||
|
||
# Simulate the race: right before claim_for_takeover writes, another
|
||
# heartbeat renews the lease. We monkey-patch claim_for_takeover to
|
||
# simulate the lease having been renewed.
|
||
original = store.claim_for_takeover
|
||
|
||
async def race_lost(run_id, *, grace_seconds, error, stop_reason=None):
|
||
# Simulate a heartbeat renewal between the read and the write
|
||
run = store._runs.get(run_id)
|
||
if run and run["status"] in ("pending", "running"):
|
||
run["lease_expires_at"] = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
return await original(run_id, grace_seconds=grace_seconds, error=error, stop_reason=stop_reason)
|
||
|
||
store.claim_for_takeover = race_lost
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
|
||
outcome = await manager.cancel("run-race")
|
||
assert outcome == CancelOutcome.requested
|
||
assert (await store.get("run-race"))["cancel_action"] == "interrupt"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_takeover_respects_grace_seconds():
|
||
"""Within the grace window, cancellation must notify rather than take over."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
# Lease expired, but only by 3s — still within the 10s grace window
|
||
just_expired = (datetime.now(UTC) - timedelta(seconds=3)).isoformat()
|
||
await store.put("run-grace", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat(), owner_worker_id="w-a", lease_expires_at=just_expired)
|
||
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
outcome = await manager.cancel("run-grace")
|
||
assert outcome == CancelOutcome.requested
|
||
assert (await store.get("run-grace"))["cancel_action"] == "interrupt"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_not_cancellable_for_store_terminal_run():
|
||
"""cancel must return not_cancellable when the store run is already in a terminal state."""
|
||
store = MemoryRunStore()
|
||
await store.put("run-done", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat())
|
||
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
outcome = await manager.cancel("run-done")
|
||
assert outcome == CancelOutcome.not_cancellable
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# HTTP-level — cancel endpoint cross-worker responses
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class _EndingCrossProcessBridge:
|
||
supports_cross_process = True
|
||
|
||
async def publish(self, run_id, event, data):
|
||
return None
|
||
|
||
async def publish_end(self, run_id):
|
||
return None
|
||
|
||
def subscribe(self, run_id, *, last_event_id=None, heartbeat_interval=15.0):
|
||
from deerflow.runtime import END_SENTINEL
|
||
|
||
async def events():
|
||
yield END_SENTINEL
|
||
|
||
return events()
|
||
|
||
async def cleanup(self, run_id, *, delay=0):
|
||
return None
|
||
|
||
|
||
def _make_cancel_test_app(mgr: RunManager, *, bridge=None):
|
||
"""Build a TestClient wired with the thread_runs router + memory bridge."""
|
||
from _router_auth_helpers import make_authed_test_app
|
||
from fastapi.testclient import TestClient
|
||
|
||
from app.gateway.routers import thread_runs
|
||
from deerflow.runtime import MemoryStreamBridge
|
||
|
||
app = make_authed_test_app()
|
||
app.include_router(thread_runs.router)
|
||
app.state.run_manager = mgr
|
||
app.state.stream_bridge = bridge or MemoryStreamBridge()
|
||
return TestClient(app, raise_server_exceptions=False)
|
||
|
||
|
||
def test_http_cancel_non_owner_valid_lease_returns_202():
|
||
"""POST /cancel must not fail solely because routing chose a non-owner."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
asyncio.run(
|
||
store.put(
|
||
"run-alive",
|
||
thread_id="t1",
|
||
status="running",
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
owner_worker_id="alive-worker",
|
||
lease_expires_at=valid_lease,
|
||
)
|
||
)
|
||
mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
client = _make_cancel_test_app(mgr)
|
||
|
||
resp = client.post("/api/threads/t1/runs/run-alive/cancel")
|
||
assert resp.status_code == 202
|
||
assert "Retry-After" not in resp.headers
|
||
|
||
# The owner remains fenced, while the cancellation request is durable.
|
||
row = asyncio.run(store.get("run-alive"))
|
||
assert row["status"] == "running"
|
||
assert row["cancel_action"] == "interrupt"
|
||
|
||
|
||
def test_http_stream_action_non_owner_without_shared_bridge_returns_202():
|
||
"""A peer cancel is accepted without subscribing to an unreachable local stream."""
|
||
store = MemoryRunStore()
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
asyncio.run(
|
||
store.put(
|
||
"run-alive-stream",
|
||
thread_id="t1",
|
||
status="running",
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
owner_worker_id="alive-worker",
|
||
lease_expires_at=valid_lease,
|
||
)
|
||
)
|
||
mgr = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
client = _make_cancel_test_app(mgr)
|
||
|
||
resp = client.post(
|
||
"/api/threads/t1/runs/run-alive-stream/stream",
|
||
params={"action": "interrupt"},
|
||
)
|
||
|
||
assert resp.status_code == 202
|
||
row = asyncio.run(store.get("run-alive-stream"))
|
||
assert row["status"] == "running"
|
||
assert row["cancel_action"] == "interrupt"
|
||
|
||
|
||
def test_http_cancel_non_owner_wait_uses_shared_bridge():
|
||
"""wait=true observes remote owner finalization through the shared bridge."""
|
||
store = MemoryRunStore()
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
asyncio.run(
|
||
store.put(
|
||
"run-alive-wait",
|
||
thread_id="t1",
|
||
status="running",
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
owner_worker_id="alive-worker",
|
||
lease_expires_at=valid_lease,
|
||
)
|
||
)
|
||
mgr = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
client = _make_cancel_test_app(mgr, bridge=_EndingCrossProcessBridge())
|
||
|
||
resp = client.post(
|
||
"/api/threads/t1/runs/run-alive-wait/cancel",
|
||
params={"action": "rollback", "wait": "true"},
|
||
)
|
||
|
||
assert resp.status_code == 204
|
||
row = asyncio.run(store.get("run-alive-wait"))
|
||
assert row["cancel_action"] == "rollback"
|
||
|
||
|
||
def test_http_stream_action_non_owner_uses_shared_bridge():
|
||
"""The SDK stop path drains the remote owner's shared stream after accept."""
|
||
store = MemoryRunStore()
|
||
valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat()
|
||
asyncio.run(
|
||
store.put(
|
||
"run-alive-shared-stream",
|
||
thread_id="t1",
|
||
status="running",
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
owner_worker_id="alive-worker",
|
||
lease_expires_at=valid_lease,
|
||
)
|
||
)
|
||
mgr = _make_manager(
|
||
store=store,
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
client = _make_cancel_test_app(mgr, bridge=_EndingCrossProcessBridge())
|
||
|
||
resp = client.post(
|
||
"/api/threads/t1/runs/run-alive-shared-stream/stream",
|
||
params={"action": "interrupt"},
|
||
)
|
||
|
||
assert resp.status_code == 200
|
||
assert "event: end" in resp.text
|
||
row = asyncio.run(store.get("run-alive-shared-stream"))
|
||
assert row["cancel_action"] == "interrupt"
|
||
|
||
|
||
def test_http_cancel_non_owner_expired_lease_returns_202_takeover():
|
||
"""POST /cancel on a non-owning worker with an expired lease must return 202 (takeover)."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 30)).isoformat()
|
||
asyncio.run(
|
||
store.put(
|
||
"run-dead",
|
||
thread_id="t1",
|
||
status="running",
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
)
|
||
)
|
||
mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
client = _make_cancel_test_app(mgr)
|
||
|
||
resp = client.post("/api/threads/t1/runs/run-dead/cancel")
|
||
assert resp.status_code == 202
|
||
|
||
# Store row must be marked error
|
||
row = asyncio.run(store.get("run-dead"))
|
||
assert row["status"] == "error"
|
||
|
||
|
||
def test_http_stream_action_interrupt_takeover_returns_202_not_hang():
|
||
"""POST /stream?action=interrupt on a dead-owner run must return 202 immediately, not hang on SSE."""
|
||
store = MemoryRunStore()
|
||
grace = 10
|
||
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 30)).isoformat()
|
||
asyncio.run(
|
||
store.put(
|
||
"run-dead-stream",
|
||
thread_id="t1",
|
||
status="running",
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
owner_worker_id="dead-worker",
|
||
lease_expires_at=expired_lease,
|
||
)
|
||
)
|
||
mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=grace))
|
||
client = _make_cancel_test_app(mgr)
|
||
|
||
# This must NOT hang — the takeover path returns 202 before reaching StreamingResponse.
|
||
resp = client.post("/api/threads/t1/runs/run-dead-stream/stream", params={"action": "interrupt"})
|
||
assert resp.status_code == 202
|
||
|
||
row = asyncio.run(store.get("run-dead-stream"))
|
||
assert row["status"] == "error"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Split-brain defences — update_status guard + heartbeat self-termination
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_update_status_rejects_terminal_row():
|
||
"""update_status must return False when the store row is already terminal
|
||
(error/success), so a late writer cannot overwrite a peer's takeover or
|
||
a completed run. interrupted is NOT terminal — the rollback path needs
|
||
``interrupted → error`` to finalize."""
|
||
store = MemoryRunStore()
|
||
# error (takeover) must stay locked
|
||
await store.put("run-err", thread_id="t1", status="error", created_at=datetime.now(UTC).isoformat())
|
||
assert await store.update_status("run-err", "success") is False
|
||
assert (await store.get("run-err"))["status"] == "error"
|
||
|
||
# success must stay locked
|
||
await store.put("run-ok", thread_id="t1", status="success", created_at=datetime.now(UTC).isoformat())
|
||
assert await store.update_status("run-ok", "error") is False
|
||
assert (await store.get("run-ok"))["status"] == "success"
|
||
|
||
# interrupted → error MUST pass (rollback finalize path)
|
||
await store.put("run-rb", thread_id="t1", status="interrupted", created_at=datetime.now(UTC).isoformat())
|
||
assert await store.update_status("run-rb", "error", error="Rolled back by user") is True
|
||
row = await store.get("run-rb")
|
||
assert row["status"] == "error"
|
||
assert row["error"] == "Rolled back by user"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_persist_status_skips_recovery_when_row_taken_over():
|
||
"""_persist_status must not recreate a row that was taken over by another worker.
|
||
|
||
When update_status returns False, the recovery path checks whether the
|
||
row still exists. A row that exists but is terminal (taken over) must
|
||
be left alone — calling put() would overwrite the takeover."""
|
||
store = MemoryRunStore()
|
||
mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
|
||
# Simulate: this worker created and started a run, but a peer took it over.
|
||
record = await mgr.create("thread-1")
|
||
await mgr.set_status(record.run_id, RunStatus.running)
|
||
# Peer takeover: directly flip the store row to error
|
||
await store.update_status(record.run_id, "error")
|
||
# Now simulate the original owner's task finishing and trying to write success
|
||
ok = await mgr._persist_status(record, RunStatus.success)
|
||
assert ok is False # skipped recovery, row already exists and is terminal
|
||
row = await store.get(record.run_id)
|
||
assert row["status"] == "error" # not overwritten
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_heartbeat_cancels_task_on_lease_loss():
|
||
"""Heartbeat must cancel the local asyncio task when update_lease returns False.
|
||
|
||
If the store row was claimed by another worker (status no longer
|
||
pending/running, or owner changed), the heartbeat tick must abort the
|
||
local task so wasted CPU is bounded to ~10s instead of the full task
|
||
lifetime."""
|
||
store = MemoryRunStore()
|
||
mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, lease_seconds=30))
|
||
|
||
# Create a run that this worker owns
|
||
record = await mgr.create("thread-1")
|
||
await mgr.set_status(record.run_id, RunStatus.running)
|
||
|
||
# Spawn a dummy task so cancel has something to stop
|
||
loop = asyncio.get_running_loop()
|
||
record.task = loop.create_task(asyncio.sleep(3600))
|
||
|
||
# Simulate takeover: directly flip the store row to error
|
||
await store.update_status(record.run_id, "error")
|
||
|
||
# Run a single heartbeat tick — it should see update_lease return False
|
||
# and cancel the task
|
||
await mgr._renew_leases()
|
||
|
||
# Let the event loop process the cancellation (task.cancel() schedules,
|
||
# doesn't await).
|
||
await asyncio.sleep(0)
|
||
assert record.ownership_lost is True
|
||
assert record.terminal_status_persisted is False
|
||
assert record.abort_event.is_set() is True
|
||
assert record.task.cancelled()
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_returns_taken_over_when_peer_claims_during_local_cancel():
|
||
"""When a peer's claim_for_takeover flips the row to error between this
|
||
worker's in-memory cancel and the guarded update_status, cancel() must
|
||
surface taken_over (not cancelled) so the client sees a status consistent
|
||
with the store."""
|
||
store = MemoryRunStore()
|
||
mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
|
||
record = await mgr.create("thread-1")
|
||
await mgr.set_status(record.run_id, RunStatus.running)
|
||
|
||
# Wrap update_status so that the first call (from cancel's _persist_status)
|
||
# is rejected as if a peer already marked the row error. This simulates
|
||
# the race: in-memory cancel succeeds, but store write is blocked.
|
||
original = store.update_status
|
||
|
||
async def race_update(run_id, status, *, error=None, stop_reason=None, **kwargs):
|
||
# Simulate peer takeover: flip to error before our write lands
|
||
run = store._runs.get(run_id)
|
||
if run and run["status"] == "running" and status == "interrupted":
|
||
run["status"] = "error"
|
||
run["error"] = "peer takeover"
|
||
run["updated_at"] = datetime.now(UTC).isoformat()
|
||
return False # our write was blocked
|
||
return await original(run_id, status, error=error, stop_reason=stop_reason, **kwargs)
|
||
|
||
store.update_status = race_update
|
||
|
||
outcome = await mgr.cancel(record.run_id)
|
||
assert outcome == CancelOutcome.taken_over
|
||
|
||
# Store row must reflect the takeover, not the local cancel
|
||
row = await store.get(record.run_id)
|
||
assert row["status"] == "error"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_action_rollback_finalizes_to_error_in_store():
|
||
"""action=rollback must end up as error in the store with the
|
||
"Rolled back by user" message preserved.
|
||
|
||
Regression guard: the update_status guard was originally
|
||
``status IN ('pending','running')`` which blocked the rollback path's
|
||
``interrupted → error`` transition — the store stayed interrupted and
|
||
the rollback message was lost.
|
||
"""
|
||
store = MemoryRunStore()
|
||
mgr = RunManager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True))
|
||
|
||
record = await mgr.create("thread-1")
|
||
await mgr.set_status(record.run_id, RunStatus.running)
|
||
|
||
# Step 1: cancel(action=rollback) flips running → interrupted
|
||
outcome = await mgr.cancel(record.run_id, action="rollback")
|
||
assert outcome == CancelOutcome.cancelled
|
||
row = await store.get(record.run_id)
|
||
assert row["status"] == "interrupted"
|
||
|
||
# Step 2: worker.py finalize path — task raises CancelledError, then
|
||
# set_status(error, "Rolled back by user"). The widened guard
|
||
# (interrupted is in the whitelist) must let this through.
|
||
await mgr.set_status(record.run_id, RunStatus.error, error="Rolled back by user")
|
||
row = await store.get(record.run_id)
|
||
assert row["status"] == "error"
|
||
assert row["error"] == "Rolled back by user"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_peer_reconciliation_fences_late_success_and_completion():
|
||
"""A stale owner cannot overwrite a peer's terminal takeover."""
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=True, grace_seconds=0)
|
||
owner = _make_manager(store=store, worker_id="worker-a", run_ownership_config=config)
|
||
peer = _make_manager(store=store, worker_id="worker-b", run_ownership_config=config)
|
||
record = await owner.create("thread-1")
|
||
await owner.set_status(record.run_id, RunStatus.running)
|
||
|
||
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
|
||
record.lease_expires_at = expired
|
||
store._runs[record.run_id]["lease_expires_at"] = expired
|
||
recovered = await peer.reconcile_orphaned_inflight_runs(error="peer takeover")
|
||
|
||
assert [recovered_record.run_id for recovered_record in recovered] == [record.run_id]
|
||
assert (await store.get(record.run_id))["status"] == "error"
|
||
|
||
await owner.set_status(record.run_id, RunStatus.success)
|
||
await owner.update_run_completion(record.run_id, status=record.status.value, total_tokens=1)
|
||
|
||
row = await store.get(record.run_id)
|
||
assert record.ownership_lost is True
|
||
assert record.status == RunStatus.error
|
||
assert row["status"] == "error"
|
||
assert row["error"] == "peer takeover"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_peer_error_takeover_is_not_mistaken_for_local_error_commit():
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=True, grace_seconds=0)
|
||
owner = _make_manager(
|
||
store=store,
|
||
worker_id="worker-a",
|
||
run_ownership_config=config,
|
||
)
|
||
peer = _make_manager(
|
||
store=store,
|
||
worker_id="worker-b",
|
||
run_ownership_config=config,
|
||
)
|
||
record = await owner.create("thread-1")
|
||
await owner.set_status(record.run_id, RunStatus.running)
|
||
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
|
||
record.lease_expires_at = expired
|
||
store._runs[record.run_id]["lease_expires_at"] = expired
|
||
|
||
recovered = await peer.reconcile_orphaned_inflight_runs(error="peer takeover")
|
||
assert [item.run_id for item in recovered] == [record.run_id]
|
||
|
||
await owner.set_status(
|
||
record.run_id,
|
||
RunStatus.error,
|
||
error="local failure",
|
||
)
|
||
|
||
row = await store.get(record.run_id)
|
||
assert row["error"] == "peer takeover"
|
||
assert record.ownership_lost is True
|
||
assert record.terminal_status_persisted is False
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_unconfirmed_success_is_fenced_when_heartbeat_is_enabled():
|
||
"""A store outage cannot turn an unconfirmed terminal write into local success."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(
|
||
store=store,
|
||
worker_id="worker-a",
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
|
||
async def fail_status_write(*_args, **_kwargs):
|
||
raise OSError("database unreachable")
|
||
|
||
store.update_status = fail_status_write
|
||
|
||
await manager.set_status(record.run_id, RunStatus.success)
|
||
|
||
row = await store.get(record.run_id)
|
||
assert record.ownership_lost is True
|
||
assert record.status == RunStatus.error
|
||
assert record.abort_event.is_set() is True
|
||
assert row["status"] == "running"
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_unconfirmed_staged_success_is_fenced_on_deferred_persistence():
|
||
"""Receipt-ordered status persistence retains the success fail-close fence."""
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(
|
||
store=store,
|
||
worker_id="worker-a",
|
||
run_ownership_config=_lease_config(heartbeat_enabled=True),
|
||
)
|
||
record = await manager.create("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
await manager.set_status(record.run_id, RunStatus.success, persist=False)
|
||
|
||
async def fail_status_write(*_args, **_kwargs):
|
||
raise OSError("database unreachable")
|
||
|
||
store.update_status = fail_status_write
|
||
|
||
persisted = await manager.persist_current_status(record.run_id)
|
||
|
||
row = await store.get(record.run_id)
|
||
assert persisted is False
|
||
assert record.ownership_lost is True
|
||
assert record.status == RunStatus.error
|
||
assert record.abort_event.is_set() is True
|
||
assert row["status"] == "running"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# cancel() claim_for_takeover False → re-read precision
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_claim_lost_to_terminal_returns_not_cancellable():
|
||
"""When cancel() reads the run as active but claim_for_takeover returns
|
||
False because the row went terminal (run finished) between the read and
|
||
the conditional UPDATE, the re-read must surface not_cancellable."""
|
||
store = MemoryRunStore()
|
||
mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=10))
|
||
|
||
# Seed as running so cancel()'s first read passes the status guard.
|
||
expired = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"run-race",
|
||
thread_id="t1",
|
||
status="running",
|
||
owner_worker_id="w-a",
|
||
lease_expires_at=expired,
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
)
|
||
|
||
# Wrap claim_for_takeover: flip the row to success just before the
|
||
# conditional UPDATE so it matches 0 rows.
|
||
original = store.claim_for_takeover
|
||
|
||
async def race_claim(run_id, *, grace_seconds, error, stop_reason=None):
|
||
store._runs[run_id]["status"] = "success"
|
||
return await original(run_id, grace_seconds=grace_seconds, error=error, stop_reason=stop_reason)
|
||
|
||
store.claim_for_takeover = race_claim
|
||
|
||
outcome = await mgr.cancel("run-race")
|
||
assert outcome == CancelOutcome.not_cancellable
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancel_claim_lost_to_takeover_returns_taken_over():
|
||
"""When cancel() reads the run as active but claim_for_takeover returns
|
||
False because another worker already took it over (row is error), the
|
||
re-read must surface taken_over."""
|
||
store = MemoryRunStore()
|
||
mgr = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True, grace_seconds=10))
|
||
|
||
expired = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
|
||
await store.put(
|
||
"run-race",
|
||
thread_id="t1",
|
||
status="running",
|
||
owner_worker_id="w-a",
|
||
lease_expires_at=expired,
|
||
created_at=datetime.now(UTC).isoformat(),
|
||
)
|
||
|
||
# Wrap claim_for_takeover: flip the row to error before the conditional
|
||
# UPDATE so it matches 0 rows (peer already took it over).
|
||
original = store.claim_for_takeover
|
||
|
||
async def race_takeover(run_id, *, grace_seconds, error, stop_reason=None):
|
||
store._runs[run_id]["status"] = "error"
|
||
store._runs[run_id]["error"] = "peer claim"
|
||
return await original(run_id, grace_seconds=grace_seconds, error=error, stop_reason=stop_reason)
|
||
|
||
store.claim_for_takeover = race_takeover
|
||
|
||
outcome = await mgr.cancel("run-race")
|
||
assert outcome == CancelOutcome.taken_over
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _compute_retry_after unit tests
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_compute_retry_after_null_lease_returns_none():
|
||
from app.gateway.routers.thread_runs import _compute_retry_after
|
||
|
||
assert _compute_retry_after(None, 10) is None
|
||
|
||
|
||
def test_compute_retry_after_unparseable_returns_none():
|
||
from app.gateway.routers.thread_runs import _compute_retry_after
|
||
|
||
assert _compute_retry_after("not-a-date", 10) is None
|
||
|
||
|
||
def test_compute_retry_after_normal():
|
||
from app.gateway.routers.thread_runs import _compute_retry_after
|
||
|
||
future = (datetime.now(UTC) + timedelta(seconds=45)).isoformat()
|
||
val = _compute_retry_after(future, 10)
|
||
assert val is not None
|
||
# lease_expires_at is ~45s from now + grace_seconds 10 = ~55, within reason
|
||
assert 40 <= val <= 65
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_cancellation_during_terminal_receipt_cannot_commit_staged_success():
|
||
config = _lease_config(heartbeat_enabled=True)
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=config)
|
||
record = await manager.create_or_reject("late-cancel")
|
||
entered, release = asyncio.Event(), asyncio.Event()
|
||
|
||
class ReceiptStore(MemoryRunEventStore):
|
||
async def put_if_absent(self, **event):
|
||
if event["event_type"] == "run.delivery":
|
||
entered.set()
|
||
await release.wait()
|
||
return await super().put_if_absent(**event)
|
||
|
||
class Agent:
|
||
async def astream(self, *args, **kwargs):
|
||
yield {"messages": []}
|
||
|
||
bridge = SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
|
||
record.task = asyncio.create_task(run_agent(bridge, manager, record, ctx=RunContext(checkpointer=None, event_store=ReceiptStore()), agent_factory=lambda **kwargs: Agent(), graph_input={}, config={}))
|
||
try:
|
||
await asyncio.wait_for(entered.wait(), 1)
|
||
assert record.status == RunStatus.success
|
||
assert await store.request_cancel(record.run_id, action="interrupt") == "interrupt"
|
||
await manager._renew_leases()
|
||
assert record.abort_event.is_set()
|
||
release.set()
|
||
await asyncio.wait_for(record.task, 1)
|
||
assert not record.ownership_lost
|
||
assert record.status == RunStatus.interrupted
|
||
assert (await store.get(record.run_id))["status"] == "interrupted"
|
||
bridge.publish_end.assert_awaited_once()
|
||
finally:
|
||
release.set()
|
||
await asyncio.wait_for(asyncio.gather(record.task, return_exceptions=True), 1)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Assembly evidence publication under lease ownership
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _evidence(record) -> dict:
|
||
return {"run_id": record.run_id, "thread_id": record.thread_id, "version": "v1"}
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_assembly_evidence_after_peer_takeover_fences_instead_of_reviving_row():
|
||
"""A heartbeat owner must not rewrite a row a peer already reclaimed."""
|
||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||
|
||
store = MemoryRunStore()
|
||
config = _lease_config(heartbeat_enabled=True, grace_seconds=1)
|
||
owner = _make_manager(store=store, run_ownership_config=config, worker_id="worker-a")
|
||
peer = _make_manager(store=store, run_ownership_config=config, worker_id="worker-b")
|
||
record = await owner.create_or_reject("thread-1")
|
||
await owner.set_status(record.run_id, RunStatus.running)
|
||
# The peer sees an expired lease that the owner still believes is live.
|
||
store._runs[record.run_id]["lease_expires_at"] = (datetime.now(UTC) - timedelta(minutes=1)).isoformat()
|
||
assert [r.run_id for r in await peer.reconcile_orphaned_inflight_runs(error="reclaimed", stop_reason=ORPHAN_RECOVERY_STOP_REASON)] == [record.run_id]
|
||
reclaimed = dict(store._runs[record.run_id])
|
||
|
||
with pytest.raises(RuntimeError, match="assembly evidence"):
|
||
await owner.record_assembly_evidence(record.run_id, _evidence(record))
|
||
|
||
assert store._runs[record.run_id] == reclaimed
|
||
assert record.ownership_lost
|
||
assert EVIDENCE_KEY not in record.metadata
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_assembly_evidence_is_merged_without_rewriting_the_owned_row():
|
||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True), worker_id="worker-a")
|
||
record = await manager.create_or_reject("thread-1", metadata={"kept": True})
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
original_put = store.put
|
||
store.put = AsyncMock(wraps=original_put)
|
||
|
||
await manager.record_assembly_evidence(record.run_id, _evidence(record))
|
||
|
||
store.put.assert_not_awaited()
|
||
|
||
row = await store.get(record.run_id)
|
||
assert row["metadata"] == {"kept": True, EVIDENCE_KEY: _evidence(record)}
|
||
assert row["status"] == "running"
|
||
assert not record.ownership_lost
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_assembly_evidence_fails_closed_on_store_without_owned_metadata_merge():
|
||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||
from deerflow.runtime.runs.store.base import RunStore
|
||
|
||
class LegacyHeartbeatStore(MemoryRunStore):
|
||
merge_owned_metadata = RunStore.merge_owned_metadata
|
||
|
||
store = LegacyHeartbeatStore()
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True), worker_id="worker-a")
|
||
record = await manager.create_or_reject("thread-1")
|
||
# A full snapshot cannot carry the owner/lease conditions.
|
||
store.put = AsyncMock()
|
||
|
||
with pytest.raises(RuntimeError, match="assembly evidence"):
|
||
await manager.record_assembly_evidence(record.run_id, _evidence(record))
|
||
|
||
store.put.assert_not_awaited()
|
||
|
||
assert EVIDENCE_KEY not in (await store.get(record.run_id))["metadata"]
|
||
assert EVIDENCE_KEY not in record.metadata
|
||
|
||
|
||
@pytest.mark.anyio
|
||
async def test_assembly_evidence_is_refused_after_local_fence():
|
||
store = MemoryRunStore()
|
||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True), worker_id="worker-a")
|
||
record = await manager.create_or_reject("thread-1")
|
||
await manager.set_status(record.run_id, RunStatus.running)
|
||
assert await manager._mark_ownership_lost(record, reason="lease expired")
|
||
store.merge_owned_metadata = AsyncMock(side_effect=AssertionError("a fenced worker must not write"))
|
||
store.put = AsyncMock(side_effect=AssertionError("a fenced worker must not write"))
|
||
|
||
with pytest.raises(RuntimeError, match="lost lease ownership"):
|
||
await manager.record_assembly_evidence(record.run_id, _evidence(record))
|