From 4117a50676b97d479839bf639d5fa943af68db1a Mon Sep 17 00:00:00 2001
From: lihongyuan99 <64824864+lihongyuan99@users.noreply.github.com>
Date: Wed, 23 Sep 2026 10:14:35 +0800
Subject: [PATCH] fix(persistence): drain the auto-create maintenance dispose
 across cancellation (#5649)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_auto_create_postgres_db opens a throwaway engine against the server's
`postgres` database and disposes it in its finally, but that await was bare.
Startup cancellation landing there abandoned a half-disposed pool that nothing
else owned — the engine is a local, and the CREATE DATABASE has already
committed under AUTOCOMMIT, so the next boot does not take this path again.

Drain the dispose through the cancellation-safe await_drained() helper the
engine's own close path already uses, and pin it with a repeated-cancellation
regression.
---
 .../harness/deerflow/persistence/engine.py    |  6 +-
 ...istence_engine_auto_create_cancellation.py | 85 +++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 backend/tests/test_persistence_engine_auto_create_cancellation.py

diff --git a/backend/packages/harness/deerflow/persistence/engine.py b/backend/packages/harness/deerflow/persistence/engine.py
index 565ced5b9..76929cd2b 100644
--- a/backend/packages/harness/deerflow/persistence/engine.py
+++ b/backend/packages/harness/deerflow/persistence/engine.py
@@ -81,7 +81,11 @@ async def _auto_create_postgres_db(url: str) -> None:
             await conn.execute(text(f'CREATE DATABASE "{db_name}"'))
         logger.info("Auto-created PostgreSQL database: %s", db_name)
     finally:
-        await maint_engine.dispose()
+        # Drain: the maintenance engine is local to this helper, so an
+        # abandoned dispose leaves its pool to the ``postgres`` database with
+        # nothing left to close it. Host cancellation must therefore be
+        # delivered only after disposal actually finishes.
+        await await_drained(maint_engine.dispose())
 
 
 async def init_engine(
diff --git a/backend/tests/test_persistence_engine_auto_create_cancellation.py b/backend/tests/test_persistence_engine_auto_create_cancellation.py
new file mode 100644
index 000000000..9b89d63c1
--- /dev/null
+++ b/backend/tests/test_persistence_engine_auto_create_cancellation.py
@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from deerflow.persistence import engine as engine_mod
+
+_TARGET_URL = "postgresql+asyncpg://deerflow:secret@127.0.0.1:5432/deerflow_target"
+
+
+class _RecordingConnection:
+    def __init__(self) -> None:
+        self.executed = asyncio.Event()
+        self.statements: list[str] = []
+
+    async def __aenter__(self) -> _RecordingConnection:
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        return None
+
+    async def execute(self, statement: object) -> None:
+        self.statements.append(str(statement))
+        self.executed.set()
+
+
+class _BlockingMaintenanceEngine:
+    """Fake maintenance engine whose ``dispose()`` can be held open."""
+
+    def __init__(self) -> None:
+        self.connection = _RecordingConnection()
+        self.dispose_started = asyncio.Event()
+        self.allow_dispose = asyncio.Event()
+        self.dispose_finished = asyncio.Event()
+
+    def connect(self) -> _RecordingConnection:
+        return self.connection
+
+    async def dispose(self) -> None:
+        self.dispose_started.set()
+        await self.allow_dispose.wait()
+        self.dispose_finished.set()
+
+
+@pytest.mark.asyncio
+async def test_auto_create_db_drains_maintenance_dispose_across_repeated_cancellation(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A cancelled first boot must still dispose the engine it opened.
+
+    ``_auto_create_postgres_db`` owns a throwaway engine against the server's
+    ``postgres`` database. Its ``finally`` dispose was a bare await, so host
+    cancellation could return from the helper with that engine's pool still
+    live — and because the engine is a local, nothing else ever disposes it.
+    """
+    maintenance_engine = _BlockingMaintenanceEngine()
+    task: asyncio.Task[None] | None = None
+
+    monkeypatch.setattr(engine_mod, "create_async_engine", lambda *args, **kwargs: maintenance_engine)
+
+    try:
+        task = asyncio.create_task(engine_mod._auto_create_postgres_db(_TARGET_URL))
+        await asyncio.wait_for(maintenance_engine.connection.executed.wait(), timeout=1)
+        assert 'CREATE DATABASE "deerflow_target"' in maintenance_engine.connection.statements[0]
+        await asyncio.wait_for(maintenance_engine.dispose_started.wait(), timeout=1)
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done(), "database auto-creation returned before the maintenance engine was disposed"
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done(), "repeated cancellation interrupted the maintenance engine dispose"
+
+        maintenance_engine.allow_dispose.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert maintenance_engine.dispose_finished.is_set()
+    finally:
+        maintenance_engine.allow_dispose.set()
+        if task is not None and not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)