From 737492d4e173edc9c78fa25eedf5753de7bb3f0e Mon Sep 17 00:00:00 2001 From: ordinary being Date: Sat, 10 Oct 2026 19:59:55 +0800 Subject: [PATCH] fix(runtime): drain staged-terminal runs on shutdown before checkpointer teardown (#6620) * fix(runtime): drain staged-terminal runs on shutdown before checkpointer teardown RunManager.shutdown() selected only pending/running records, but with an event store a run stages its terminal status in memory and keeps finalizing: journal flush, delivery receipt, workspace scan, the duration checkpoint write and the deferred terminal commit, all while terminal_commit_pending is set. A Gateway shutdown landing in that window did not await the run, so langgraph_runtime could close the checkpointer before the duration write landed - the #3373 class the drain exists to prevent. Drain those records too, using the existing _awaits_terminal_commit state that _renew_leases already relies on; await them without cancelling (a cancel would skip their terminal tail) and never overwrite their staged terminal status with interrupted. * fix(runtime): surface shutdown terminal finalizer failures Cover timed-out staged terminal runs for both commit barriers and terminal outcomes. Log failed finalizers with their run ID, staged status, and traceback while preserving status and completing the drain. --------- Co-authored-by: Willem Jiang --- .../harness/deerflow/runtime/AGENTS.md | 8 +- .../harness/deerflow/runtime/runs/manager.py | 28 ++- .../tests/test_gateway_run_drain_shutdown.py | 164 ++++++++++++++++++ 3 files changed, 197 insertions(+), 3 deletions(-) diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index a699fe54e9..d2e0c07949 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -77,7 +77,13 @@ status locally and commits it only after finalization's receipt and duration writes. `RunRecord.terminal_commit_pending` keeps `_renew_leases()` renewing that still-active row until the commit is attempted; a renewal rejected by the worker's own commit is confirmed by re-reading the row, while a peer claim fences the run. -Never select runs for renewal by local status alone. +Never select runs for renewal by local status alone. Shutdown's drain treats the +same state as active: such runs are awaited (never cancelled — cancellation would +skip their terminal tail) so the duration checkpoint and terminal commit land +before checkpointer teardown. A drain timeout leaves staged statuses and commit +barriers intact without persisting `interrupted`. A finalizer that raises is +logged with its run ID, staged status and exception while the other runs continue +draining. Coverage: `tests/test_gateway_run_drain_shutdown.py`. **Deferred-tool promotion event deduplication** (`runtime/journal.py`): one `RunJournal` owns the lead graph's run-scoped atomic promotion claim. Parallel diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 7af7170367..589b29d715 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -2510,7 +2510,12 @@ class RunManager: resources are still open. Only runs that do **not** settle on their own are marked ``interrupted`` — a run that completes (e.g. ``success``) during the drain keeps its real terminal status instead of being - blanket-overwritten. The whole drain, including the trailing status + blanket-overwritten. Runs whose terminal status is staged in memory with + a deferred commit are still writing their final duration checkpoint and + then commit that status, so they are awaited in the same bounded drain — + without cancellation, which would skip their terminal tail (#5542), + exactly as ``_renew_leases`` keeps them active through + ``_awaits_terminal_commit``. The whole drain, including the trailing status persistence, is bounded by ``timeout`` so a run stuck in cleanup (or a slow store under DB pressure) cannot hang worker shutdown — the precondition for the signal-reentrancy deadlock guarded by @@ -2522,12 +2527,20 @@ class RunManager: async with self._lock: inflight = [record for record in self._runs.values() if record.status in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()] + # A run whose terminal status is staged in memory (deferred commit) + # keeps writing the duration checkpoint before committing that + # status. Drain it too — but never cancel it — so the write lands + # while the checkpointer is still open. ``_renew_leases`` already + # treats these runs as active through ``_awaits_terminal_commit``. + finalizing = [record for record in self._runs.values() if self._awaits_terminal_commit(record) and record.status not in (RunStatus.pending, RunStatus.running) and record.task is not None and not record.task.done()] + finalizing_ids = {record.run_id for record in finalizing} for record in inflight: record.abort_action = "interrupt" record.abort_event.set() record.task.cancel() # type: ignore[union-attr] # filtered above # Status is decided AFTER the drain (below), not here: a run that # completes on its own during the drain must keep its real status. + inflight += finalizing await self.stop_heartbeat(timeout=max(0.0, deadline - loop.time())) @@ -2548,7 +2561,18 @@ class RunManager: if task not in pending and not task.cancelled(): # Completed on its own — retrieve any surfaced exception so it # is not reported as "never retrieved", and keep its status. - task.exception() # type: ignore[union-attr] # done & not cancelled + error = task.exception() # type: ignore[union-attr] # done & not cancelled + if error is not None and record.run_id in finalizing_ids: + logger.warning( + "Run %s failed during terminal finalization on shutdown; staged status %s may not be committed", + record.run_id, + record.status.value, + exc_info=(type(error), error, error.__traceback__), + ) + continue + if record.run_id in finalizing_ids: + # The staged terminal status is committed by the run's own + # finalizer; shutdown must not overwrite it with interrupted. continue if record.status in (RunStatus.pending, RunStatus.running): record.status = RunStatus.interrupted diff --git a/backend/tests/test_gateway_run_drain_shutdown.py b/backend/tests/test_gateway_run_drain_shutdown.py index 22a72b4906..e1af57e77c 100644 --- a/backend/tests/test_gateway_run_drain_shutdown.py +++ b/backend/tests/test_gateway_run_drain_shutdown.py @@ -22,15 +22,18 @@ precondition for the signal-reentrancy deadlock guarded by from __future__ import annotations import asyncio +import logging import operator from contextlib import asynccontextmanager, suppress from types import SimpleNamespace from typing import Annotated, TypedDict +from unittest.mock import AsyncMock import pytest from langgraph.checkpoint.memory import InMemorySaver from deerflow.runtime import RunManager, RunStatus +from deerflow.runtime.runs.store.memory import MemoryRunStore # Module-level so langgraph's get_type_hints (which resolves annotations against @@ -401,6 +404,167 @@ async def test_shutdown_preserves_status_of_run_completed_during_drain(): await record.task +@pytest.mark.asyncio +async def test_shutdown_awaits_staged_terminal_finalization_without_cancelling(): + """shutdown() must drain runs whose terminal status is staged in memory. + + With an event store the worker stages its terminal status (in memory, + ``terminal_commit_pending=True``) and keeps finalizing: journal flush, + delivery receipt, workspace scan, the duration checkpoint write and the + deferred terminal commit. ``langgraph_runtime`` closes the checkpointer + right after ``shutdown()``, so a finalizing run that is not drained can + still write its duration checkpoint against a closed pool (#3373 class). + The run must be awaited — never cancelled, which would skip the terminal + tail (#5542) — and its staged status must survive the drain. + """ + rm = RunManager() + record = await rm.create("t-finalizing") + await rm.set_status(record.run_id, RunStatus.running) + + finalizer_started = asyncio.Event() + allow_finish = asyncio.Event() + + async def finalizer() -> None: + finalizer_started.set() + await allow_finish.wait() + record.terminal_commit_pending = False + + record.task = asyncio.create_task(finalizer()) + try: + await asyncio.wait_for(finalizer_started.wait(), timeout=1.0) + # The worker stages the terminal status before the deferred commit. + record.status = RunStatus.success + record.terminal_commit_pending = True + + shutdown_task = asyncio.create_task(rm.shutdown(timeout=5.0)) + # The idle-manager shutdown path has no suspension points, so one + # scheduler turn is enough for the unfixed code to run to completion; + # the fixed path parks in ``asyncio.wait`` on the finalizer. + for _ in range(3): + await asyncio.sleep(0) + assert not shutdown_task.done(), "shutdown() returned without awaiting the staged-terminal finalizer" + + allow_finish.set() + await asyncio.wait_for(shutdown_task, timeout=5.0) + + assert not record.task.cancelled(), "shutdown() cancelled a staged-terminal finalizer" + assert record.terminal_commit_pending is False + assert record.status == RunStatus.success, f"shutdown overwrote the staged terminal status: {record.status}" + finally: + if not record.task.done(): + record.task.cancel() + with suppress(asyncio.CancelledError): + await record.task + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [RunStatus.success, RunStatus.error]) +@pytest.mark.parametrize("barrier", ["terminal_commit_pending", "scheduled_goal_cleanup_pending"]) +async def test_shutdown_preserves_staged_terminal_status_after_drain_timeout(monkeypatch, caplog, status, barrier): + """A blocked terminal tail must survive shutdown's bounded drain untouched.""" + store = MemoryRunStore() + rm = RunManager(store=store) + record = await rm.create("t-finalizing-timeout") + await rm.set_status(record.run_id, RunStatus.running) + await rm.set_status(record.run_id, status, persist=False) + setattr(record, barrier, True) + persist_status = AsyncMock(wraps=rm._persist_status) + monkeypatch.setattr(rm, "_persist_status", persist_status) + + started = asyncio.Event() + allow_finish = asyncio.Event() + + async def finalizer() -> None: + started.set() + await allow_finish.wait() + + record.task = asyncio.create_task(finalizer()) + try: + await asyncio.wait_for(started.wait(), timeout=1.0) + with caplog.at_level(logging.WARNING, logger="deerflow.runtime.runs.manager"): + await asyncio.wait_for(rm.shutdown(timeout=0.01), timeout=1.0) + + assert not record.task.done() + assert record.task.cancelling() == 0 + assert not record.abort_event.is_set() + assert record.status == status + assert getattr(record, barrier) is True + persist_status.assert_not_awaited() + assert (await store.get(record.run_id))["status"] == "running" + assert "run task(s) still active" in caplog.text + # Without the finalizing skip, shutdown incorrectly queues this staged + # record for interrupted persistence, even though the budget is spent. + assert "before persisting" not in caplog.text + finally: + allow_finish.set() + await asyncio.gather(record.task, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [RunStatus.success, RunStatus.error]) +async def test_shutdown_warns_when_staged_terminal_finalization_fails(monkeypatch, caplog, status): + """Surface a lost terminal tail without changing its status or aborting drain.""" + store = MemoryRunStore() + rm = RunManager(store=store) + failed = await rm.create("t-finalizing-failed") + healthy = await rm.create("t-finalizing-healthy") + for record in (failed, healthy): + await rm.set_status(record.run_id, RunStatus.running) + await rm.set_status(record.run_id, status, persist=False) + record.terminal_commit_pending = True + persist_status = AsyncMock(wraps=rm._persist_status) + monkeypatch.setattr(rm, "_persist_status", persist_status) + + drain_started = asyncio.Event() + allow_finish = asyncio.Event() + failure = RuntimeError("terminal commit failed") + stop_heartbeat = rm.stop_heartbeat + + async def observed_stop_heartbeat(*, timeout): + await stop_heartbeat(timeout=timeout) + drain_started.set() + + monkeypatch.setattr(rm, "stop_heartbeat", observed_stop_heartbeat) + + async def finalizer(record) -> None: + await allow_finish.wait() + if record is failed: + raise failure + record.terminal_commit_pending = False + + failed.task = asyncio.create_task(finalizer(failed)) + healthy.task = asyncio.create_task(finalizer(healthy)) + shutdown_task = None + try: + with caplog.at_level(logging.WARNING, logger="deerflow.runtime.runs.manager"): + shutdown_task = asyncio.create_task(rm.shutdown(timeout=1.0)) + await asyncio.wait_for(drain_started.wait(), timeout=1.0) + allow_finish.set() + await asyncio.wait_for(shutdown_task, timeout=1.0) + + assert failed.task.exception() is failure + assert failed.status == status + assert failed.terminal_commit_pending is True + assert (await store.get(failed.run_id))["status"] == "running" + assert healthy.task.done() and not healthy.task.cancelled() + assert healthy.terminal_commit_pending is False + assert healthy.status == status + persist_status.assert_not_awaited() + warnings = [entry for entry in caplog.records if entry.levelno == logging.WARNING] + assert len(warnings) == 1 + warning = warnings[0] + assert failed.run_id in warning.getMessage() + assert healthy.run_id not in warning.getMessage() + assert status.value in warning.getMessage() + assert warning.exc_info is not None and warning.exc_info[1] is failure + finally: + allow_finish.set() + tasks = [failed.task, healthy.task] + if shutdown_task is not None: + tasks.append(shutdown_task) + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.asyncio async def test_shutdown_surfaces_failed_interrupted_persist(caplog): """A failed interrupted-status persist during the drain must be surfaced (with