From a7dd4849ebe985fd2a48a73354b1389484b6e959 Mon Sep 17 00:00:00 2001 From: ordinary being Date: Sun, 11 Oct 2026 18:30:09 +0800 Subject: [PATCH] fix(runtime): mark cancel takeovers with the orphan-recovery stop reason (#6674) * fix(runtime): mark cancel takeovers with the orphan-recovery stop reason A cross-worker cancel that takes over an expired-lease run writes error but left stop_reason unset. `orphan_recovered` is the only signal that lets a store-only SSE/`wait` consumer synthesize END after a heartbeat, because the producer is known to be gone. Without it, a subscriber that outlives the crashed owner waits forever for an END no surviving worker publishes; periodic reconciliation ignores the row because it is already terminal. Mark the takeover like periodic orphan recovery: claim_for_takeover already accepts the stop reason and the reconcile path already passes it. The new tests pin the takeover row's marker and prove a real takeover wakes a waiting consumer; three existing takeover test doubles now forward the store argument. * docs(gateway): keep guidance within the 48KB budget The takeover stop-reason sentence pushed backend/app/gateway/AGENTS.md over its 48KB hard limit, failing agent-guidance and the backend guidance-budget test. Trim the sentence; the rationale stays in the manager.py comment. * fix(gateway): complete streams after scheduler orphan recovery --------- Co-authored-by: Willem Jiang --- README.md | 4 + backend/app/gateway/AGENTS.md | 4 +- backend/app/gateway/services.py | 7 +- .../harness/deerflow/runtime/runs/manager.py | 6 + .../tests/test_multi_worker_run_ownership.py | 35 ++++- .../tests/test_wait_disconnect_handling.py | 146 +++++++++++++++++- 6 files changed, 189 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 496955de3c..c7fcc553e9 100644 --- a/README.md +++ b/README.md @@ -543,6 +543,10 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser > > During SSE replay-gap recovery, the Web UI preserves a reconnect pointer claimed by another run while durable state is loading. An older recovery cannot overwrite that pointer when it rejoins its own stream. Even a stale foreign pointer is preserved, so the next reload may first check that run before active-run recovery rejoins, adding a reconnect cycle. > +> Store-only SSE and `/wait` consumers can finish after a heartbeat when run-manager +> or scheduled-task reconciliation recovers a dead worker's run, even if no END +> marker was published. Ordinary completions still wait for END to preserve tail events. +> > In single-process JSONL deployments, cancelling an admitted event-store mutation waits for its background file I/O, rollback, and bookkeeping to settle before releasing the thread write lock. This prevents an older cancelled write from recreating deleted records or rolling back a later successful write. Cancellation can therefore wait on slow storage; it does not stop an in-flight filesystem operation. Callers still waiting to acquire the lock can cancel without starting a mutation. A batch spanning multiple threads drains its current thread group before propagating cancellation; subsequent thread groups do not start. > > After a run publishes its terminal stream marker, its process-local `RunRecord` remains available for the existing five-minute grace period before cleanup; durable run history remains available through `RunStore`, while the stream bridge retains its delivery tail on its separate cleanup schedule. diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index 51d84af46e..031546345e 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -193,7 +193,7 @@ or `deployment.multi_instance: true`. - Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window. - `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `requested` (the non-owning worker durably recorded the first cancellation action for the live owner), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (legacy/custom store lacks the durable request primitive — caller retains the safe 409 + `Retry-After` fallback), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions. - Interrupt/rollback admission registers the replacement before its best-effort persistence of locally interrupted predecessors. If the admitting caller is cancelled during that post-registration await, `RunManager` drains a shielded replacement cleanup before propagating `CancelledError`, including across repeated cancellation. The cleanup normally persists `interrupted`; if that best-effort transition fails, it retries the active-to-`interrupted` store transition strictly and verifies the result with the replacement's captured owner identity. A concurrent peer terminal transition wins and is synchronized back into the local record rather than being overwritten or deleted. -- Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run records `runs.cancel_action` / `cancel_requested_at` while the owner's lease is live; the first action wins even if a retry later lands on the owner. `RunStore.request_cancel()` and owner completion through `finalize_if_not_cancelled()` are competing active-row CAS operations, so an accepted cancel cannot be overwritten by a later success. `RunStore.renew_lease()` renews and observes the request atomically in the SQL implementation. The owner then executes the normal process-local interrupt/rollback and terminal stream path without transferring the lease. An expired owner is still taken over and marked `error`. `wait=true` and cancel-then-stream use the shared bridge to observe owner finalization; a non-standard process-local bridge returns accepted 202 instead of subscribing to an unreachable stream. In single-worker mode (heartbeat off), store-only runs still return 409. +- Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run records `runs.cancel_action` / `cancel_requested_at` while the owner's lease is live; the first action wins even if a retry later lands on the owner. `RunStore.request_cancel()` and owner completion through `finalize_if_not_cancelled()` are competing active-row CAS operations, so an accepted cancel cannot be overwritten by a later success. `RunStore.renew_lease()` renews and observes the request atomically in the SQL implementation. The owner then executes the normal process-local interrupt/rollback and terminal stream path without transferring the lease. An expired owner is taken over with the `orphan_recovered` stop reason so cross-worker consumers can synthesize END. `wait=true` and cancel-then-stream use the shared bridge to observe owner finalization; a non-standard process-local bridge returns accepted 202 instead of subscribing to an unreachable stream. In single-worker mode (heartbeat off), store-only runs still return 409. - A local worker's `RunRecord.lease_expires_at` is the last durably confirmed ownership deadline. `_renew_leases()` bounds each renewal attempt by that deadline: transient store exceptions remain retryable while it is valid, but an exception or blocked call that reaches expiry sets the process-local `ownership_lost` fence, raises `abort_event`, and cancels the run task. Successful renewals collect durable cancellation actions; after all local renewals have been attempted, heartbeat only signals the corresponding process-local tasks, leaving status writes and rollback cleanup to the worker finalization path. Fenced workers do not perform subsequent journal/delivery-receipt, progress/completion/status, checkpoint/thread-metadata, or `on_run_completed` writes; the peer recovery path owns the terminal receipt. `RunStore.update_run_completion()` also refuses to replace a different terminal status, closing the peer-takeover/late-finalization race. `grace_seconds` delays peer reclamation for clock skew but is not extra execution time for an owner that can no longer confirm its lease. Already-committed remote tool side effects remain outside this local cancellation boundary. - Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active. - Run admission and independent writes are first-class thread operations. `runs.operation_kind` distinguishes `run` from `checkpoint_write`, `artifact_write`, `artifact_archive`, `branch`, and `delete`; every active kind shares the durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds. @@ -202,7 +202,7 @@ or `deployment.multi_instance: true`. - `DELETE /api/threads/{id}` holds a durable `delete` reservation for the whole cleanup and removes thread filesystem data (off-loop, drained on cancel), checkpoints, owner-scoped historical runs, run events, feedback and thread metadata best-effort. Historical-run cleanup removes only `operation_kind="run"` rows, so internal thread-operation rows — including the reservation protecting that request — stay durable until `RunManager.reserve_thread_operation()` exits. This cleans persisted rows; preventing an already-admitted write from re-creating state for a deleted thread is a separate thread-incarnation contract, not part of this cleanup. - Creator SSE and `/wait` apply `on_disconnect` only to confirmed request disconnect or cancellation; bridge failures propagate without cancelling the run. `/wait` consumes bridge END; errors return SDK `__error__`, never stale state. - Memory and Redis `StreamBridge` implementations retain only `stream_bridge.queue_maxsize` data events. A syntactically valid `Last-Event-ID` older than the retained watermark, or a live subscriber that falls behind it, yields `StreamGap` before any partial replay. `sse_consumer` maps that control item to an id-less SSE `gap` payload (`stream_replay_gap`) and intentionally leaves the run active; internal `/wait` consumers resume from its latest retained ID because they only need terminal completion. Redis checks bounds plus the non-blocking read in one transaction, using blocking `XREAD` only as a wake-up before repeating the atomic snapshot. For a no-cursor subscriber that established a wait on an empty stream, the first wake response remains provisional until that next snapshot verifies its tail is still retained; this closes the pre-first-delivery trimming window without changing malformed-cursor live tailing. The correctness tradeoff is one three-command snapshot pipeline per poll plus the blocking wake round trip while idle. Malformed cursor behavior remains backend-specific. Memory treats a syntactically numeric cursor below its watermark conservatively as a gap even when the evicted timestamp can no longer be verified; unknown ids at or above the watermark retain the legacy replay-from-earliest policy. -- Redis `StreamBridge` keys use a rolling retained-buffer TTL (`stream_bridge.stream_ttl_seconds`, refreshed on `publish()` / `publish_end()`) as a leak safety net, not as a run timeout. Startup and lease-driven periodic orphan recovery share one Gateway stream-terminalization path: after `RunManager` durably marks a run `error` with `stop_reason=orphan_recovered`, Gateway publishes `END_SENTINEL` and schedules stream cleanup. The periodic store scan, per-row status writes, and Gateway callback run as one supervised single-flight task, so a slow pass is skipped at the next interval instead of piling up or pausing the sole lease-renewal loop. Store retries have bounded attempts/backoff; an individual operation still relies on the database driver/pool timeout. `RunManager.shutdown()` gives active user runs priority within its shared deadline, then drains or cancels orphan recovery. Gateway tracks delayed recovered-stream cleanups and converts unfinished delays to immediate deletes before closing the bridge; the Redis TTL remains the outage safety net. Only startup recovery, before the runtime yields to requests, projects the latest affected thread to `error`; periodic recovery deliberately avoids that non-atomic projection because `ThreadMetaStore` has no `latest_run_id` conditional-update contract. Store-only SSE and `/wait` consumers wait for the bridge's real END marker after an ordinary durable terminal status, because status persistence can precede tail events. The explicit `orphan_recovered` signal is the only heartbeat fallback: its publisher is known to be gone, so it supplies the liveness boundary if END publication fails or the retained key expires. Malformed `Last-Event-ID` reconnect values live-tail new Redis events rather than replaying the retained buffer. Keep cross-component recovery orchestration in Gateway through the generic `RunManager.on_orphans_recovered` callback; do not introduce a harness-to-app dependency. Callback failure warnings include every recovered `run_id` so operators can identify rows whose Gateway-side terminalization needs inspection. +- Redis `StreamBridge` keys use a rolling retained-buffer TTL (`stream_bridge.stream_ttl_seconds`, refreshed on `publish()` / `publish_end()`) as a leak safety net, not as a run timeout. Startup and lease-driven periodic orphan recovery share one Gateway stream-terminalization path: after `RunManager` durably marks a run `error` with `stop_reason=orphan_recovered`, Gateway publishes `END_SENTINEL` and schedules stream cleanup. The periodic store scan, per-row status writes, and Gateway callback run as one supervised single-flight task, so a slow pass is skipped at the next interval instead of piling up or pausing the sole lease-renewal loop. Store retries have bounded attempts/backoff; an individual operation still relies on the database driver/pool timeout. `RunManager.shutdown()` gives active user runs priority within its shared deadline, then drains or cancels orphan recovery. Gateway tracks delayed recovered-stream cleanups and converts unfinished delays to immediate deletes before closing the bridge; the Redis TTL remains the outage safety net. Only startup recovery, before the runtime yields to requests, projects the latest affected thread to `error`; periodic recovery deliberately avoids that non-atomic projection because `ThreadMetaStore` has no `latest_run_id` conditional-update contract. Store-only SSE and `/wait` consumers wait for the bridge's real END marker after an ordinary durable terminal status, because status persistence can precede tail events. Terminal `orphan_recovered` and `scheduled_task_orphan_recovered` rows permit heartbeat fallback because their producer is gone; ordinary terminal rows still await END. Malformed `Last-Event-ID` reconnect values live-tail new Redis events rather than replaying the retained buffer. Keep cross-component recovery orchestration in Gateway through the generic `RunManager.on_orphans_recovered` callback; do not introduce a harness-to-app dependency. Callback failure warnings include every recovered `run_id` so operators can identify rows whose Gateway-side terminalization needs inspection. - Thread-scoped create, stream and wait share an optional owner/thread-scoped `Idempotency-Key`; never pass the raw external key to `RunManager`'s process-wide index. Stores resolve an existing key before applying `interrupt`/`rollback`, so a retry reuses rather than interrupts its run. Reuse requires matching input, assistant and conversation references; keyed resume stores a versioned strict typed-JSON SHA-256 (`true`, `1` and `1.0` differ) in the private `runs.idempotency_request_json` column, never public `kwargs`, while ordinary input needs no extra identity. New readers accept the exact old-writer resume identity shape from `kwargs` and every response still strips it; truly identity-less legacy resume fails closed. Validate strict resume JSON with or without a key, but persist the identity only for keyed resume admission. Missing headers stay non-idempotent; stateless `/api/runs/*` is outside this contract because requests without a thread create one before admission. - For a rolling upgrade, route every keyed resume submission/retry on create/stream/wait only to upgraded workers; otherwise pause that traffic or replace old workers first. Old workers ignore the private identity and can reuse `approve` for `deny` with null input. Old-reader schema compatibility does not protect admission. Upgraded workers reject identity-less legacy resumes with 409. Rollout details: `README.md`, `backend/docs/API.md`. - Reused `/wait` returns durable status/error, never the latest thread checkpoint; `task is None` is not completion, and store-backed fields must refresh after waiting. Capture reuse before awaiting because the cached record's `idempotency_reused` flag is sticky. A creating `/runs/stream` retry whose terminal stream is gone opts into `gap`/`stream_replay_gap` via `emit_gap_on_missing_stream`, while observer joins and default/stateless consumers emit `end`; never infer this from disconnect policy or the sticky flag. A reused running `store_only` record on a process-local bridge returns 409 without `Retry-After`. diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index f20613606d..927b18c646 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -306,13 +306,14 @@ async def _orphan_recovery_observed_after_heartbeat( A normal terminal status is not sufficient: the producer persists status before publishing its final error/data frames and END. Orphan recovery is different because the producer is known to be gone and the durable - ``stop_reason`` is written atomically with the terminal status. Only that - explicit signal may synthesize END after a heartbeat. + ``stop_reason`` is written atomically with the terminal status by run or + scheduled-task reconciliation. Only those explicit recovery signals may + synthesize END after a heartbeat. """ if not record.store_only: return False refreshed = await run_mgr.get(record.run_id, user_id=record.user_id) - return refreshed is not None and _run_is_terminal(refreshed) and refreshed.stop_reason == ORPHAN_RECOVERY_STOP_REASON + return refreshed is not None and _run_is_terminal(refreshed) and refreshed.stop_reason in (ORPHAN_RECOVERY_STOP_REASON, "scheduled_task_orphan_recovered") # --------------------------------------------------------------------------- diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 6a493e449f..c1c5558298 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -1506,6 +1506,12 @@ class RunManager: run_id, grace_seconds=grace_seconds, error=take_over_msg, + # The owner is presumed dead, so a cross-worker consumer's + # only liveness edge is the durable row. Mark the takeover + # like periodic orphan recovery; without this stop reason + # SSE/`wait` subscribers that outlive the crashed owner + # loop on heartbeats forever (their fallback requires it). + stop_reason=ORPHAN_RECOVERY_STOP_REASON, ), ) except Exception: diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py index 5bda66218c..645c702224 100644 --- a/backend/tests/test_multi_worker_run_ownership.py +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -1966,6 +1966,29 @@ async def test_cancel_takeover_from_crashed_worker(): 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.""" @@ -2172,12 +2195,12 @@ async def test_cancel_takeover_race_owner_renewed_lease(): # simulate the lease having been renewed. original = store.claim_for_takeover - async def race_lost(run_id, *, grace_seconds, error): + 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) + 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)) @@ -2690,9 +2713,9 @@ async def test_cancel_claim_lost_to_terminal_returns_not_cancellable(): # conditional UPDATE so it matches 0 rows. original = store.claim_for_takeover - async def race_claim(run_id, *, grace_seconds, error): + 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) + return await original(run_id, grace_seconds=grace_seconds, error=error, stop_reason=stop_reason) store.claim_for_takeover = race_claim @@ -2722,10 +2745,10 @@ async def test_cancel_claim_lost_to_takeover_returns_taken_over(): # 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): + 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) + return await original(run_id, grace_seconds=grace_seconds, error=error, stop_reason=stop_reason) store.claim_for_takeover = race_takeover diff --git a/backend/tests/test_wait_disconnect_handling.py b/backend/tests/test_wait_disconnect_handling.py index 566c5fa275..e8e19f013a 100644 --- a/backend/tests/test_wait_disconnect_handling.py +++ b/backend/tests/test_wait_disconnect_handling.py @@ -12,8 +12,8 @@ it subscribes to the stream bridge until ``END_SENTINEL``, polls ``on_disconnect`` mode by cancelling the background run on real client disconnect. Store-only consumers wait for the bridge's real terminal marker instead of treating an ordinary durable terminal status as proof that all tail -events have already been published. A durable ``orphan_recovered`` stop reason -provides the narrow heartbeat fallback when the publisher is known to be gone. +events have already been published. Explicit run or scheduled-task orphan-recovery +stop reasons provide the narrow heartbeat fallback when the publisher is gone. """ from __future__ import annotations @@ -595,6 +595,148 @@ class TestWaitForRunCompletion: asyncio.run(run()) + def test_wait_completes_after_cancel_takeover_of_crashed_owner(self) -> None: + """A cancel takeover must wake a waiting consumer; no producer survives. + + Regression: the takeover path wrote ``error`` without the + orphan-recovery stop reason, so ``/wait`` (and SSE joins) kept looping + on heartbeats waiting for an END the dead owner can no longer publish. + """ + from datetime import UTC, datetime, timedelta + + from app.gateway.services import wait_for_run_completion + from deerflow.config.run_ownership_config import RunOwnershipConfig + from deerflow.runtime.runs.store.memory import MemoryRunStore + + async def run() -> None: + store = MemoryRunStore() + grace = 10 + expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat() + await store.put( + "takeover-wait", + thread_id=THREAD_ID, + status="running", + created_at=datetime.now(UTC).isoformat(), + owner_worker_id="dead-worker", + lease_expires_at=expired_lease, + ) + mgr = RunManager( + store=store, + run_ownership_config=RunOwnershipConfig(lease_seconds=30, grace_seconds=grace, heartbeat_enabled=True), + ) + record = await mgr.get("takeover-wait") + assert record is not None + assert record.store_only is True + + outcome = await mgr.cancel("takeover-wait", action="interrupt") + assert outcome == CancelOutcome.taken_over + + bridge = _FastHeartbeatBridge() + completed = await asyncio.wait_for( + wait_for_run_completion(bridge, record, _FakeRequest(), mgr), + timeout=1.0, + ) + + assert completed is True + + asyncio.run(run()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("consumer_kind", ["sse", "wait"]) +@pytest.mark.parametrize("recovery_path", ["occurrence", "once"]) +async def test_scheduler_recovery_completes_subscribed_consumers(tmp_path, consumer_kind, recovery_path) -> None: + """Both scheduler recovery paths terminate already-subscribed consumers.""" + from datetime import UTC, datetime, timedelta + + from app.gateway.services import sse_consumer, wait_for_run_completion + from deerflow.config.database_config import DatabaseConfig + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config + from deerflow.persistence.run import RunRepository + from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository + from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository + from deerflow.runtime import RunManager, RunStatus + + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + consumer_task = None + try: + session_factory = get_session_factory() + assert session_factory is not None + task_repo = ScheduledTaskRepository(session_factory) + occurrence_repo = ScheduledTaskRunRepository(session_factory) + run_repo = RunRepository(session_factory) + now = datetime.now(UTC) + await task_repo.create( + task_id="orphan-task", + user_id="user-1", + thread_id="orphan-thread", + context_mode="reuse_thread", + assistant_id="lead_agent", + title="Recovered task", + prompt="test", + schedule_type="once" if recovery_path == "once" else "cron", + schedule_spec={"run_at": now.isoformat()} if recovery_path == "once" else {"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await occurrence_repo.create( + run_record_id="orphan-occurrence", + task_id="orphan-task", + thread_id="orphan-thread", + scheduled_for=now, + trigger="scheduled", + status="running", + ) + await occurrence_repo.update_status("orphan-occurrence", status="running", run_id="orphan-run") + if recovery_path == "once": + await task_repo.update("orphan-task", user_id="user-1", updates={"status": "running"}) + await run_repo.put( + "orphan-run", + thread_id="orphan-thread", + user_id="user-1", + status="running", + owner_worker_id="dead-worker", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + manager = RunManager(store=run_repo) + record = await manager.get("orphan-run", user_id="user-1") + assert record is not None and record.store_only + assert record.status == RunStatus.running + subscribed = asyncio.Event() + + class SubscribedBridge(_FastHeartbeatBridge): + def subscribe(self, run_id, **kwargs): + subscribed.set() + return super().subscribe(run_id, **kwargs) + + bridge = SubscribedBridge() + await bridge.publish(record.run_id, "values", {"step": 1}) + + async def consume(): + if consumer_kind == "wait": + return await wait_for_run_completion(bridge, record, _FakeRequest(), manager) + return [frame async for frame in sse_consumer(bridge, record, _FakeRequest(), manager)] + + consumer_task = asyncio.create_task(consume()) + await asyncio.wait_for(subscribed.wait(), timeout=1.0) + assert not consumer_task.done() + recover = task_repo.reconcile_stuck_once_tasks if recovery_path == "once" else occurrence_repo.reconcile_active_runs + await recover(error="crashed owner", now=now) + recovered = await run_repo.get(record.run_id, user_id="user-1") + assert recovered is not None + assert recovered["status"] == "error" + assert recovered["stop_reason"] == "scheduled_task_orphan_recovered" + result = await asyncio.wait_for(consumer_task, timeout=1.0) + if consumer_kind == "wait": + assert result is True + else: + assert result[-1] == "event: end\ndata: null\n\n" + finally: + if consumer_task is not None and not consumer_task.done(): + consumer_task.cancel() + await asyncio.gather(consumer_task, return_exceptions=True) + await close_engine() + @pytest.mark.parametrize("route_name", ["stateless_wait", "thread_wait", "cancel_wait", "cancel_stream_wait"]) def test_wait_routes_propagate_bridge_failures(route_name: str) -> None: