mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
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 <willem.jiang@gmail.com>
This commit is contained in:
parent
29a958100e
commit
a7dd4849eb
@ -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.
|
||||
|
||||
@ -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`.
|
||||
|
||||
@ -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")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user