From 6b3d64cccb9cd8fb6733fd453b8e911052bbe043 Mon Sep 17 00:00:00 2001
From: yetuge <182123929+yetuge@users.noreply.github.com>
Date: Thu, 24 Sep 2026 17:34:06 +0800
Subject: [PATCH] fix(runtime): drain postgres schema-setup connection exit
 across cancellation (#5812)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_ensure_postgres_schema_with_pool runs inside the drained pool context but
acquired its connection 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 to the pool and the pool's own
drained teardown closes with a checked-out connection — the backend-ownership
invariant documented in runtime/AGENTS.md (same class as the provider-context
fix that introduced drained_async_context).

Use drained_async_context for the schema-setup connection so the connection
returns to the pool before the pool context unwinds.
---
 .../runtime/checkpointer/async_provider.py    |  5 +-
 ..._runtime_db_provider_close_cancellation.py | 69 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

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"