From cd2633725b66fc9faa9711cdf8f4911ef873d3d9 Mon Sep 17 00:00:00 2001 From: RongJie G <111257566+CorgiBoyG@users.noreply.github.com> Date: Sun, 6 Sep 2026 08:39:26 +0800 Subject: [PATCH] fix(runtime): finish terminal signaling after hook cancellation (#5191) * fix(runtime): finish terminal signaling after hook cancellation * fix(runtime): shield task-stop observer fan-out --------- Co-authored-by: CorgiBoyG --- .../harness/deerflow/runtime/AGENTS.md | 2 +- .../harness/deerflow/runtime/runs/worker.py | 63 +++++++++- .../tests/test_extension_task_lifecycle.py | 110 ++++++++++++++++++ backend/tests/test_run_worker_rollback.py | 2 + 4 files changed, 170 insertions(+), 7 deletions(-) diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index 84633b6c2..4da7ab93e 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -149,7 +149,7 @@ The first `RunManager.list_by_thread()` hydration page uses a 100-row floor or the number of required IDs, whichever is larger; missing exact runs use targeted `get()` calls. -**Terminal run cleanup explicitly breaks graph-scoped references while preserving the existing `RunRecord` grace period.** Every `agent.astream()` iterator is closed in `_stream_once`, including abort/exception/early-break paths. A close failure after an abort is warning-only and cannot replace the user-requested `interrupted` outcome; normal-completion close failures still surface, and an in-flight stream exception remains authoritative over a secondary close failure. Journal construction and cancellable preflight work (including MCP task projection and the prior-finalization wait) live inside the worker's guarded body, so cancellation before agent startup still terminalizes the run and closes its stream. `run_agent()` wraps the complete terminal-finalization sequence in an outer teardown guard, so cancellation or failure from any terminal-stage await cannot skip `RunJournal.close()`, removal of the journal, `__pregel_runtime`, and internal runtime-context values from every runnable config, or release of local graph/payload references. That guard schedules bridge cleanup, run-record cleanup, and cyclic GC even when interruption happens before the terminal stream marker or terminal publication itself fails, so neither a cancelled observer nor a delivery-backend outage can strand process-local run state. `RunJournal.flush()` clears its `_pending_progress_task` after awaiting or cancelling it; ordinary `close()` detaches the event store/progress reporter and clears callback bookkeeping only after that flush succeeds, preserving the buffer for retry on a transient store failure. A fenced worker instead calls `close(flush=False)`, which cancels pending journal work and detaches without initiating another event-store write after lease ownership is lost; its final detach runs even if a second cancellation interrupts pending-task shutdown. `RunManager.cleanup(run_id)` retains the process-local `RunRecord`, completed task, and request payload for its default 300-second local join/status window before releasing them. Durable history remains in `RunStore`; `StreamBridge` data keeps its separate 60-second late-subscriber window, and both cleanup coroutines run in a fresh empty `contextvars.Context`. A contextless full cyclic-GC pass, coalesced to at most once every 10 seconds and dispatched through the default executor, bounds the lifetime of unreachable LangGraph callback/loop cycles without synchronously walking the heap in the event-loop timer; passes taking at least 100 ms are logged at INFO because CPython GC may still impose interpreter-level pauses. +**Terminal run cleanup explicitly breaks graph-scoped references while preserving the existing `RunRecord` grace period.** Every `agent.astream()` iterator is closed in `_stream_once`, including abort/exception/early-break paths. A close failure after an abort is warning-only and cannot replace the user-requested `interrupted` outcome; normal-completion close failures still surface, and an in-flight stream exception remains authoritative over a secondary close failure. Journal construction and cancellable preflight work (including MCP task projection and the prior-finalization wait) live inside the worker's guarded body, so cancellation before agent startup still terminalizes the run and closes its stream. `run_agent()` wraps the complete terminal-finalization sequence in an outer teardown guard, so cancellation or failure from any terminal-stage await cannot skip `RunJournal.close()`, removal of the journal, `__pregel_runtime`, and internal runtime-context values from every runnable config, or release of local graph/payload references. That guard schedules bridge cleanup, run-record cleanup, and cyclic GC even when interruption happens before the terminal stream marker or terminal publication itself fails, so neither a cancelled observer nor a delivery-backend outage can strand process-local run state. A non-`Exception` `BaseException` caught while awaiting the completion hook or task-stop notification (including host-task cancellation) is deferred through the ordinary remaining finalization, with the first interruption preserved and every caught host-task `CancelledError` balanced by calling `Task.uncancel()` until the current task’s cumulative cancellation count is clear. Task-stop fan-out runs in one child task and every host wait uses `shield`, so repeated cancellation of the worker cannot cancel that fan-out or skip later observers; the worker keeps awaiting the same child task. A rogue observer that raises its own `CancelledError` remains contained by the extension dispatcher and distinguishable from host cancellation. This guarantee applies only to cancellation caught during those hook stages: clearing the finalizing barrier and publishing END remain direct awaits, so another cancellation in the subsequent critical tail retains forceful-termination semantics instead of creating an unbounded shield. If that tail completes without another interruption, the first deferred interruption is re-raised after END; a barrier-clear failure prevents END publication, while an END failure is raised after the barrier is clear. `RunJournal.flush()` clears its `_pending_progress_task` after awaiting or cancelling it; ordinary `close()` detaches the event store/progress reporter and clears callback bookkeeping only after that flush succeeds, preserving the buffer for retry on a transient store failure. A fenced worker instead calls `close(flush=False)`, which cancels pending journal work and detaches without initiating another event-store write after lease ownership is lost; its final detach runs even if a second cancellation interrupts pending-task shutdown. `RunManager.cleanup(run_id)` retains the process-local `RunRecord`, completed task, and request payload for its default 300-second local join/status window before releasing them. Durable history remains in `RunStore`; `StreamBridge` data keeps its separate 60-second late-subscriber window, and both cleanup coroutines run in a fresh empty `contextvars.Context`. A contextless full cyclic-GC pass, coalesced to at most once every 10 seconds and dispatched through the default executor, bounds the lifetime of unreachable LangGraph callback/loop cycles without synchronously walking the heap in the event-loop timer; passes taking at least 100 ms are logged at INFO because CPython GC may still impose interpreter-level pauses. **Where things live**: - `runtime/checkpoint_mode.py` — mode + snapshot-frequency freeze, marker injection, delta detection, compatibility gate, both error types diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 1685fd8cc..f4ff71f93 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -728,6 +728,36 @@ def _bind_trace_id(config: dict[str, Any], runtime_ctx: dict[str, Any]) -> str: return trace_id +def _defer_finalization_interrupt( + deferred: BaseException | None, + interrupt: BaseException, +) -> BaseException: + """Preserve the first interrupt while allowing terminal awaits to finish.""" + if isinstance(interrupt, asyncio.CancelledError): + task = asyncio.current_task() + if task is not None: + while task.cancelling(): + task.uncancel() + return deferred if deferred is not None else interrupt + + +async def _await_task_stop_after_host_cancellation( + task: asyncio.Task[None], + deferred: BaseException | None, +) -> BaseException | None: + """Wait for one task-stop fan-out despite repeated host cancellation.""" + while True: + try: + await asyncio.shield(task) + return deferred + except asyncio.CancelledError as exc: + host = asyncio.current_task() + if host is None or not host.cancelling(): + # The fan-out task itself was cancelled rather than the host. + raise + deferred = _defer_finalization_interrupt(deferred, exc) + + async def run_agent( bridge: StreamBridge, run_manager: RunManager, @@ -768,7 +798,7 @@ async def run_agent( extensions = ctx.extensions if ctx.extensions is not None else get_loaded_extensions() task_store: ExtensionData | None = None task_info: TaskInfo | None = None - deferred_stop_interrupt: BaseException | None = None + deferred_finalization_interrupt: BaseException | None = None pre_run_checkpoint_id: str | None = None pre_run_workspace_snapshot: WorkspaceSnapshot | None = None workspace_changes_user_id: str | None = None @@ -1534,12 +1564,23 @@ async def run_agent( await ctx.on_run_completed(record) except Exception: logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True) + except BaseException as exc: + # A terminal hook must not leave replacement runs blocked or + # stream consumers waiting indefinitely. + deferred_finalization_interrupt = _defer_finalization_interrupt( + deferred_finalization_interrupt, + exc, + ) + logger.warning( + "Run completion hook interrupted for %s; completing finalization first", + run_id, + ) if task_info is not None and task_store is not None: # Keep the finalizing barrier held until stop observers finish, so # a same-thread replacement cannot overlap this task's lifecycle. - try: - await notify_task_stop( + task_stop = asyncio.create_task( + notify_task_stop( extensions, task_store, task_info, @@ -1548,6 +1589,13 @@ async def run_agent( succeeded=record.status == RunStatus.success, ), timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS, + ), + name=f"extension-task-stop-{run_id}", + ) + try: + deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation( + task_stop, + deferred_finalization_interrupt, ) except Exception: logger.warning( @@ -1558,7 +1606,10 @@ async def run_agent( except BaseException as exc: # Cancellation here must not strand the finalizing barrier or # leave stream consumers waiting for the end frame. - deferred_stop_interrupt = exc + deferred_finalization_interrupt = _defer_finalization_interrupt( + deferred_finalization_interrupt, + exc, + ) logger.warning( "Extension task-stop notification interrupted for run %s; completing cleanup first", run_id, @@ -1568,8 +1619,8 @@ async def run_agent( await bridge.publish_end(run_id) - if deferred_stop_interrupt is not None: - raise deferred_stop_interrupt + if deferred_finalization_interrupt is not None: + raise deferred_finalization_interrupt finally: try: if journal is not None: diff --git a/backend/tests/test_extension_task_lifecycle.py b/backend/tests/test_extension_task_lifecycle.py index 7e227f4a3..15b4d3b89 100644 --- a/backend/tests/test_extension_task_lifecycle.py +++ b/backend/tests/test_extension_task_lifecycle.py @@ -359,6 +359,116 @@ async def test_lead_stop_interrupt_is_deferred_until_final_cleanup(): bridge.publish_end.assert_awaited_once_with(record.run_id) +@pytest.mark.asyncio +async def test_completion_cancellation_does_not_skip_survivor_after_rogue_stop_cancel(): + completion_hook_entered = asyncio.Event() + + async def _block_completion(_record): + completion_hook_entered.set() + await asyncio.wait_for(asyncio.Event().wait(), timeout=1) + + class _Rogue: + async def on_task_stop(self, app_store, task_store, info, outcome): + raise asyncio.CancelledError() + + survivor = _RunRecorder() + manager = RunManager() + record = await manager.create("thread-completion-cancel-rogue-stop") + bridge = _bridge() + task = asyncio.create_task( + run_agent( + bridge, + manager, + record, + ctx=RunContext( + checkpointer=InMemorySaver(), + extensions=_extensions(_Rogue(), survivor), + on_run_completed=_block_completion, + ), + agent_factory=lambda *, config: _OkAgent(), + graph_input={}, + config={}, + ) + ) + + await asyncio.wait_for(completion_hook_entered.wait(), timeout=1) + task.cancel("first completion cancellation") + + with pytest.raises(asyncio.CancelledError, match="first completion cancellation"): + await asyncio.wait_for(task, timeout=1) + + assert survivor.events[-1] == ("stop", record.run_id, "completed") + assert task.cancelling() == 0 + assert record.finalizing is False + bridge.publish_end.assert_awaited_once_with(record.run_id) + + +@pytest.mark.asyncio +async def test_second_real_stop_cancellation_finishes_all_observers_and_cleanup(): + class _CleanupTrackingRunManager(RunManager): + def __init__(self) -> None: + super().__init__() + self.cleanup_calls: list[tuple[str, float]] = [] + + async def cleanup(self, run_id: str, *, delay: float = 300) -> None: + self.cleanup_calls.append((run_id, delay)) + + completion_hook_entered = asyncio.Event() + first_stop_entered = asyncio.Event() + release_first_stop = asyncio.Event() + + async def _block_completion(_record): + completion_hook_entered.set() + await asyncio.Event().wait() + + class _FirstStop(_RunRecorder): + async def on_task_stop(self, app_store, task_store, info, outcome): + await super().on_task_stop(app_store, task_store, info, outcome) + first_stop_entered.set() + await release_first_stop.wait() + + first = _FirstStop() + second = _RunRecorder() + manager = _CleanupTrackingRunManager() + record = await manager.create("thread-double-finalization-cancel") + bridge = _bridge() + task = asyncio.create_task( + run_agent( + bridge, + manager, + record, + ctx=RunContext( + checkpointer=InMemorySaver(), + extensions=_extensions(first, second), + on_run_completed=_block_completion, + ), + agent_factory=lambda *, config: _OkAgent(), + graph_input={}, + config={}, + ) + ) + + await asyncio.wait_for(completion_hook_entered.wait(), timeout=1) + task.cancel("first completion cancellation") + await asyncio.wait_for(first_stop_entered.wait(), timeout=1) + task.cancel("second task-stop cancellation") + await asyncio.sleep(0) + release_first_stop.set() + + with pytest.raises(asyncio.CancelledError, match="first completion cancellation"): + await asyncio.wait_for(task, timeout=1) + await asyncio.sleep(0) + + expected = [("stop", record.run_id, "completed")] + assert first.events[-1:] == expected + assert second.events[-1:] == expected + assert task.cancelling() == 0 + assert record.finalizing is False + bridge.publish_end.assert_awaited_once_with(record.run_id) + bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60) + assert manager.cleanup_calls == [(record.run_id, 300)] + + @pytest.mark.asyncio async def test_contributor_raising_cancellederror_cannot_interrupt_run_cleanup(): # Fail-open is decided by origin, not base class: a contributor that lets a diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 0f0783be6..77bbaed10 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -1031,6 +1031,8 @@ async def test_run_agent_schedules_terminal_cleanup_when_completion_hook_is_canc assert journal not in config["callbacks"] assert journal._closed is True assert journal._store is None + assert record.finalizing is False + bridge.publish_end.assert_awaited_once_with(record.run_id) bridge.cleanup.assert_awaited_once_with(record.run_id, delay=60) assert run_manager.cleanup_calls == [(record.run_id, 300)] schedule_collection.assert_called_once_with()