fix(persistence): drain alembic workers across cancellation (#5631)

* fix(persistence): drain alembic workers across cancellation

* docs(persistence): clarify alembic cancellation ownership

* docs(persistence): align alembic helper ownership contract

* fix(persistence): restore docstring underline

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

@@ -3,3 +3,5 @@
 Postgres bootstrap owns its session-scoped advisory lock until `pg_advisory_unlock` completes. Drain that unlock across host cancellation before leaving the SQLAlchemy connection context; repeated cancellation must not return a pooled session while it still holds the bootstrap mutex. Ordinary database errors remain best-effort and are logged.
 
 When `database.postgres_schema` is configured, both async ORM connections and the synchronous SQLAlchemy connections used by DB-backed custom agents and managed subagents must use the same `search_path`; preserve this invariant when adding another persistence entry point.
+
+Alembic stamp/upgrade workers started inside `bootstrap_schema()` remain owned by the bootstrap critical section until the worker finishes. Drain those `asyncio.to_thread()` calls across host cancellation before releasing the in-process SQLite bootstrap lock or PostgreSQL advisory lock; otherwise another bootstrap can overlap a still-running migration worker.

backend/packages/harness/deerflow/persistence/bootstrap.py

@@ -507,12 +507,12 @@ def _run_baseline_create_all_sync(sync_conn: Any) -> None:
 
 
 def _stamp(cfg: AlembicConfig, revision: str) -> None:
-    """Synchronous alembic stamp; callers must wrap in ``asyncio.to_thread``."""
+    """Synchronous alembic stamp; callers must drain via ``await_drained(asyncio.to_thread(...))``."""
     alembic_command.stamp(cfg, revision)
 
 
 def _upgrade(cfg: AlembicConfig, revision: str) -> None:
-    """Synchronous alembic upgrade; callers must wrap in ``asyncio.to_thread``."""
+    """Synchronous alembic upgrade; callers must drain via ``await_drained(asyncio.to_thread(...))``."""
     alembic_command.upgrade(cfg, revision)
 
 
@@ -611,7 +611,8 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema
 
     Branch dispatch is documented at module top. ``alembic.command.stamp`` and
     ``alembic.command.upgrade`` are synchronous and would block the event
-    loop; both are wrapped in ``asyncio.to_thread``.
+    loop; both are wrapped in ``await_drained(asyncio.to_thread(...))`` so the
+    worker finishes before the bootstrap lock is released on cancellation.
 
     *postgres_schema*, when set, is forwarded to the alembic config so the
     alembic-spawned engine pins its ``search_path`` to that schema. The target
@@ -631,7 +632,7 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema
             logger.info("bootstrap: branch=empty -> create_all + stamp head (%s)", head)
             async with engine.begin() as conn:
                 await conn.run_sync(_run_create_all_sync)
-            await asyncio.to_thread(_stamp, cfg, head)
+            await await_drained(asyncio.to_thread(_stamp, cfg, head))
 
         elif decision == "legacy":
             logger.info(
@@ -651,8 +652,8 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema
             # columns those revisions would add.
             async with engine.begin() as conn:
                 await conn.run_sync(_run_baseline_create_all_sync)
-            await asyncio.to_thread(_stamp, cfg, _BASELINE_REVISION)
+            await await_drained(asyncio.to_thread(_stamp, cfg, _BASELINE_REVISION))
-            await asyncio.to_thread(_upgrade, cfg, "head")
+            await await_drained(asyncio.to_thread(_upgrade, cfg, "head"))
 
         elif decision == "versioned":
             # The same revision id once named a different out-of-tree schema.
@@ -669,7 +670,7 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema
                     head,
                 )
                 try:
-                    await asyncio.to_thread(_upgrade, cfg, "head")
+                    await await_drained(asyncio.to_thread(_upgrade, cfg, "head"))
                 except CommandError:
                     # SQLite has no cross-process bootstrap mutex. Another
                     # process may advance 0018 to the reviewed 0019 after this

backend/tests/test_persistence_bootstrap_concurrency.py

@@ -15,6 +15,7 @@ idempotent revision helpers.
 from __future__ import annotations
 
 import asyncio
+import threading
 from pathlib import Path
 
 import pytest
@@ -148,3 +149,98 @@ async def test_slow_upgrade_does_not_corrupt_concurrent_state(monkeypatch, tmp_p
         assert await _alembic_version(engine) == HEAD
     finally:
         await engine.dispose()
+
+
+async def test_cancelled_bootstrap_keeps_sqlite_lock_until_alembic_worker_finishes(monkeypatch, tmp_path: Path) -> None:
+    """Cancellation must not detach an Alembic worker from its bootstrap lock."""
+    engine = create_async_engine(_url(tmp_path))
+    stamp_started = threading.Event()
+    allow_stamp = threading.Event()
+    second_reflect_started = asyncio.Event()
+    reflect_calls = 0
+    original_reflect = bootstrap_mod._reflect_state
+
+    def blocking_stamp(_cfg, _revision: str) -> None:
+        stamp_started.set()
+        assert allow_stamp.wait(5), "test did not release the blocked stamp worker"
+
+    def recording_reflect(sync_conn):
+        nonlocal reflect_calls
+        reflect_calls += 1
+        if reflect_calls >= 2:
+            second_reflect_started.set()
+        return original_reflect(sync_conn)
+
+    monkeypatch.setattr(bootstrap_mod, "_stamp", blocking_stamp)
+    monkeypatch.setattr(bootstrap_mod, "_upgrade", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(bootstrap_mod, "_reflect_state", recording_reflect)
+
+    first = asyncio.create_task(bootstrap_schema(engine, backend="sqlite"))
+    second: asyncio.Task[None] | None = None
+    try:
+        assert await asyncio.to_thread(stamp_started.wait, 2), "bootstrap did not reach the Alembic stamp worker"
+
+        first.cancel()
+        await asyncio.sleep(0)
+        first.cancel()
+        for _ in range(10):
+            await asyncio.sleep(0)
+
+        assert not first.done(), "cancelled bootstrap returned while its Alembic worker was still running"
+
+        second = asyncio.create_task(bootstrap_schema(engine, backend="sqlite"))
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert not second_reflect_started.is_set(), "cancelled bootstrap released the SQLite mutex before its worker finished"
+    finally:
+        allow_stamp.set()
+        await asyncio.gather(first, *(task for task in (second,) if task is not None), return_exceptions=True)
+        await engine.dispose()
+
+
+async def test_cancelled_versioned_bootstrap_keeps_sqlite_lock_until_upgrade_worker_finishes(monkeypatch, tmp_path: Path) -> None:
+    """The versioned upgrade path must retain the same lock while its worker drains."""
+    engine = create_async_engine(_url(tmp_path))
+    await bootstrap_schema(engine, backend="sqlite")
+
+    upgrade_started = threading.Event()
+    allow_upgrade = threading.Event()
+    second_reflect_started = asyncio.Event()
+    reflect_calls = 0
+    original_reflect = bootstrap_mod._reflect_state
+
+    def blocking_upgrade(_cfg, _revision: str) -> None:
+        upgrade_started.set()
+        assert allow_upgrade.wait(5), "test did not release the blocked upgrade worker"
+
+    def recording_reflect(sync_conn):
+        nonlocal reflect_calls
+        reflect_calls += 1
+        if reflect_calls >= 2:
+            second_reflect_started.set()
+        return original_reflect(sync_conn)
+
+    monkeypatch.setattr(bootstrap_mod, "_upgrade", blocking_upgrade)
+    monkeypatch.setattr(bootstrap_mod, "_reflect_state", recording_reflect)
+
+    first = asyncio.create_task(bootstrap_schema(engine, backend="sqlite"))
+    second: asyncio.Task[None] | None = None
+    try:
+        assert await asyncio.to_thread(upgrade_started.wait, 2), "bootstrap did not reach the Alembic upgrade worker"
+
+        first.cancel()
+        await asyncio.sleep(0)
+        first.cancel()
+        for _ in range(10):
+            await asyncio.sleep(0)
+
+        assert not first.done(), "cancelled bootstrap returned while its Alembic upgrade worker was still running"
+
+        second = asyncio.create_task(bootstrap_schema(engine, backend="sqlite"))
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert not second_reflect_started.is_set(), "cancelled bootstrap released the SQLite mutex before upgrade finished"
+    finally:
+        allow_upgrade.set()
+        await asyncio.gather(first, *(task for task in (second,) if task is not None), return_exceptions=True)
+        await engine.dispose()