From e1430c832c9ef7c7c6418a3e21923e2a96e9b89f Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Sat, 24 Oct 2026 11:36:06 +0800 Subject: [PATCH] fix(extensions): make notify dispatch admission atomic with shutdown (#7383) * fix(extensions): fence notify admission against shutdown * test(extensions): cover notify admission shutdown race * docs(extensions): record atomic notification admission * style(tests): format extension notify shutdown regression --------- Co-authored-by: Willem Jiang --- .../harness/deerflow/extensions/AGENTS.md | 6 +- .../harness/deerflow/extensions/notify.py | 81 +++++++++++------- ..._extension_notify_shutdown_cancellation.py | 82 +++++++++++++++++++ 3 files changed, 137 insertions(+), 32 deletions(-) diff --git a/backend/packages/harness/deerflow/extensions/AGENTS.md b/backend/packages/harness/deerflow/extensions/AGENTS.md index 22a06a1521..121eee95a3 100644 --- a/backend/packages/harness/deerflow/extensions/AGENTS.md +++ b/backend/packages/harness/deerflow/extensions/AGENTS.md @@ -284,7 +284,11 @@ snapshot immutable in fact rather than only by declaration. Gateway registers on hooks and async system observations are dispatched to that loop even when the caller is a subagent's isolated loop, while synchronous system callbacks submit fire-and-forget work there. Shutdown stops accepting detached observations before the memory shutdown flush and -resets the loop only after in-flight run/subagent drain ordering is complete. +resets the loop only after in-flight run/subagent drain ordering is complete. The +synchronous dispatch admission check, submission, and pending-future registration share +a process-local lock with suspension; draining snapshots the set only after suspension +acquires that lock, so an accepted observation cannot be registered after the shutdown +snapshot. Completion callbacks take the same lock before removing futures. `ContextCompactionObserver` reports the one moment a lossy context transform can still be described: `DeerFlowSummarizationMiddleware.compact_state()` / `acompact_state()` hash each diff --git a/backend/packages/harness/deerflow/extensions/notify.py b/backend/packages/harness/deerflow/extensions/notify.py index 57785e9a52..c652c7a32f 100644 --- a/backend/packages/harness/deerflow/extensions/notify.py +++ b/backend/packages/harness/deerflow/extensions/notify.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import concurrent.futures import logging +import threading import time from collections.abc import Awaitable, Callable, Coroutine, Mapping from typing import Any @@ -204,14 +205,17 @@ _notify_loop: asyncio.AbstractEventLoop | None = None _pending_dispatches: set[concurrent.futures.Future[None]] = set() _warned_no_loop = False _system_observations_enabled = True +# Serializes detached-dispatch admission with shutdown suspension and snapshots. +_dispatch_state_lock = threading.Lock() def set_extension_notify_loop(loop: asyncio.AbstractEventLoop | None) -> None: """Bind extension notifications to the loop that owns extension resources.""" global _notify_loop, _system_observations_enabled, _warned_no_loop - _notify_loop = loop - _system_observations_enabled = True - _warned_no_loop = False + with _dispatch_state_lock: + _notify_loop = loop + _system_observations_enabled = True + _warned_no_loop = False async def drain_extension_notify_dispatches(*, timeout: float = 5.0) -> None: @@ -222,7 +226,10 @@ async def drain_extension_notify_dispatches(*, timeout: float = 5.0) -> None: extension observer from holding Gateway teardown forever. """ suspend_extension_system_observations() - pending = tuple(_pending_dispatches) + # Once suspension owns the same lock as dispatch admission, no new future + # can be registered; completed futures may still be removed by callbacks. + with _dispatch_state_lock: + pending = tuple(_pending_dispatches) if not pending: return @@ -246,17 +253,19 @@ async def drain_extension_notify_dispatches(*, timeout: float = 5.0) -> None: def reset_extension_notify_loop() -> None: """Remove the process-wide loop binding during host shutdown or tests.""" global _notify_loop, _system_observations_enabled, _warned_no_loop - _notify_loop = None - _system_observations_enabled = True - _warned_no_loop = False - _pending_dispatches.clear() + with _dispatch_state_lock: + _notify_loop = None + _system_observations_enabled = True + _warned_no_loop = False + _pending_dispatches.clear() def suspend_extension_system_observations() -> None: """Drop new fire-and-forget observations while awaited hooks still drain.""" global _system_observations_enabled - if _notify_loop is not None: - _system_observations_enabled = False + with _dispatch_state_lock: + if _notify_loop is not None: + _system_observations_enabled = False async def _notify_each_on_extension_loop( @@ -472,6 +481,11 @@ async def observe_system_model_call( return response +def _discard_pending_dispatch(future: concurrent.futures.Future[None]) -> None: + with _dispatch_state_lock: + _pending_dispatches.discard(future) + + def dispatch_system_model_observation( coro: Coroutine[Any, Any, None], what: str, @@ -479,31 +493,36 @@ def dispatch_system_model_observation( """Submit a synchronous call site's observation to the registered loop.""" global _warned_no_loop - loop = _notify_loop submitted = False try: - if not _system_observations_enabled: - return False - if loop is None or not loop.is_running(): - if not _warned_no_loop: - _warned_no_loop = True - logger.warning( - "No running loop registered for extension observations; %s and later ones are dropped", + with _dispatch_state_lock: + if not _system_observations_enabled: + return False + loop = _notify_loop + if loop is None or not loop.is_running(): + if not _warned_no_loop: + _warned_no_loop = True + logger.warning( + "No running loop registered for extension observations; %s and later ones are dropped", + what, + ) + return False + try: + future = asyncio.run_coroutine_threadsafe(coro, loop) + except Exception: + logger.debug( + "Could not dispatch %s to the extension notify loop", what, + exc_info=True, ) - return False - try: - future = asyncio.run_coroutine_threadsafe(coro, loop) - except Exception: - logger.debug( - "Could not dispatch %s to the extension notify loop", - what, - exc_info=True, - ) - return False - _pending_dispatches.add(future) - future.add_done_callback(_pending_dispatches.discard) - submitted = True + return False + # Registration is part of admission: shutdown cannot suspend and + # snapshot between the enabled check and this insertion. + _pending_dispatches.add(future) + submitted = True + # Register outside the lock because add_done_callback may run inline for + # an already-completed Future; the callback acquires the same lock. + future.add_done_callback(_discard_pending_dispatch) return True finally: if not submitted: diff --git a/backend/tests/test_extension_notify_shutdown_cancellation.py b/backend/tests/test_extension_notify_shutdown_cancellation.py index 31a7ff4a71..9fa627eded 100644 --- a/backend/tests/test_extension_notify_shutdown_cancellation.py +++ b/backend/tests/test_extension_notify_shutdown_cancellation.py @@ -1,14 +1,19 @@ from __future__ import annotations import asyncio +import threading from concurrent.futures import Future import pytest +import deerflow.extensions.notify as notify_module from deerflow.extensions.notify import ( _pending_dispatches, + dispatch_system_model_observation, drain_extension_notify_dispatches, reset_extension_notify_loop, + set_extension_notify_loop, + suspend_extension_system_observations, ) @@ -62,3 +67,80 @@ async def test_extension_notify_shutdown_drain_is_bounded() -> None: if not pending.done(): pending.cancel() reset_extension_notify_loop() + + +def test_shutdown_snapshot_includes_dispatch_already_admitted(monkeypatch) -> None: + """Shutdown must wait for a dispatch that passed admission before suspension.""" + reset_extension_notify_loop() + submit_entered = threading.Event() + release_submit = threading.Event() + suspend_entered = threading.Event() + dispatch_done = threading.Event() + drain_done = threading.Event() + pending: Future[None] = Future() + + class _Loop: + @staticmethod + def is_running() -> bool: + return True + + def _submit(_coro, _loop): + submit_entered.set() + if not release_submit.wait(timeout=2): + raise TimeoutError("test did not release the simulated submission") + return pending + + original_suspend = suspend_extension_system_observations + + def _tracked_suspend() -> None: + suspend_entered.set() + original_suspend() + + monkeypatch.setattr(notify_module.asyncio, "run_coroutine_threadsafe", _submit) + monkeypatch.setattr(notify_module, "suspend_extension_system_observations", _tracked_suspend) + set_extension_notify_loop(_Loop()) + + async def _observation() -> None: + return None + + def _dispatch() -> None: + coro = _observation() + try: + accepted = dispatch_system_model_observation(coro, "shutdown-race") + if accepted: + # The fake submitter deliberately does not schedule this coroutine. + coro.close() + finally: + dispatch_done.set() + + def _drain() -> None: + try: + asyncio.run(drain_extension_notify_dispatches(timeout=1)) + finally: + drain_done.set() + + dispatcher = threading.Thread(target=_dispatch, daemon=True) + shutdown = threading.Thread(target=_drain, daemon=True) + try: + dispatcher.start() + assert submit_entered.wait(timeout=1) + shutdown.start() + assert suspend_entered.wait(timeout=1) + assert not drain_done.wait(timeout=0.05), "shutdown returned while an admitted dispatch had not been registered" + + release_submit.set() + assert dispatch_done.wait(timeout=1) + assert pending in _pending_dispatches + pending.set_result(None) + assert drain_done.wait(timeout=1) + assert pending not in _pending_dispatches + finally: + release_submit.set() + if not pending.done(): + pending.set_result(None) + dispatcher.join(timeout=1) + shutdown.join(timeout=1) + reset_extension_notify_loop() + + assert not dispatcher.is_alive() + assert not shutdown.is_alive()