fix(extensions): drain pending notify dispatches on shutdown (#6357)

* fix(extensions): drain pending notify dispatches on shutdown

* fix(gateway): await extension notify drain before reset

* test(extensions): retain pending notify dispatches on shutdown

* fix(extensions): bound and close notify shutdown drain

* test(extensions): cover bounded notify shutdown drain

backend/app/gateway/deps.py

@@ -535,6 +535,7 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
         # startup-failure and cancellation path below.
         try:
             from deerflow.extensions.notify import (
+                drain_extension_notify_dispatches,
                 reset_extension_notify_loop,
                 set_extension_notify_loop,
             )
@@ -544,16 +545,24 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
             logger.exception("Failed to register the extension notify loop; sync observations will be dropped")
         else:
 
-            def reset_notify_loop_safely() -> None:
+            async def reset_notify_loop_safely() -> None:
                 try:
-                    reset_extension_notify_loop()
+                    await drain_extension_notify_dispatches()
                 except Exception:
                     logger.debug(
-                        "Failed to reset the extension notify loop (non-fatal)",
+                        "Failed to drain pending extension notifications (non-fatal)",
                         exc_info=True,
                     )
+                finally:
+                    try:
+                        reset_extension_notify_loop()
+                    except Exception:
+                        logger.debug(
+                            "Failed to reset the extension notify loop (non-fatal)",
+                            exc_info=True,
+                        )
 
-            stack.callback(reset_notify_loop_safely)
+            stack.push_async_callback(reset_notify_loop_safely)
 
         config = startup_config
         app.state.checkpoint_channel_mode = freeze_checkpoint_channel_mode(config.database.checkpoint_channel_mode)

backend/packages/harness/deerflow/extensions/notify.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import logging
 import time
 from collections.abc import Awaitable, Callable, Coroutine, Mapping
@@ -20,6 +21,7 @@ from deerflow_extension_api import (
 )
 
 from deerflow.extensions.registry import LoadedExtensions
+from deerflow.utils.file_io import await_drained
 
 logger = logging.getLogger(__name__)
 
@@ -182,7 +184,7 @@ async def _notify_each(
 # loops, but extension resources must always be touched on the loop where they
 # were started.
 _notify_loop: asyncio.AbstractEventLoop | None = None
-_pending_dispatches: set[asyncio.Future[Any]] = set()
+_pending_dispatches: set[concurrent.futures.Future[None]] = set()
 _warned_no_loop = False
 _system_observations_enabled = True
 
@@ -195,6 +197,35 @@ def set_extension_notify_loop(loop: asyncio.AbstractEventLoop | None) -> None:
     _warned_no_loop = False
 
 
+async def drain_extension_notify_dispatches(*, timeout: float = 5.0) -> None:
+    """Boundedly drain already-submitted fire-and-forget observations.
+
+    New detached observations are suspended before the snapshot so shutdown
+    cannot open a submit-after-snapshot window. The bounded wait keeps a wedged
+    extension observer from holding Gateway teardown forever.
+    """
+    suspend_extension_system_observations()
+    pending = tuple(_pending_dispatches)
+    if not pending:
+        return
+
+    async def _wait() -> None:
+        try:
+            async with asyncio.timeout(timeout):
+                await asyncio.gather(
+                    *(asyncio.wrap_future(future) for future in pending),
+                    return_exceptions=True,
+                )
+        except TimeoutError:
+            logger.warning(
+                "Timed out after %.1fs draining %d pending extension observation(s)",
+                timeout,
+                len(pending),
+            )
+
+    await await_drained(_wait())
+
+
 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

backend/tests/test_extension_notify_shutdown_cancellation.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import asyncio
+from concurrent.futures import Future
+
+import pytest
+
+from deerflow.extensions.notify import (
+    _pending_dispatches,
+    drain_extension_notify_dispatches,
+    reset_extension_notify_loop,
+)
+
+
+@pytest.mark.asyncio
+async def test_extension_notify_shutdown_drains_pending_dispatch_across_repeated_cancellation() -> None:
+    reset_extension_notify_loop()
+    pending: Future[None] = Future()
+    _pending_dispatches.add(pending)
+    pending.add_done_callback(_pending_dispatches.discard)
+
+    task = asyncio.create_task(drain_extension_notify_dispatches())
+    try:
+        await asyncio.sleep(0)
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done()
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done()
+
+        pending.set_result(None)
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        if not pending.done():
+            pending.set_result(None)
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+        reset_extension_notify_loop()
+
+    assert pending not in _pending_dispatches
+
+
+@pytest.mark.asyncio
+async def test_extension_notify_shutdown_drain_is_bounded() -> None:
+    reset_extension_notify_loop()
+    pending: Future[None] = Future()
+    _pending_dispatches.add(pending)
+    pending.add_done_callback(_pending_dispatches.discard)
+
+    try:
+        await asyncio.wait_for(
+            drain_extension_notify_dispatches(timeout=0.01),
+            timeout=0.5,
+        )
+    finally:
+        if not pending.done():
+            pending.cancel()
+        reset_extension_notify_loop()