diff --git a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py index 72a2286a56..0e73dff1a5 100644 --- a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py +++ b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py @@ -82,7 +82,10 @@ async def _ensure_postgres_schema_with_pool(pool, schema: str) -> None: statement = create_schema_sql(schema) if statement is None: return - async with pool.connection() as conn: + # Drain the connection's exit like the pool context around it: caller + # cancellation must not leave the connection checked out while the pool's + # own teardown runs (backend-ownership invariant in runtime/AGENTS.md). + async with drained_async_context(pool.connection()) as conn: await conn.execute(statement) diff --git a/backend/tests/test_runtime_db_provider_close_cancellation.py b/backend/tests/test_runtime_db_provider_close_cancellation.py index b1505adc66..ef87b098ee 100644 --- a/backend/tests/test_runtime_db_provider_close_cancellation.py +++ b/backend/tests/test_runtime_db_provider_close_cancellation.py @@ -178,3 +178,72 @@ async def test_postgres_store_context_exit_drains_across_repeated_cancellation(m ) await _assert_context_exit_is_drained(store_provider._async_store(config), resource) + + +class _PoolConnection(_BlockingAsyncContext): + """A pool connection whose exit blocks until the test allows it.""" + + def __init__(self) -> None: + super().__init__() + self.executed: list[str] = [] + + async def execute(self, statement: str) -> None: + self.executed.append(statement) + + +class _FakePool: + """A pool that records being closed while a connection is checked out.""" + + def __init__(self, connection: _PoolConnection) -> None: + self._connection = connection + self.closed_with_checked_out_connection = False + + def connection(self) -> _PoolConnection: + return self._connection + + async def __aenter__(self) -> "_FakePool": + return self + + async def __aexit__(self, exc_type, exc, tb) -> bool: + if not self._connection.exit_finished.is_set(): + self.closed_with_checked_out_connection = True + return False + + +@pytest.mark.asyncio +async def test_ensure_postgres_schema_drains_pool_connection_exit(monkeypatch: pytest.MonkeyPatch) -> None: + """The schema-setup connection must return to the pool before the pool drains. + + ``_ensure_postgres_schema_with_pool`` runs inside the drained pool context. + With a raw ``async with pool.connection()``, caller cancellation arriving + while the connection's ``__aexit__`` is blocked interrupts that exit, so the + connection is never returned and the pool's own drained teardown closes + with a checked-out connection — the backend-ownership invariant in + ``runtime/AGENTS.md``. + """ + monkeypatch.setattr(checkpointer_provider, "create_schema_sql", lambda schema: f'CREATE SCHEMA "{schema}"') + connection = _PoolConnection() + pool = _FakePool(connection) + + async def owner() -> None: + async with checkpointer_provider.drained_async_context(pool): + await checkpointer_provider._ensure_postgres_schema_with_pool(pool, "custom") + + task = asyncio.create_task(owner()) + await asyncio.wait_for(connection.exit_started.wait(), timeout=1) + assert connection.executed == ['CREATE SCHEMA "custom"'] + + task.cancel() + for _ in range(5): + await asyncio.sleep(0) + assert not task.done(), "pool teardown returned before the connection exit finished" + + task.cancel() + for _ in range(5): + await asyncio.sleep(0) + assert not task.done(), "repeated cancellation interrupted the connection exit" + + connection.allow_exit.set() + with pytest.raises(asyncio.CancelledError): + await task + assert not pool.closed_with_checked_out_connection, "pool closed with a checked-out connection"