From c7538cfb35553aa20292f4d8636304f16ff14040 Mon Sep 17 00:00:00 2001 From: H Haidong <161136270+toforeverlove@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:34:20 +0800 Subject: [PATCH] fix(runs): terminate orphaned streams after lease recovery (#4420) * fix(runs): terminate orphaned streams after lease recovery * fix(runs): include recovered ids in callback warnings * fix(runs): harden orphan recovery lifecycle --- README.md | 2 +- backend/AGENTS.md | 2 +- backend/app/gateway/deps.py | 132 ++++++++- backend/app/gateway/services.py | 23 +- backend/docs/STREAMING.md | 2 +- .../harness/deerflow/persistence/run/sql.py | 10 +- .../harness/deerflow/runtime/__init__.py | 4 +- .../harness/deerflow/runtime/runs/__init__.py | 4 +- .../harness/deerflow/runtime/runs/manager.py | 110 ++++++-- .../deerflow/runtime/runs/store/base.py | 4 +- .../deerflow/runtime/runs/store/memory.py | 3 + backend/tests/test_gateway_run_recovery.py | 120 +++++++- .../tests/test_multi_worker_run_ownership.py | 267 +++++++++++++++++- backend/tests/test_run_manager.py | 21 +- .../tests/test_wait_disconnect_handling.py | 165 ++++++++++- 15 files changed, 820 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index 0a0f5d542..79a5a7829 100644 --- a/README.md +++ b/README.md @@ -275,7 +275,7 @@ Browser login uses `HttpOnly` session cookies. The login page offers a "keep me DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser-facing scheme and origin behind a proxy. The bundled nginx sets `X-Forwarded-Proto`, but preserves an upstream HTTPS value and does not overwrite every forwarded header. Configure the outer trusted proxy to replace or strip client-supplied forwarding headers before traffic reaches DeerFlow. > [!IMPORTANT] -> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). The Redis stream bridge (`stream_bridge.type: redis`) shares SSE delivery and `Last-Event-ID` replay across workers, with a rolling retained-buffer TTL (`stream_ttl_seconds`) as a cleanup safety net. Malformed reconnect IDs live-tail new events instead of replaying the retained buffer. It does not make run cancellation, request de-duplication, or IM channel state fully cross-worker by itself; use single-worker Gateway or explicit sticky routing/ownership before raising `GATEWAY_WORKERS`. +> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), and `run_ownership.heartbeat_enabled: true`. The bridge shares SSE delivery and `Last-Event-ID` replay across workers, while lease reconciliation marks runs from dead workers as errors, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE and `/wait` consumers also refresh durable status on heartbeats as a fallback if terminal publication fails. The rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. IM channel state and other process-local services still need their own multi-worker coordination. See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 990b2dfc2..e84534db8 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -444,7 +444,7 @@ metadata only. - Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run can take over (mark `error`) when the owner's lease has expired past the grace window; otherwise it fails with 409 + `Retry-After`. In single-worker mode (heartbeat off), store-only runs still return 409. - 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. - `POST /wait` (both thread-scoped and `/api/runs/wait`) drains the stream bridge via `wait_for_run_completion()` instead of bare `await record.task`, so it honours the run's `on_disconnect` setting and cancels the background run on real client disconnect rather than returning a stale checkpoint (issue #3265). -- 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 orphan recovery publishes `END_SENTINEL` and schedules stream cleanup for recovered runs; malformed `Last-Event-ID` reconnect values live-tail new Redis events rather than replaying the retained buffer. Do not broaden this into a shared-database multi-pod reaper without adding worker ownership/liveness first. +- 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. - Thread-scoped run creation accepts `checkpoint` / `checkpoint_id`; Gateway validates the checkpoint belongs to the request thread before writing `checkpoint_id` / `checkpoint_ns` into `config.configurable` for LangGraph branching. - Thread-scoped Gateway runs evaluate an active `ThreadState.goal` after the visible turn completes. `runtime/goal.py` asks a non-thinking evaluator model to judge only visible conversation evidence and return a typed blocker; the evaluator model is created once per run and reused across hidden continuation checks. The evaluator runs after the graph root's tracing scope has already closed, so `create_goal_evaluator_model`/`evaluate_goal_completion` attach their own model-level tracing callbacks (`attach_tracing=True`) and inject Langfuse trace metadata (`thread_id`/`user_id`/`deerflow_trace_id`) directly onto the `ainvoke` call — the same standalone-caller pattern as `oneshot_llm.run_oneshot_llm` and `MemoryUpdater` (see Tracing System below). Satisfied goals are cleared; every non-satisfied evaluation — continuable or stand-down — is persisted with `last_evaluation` (the blocker, reason, and evidence summary; outcomes that stop the loop additionally record a `stand_down_reason` for observability), but only `goal_not_met_yet` evaluations are streamed as hidden `HumanMessage` continuations, and only when a durable assistant end-of-turn checkpoint exists, the run has not been aborted, the thread did not change during evaluation, and the no-progress breaker has not fired. The continuation cap is 8 — a hard maximum in the `0`–`8` range; callers requesting more are clamped (`set_goal`/TUI) or rejected with 422 (`PUT /goal`). The no-progress breaker keys on the latest visible assistant evidence (not the evaluator's free-text reason, which an LLM rewords every turn), so two consecutive continuations that add no new visible assistant output stop the loop after 2 attempts. Model-response cleanup helpers such as think-block stripping and code-fence stripping live in `deerflow.utils.llm_text` so `runtime/goal.py` and Gateway suggestion parsing share the same JSON-prep behavior. - Run event stream changes must keep producer code, `deerflow/constants.py`, `runtime/events/catalog.py`, `contracts/run_event_stream_contract.json`, `backend/docs/RUN_EVENT_STREAM.md`, and `tests/test_run_event_stream_contract.py` in sync. The dependency-free constants module owns the persisted envelope limits (`event_type` 32 characters, `category` 16) and cross-layer workspace event identity; the catalog owns validated runtime definitions and categories. Dynamic middleware tags are limited to 21 characters after the `middleware:` prefix. The JSON contract owns payload schemas, backend-specific storage semantics, legacy aliases, and compatibility rules; conformance tests require both views and all producer groups to agree. `run.end.content` remains opaque and may retain nested Python values in memory while JSONL/database stores stringify non-JSON nested values, so consumers must not assume backend-identical nested output representations. diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index e11c75a53..9946f5c3f 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -30,7 +30,7 @@ from langgraph.types import Checkpointer from deerflow.community.browser_automation.session import browser_multi_worker_error from deerflow.config.app_config import AppConfig, get_app_config from deerflow.persistence.feedback import FeedbackRepository -from deerflow.runtime import RunContext, RunManager, StreamBridge +from deerflow.runtime import ORPHAN_RECOVERY_STOP_REASON, STARTUP_ORPHAN_RECOVERY_ERROR, RunContext, RunManager, StreamBridge from deerflow.runtime.events.store.base import RunEventStore from deerflow.runtime.runs.store.base import RunStore @@ -163,8 +163,10 @@ async def _publish_recovered_run_stream_end( recovered_runs: list[RunRecord], *, cleanup_delay: float = 60.0, -) -> None: - """Terminate retained streams for runs recovered as orphaned at startup.""" + on_cleanup_scheduled: Callable[[str, asyncio.Task[None]], None] | None = None, +) -> list[tuple[str, asyncio.Task[None]]]: + """Terminate retained streams for runs recovered as orphaned.""" + cleanup_tasks: list[tuple[str, asyncio.Task[None]]] = [] for record in recovered_runs: stream_exists = getattr(bridge, "stream_exists", None) if stream_exists is not None: @@ -185,6 +187,10 @@ async def _publish_recovered_run_stream_end( continue task = asyncio.create_task(bridge.cleanup(record.run_id, delay=cleanup_delay)) task.add_done_callback(lambda task, run_id=record.run_id: _log_recovered_stream_cleanup_result(task, run_id)) + cleanup_tasks.append((record.run_id, task)) + if on_cleanup_scheduled is not None: + on_cleanup_scheduled(record.run_id, task) + return cleanup_tasks def _log_recovered_stream_cleanup_result(task: asyncio.Task[None], run_id: str) -> None: @@ -196,6 +202,45 @@ def _log_recovered_stream_cleanup_result(task: asyncio.Task[None], run_id: str) logger.warning("Failed to clean up recovered run stream for %s", run_id, exc_info=True) +async def _flush_recovered_stream_cleanups( + bridge: StreamBridge, + cleanup_tasks: dict[asyncio.Task[None], str], + *, + timeout: float = 1.0, +) -> None: + """Cancel delayed cleanups and delete their streams before bridge shutdown.""" + pending = [(task, run_id) for task, run_id in cleanup_tasks.items() if not task.done()] + if not pending: + return + for task, _run_id in pending: + task.cancel() + await asyncio.gather(*(task for task, _run_id in pending), return_exceptions=True) + + run_ids = list(dict.fromkeys(run_id for _task, run_id in pending)) + try: + results = await asyncio.wait_for( + asyncio.gather( + *(bridge.cleanup(run_id, delay=0) for run_id in run_ids), + return_exceptions=True, + ), + timeout=max(0.0, timeout), + ) + except TimeoutError: + logger.warning( + "Immediate recovered stream cleanup exceeded %.1fs for run_ids=%s; bridge TTL remains the final safety net", + timeout, + run_ids, + ) + else: + for run_id, result in zip(run_ids, results): + if isinstance(result, BaseException): + logger.warning( + "Failed to immediately clean up recovered run stream for %s: %r", + run_id, + result, + ) + + if TYPE_CHECKING: from app.gateway.auth.local_provider import LocalAuthProvider from app.gateway.auth.repositories.sqlite import SQLiteUserRepository @@ -206,12 +251,19 @@ if TYPE_CHECKING: T = TypeVar("T") -async def _mark_latest_recovered_threads_error( +async def _mark_latest_startup_recovered_threads_error( run_manager: RunManager, thread_store: ThreadMetaStore, recovered_runs: list[RunRecord], ) -> None: - """Mark thread status as error only when its newest run was recovered.""" + """Project startup recovery before request-serving concurrency begins. + + This helper must remain on the pre-``yield`` startup path. ``ThreadMetaStore`` + has no ``latest_run_id`` column, so it cannot express an atomic conditional + update keyed by the recovered run. Periodic recovery deliberately skips this + projection; moving this helper after request serving starts would reintroduce + a read/update race with newer runs. + """ recovered_by_thread: dict[str, set[str]] = {} for record in recovered_runs: recovered_by_thread.setdefault(record.thread_id, set()).add(record.run_id) @@ -230,6 +282,22 @@ async def _mark_latest_recovered_threads_error( logger.warning("Failed to mark thread %s as error during run reconciliation", thread_id, exc_info=True) +async def _terminalize_recovered_runs( + bridge: StreamBridge, + recovered_runs: list[RunRecord], + *, + cleanup_delay: float, + on_cleanup_scheduled: Callable[[str, asyncio.Task[None]], None] | None = None, +) -> list[tuple[str, asyncio.Task[None]]]: + """Publish terminal markers and schedule retained-stream cleanup.""" + return await _publish_recovered_run_stream_end( + bridge, + recovered_runs, + cleanup_delay=cleanup_delay, + on_cleanup_scheduled=on_cleanup_scheduled, + ) + + def get_config() -> AppConfig: """Return the freshest ``AppConfig`` for the current request. @@ -358,9 +426,29 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen # RunManager with store backing for persistence run_ownership_config = getattr(config, "run_ownership", None) + sb_config = getattr(config, "stream_bridge", None) + cleanup_delay = getattr(sb_config, "recovered_stream_cleanup_delay_seconds", 60.0) if sb_config else 60.0 + recovered_stream_cleanup_tasks: dict[asyncio.Task[None], str] = {} + + def track_recovered_stream_cleanup( + run_id: str, + task: asyncio.Task[None], + ) -> None: + recovered_stream_cleanup_tasks[task] = run_id + task.add_done_callback(lambda completed: recovered_stream_cleanup_tasks.pop(completed, None)) + + async def terminalize_recovered_runs(recovered_runs: list[RunRecord]) -> None: + await _terminalize_recovered_runs( + app.state.stream_bridge, + recovered_runs, + cleanup_delay=cleanup_delay, + on_cleanup_scheduled=track_recovered_stream_cleanup, + ) + app.state.run_manager = RunManager( store=app.state.run_store, run_ownership_config=run_ownership_config, + on_orphans_recovered=terminalize_recovered_runs, ) # Startup recovery: mark inflight runs whose lease has expired as error. # In single-worker mode (SQLite / backend=memory), no run has a lease, so @@ -370,13 +458,21 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen from deerflow.utils.time import now_iso recovered_runs = await app.state.run_manager.reconcile_orphaned_inflight_runs( - error="Gateway restarted before this run reached a durable final state.", + error=STARTUP_ORPHAN_RECOVERY_ERROR, before=now_iso(), + stop_reason=ORPHAN_RECOVERY_STOP_REASON, + ) + await _terminalize_recovered_runs( + app.state.stream_bridge, + recovered_runs, + cleanup_delay=cleanup_delay, + on_cleanup_scheduled=track_recovered_stream_cleanup, + ) + await _mark_latest_startup_recovered_threads_error( + app.state.run_manager, + app.state.thread_store, + recovered_runs, ) - sb_config = getattr(config, "stream_bridge", None) - cleanup_delay = getattr(sb_config, "recovered_stream_cleanup_delay_seconds", 60.0) if sb_config else 60.0 - await _publish_recovered_run_stream_end(app.state.stream_bridge, recovered_runs, cleanup_delay=cleanup_delay) - await _mark_latest_recovered_threads_error(app.state.run_manager, app.state.thread_store, recovered_runs) # Start the lease heartbeat if enabled (multi-worker deployments). await app.state.run_manager.start_heartbeat() @@ -391,7 +487,21 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen # raises PoolClosed (issue #3373). run_manager = getattr(app.state, "run_manager", None) if run_manager is not None: - await _drain_inflight_runs(run_manager) + shutdown_deadline = asyncio.get_running_loop().time() + _RUN_DRAIN_TIMEOUT_SECONDS + try: + await _drain_inflight_runs(run_manager) + finally: + await _flush_recovered_stream_cleanups( + app.state.stream_bridge, + recovered_stream_cleanup_tasks, + timeout=min( + 1.0, + max( + 0.0, + shutdown_deadline - asyncio.get_running_loop().time(), + ), + ), + ) await close_engine() diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 44ef5e8f2..e125d4eae 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -36,6 +36,7 @@ from deerflow.config.app_config import get_app_config from deerflow.runtime import ( END_SENTINEL, HEARTBEAT_SENTINEL, + ORPHAN_RECOVERY_STOP_REASON, CheckpointStateAccessor, ConflictError, DisconnectMode, @@ -122,6 +123,24 @@ async def _terminal_record_stream_missing(bridge: StreamBridge, record: RunRecor return False +async def _orphan_recovery_observed_after_heartbeat( + record: RunRecord, + run_mgr: RunManager, +) -> bool: + """Return whether durable orphan recovery is the consumer's liveness edge. + + 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. + """ + 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 + + # --------------------------------------------------------------------------- # Input / config helpers # --------------------------------------------------------------------------- @@ -1111,7 +1130,7 @@ async def sse_consumer( break if entry is HEARTBEAT_SENTINEL: - if await _terminal_record_stream_missing(bridge, record): + if await _orphan_recovery_observed_after_heartbeat(record, run_mgr): yield format_sse("end", None) return yield ": heartbeat\n\n" @@ -1174,7 +1193,7 @@ async def wait_for_run_completion( if entry is END_SENTINEL: completed = True return True - if entry is HEARTBEAT_SENTINEL and await _terminal_record_stream_missing(bridge, record): + if entry is HEARTBEAT_SENTINEL and await _orphan_recovery_observed_after_heartbeat(record, run_mgr): completed = True return True if await request.is_disconnected(): diff --git a/backend/docs/STREAMING.md b/backend/docs/STREAMING.md index 7b0039b96..77b5e2925 100644 --- a/backend/docs/STREAMING.md +++ b/backend/docs/STREAMING.md @@ -143,7 +143,7 @@ sequenceDiagram 关键组件： - `runtime/runs/worker.py::run_agent` — 在 `asyncio.Task` 里跑 `agent.astream()`，把每个 chunk 通过 `serialize(chunk, mode=mode)` 转成 JSON，再 `bridge.publish()`。 -- `runtime/stream_bridge` — 抽象 Queue。`publish/subscribe` 解耦生产者和消费者，支持 `Last-Event-ID` 重连、心跳、多订阅者 fan-out。Redis backend 会在每次 `publish()` / `publish_end()` 刷新 retained stream key TTL；启动恢复将 orphan run 标记为 error 后也会发布 `END_SENTINEL`，避免重连的 SSE 客户端只收到心跳。注意：TTL 是 Redis 内存安全网（防止 key 泄漏），不是 subscriber 终止机制——如果 worker 和 gateway 同时挂掉，已连接的 SSE 客户端在 TTL 过期后仍无法收到 END 信号，需要依赖客户端侧超时。完整的跨 pod subscriber 终止需要 worker 存活检测（liveness），当前版本不包含此功能。 +- `runtime/stream_bridge` — 抽象 Queue。`publish/subscribe` 解耦生产者和消费者，支持 `Last-Event-ID` 重连、心跳、多订阅者 fan-out。Redis backend 会在每次 `publish()` / `publish_end()` 刷新 retained stream key TTL；启动恢复与基于 worker lease 的周期恢复共用 Gateway stream terminalization 路径：`RunManager` 先将 orphan run 持久化为 `error` 并写入显式的 `stop_reason=orphan_recovered`，随后 Gateway 发布 `END_SENTINEL` 并安排 stream cleanup。周期扫描、逐行状态写入和 Gateway callback 作为一个受监督的 single-flight 后台 task 执行；慢任务不会堆积，也不会阻塞唯一的 lease heartbeat。shutdown 优先收敛活跃 run，再处理恢复 task；尚未执行的延迟 stream cleanup 会改为立即删除。只有 runtime `yield` 前、无并发请求的启动恢复会把最新受影响 thread 标记为 error；周期恢复不做非原子的 thread 投影。store-only SSE 与 `/wait` consumer 不能把普通 durable terminal status 当成流已完成，否则可能跳过延迟发布的 error 等尾部事件；只有 `orphan_recovered` 信号能在 heartbeat 时触发 END fallback，因为此时 producer 已被确认失联。TTL 仍是 Redis 内存和故障安全网，不是正常的 subscriber 终止机制。 - `app/gateway/services.py::sse_consumer` — 从 bridge 订阅，格式化为 SSE wire 帧。 - `runtime/serialization.py::serialize` — mode-aware 序列化；`messages` mode 下 `serialize_messages_tuple` 把 `(chunk, metadata)` 转成 `[chunk.model_dump(), metadata]`。 diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index c84e56b3f..97affff83 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -461,8 +461,16 @@ class RunRepository(RunStore): *, grace_seconds: int, error: str, + stop_reason: str | None = None, ) -> bool: cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds) + values: dict[str, Any] = { + "status": "error", + "error": error, + "updated_at": datetime.now(UTC), + } + if stop_reason is not None: + values["stop_reason"] = stop_reason async with self._sf() as session: result = await session.execute( update(RunRow) @@ -471,7 +479,7 @@ class RunRepository(RunStore): RunRow.status.in_(("pending", "running")), _lease_expired_or_null(RunRow.lease_expires_at, cutoff), ) - .values(status="error", error=error, updated_at=datetime.now(UTC)) + .values(**values) ) await session.commit() return result.rowcount != 0 diff --git a/backend/packages/harness/deerflow/runtime/__init__.py b/backend/packages/harness/deerflow/runtime/__init__.py index 2940766a5..c47eca5b1 100644 --- a/backend/packages/harness/deerflow/runtime/__init__.py +++ b/backend/packages/harness/deerflow/runtime/__init__.py @@ -7,7 +7,7 @@ directly from ``deerflow.runtime``. from .checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer -from .runs import CancelOutcome, ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent +from .runs import ORPHAN_RECOVERY_STOP_REASON, STARTUP_ORPHAN_RECOVERY_ERROR, CancelOutcome, ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent from .serialization import serialize, serialize_channel_values, serialize_channel_values_for_api, serialize_lc_object, serialize_messages_tuple, strip_data_url_image_blocks from .store import get_store, make_store, reset_store, store_context @@ -29,10 +29,12 @@ __all__ = [ "CancelOutcome", "ConflictError", "DisconnectMode", + "ORPHAN_RECOVERY_STOP_REASON", "RunContext", "RunManager", "RunRecord", "RunStatus", + "STARTUP_ORPHAN_RECOVERY_ERROR", "UnsupportedStrategyError", "run_agent", # serialization diff --git a/backend/packages/harness/deerflow/runtime/runs/__init__.py b/backend/packages/harness/deerflow/runtime/runs/__init__.py index 5d2a2661f..f78ae6be0 100644 --- a/backend/packages/harness/deerflow/runtime/runs/__init__.py +++ b/backend/packages/harness/deerflow/runtime/runs/__init__.py @@ -1,6 +1,6 @@ """Run lifecycle management for LangGraph Platform API compatibility.""" -from .manager import CancelOutcome, ConflictError, RunManager, RunRecord, UnsupportedStrategyError +from .manager import ORPHAN_RECOVERY_STOP_REASON, STARTUP_ORPHAN_RECOVERY_ERROR, CancelOutcome, ConflictError, RunManager, RunRecord, UnsupportedStrategyError from .schemas import DisconnectMode, RunStatus from .worker import RunContext, run_agent @@ -8,10 +8,12 @@ __all__ = [ "CancelOutcome", "ConflictError", "DisconnectMode", + "ORPHAN_RECOVERY_STOP_REASON", "RunContext", "RunManager", "RunRecord", "RunStatus", + "STARTUP_ORPHAN_RECOVERY_ERROR", "UnsupportedStrategyError", "run_agent", ] diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 0a8b6aad6..04a5d335c 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -27,6 +27,10 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +ORPHAN_RECOVERY_STOP_REASON = "orphan_recovered" +STARTUP_ORPHAN_RECOVERY_ERROR = "Gateway restarted before this run reached a durable final state." +LEASE_ORPHAN_RECOVERY_ERROR = "Run lease expired — owning worker is unreachable." + _RETRYABLE_SQLITE_MESSAGES = ( "database is locked", "database table is locked", @@ -184,6 +188,9 @@ class RunRecord: stop_reason: str | None = None +OrphanRecoveryCallback = Callable[[list[RunRecord]], Awaitable[None]] + + class RunManager: """In-memory run registry with optional persistent RunStore backing. @@ -199,6 +206,7 @@ class RunManager: persistence_retry_policy: PersistenceRetryPolicy | None = None, worker_id: str | None = None, run_ownership_config: RunOwnershipConfig | None = None, + on_orphans_recovered: OrphanRecoveryCallback | None = None, ) -> None: self._runs: dict[str, RunRecord] = {} # Secondary index: thread_id -> insertion-ordered run_id set (a dict is @@ -211,8 +219,10 @@ class RunManager: self._persistence_retry_policy = persistence_retry_policy or PersistenceRetryPolicy() self._worker_id = worker_id or _generate_worker_id() self._run_ownership_config = run_ownership_config + self._on_orphans_recovered = on_orphans_recovered self._heartbeat_task: asyncio.Task | None = None self._heartbeat_stop: asyncio.Event | None = None + self._orphan_recovery_task: asyncio.Task[None] | None = None def _index_run_locked(self, record: RunRecord) -> None: """Register *record* in the thread index. Caller must hold ``self._lock``.""" @@ -580,7 +590,12 @@ class RunManager: if self._store is None: return sorted(memory_records, key=lambda r: r.created_at, reverse=True)[:limit] records_by_id = {record.run_id: record for record in memory_records} - store_limit = max(0, limit - len(memory_records)) + # Query enough rows to cover both the requested page and every possible + # in-memory/store duplicate. Local records can be older than persisted + # rows, so subtracting them from the store limit can hide the actual + # newest run before the merge; querying only ``limit`` can still lose a + # distinct row when that page is occupied by duplicate local records. + store_limit = limit + len(memory_records) try: rows = await self._store.list_by_thread(thread_id, user_id=user_id, limit=store_limit) except Exception: @@ -1090,6 +1105,7 @@ class RunManager: *, error: str, before: str | None = None, + stop_reason: str | None = None, ) -> list[RunRecord]: """Mark persisted active runs as failed when their lease has expired. @@ -1138,6 +1154,7 @@ class RunManager: record.run_id, grace_seconds=grace_seconds, error=error, + stop_reason=stop_reason, ), ) except Exception: @@ -1151,6 +1168,7 @@ class RunManager: continue record.status = RunStatus.error record.error = error + record.stop_reason = stop_reason record.updated_at = now recovered.append(record) @@ -1215,21 +1233,21 @@ class RunManager: self._heartbeat_task = task logger.info("Run lease heartbeat started for worker %s", self._worker_id) - async def stop_heartbeat(self) -> None: - """Stop the background heartbeat task.""" + async def stop_heartbeat(self, *, timeout: float = 5.0) -> None: + """Stop the background heartbeat task within ``timeout`` seconds.""" if self._heartbeat_stop is not None: self._heartbeat_stop.set() if self._heartbeat_task is not None and not self._heartbeat_task.done(): - try: - await asyncio.wait_for(self._heartbeat_task, timeout=5.0) - except TimeoutError: + _, pending = await asyncio.wait( + (self._heartbeat_task,), + timeout=max(0.0, timeout), + ) + if pending: self._heartbeat_task.cancel() try: await self._heartbeat_task except asyncio.CancelledError: pass - except asyncio.CancelledError: - pass self._heartbeat_task = None self._heartbeat_stop = None logger.info("Run lease heartbeat stopped for worker %s", self._worker_id) @@ -1273,10 +1291,7 @@ class RunManager: # replacement starts before the lease expires, and the # startup pass skips the still-valid lease. if cycle % 3 == 0: - try: - await self._reconcile_orphans_periodic() - except Exception: - logger.warning("Periodic orphan reconciliation failed", exc_info=True) + self._schedule_orphan_reconciliation() async def _renew_leases(self) -> None: """Renew the lease on every locally-owned active run.""" @@ -1339,22 +1354,74 @@ class RunManager: async def _reconcile_orphans_periodic(self) -> None: """Sweep for expired leases owned by dead peers. - Called from ``_heartbeat_loop`` every ``lease_seconds``. Startup - reconciliation handles the initial sweep; this periodic pass - catches orphans whose lease expires between restarts. + Scheduled as a single-flight background task by ``_heartbeat_loop``. + This keeps both the store scan/status writes and the Gateway callback + off the lease-renewal loop. Startup reconciliation handles the initial + sweep; this periodic pass catches orphans whose lease expires between + restarts. """ - error_msg = "Run lease expired — owning worker is unreachable." - recovered = await self.reconcile_orphaned_inflight_runs(error=error_msg) + recovered = await self.reconcile_orphaned_inflight_runs( + error=LEASE_ORPHAN_RECOVERY_ERROR, + stop_reason=ORPHAN_RECOVERY_STOP_REASON, + ) if recovered: logger.warning( "Periodic reconciliation recovered %d orphaned run(s) as error", len(recovered), ) + if self._on_orphans_recovered is not None: + try: + await self._on_orphans_recovered(recovered) + except Exception: + logger.warning( + "Periodic orphan recovery callback failed for %d run(s): run_ids=%s", + len(recovered), + [record.run_id for record in recovered], + exc_info=True, + ) + + def _schedule_orphan_reconciliation(self) -> None: + """Start one supervised recovery pass unless one is already running.""" + task = self._orphan_recovery_task + if task is not None and not task.done(): + logger.debug("Skipping periodic orphan reconciliation: previous pass is still running") + return + task = asyncio.create_task(self._reconcile_orphans_periodic()) + task.set_name("deerflow-periodic-orphan-recovery") + self._orphan_recovery_task = task + task.add_done_callback(self._orphan_reconciliation_done) + + def _orphan_reconciliation_done(self, task: asyncio.Task[None]) -> None: + """Clear and inspect the supervised single-flight recovery task.""" + if self._orphan_recovery_task is task: + self._orphan_recovery_task = None + if task.cancelled(): + return + try: + task.result() + except Exception: + logger.warning("Periodic orphan reconciliation failed", exc_info=True) + + async def _drain_orphan_recovery_task(self, *, timeout: float) -> None: + """Boundedly await the supervised recovery pass during shutdown.""" + task = self._orphan_recovery_task + if task is None or task.done(): + return + _, pending = await asyncio.wait((task,), timeout=max(0.0, timeout)) + if pending: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + logger.warning( + "Orphan recovery drain exceeded %.1fs on shutdown; cancelled the active pass", + timeout, + ) async def shutdown(self, *, timeout: float = 5.0) -> None: """Cancel and bounded-await all in-flight runs on process shutdown. - Stops the lease heartbeat first so no renewal races against the drain. + Signals active runs first so their cancellation/cleanup can overlap a + bounded heartbeat stop. The heartbeat may perform one final benign + lease renewal before it observes the stop event. Chat runs execute in fire-and-forget background ``asyncio`` tasks that write checkpoints through a shared checkpointer. On shutdown the @@ -1380,7 +1447,6 @@ class RunManager: ``app.gateway.app._SHUTDOWN_HOOK_TIMEOUT_SECONDS``. Runs still active after ``timeout`` are logged and may still race teardown. """ - await self.stop_heartbeat() loop = asyncio.get_running_loop() deadline = loop.time() + timeout @@ -1393,11 +1459,14 @@ class RunManager: # Status is decided AFTER the drain (below), not here: a run that # completes on its own during the drain must keep its real status. + await self.stop_heartbeat(timeout=max(0.0, deadline - loop.time())) + if not inflight: + await self._drain_orphan_recovery_task(timeout=max(0.0, deadline - loop.time())) return tasks = [record.task for record in inflight] - _, pending = await asyncio.wait(tasks, timeout=timeout) + _, pending = await asyncio.wait(tasks, timeout=max(0.0, deadline - loop.time())) # Only mark/persist ``interrupted`` for runs that did not settle on their # own (still pending after the timeout, or ended cancelled). A run that @@ -1445,6 +1514,7 @@ class RunManager: if pending: logger.warning("Run drain exceeded %.1fs on shutdown; %d run task(s) still active and may race checkpointer teardown", timeout, len(pending)) logger.info("Drained %d in-flight run(s) on shutdown (%d settled within %.1fs)", len(inflight), len(inflight) - len(pending), timeout) + await self._drain_orphan_recovery_task(timeout=max(0.0, deadline - loop.time())) class CancelOutcome(StrEnum): diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index cd6e9eae0..017bf374b 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -188,13 +188,15 @@ class RunStore(abc.ABC): *, grace_seconds: int, error: str, + stop_reason: str | None = None, ) -> bool: """Atomically mark an expired-lease active run as ``error``. Only rows whose lease has expired past *grace_seconds* (or whose lease is NULL — pre-ownership data) are updated. The conditional WHERE closes the race between the caller's stale read of the lease - and a concurrent heartbeat renewal by the owning worker. + and a concurrent heartbeat renewal by the owning worker. When + provided, *stop_reason* is persisted in the same atomic update. Returns ``False`` when: - the run is no longer ``pending`` / ``running``, diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index a91a958f4..4f7259981 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -227,6 +227,7 @@ class MemoryRunStore(RunStore): *, grace_seconds: int, error: str, + stop_reason: str | None = None, ) -> bool: from deerflow.utils.time import is_lease_expired @@ -240,6 +241,8 @@ class MemoryRunStore(RunStore): return False run["status"] = "error" run["error"] = error + if stop_reason is not None: + run["stop_reason"] = stop_reason run["updated_at"] = datetime.now(UTC).isoformat() return True diff --git a/backend/tests/test_gateway_run_recovery.py b/backend/tests/test_gateway_run_recovery.py index 921162a48..1b23234ee 100644 --- a/backend/tests/test_gateway_run_recovery.py +++ b/backend/tests/test_gateway_run_recovery.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio from contextlib import asynccontextmanager +from datetime import UTC, datetime, timedelta from types import SimpleNamespace import anyio @@ -11,10 +13,13 @@ from fastapi import FastAPI import deerflow.runtime as runtime_module from app.gateway import deps as gateway_deps +from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.persistence import engine as engine_module from deerflow.persistence import thread_meta as thread_meta_module +from deerflow.runtime import END_SENTINEL, MemoryStreamBridge, RunManager from deerflow.runtime.checkpointer import async_provider as checkpointer_module from deerflow.runtime.events import store as event_store_module +from deerflow.runtime.runs.store.memory import MemoryRunStore @asynccontextmanager @@ -29,16 +34,23 @@ class _FakeRunManager: recovered_runs = [SimpleNamespace(run_id="run-1", thread_id="thread-1")] latest_by_thread: dict[str, list[SimpleNamespace]] = {} - def __init__(self, *, store, run_ownership_config=None): + def __init__(self, *, store, run_ownership_config=None, on_orphans_recovered=None): self.store = store self.run_ownership_config = run_ownership_config + self.on_orphans_recovered = on_orphans_recovered self.reconcile_calls: list[dict] = [] self.list_by_thread_calls: list[dict] = [] self.shutdown_calls: int = 0 _FakeRunManager.instances.append(self) - async def reconcile_orphaned_inflight_runs(self, *, error: str, before: str | None = None): - self.reconcile_calls.append({"error": error, "before": before}) + async def reconcile_orphaned_inflight_runs( + self, + *, + error: str, + before: str | None = None, + stop_reason: str | None = None, + ): + self.reconcile_calls.append({"error": error, "before": before, "stop_reason": stop_reason}) return self.recovered_runs async def list_by_thread(self, thread_id: str, *, user_id=None, limit: int = 100): @@ -81,6 +93,35 @@ class _FakeStreamBridge: self.cleanup_calls.append((run_id, delay)) +class _RetainedMemoryStreamBridge(MemoryStreamBridge): + """Memory bridge that records cleanup without deleting test history.""" + + def __init__(self) -> None: + super().__init__() + self.cleanup_calls: list[tuple[str, float]] = [] + + async def cleanup(self, run_id: str, *, delay: float = 0) -> None: + self.cleanup_calls.append((run_id, delay)) + + +class _DelayedCleanupStreamBridge(_FakeStreamBridge): + def __init__(self) -> None: + super().__init__(existing_streams={"run-1"}) + self.delayed_cleanup_started = asyncio.Event() + self.delayed_cleanup_cancelled = asyncio.Event() + + async def cleanup(self, run_id: str, *, delay: float = 0) -> None: + self.cleanup_calls.append((run_id, delay)) + if delay <= 0: + return + self.delayed_cleanup_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.delayed_cleanup_cancelled.set() + raise + + @pytest.mark.anyio async def test_recovered_run_stream_end_skips_expired_stream(): """Startup recovery should not recreate an already-expired retained stream.""" @@ -95,6 +136,78 @@ async def test_recovered_run_stream_end_skips_expired_stream(): assert stream_bridge.cleanup_calls == [] +@pytest.mark.anyio +async def test_shutdown_flushes_delayed_recovered_stream_cleanup_immediately(): + """Bridge shutdown must not abandon a delayed cleanup until the stream TTL.""" + stream_bridge = _DelayedCleanupStreamBridge() + cleanups = await gateway_deps._publish_recovered_run_stream_end( + stream_bridge, + [SimpleNamespace(run_id="run-1", thread_id="thread-1")], + cleanup_delay=60.0, + ) + cleanup_tasks = {task: run_id for run_id, task in cleanups} + await asyncio.wait_for(stream_bridge.delayed_cleanup_started.wait(), timeout=0.5) + + await gateway_deps._flush_recovered_stream_cleanups( + stream_bridge, + cleanup_tasks, + ) + + assert stream_bridge.delayed_cleanup_cancelled.is_set() + assert stream_bridge.cleanup_calls == [("run-1", 60.0), ("run-1", 0)] + + +@pytest.mark.anyio +async def test_periodic_recovery_terminalizes_stream_without_thread_projection(): + """Periodic recovery must close streams without racing thread projection.""" + store = MemoryRunStore() + stream_bridge = _RetainedMemoryStreamBridge() + thread_store = _FakeThreadStore() + expired = (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, + created_at=expired, + ) + await stream_bridge.publish("periodic-orphan", "values", {"step": 1}) + + async def terminalize(recovered_runs): + await gateway_deps._terminalize_recovered_runs( + stream_bridge, + recovered_runs, + cleanup_delay=60.0, + ) + + manager = RunManager( + store=store, + worker_id="live-worker", + run_ownership_config=RunOwnershipConfig( + heartbeat_enabled=True, + lease_seconds=30, + grace_seconds=10, + ), + on_orphans_recovered=terminalize, + ) + + await manager._reconcile_orphans_periodic() + await anyio.sleep(0) + received = [ + entry + async for entry in stream_bridge.subscribe( + "periodic-orphan", + heartbeat_interval=0.01, + ) + ] + await anyio.sleep(0) + + assert received[-1] is END_SENTINEL + assert stream_bridge.cleanup_calls == [("periodic-orphan", 60.0)] + assert thread_store.status_updates == [] + + @pytest.mark.anyio async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch): """SQLite startup should recover stale active runs before serving requests.""" @@ -133,6 +246,7 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch): assert len(_FakeRunManager.instances) == 1 assert _FakeRunManager.instances[0].reconcile_calls assert _FakeRunManager.instances[0].reconcile_calls[0]["error"] + assert _FakeRunManager.instances[0].reconcile_calls[0]["stop_reason"] == runtime_module.ORPHAN_RECOVERY_STOP_REASON assert _FakeRunManager.instances[0].list_by_thread_calls == [{"thread_id": "thread-1", "user_id": None, "limit": 1}] assert thread_store.status_updates == [("thread-1", "error", None)] assert stream_bridge.publish_end_calls == ["run-1"] diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py index 7f298c663..8d0939dcf 100644 --- a/backend/tests/test_multi_worker_run_ownership.py +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -5,6 +5,7 @@ Coverage: - create_or_reject with interrupt strategy claims and cancels old runs - create_run_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 - GATEWAY_WORKERS=1 + heartbeat_enabled=false behaviour unchanged @@ -19,7 +20,7 @@ from unittest.mock import AsyncMock import pytest from deerflow.config.run_ownership_config import RunOwnershipConfig -from deerflow.runtime import RunManager, RunStatus +from deerflow.runtime import ORPHAN_RECOVERY_STOP_REASON, RunManager, RunStatus from deerflow.runtime.runs.manager import CancelOutcome, ConflictError, _generate_worker_id from deerflow.runtime.runs.store.memory import MemoryRunStore @@ -389,6 +390,262 @@ async def test_reconciliation_returns_empty_when_no_orphaned_runs(): 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 # --------------------------------------------------------------------------- @@ -1023,13 +1280,19 @@ async def test_claim_for_takeover_succeeds_with_expired_lease(): 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") + 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 diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index 9a3433749..ac9f0e8eb 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -68,7 +68,7 @@ class FailingTakeoverRunStore(MemoryRunStore): super().__init__() self.takeover_attempts = 0 - async def claim_for_takeover(self, run_id, *, grace_seconds, error): + async def claim_for_takeover(self, run_id, *, grace_seconds, error, stop_reason=None): self.takeover_attempts += 1 raise sqlite3.OperationalError("database is locked") @@ -576,6 +576,25 @@ async def test_list_by_thread_merges_store_runs_newest_first(): assert runs[0] is memory_record +@pytest.mark.anyio +async def test_list_by_thread_limit_does_not_let_old_memory_hide_new_store_run(): + """A local row must not consume the store query's newest-run limit.""" + store = MemoryRunStore() + manager = RunManager(store=store) + old_memory = await manager.create("thread-1") + old_memory.created_at = "2026-01-01T00:00:00+00:00" + await store.put( + "new-store", + thread_id="thread-1", + status="success", + created_at="2026-01-02T00:00:00+00:00", + ) + + runs = await manager.list_by_thread("thread-1", limit=1) + + assert [run.run_id for run in runs] == ["new-store"] + + @pytest.mark.anyio async def test_create_defaults(manager: RunManager): """Create with no optional args should use defaults.""" diff --git a/backend/tests/test_wait_disconnect_handling.py b/backend/tests/test_wait_disconnect_handling.py index a32394d20..dcfec4a03 100644 --- a/backend/tests/test_wait_disconnect_handling.py +++ b/backend/tests/test_wait_disconnect_handling.py @@ -1,4 +1,4 @@ -"""Regression tests for issue #3265. +"""Regression tests for issues #3265 and #3932. The non-streaming ``/wait`` endpoints used to ``await record.task`` with no disconnect handling and silently swallow ``CancelledError``. When a long @@ -10,7 +10,10 @@ The fix introduces ``wait_for_run_completion`` in ``app.gateway.services``: it subscribes to the stream bridge until ``END_SENTINEL``, polls ``request.is_disconnected()`` on every wake-up, and honours the record's ``on_disconnect`` mode by cancelling the background run on real client -disconnect. +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. """ from __future__ import annotations @@ -19,7 +22,7 @@ import asyncio from dataclasses import dataclass, field from typing import Any -from deerflow.runtime import RunManager, RunRecord, RunStatus +from deerflow.runtime import ORPHAN_RECOVERY_STOP_REASON, RunManager, RunRecord, RunStatus from deerflow.runtime.runs.schemas import DisconnectMode from deerflow.runtime.stream_bridge.memory import MemoryStreamBridge @@ -69,6 +72,17 @@ class _MissingStreamBridge: return None +class _FastHeartbeatBridge(MemoryStreamBridge): + """Memory bridge with a short heartbeat for durable-status refresh tests.""" + + def subscribe(self, run_id, *, last_event_id=None, heartbeat_interval=15.0): + return super().subscribe( + run_id, + last_event_id=last_event_id, + heartbeat_interval=0.01, + ) + + async def _create_running_record(mgr: RunManager, *, on_disconnect: DisconnectMode) -> Any: record = await mgr.create_or_reject( THREAD_ID, @@ -249,3 +263,148 @@ class TestWaitForRunCompletion: assert bridge.subscribed is False asyncio.run(run()) + + def test_sse_consumer_preserves_tail_events_after_durable_terminal_status(self) -> None: + """A durable terminal row must not overtake delayed error and END events.""" + from app.gateway.services import sse_consumer + from deerflow.runtime.runs.store.memory import MemoryRunStore + + async def run() -> None: + store = MemoryRunStore() + await store.put( + "periodic-orphan", + thread_id=THREAD_ID, + status="running", + ) + mgr = RunManager(store=store) + record = await mgr.get("periodic-orphan") + assert record is not None + assert record.store_only is True + bridge = _FastHeartbeatBridge() + await bridge.publish(record.run_id, "values", {"step": 1}) + request = _FakeRequest() + consumer = sse_consumer(bridge, record, request, mgr) + + first_frame = await anext(consumer) + assert first_frame.startswith("event: values\n") + + await store.update_status(record.run_id, "error", error="lease expired") + + async def publish_tail() -> None: + await asyncio.sleep(0.05) + await bridge.publish(record.run_id, "error", {"message": "late error"}) + await bridge.publish_end(record.run_id) + + publisher = asyncio.create_task(publish_tail()) + tail_frames = [frame async for frame in consumer] + await publisher + + error_index = next(index for index, frame in enumerate(tail_frames) if frame.startswith("event: error\n")) + end_index = next(index for index, frame in enumerate(tail_frames) if frame.startswith("event: end\n")) + assert error_index < end_index + assert record.status == RunStatus.running + + asyncio.run(run()) + + def test_wait_preserves_tail_events_after_durable_terminal_status(self) -> None: + """The wait path must remain blocked until the real END is published.""" + from app.gateway.services import wait_for_run_completion + from deerflow.runtime.runs.store.memory import MemoryRunStore + + async def run() -> None: + store = MemoryRunStore() + await store.put( + "periodic-orphan", + thread_id=THREAD_ID, + status="running", + ) + mgr = RunManager(store=store) + record = await mgr.get("periodic-orphan") + assert record is not None + assert record.store_only is True + bridge = _FastHeartbeatBridge() + await bridge.publish(record.run_id, "values", {"step": 1}) + await store.update_status(record.run_id, "error", error="lease expired") + + wait_task = asyncio.create_task(wait_for_run_completion(bridge, record, _FakeRequest(), mgr)) + await asyncio.sleep(0.05) + assert wait_task.done() is False + + await bridge.publish(record.run_id, "error", {"message": "late error"}) + await asyncio.sleep(0) + assert wait_task.done() is False + + await bridge.publish_end(record.run_id) + completed = await asyncio.wait_for(wait_task, timeout=1.0) + + assert completed is True + assert record.status == RunStatus.running + + asyncio.run(run()) + + def test_sse_consumer_uses_explicit_orphan_recovery_liveness_boundary( + self, + ) -> None: + """A recovered orphan may synthesize END when its publisher is gone.""" + from app.gateway.services import sse_consumer + from deerflow.runtime.runs.store.memory import MemoryRunStore + + async def run() -> None: + store = MemoryRunStore() + await store.put( + "periodic-orphan", + thread_id=THREAD_ID, + status="running", + ) + mgr = RunManager(store=store) + record = await mgr.get("periodic-orphan") + assert record is not None + bridge = _FastHeartbeatBridge() + await bridge.publish(record.run_id, "values", {"step": 1}) + consumer = sse_consumer(bridge, record, _FakeRequest(), mgr) + assert (await anext(consumer)).startswith("event: values\n") + + await store.update_status( + record.run_id, + "error", + error="lease expired", + stop_reason=ORPHAN_RECOVERY_STOP_REASON, + ) + + end_frame = await asyncio.wait_for(anext(consumer), timeout=1.0) + assert end_frame == "event: end\ndata: null\n\n" + + asyncio.run(run()) + + def test_wait_uses_explicit_orphan_recovery_liveness_boundary(self) -> None: + """The non-streaming consumer shares the recovered-orphan boundary.""" + from app.gateway.services import wait_for_run_completion + from deerflow.runtime.runs.store.memory import MemoryRunStore + + async def run() -> None: + store = MemoryRunStore() + await store.put( + "periodic-orphan", + thread_id=THREAD_ID, + status="running", + ) + mgr = RunManager(store=store) + record = await mgr.get("periodic-orphan") + assert record is not None + bridge = _FastHeartbeatBridge() + await bridge.publish(record.run_id, "values", {"step": 1}) + await store.update_status( + record.run_id, + "error", + error="lease expired", + stop_reason=ORPHAN_RECOVERY_STOP_REASON, + ) + + completed = await asyncio.wait_for( + wait_for_run_completion(bridge, record, _FakeRequest(), mgr), + timeout=1.0, + ) + + assert completed is True + + asyncio.run(run())