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 <willem.jiang@gmail.com>

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

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

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