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

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

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:

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()