fix(runtime): drain DB provider context exits across cancellation (#5757)

* fix(runtime): drain DB provider context exits

* test(runtime): pin drained async context semantics

* style(runtime): format drained context tests

* fix(runtime): preserve drained context exit failures

* style(runtime): format exit failure regression

backend/packages/harness/deerflow/runtime/AGENTS.md

@@ -1,6 +1,6 @@
 ### Stream Bridge Heartbeats
 
-Memory and Redis bridges keep the startup-only `stream_bridge.heartbeat_interval_seconds` default on the instance; explicit `subscribe(..., heartbeat_interval=...)` overrides it per subscription. Provider context managers own cache/bridge backends through exit: drain `aclose()` / `close()` across caller cancellation before propagating cancellation.
+Memory/Redis bridges keep startup-only `stream_bridge.heartbeat_interval_seconds`; explicit `subscribe(..., heartbeat_interval=...)` overrides it. Provider contexts retain backend ownership through exit: drain cache/bridge `aclose()`/`close()` and SQLite/PostgreSQL checkpointer/Store `__aexit__` across cancellation before it propagates.
 
 ### Checkpoint Channel Modes (`full` / `delta`)
 

backend/packages/harness/deerflow/runtime/cancellation.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncIterator, Awaitable
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from typing import TypeVar
 
 T = TypeVar("T")
@@ -22,3 +24,40 @@ async def wait_for_task_until(  # noqa: UP047
         if task in done:
             return True
     return True
+
+
+async def _drain_context_exit(awaitable: Awaitable[bool | None]) -> bool | None:
+    """Drain an async context exit without hiding a teardown failure.
+
+    If caller cancellation arrives while __aexit__ is running, keep waiting
+    through repeated cancellation. A later exit failure replaces/chains the
+    cancellation, matching normal context-manager semantics; only a successful
+    exit re-raises the caller cancellation.
+    """
+    task = asyncio.ensure_future(awaitable)
+    try:
+        return await asyncio.shield(task)
+    except asyncio.CancelledError:
+        while not task.done():
+            try:
+                await asyncio.wait({task})
+            except asyncio.CancelledError:
+                continue
+        task.result()
+        raise
+
+
+@asynccontextmanager
+async def drained_async_context[T](
+    manager: AbstractAsyncContextManager[T],
+) -> AsyncIterator[T]:
+    """Keep an entered async context owned until its exit fully settles."""
+    value = await manager.__aenter__()
+    try:
+        yield value
+    except BaseException as exc:
+        suppressed = await _drain_context_exit(manager.__aexit__(type(exc), exc, exc.__traceback__))
+        if not suppressed:
+            raise
+    else:
+        await _drain_context_exit(manager.__aexit__(None, None, None))

backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py

@@ -26,6 +26,7 @@ from langgraph.types import Checkpointer
 
 from deerflow.config.app_config import AppConfig, get_app_config
 from deerflow.persistence.postgres_schema import create_schema_sql, dsn_with_search_path, normalize_libpq_dsn
+from deerflow.runtime.cancellation import drained_async_context
 from deerflow.runtime.checkpointer.provider import (
     POSTGRES_CONN_REQUIRED,
     POSTGRES_INSTALL,
@@ -121,7 +122,7 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
             raise ImportError(SQLITE_INSTALL) from exc
 
         conn_str = await asyncio.to_thread(_prepare_sqlite_checkpointer_path, config.connection_string or "store.db")
-        async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
+        async with drained_async_context(AsyncSqliteSaver.from_conn_string(conn_str)) as saver:
             await saver.setup()
             yield saver
         return
@@ -132,7 +133,7 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
 
         AsyncPostgresSaver, _ = _ensure_postgres_imports()
         pool = _build_postgres_pool(config.connection_string, config.postgres_schema)
-        async with pool:
+        async with drained_async_context(pool):
             await _ensure_postgres_schema_with_pool(pool, config.postgres_schema)
             saver = AsyncPostgresSaver(conn=pool)
             await saver.setup()
@@ -163,7 +164,7 @@ async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpoi
             raise ImportError(SQLITE_INSTALL) from exc
 
         conn_str = await asyncio.to_thread(_prepare_database_sqlite_checkpointer_path, db_config)
-        async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
+        async with drained_async_context(AsyncSqliteSaver.from_conn_string(conn_str)) as saver:
             await saver.setup()
             yield saver
         return
@@ -174,7 +175,7 @@ async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpoi
 
         AsyncPostgresSaver, _ = _ensure_postgres_imports()
         pool = _build_postgres_pool(db_config.postgres_url, db_config.postgres_schema)
-        async with pool:
+        async with drained_async_context(pool):
             await _ensure_postgres_schema_with_pool(pool, db_config.postgres_schema)
             saver = AsyncPostgresSaver(conn=pool)
             await saver.setup()

backend/packages/harness/deerflow/runtime/store/async_provider.py

@@ -26,6 +26,7 @@ from langgraph.store.base import BaseStore
 
 from deerflow.config.app_config import AppConfig, get_app_config
 from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema_async
+from deerflow.runtime.cancellation import drained_async_context
 from deerflow.runtime.store.provider import (
     POSTGRES_CONN_REQUIRED,
     POSTGRES_STORE_INSTALL,
@@ -71,7 +72,7 @@ async def _async_store(config) -> AsyncIterator[BaseStore]:
         conn_str = resolve_sqlite_conn_str(config.connection_string or "store.db")
         await asyncio.to_thread(ensure_sqlite_parent_dir, conn_str)
 
-        async with AsyncSqliteStore.from_conn_string(conn_str) as store:
+        async with drained_async_context(AsyncSqliteStore.from_conn_string(conn_str)) as store:
             await store.setup()
             logger.info("Store: using AsyncSqliteStore (%s)", conn_str)
             yield store
@@ -88,7 +89,7 @@ async def _async_store(config) -> AsyncIterator[BaseStore]:
 
         await _ensure_postgres_schema(config.connection_string, config.postgres_schema)
         conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
-        async with AsyncPostgresStore.from_conn_string(conn_string) as store:
+        async with drained_async_context(AsyncPostgresStore.from_conn_string(conn_string)) as store:
             await store.setup()
             logger.info("Store: using AsyncPostgresStore")
             yield store

backend/tests/test_runtime_cancellation.py

@@ -4,7 +4,7 @@ from types import SimpleNamespace
 import pytest
 
 import deerflow.runtime.cancellation as cancellation
-from deerflow.runtime.cancellation import wait_for_task_until
+from deerflow.runtime.cancellation import drained_async_context, wait_for_task_until
 
 
 @pytest.mark.anyio
@@ -89,3 +89,137 @@ async def test_wait_for_task_until_repeated_cancellation_keeps_original_deadline
     assert child.done() is False
     event.set()
     await child
+
+
+class _RecordingAsyncContext:
+    def __init__(self, *, suppress: bool = False, exit_error: Exception | None = None) -> None:
+        self.suppress = suppress
+        self.exit_error = exit_error
+        self.exit_args = None
+
+    async def __aenter__(self):
+        return "resource"
+
+    async def __aexit__(self, exc_type, exc, tb) -> bool:
+        self.exit_args = (exc_type, exc, tb)
+        if self.exit_error is not None:
+            raise self.exit_error
+        return self.suppress
+
+
+@pytest.mark.anyio
+async def test_drained_async_context_suppresses_when_exit_returns_true():
+    manager = _RecordingAsyncContext(suppress=True)
+
+    async with drained_async_context(manager):
+        raise asyncio.CancelledError()
+
+    assert manager.exit_args is not None
+    assert manager.exit_args[0] is asyncio.CancelledError
+
+
+@pytest.mark.anyio
+async def test_drained_async_context_reraises_when_exit_returns_false():
+    manager = _RecordingAsyncContext()
+
+    with pytest.raises(ValueError, match="body failed"):
+        async with drained_async_context(manager):
+            raise ValueError("body failed")
+
+    assert manager.exit_args is not None
+    assert manager.exit_args[0] is ValueError
+
+
+@pytest.mark.anyio
+async def test_drained_async_context_exit_error_replaces_body_error():
+    manager = _RecordingAsyncContext(exit_error=RuntimeError("exit failed"))
+
+    with pytest.raises(RuntimeError, match="exit failed") as exc_info:
+        async with drained_async_context(manager):
+            raise ValueError("body failed")
+
+    assert isinstance(exc_info.value.__context__, ValueError)
+
+
+@pytest.mark.anyio
+async def test_drained_async_context_drains_repeated_cancellation_during_exit():
+    exit_started = asyncio.Event()
+    allow_exit = asyncio.Event()
+    exit_finished = asyncio.Event()
+    entered = asyncio.Event()
+
+    class _BlockingContext:
+        async def __aenter__(self):
+            return "resource"
+
+        async def __aexit__(self, _exc_type, _exc, _tb) -> bool:
+            exit_started.set()
+            await allow_exit.wait()
+            exit_finished.set()
+            return False
+
+    async def owner() -> None:
+        async with drained_async_context(_BlockingContext()):
+            entered.set()
+            await asyncio.Event().wait()
+
+    task = asyncio.create_task(owner())
+    try:
+        await asyncio.wait_for(entered.wait(), timeout=1)
+        task.cancel()
+        await asyncio.wait_for(exit_started.wait(), timeout=1)
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done()
+
+        allow_exit.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert exit_finished.is_set()
+    finally:
+        allow_exit.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
+@pytest.mark.anyio
+async def test_drained_async_context_exit_failure_replaces_cancellation_during_exit():
+    exit_started = asyncio.Event()
+    allow_exit = asyncio.Event()
+    entered = asyncio.Event()
+    leave = asyncio.Event()
+
+    class _FailingExitContext:
+        async def __aenter__(self):
+            return "resource"
+
+        async def __aexit__(self, _exc_type, _exc, _tb) -> bool:
+            exit_started.set()
+            await allow_exit.wait()
+            raise RuntimeError("exit failed")
+
+    async def owner() -> None:
+        async with drained_async_context(_FailingExitContext()):
+            entered.set()
+            await leave.wait()
+
+    task = asyncio.create_task(owner())
+    try:
+        await asyncio.wait_for(entered.wait(), timeout=1)
+        leave.set()
+        await asyncio.wait_for(exit_started.wait(), timeout=1)
+        task.cancel()
+        await asyncio.sleep(0)
+        assert not task.done()
+
+        allow_exit.set()
+        with pytest.raises(RuntimeError, match="exit failed") as exc_info:
+            await task
+        assert isinstance(exc_info.value.__context__, asyncio.CancelledError)
+    finally:
+        allow_exit.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)

backend/tests/test_runtime_db_provider_close_cancellation.py

@@ -0,0 +1,180 @@
+import asyncio
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+
+from deerflow.runtime.checkpointer import async_provider as checkpointer_provider
+from deerflow.runtime.store import async_provider as store_provider
+
+
+class _BlockingAsyncContext:
+    def __init__(self, value=None) -> None:
+        self.value = self if value is None else value
+        self.exit_started = asyncio.Event()
+        self.allow_exit = asyncio.Event()
+        self.exit_finished = asyncio.Event()
+
+    async def __aenter__(self):
+        return self.value
+
+    async def __aexit__(self, exc_type, exc, tb) -> bool:
+        self.exit_started.set()
+        await self.allow_exit.wait()
+        self.exit_finished.set()
+        return False
+
+
+class _SetupResource:
+    serde = object()
+
+    async def setup(self) -> None:
+        return None
+
+
+def _install_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs) -> None:
+    module = ModuleType(name)
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    monkeypatch.setitem(sys.modules, name, module)
+
+
+async def _assert_context_exit_is_drained(cm, resource: _BlockingAsyncContext) -> None:
+    entered = asyncio.Event()
+    leave = asyncio.Event()
+
+    async def owner() -> None:
+        async with cm:
+            entered.set()
+            await leave.wait()
+
+    task: asyncio.Task[None] | None = None
+    try:
+        task = asyncio.create_task(owner())
+        await asyncio.wait_for(entered.wait(), timeout=1)
+        leave.set()
+        await asyncio.wait_for(resource.exit_started.wait(), timeout=1)
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done(), "provider teardown returned before backend __aexit__ finished"
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done(), "repeated cancellation interrupted backend __aexit__"
+
+        resource.allow_exit.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+        assert resource.exit_finished.is_set()
+    finally:
+        resource.allow_exit.set()
+        if task is not None and not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("unified", [False, True])
+async def test_sqlite_checkpointer_context_exit_drains_across_repeated_cancellation(monkeypatch: pytest.MonkeyPatch, unified: bool) -> None:
+    resource = _BlockingAsyncContext(_SetupResource())
+
+    class _FakeAsyncSqliteSaver:
+        @classmethod
+        def from_conn_string(cls, _conn_str: str):
+            return resource
+
+    _install_module(
+        monkeypatch,
+        "langgraph.checkpoint.sqlite.aio",
+        AsyncSqliteSaver=_FakeAsyncSqliteSaver,
+    )
+
+    if unified:
+        config = SimpleNamespace(backend="sqlite", checkpointer_sqlite_path=":memory:")
+        cm = checkpointer_provider._async_checkpointer_from_database(config)
+    else:
+        config = SimpleNamespace(type="sqlite", connection_string=":memory:")
+        cm = checkpointer_provider._async_checkpointer(config)
+
+    await _assert_context_exit_is_drained(cm, resource)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("unified", [False, True])
+async def test_postgres_checkpointer_context_exit_drains_across_repeated_cancellation(monkeypatch: pytest.MonkeyPatch, unified: bool) -> None:
+    resource = _BlockingAsyncContext()
+
+    class _FakeAsyncPostgresSaver(_SetupResource):
+        def __init__(self, *, conn) -> None:
+            assert conn is resource
+
+    monkeypatch.setattr(checkpointer_provider, "_build_postgres_pool", lambda *_args, **_kwargs: resource)
+    monkeypatch.setattr(
+        checkpointer_provider,
+        "_ensure_postgres_imports",
+        lambda: (_FakeAsyncPostgresSaver, object()),
+    )
+
+    async def _no_schema(_pool, _schema: str) -> None:
+        return None
+
+    monkeypatch.setattr(checkpointer_provider, "_ensure_postgres_schema_with_pool", _no_schema)
+
+    if unified:
+        config = SimpleNamespace(backend="postgres", postgres_url="postgresql://example/db", postgres_schema="")
+        cm = checkpointer_provider._async_checkpointer_from_database(config)
+    else:
+        config = SimpleNamespace(type="postgres", connection_string="postgresql://example/db", postgres_schema="")
+        cm = checkpointer_provider._async_checkpointer(config)
+
+    await _assert_context_exit_is_drained(cm, resource)
+
+
+@pytest.mark.asyncio
+async def test_sqlite_store_context_exit_drains_across_repeated_cancellation(monkeypatch: pytest.MonkeyPatch) -> None:
+    resource = _BlockingAsyncContext(_SetupResource())
+
+    class _FakeAsyncSqliteStore:
+        @classmethod
+        def from_conn_string(cls, _conn_str: str):
+            return resource
+
+    _install_module(
+        monkeypatch,
+        "langgraph.store.sqlite.aio",
+        AsyncSqliteStore=_FakeAsyncSqliteStore,
+    )
+    config = SimpleNamespace(type="sqlite", connection_string=":memory:")
+
+    await _assert_context_exit_is_drained(store_provider._async_store(config), resource)
+
+
+@pytest.mark.asyncio
+async def test_postgres_store_context_exit_drains_across_repeated_cancellation(monkeypatch: pytest.MonkeyPatch) -> None:
+    resource = _BlockingAsyncContext(_SetupResource())
+
+    class _FakeAsyncPostgresStore:
+        @classmethod
+        def from_conn_string(cls, _conn_str: str):
+            return resource
+
+    _install_module(
+        monkeypatch,
+        "langgraph.store.postgres.aio",
+        AsyncPostgresStore=_FakeAsyncPostgresStore,
+    )
+
+    async def _no_schema(_conn_string: str, _schema: str) -> None:
+        return None
+
+    monkeypatch.setattr(store_provider, "_ensure_postgres_schema", _no_schema)
+    config = SimpleNamespace(
+        type="postgres",
+        connection_string="postgresql://example/db",
+        postgres_schema="",
+    )
+
+    await _assert_context_exit_is_drained(store_provider._async_store(config), resource)