diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index d24ae7845f..507c24f420 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/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`) diff --git a/backend/packages/harness/deerflow/runtime/cancellation.py b/backend/packages/harness/deerflow/runtime/cancellation.py index 6e2cd99e89..6dd3d44e3a 100644 --- a/backend/packages/harness/deerflow/runtime/cancellation.py +++ b/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)) diff --git a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py index 125373feb4..72a2286a56 100644 --- a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py +++ b/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() diff --git a/backend/packages/harness/deerflow/runtime/store/async_provider.py b/backend/packages/harness/deerflow/runtime/store/async_provider.py index 4ef4967f1b..bddfe7a6d4 100644 --- a/backend/packages/harness/deerflow/runtime/store/async_provider.py +++ b/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 diff --git a/backend/tests/test_runtime_cancellation.py b/backend/tests/test_runtime_cancellation.py index 929540cfa0..7a1cb55e2c 100644 --- a/backend/tests/test_runtime_cancellation.py +++ b/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) diff --git a/backend/tests/test_runtime_db_provider_close_cancellation.py b/backend/tests/test_runtime_db_provider_close_cancellation.py new file mode 100644 index 0000000000..b1505adc66 --- /dev/null +++ b/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)