From 99367100fbdddeaf7cec29aebbc9170cb5653def Mon Sep 17 00:00:00 2001 From: RongJie G <111257566+CorgiBoyG@users.noreply.github.com> Date: Mon, 7 Sep 2026 15:27:00 +0800 Subject: [PATCH] fix(persistence): preserve rollback across the incarnation migration (#5219) * fix(persistence): tolerate thread incarnation migration * docs(persistence): pin forward revision contract --------- Co-authored-by: CorgiBoyG --- .../harness/deerflow/persistence/bootstrap.py | 94 ++++- .../deerflow/persistence/migrations/AGENTS.md | 29 +- .../blocking_io/test_persistence_bootstrap.py | 9 +- ...est_persistence_forward_revision_compat.py | 357 ++++++++++++++++++ 4 files changed, 476 insertions(+), 13 deletions(-) create mode 100644 backend/tests/test_persistence_forward_revision_compat.py diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py index be5e5919a..b1a8861dd 100644 --- a/backend/packages/harness/deerflow/persistence/bootstrap.py +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -13,11 +13,13 @@ Combines two ideas: Three-branch decision (see ``_decide_state``) --------------------------------------------- -| DB state | Action | -|---------------------------------------|-----------------------------------------| -| empty (no DeerFlow tables) | ``create_all`` + ``alembic stamp head`` | -| legacy (DeerFlow tables, no alembic) | ``create_all`` (baseline tables only, as backfill) + ``stamp 0001_baseline`` + ``upgrade head`` | -| versioned (``alembic_version`` row) | ``alembic upgrade head`` | +| DB state | Action | +|-----------------------------------------------|-----------------------------------------| +| empty (no DeerFlow tables) | ``create_all`` + ``alembic stamp head`` | +| legacy (DeerFlow tables, no alembic) | ``create_all`` (baseline tables only, as backfill) + ``stamp 0001_baseline`` + ``upgrade head`` | +| versioned (one locally known revision) | ``alembic upgrade head`` | +| reviewed forward-compatible revision 0019 | warn and skip migration | +| unknown, empty, or multiple revision rows | refuse to start | The legacy branch handles pre-alembic databases that already have at least one DeerFlow-owned table. ``create_all`` runs first because stamping at @@ -70,6 +72,11 @@ best-effort; multi-instance deployments should use Postgres. * **Idempotent revisions -- retry fallback.** Column revisions use the helpers in ``migrations/_helpers.py`` so repeated post-baseline changes, manual ALTERs, or retries after SQLite lock contention do not duplicate work. + During the 0018-to-0019 compatibility window, an old SQLite process also + re-reads ``alembic_version`` after an Alembic ``CommandError``. It recovers + only when another process advanced the file to the explicitly reviewed 0019 + and 0019 is still absent from the local migration tree; every other migration + failure remains fatal. ``alembic upgrade head`` on a DB already at head is a no-op by alembic's own semantics, so the second-N-th actor simply observes head and exits. @@ -87,6 +94,7 @@ from typing import Any from alembic import command as alembic_command from alembic.config import Config as AlembicConfig from alembic.script import ScriptDirectory +from alembic.util.exc import CommandError from sqlalchemy import inspect as sa_inspect from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncEngine @@ -99,6 +107,17 @@ _MIGRATIONS_DIR = Path(__file__).resolve().parent / "migrations" # Cached migration head, computed once per process from the disk script tree. _HEAD_REVISION: str | None = None +_KNOWN_REVISIONS: frozenset[str] | None = None + +# One additive revision may be present when an older Gateway starts during the +# thread-incarnation rollout. This allowlist was reviewed only for revision +# ``0019_thread_incarnations`` adding nullable VARCHAR(32) columns +# ``threads_meta.incarnation`` and ``mcp_tasks.thread_incarnation`` without a +# server default, table, index, constraint, or data backfill. The owning 0019 +# change must cross-pin this revision id and schema shape in tests. Amending +# that DDL requires re-auditing old-repository reads and writes before this +# exception remains valid. +_FORWARD_COMPATIBLE_REVISION = "0019_thread_incarnations" # Baseline (stamp target for legacy DBs). Pinned here so the bootstrap layer # fails loudly if the baseline revision is ever renamed without updating the @@ -270,6 +289,34 @@ def _get_head_revision() -> str: return _HEAD_REVISION +def _get_known_revisions() -> frozenset[str]: + """Return every revision id available in the local migration tree.""" + global _KNOWN_REVISIONS + if _KNOWN_REVISIONS is None: + cfg = AlembicConfig() + cfg.set_main_option("script_location", str(_MIGRATIONS_DIR)) + script = ScriptDirectory.from_config(cfg) + _KNOWN_REVISIONS = frozenset(revision.revision for revision in script.walk_revisions()) + return _KNOWN_REVISIONS + + +def _get_revision_metadata() -> tuple[str, frozenset[str]]: + """Load the local head and revision set off the event loop.""" + return _get_head_revision(), _get_known_revisions() + + +async def _read_database_revision(conn: Any) -> str: + """Read and validate the database's single alembic revision row.""" + result = await conn.execute(text("SELECT version_num FROM alembic_version")) + rows = list(result.scalars()) + if len(rows) != 1: + raise RuntimeError(f"bootstrap: expected exactly one alembic_version row, found {len(rows)}") + revision = rows[0] + if not isinstance(revision, str) or not revision: + raise RuntimeError("bootstrap: alembic_version contains an empty revision") + return revision + + def _reflect_state(sync_conn: Any) -> dict[str, bool]: """Inspect *sync_conn* (sync connection inside ``run_sync``) and return: @@ -495,12 +542,13 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema schema must already exist (``init_engine`` issues ``CREATE SCHEMA`` before calling this). Ignored for non-postgres backends. """ - head = _get_head_revision() + head, known_revisions = await asyncio.to_thread(_get_revision_metadata) cfg = _get_alembic_config(engine, postgres_schema=postgres_schema if backend == "postgres" else "") async with _bootstrap_lock(engine, backend=backend): async with engine.connect() as conn: state = await conn.run_sync(_reflect_state) + database_revision = await _read_database_revision(conn) if state["has_alembic_version"] else None decision = _decide_state(state) if decision == "empty": @@ -531,8 +579,38 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema await asyncio.to_thread(_upgrade, cfg, "head") elif decision == "versioned": - logger.info("bootstrap: branch=versioned -> upgrade head (%s)", head) - await asyncio.to_thread(_upgrade, cfg, "head") + if database_revision in known_revisions: + logger.info( + "bootstrap: branch=versioned revision=%s -> upgrade head (%s)", + database_revision, + head, + ) + try: + await 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 + # process reads the version but before Alembic starts. + # Do not apply this recovery once 0019 belongs to the local + # tree: a new binary's migration failure must stay fatal. + if backend != "sqlite" or _FORWARD_COMPATIBLE_REVISION in known_revisions: + raise + async with engine.connect() as conn: + current_revision = await _read_database_revision(conn) + if current_revision != _FORWARD_COMPATIBLE_REVISION: + raise + logger.warning( + "bootstrap: database advanced concurrently to explicitly forward-compatible revision %s; skipping the stale local upgrade", + current_revision, + ) + elif database_revision == _FORWARD_COMPATIBLE_REVISION: + logger.warning( + "bootstrap: database revision %s is newer than local head %s but is explicitly forward-compatible; skipping migration", + database_revision, + head, + ) + else: + raise RuntimeError(f"bootstrap: database revision {database_revision!r} is not known to this build (local head {head!r}); refusing to start") else: # pragma: no cover -- defensive raise RuntimeError(f"bootstrap: unhandled decision {decision!r}") diff --git a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md index 7597e7a29..ba7e05057 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md @@ -10,12 +10,37 @@ DeerFlow's application tables (`runs`, `threads_meta`, `feedback`, `users`, `run |-------------------------------------------|-----------------------------------------| | empty (no DeerFlow tables) | `create_all` + `alembic stamp head` | | legacy (DeerFlow tables, no `alembic_version`) | `create_all` (baseline tables only, backfill) + `alembic stamp 0001_baseline` + `upgrade head` | -| versioned (`alembic_version` row exists) | `alembic upgrade head` | +| versioned (one locally known `alembic_version` row) | `alembic upgrade head` | +| forward-compatible (`0019_thread_incarnations`) | warn and skip migration | +| unknown revision, empty version table, or multiple version rows | fail closed and refuse to start | The legacy branch handles pre-alembic databases that already have at least one DeerFlow-owned table. `create_all` runs first because stamping at `0001_baseline` makes alembic skip the baseline's own `create_table` DDL on the subsequent upgrade — so any baseline table introduced into `Base.metadata` after the user's DB was first provisioned (e.g. the `channel_*` tables from PR #1930 for users upgrading across multiple releases) would otherwise never be created, and the first request hitting that table would 500 with `no such table`. The backfill is **restricted to `_BASELINE_TABLE_NAMES`** so it does not also create tables that future revisions introduce — those revisions' own `op.create_table` would otherwise fail with `relation already exists`. A guard test pins `_BASELINE_TABLE_NAMES` against `0001_baseline.upgrade()`'s actual output, so editing 0001 to add or remove a table forces a matching update to the constant. Column-level shape (pre-#3658 vs post-#3658 vs manual-ALTER for `token_usage_by_model`) is answered by each `versions/*.py` revision via the idempotent helpers in `migrations/_helpers.py` (`safe_add_column` / `safe_drop_column`) which no-op when the change is already present and `logger.warning` on shape drift. **Adding a new ORM column / table only requires a new revision file — no edit to `bootstrap.py` is needed** *unless* the new revision adds a new baseline table (rare; only happens when a new model is part of the baseline rather than introduced by its own revision). The empty-DB path keeps using `create_all` because `Base.metadata` is the only authoritative schema source — `create_all` renders both SQLite (JSON, type affinity) and Postgres (JSONB, partial indexes) correctly without anyone having to keep a hand-written baseline in lockstep. `0001_baseline.upgrade()` is therefore almost never executed in practice; it exists as a stamp target + chain root. +**Rolling forward compatibility**: this migration tree deliberately remains at +`0018_oauth_identity_pg_partial`, but an older Gateway may briefly share a +database with the expand-only `0019_thread_incarnations` deployment. Bootstrap +reads `alembic_version` while holding its backend lock and accepts exactly one +row. A locally known revision follows the normal upgrade path. The one unknown +revision `0019_thread_incarnations` is explicitly allowlisted: bootstrap logs a +warning and leaves the newer schema untouched. Any other unknown revision, an +empty version table, or multiple version rows fails closed. Do not broaden the +allowlist without proving that old repositories can read, insert, and update +through the newer schema; nullable additive columns are covered by +`tests/test_persistence_forward_revision_compat.py`. This exception is reviewed +only for the expand-only 0019 shape: nullable VARCHAR(32) +`threads_meta.incarnation` and `mcp_tasks.thread_incarnation` columns with no +server default, table, index, constraint, or data backfill. The 0019 migration +must cross-pin its revision id and schema shape against the bootstrap contract; +amending that DDL requires a fresh old-repository compatibility audit. Because SQLite has no +cross-process bootstrap mutex, an old process may read 0018 immediately before +another process commits 0019. If its now-stale Alembic upgrade fails, bootstrap +re-reads the version and recovers only for the exact allowlisted 0019 while that +revision remains absent from the local migration tree. Do not generalize this +recovery or apply it to a binary that owns 0019; its migration failures must +remain fatal. + **Concurrency safety**: Postgres uses `pg_advisory_lock` to serialise concurrent Gateway instances. SQLite uses a per-engine `asyncio.Lock` for same-process startup and is best-effort across processes via SQLite's file-level write lock + `PRAGMA busy_timeout`; multi-instance deployments should use Postgres. Column revisions in `versions/` additionally use idempotent helpers (`_helpers.py::safe_add_column`, `safe_drop_column`) so repeated post-baseline changes and retries are no-ops when the change is already present. **Authoring a new revision**: @@ -88,6 +113,6 @@ on installs that never enabled it. The convention is: - `migrations/versions/0016_subagent_batches.py` — creates durable native-subagent batch and item tables, including owner/submission idempotency, item identity, lease/recovery state, and result fields - `migrations/versions/0017_personal_access_tokens.py` — creates the personal access token table for programmatic API access - `migrations/versions/0018_oauth_identity_pg_partial.py` — converts `idx_users_oauth_identity` to a partial index on Postgres (`postgresql_where`), matching what `UserRow.__table_args__` already builds via `create_all`; `0001_baseline` never applied the predicate on Postgres, so every `alembic upgrade head`-provisioned deployment carried a full index until this revision. Postgres-only, idempotent (checks `pg_index.indpred` directly), no-op on SQLite (already partial via `sqlite_where`) and on a DB where the index doesn't exist yet. Originally generated as 0017 and renumbered to 0018 after 0017_personal_access_tokens merged first and kept that slot -- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking +- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch provisioning decision, locked revision validation, and the narrow 0019 forward-compatibility exception - `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()` - Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps) diff --git a/backend/tests/blocking_io/test_persistence_bootstrap.py b/backend/tests/blocking_io/test_persistence_bootstrap.py index ff81d1be2..106a52750 100644 --- a/backend/tests/blocking_io/test_persistence_bootstrap.py +++ b/backend/tests/blocking_io/test_persistence_bootstrap.py @@ -1,5 +1,5 @@ -"""Regression: ``bootstrap_schema`` offloads ``alembic.command.stamp`` / -``alembic.command.upgrade`` via ``asyncio.to_thread``. +"""Regression: ``bootstrap_schema`` offloads migration-tree parsing and +``alembic.command.stamp`` / ``alembic.command.upgrade`` via ``asyncio.to_thread``. The alembic commands are synchronous: they open their own engine and execute DDL. Calling them directly on the FastAPI lifespan event loop would block -- @@ -59,17 +59,20 @@ async def test_bootstrap_offloads_alembic_stamp_and_upgrade(monkeypatch: pytest. monkeypatch.setattr(bootstrap_mod.asyncio, "to_thread", spy_to_thread) - # Use a real SQLite DB so alembic actually runs stamp + upgrade. + # Use a real SQLite DB so alembic actually runs metadata parsing, stamp, + # and upgrade. db_path = tmp_path / "spy.db" engine = create_async_engine(f"sqlite+aiosqlite:///{db_path.as_posix()}") try: # Empty branch -> create_all + stamp head. ``_stamp`` must be offloaded. await bootstrap_mod.bootstrap_schema(engine, backend="sqlite") + assert "_get_revision_metadata" in seen, f"migration metadata parsing not offloaded; saw: {seen}" assert "_stamp" in seen, f"_stamp not offloaded; saw: {seen}" # Re-run -> versioned branch -> upgrade head (no-op at head). ``_upgrade`` must be offloaded. seen.clear() await bootstrap_mod.bootstrap_schema(engine, backend="sqlite") + assert "_get_revision_metadata" in seen, f"migration metadata parsing not offloaded; saw: {seen}" assert "_upgrade" in seen, f"_upgrade not offloaded; saw: {seen}" finally: await engine.dispose() diff --git a/backend/tests/test_persistence_forward_revision_compat.py b/backend/tests/test_persistence_forward_revision_compat.py new file mode 100644 index 000000000..f5ca2e159 --- /dev/null +++ b/backend/tests/test_persistence_forward_revision_compat.py @@ -0,0 +1,357 @@ +"""Forward-compatibility tests for an old Gateway against migration 0019.""" + +from __future__ import annotations + +import asyncio +import os +import threading +import uuid +from datetime import UTC, datetime, timedelta +from pathlib import Path +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import pytest +import sqlalchemy as sa +from alembic.util.exc import CommandError +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +import deerflow.persistence.models # noqa: F401 +from deerflow.config.database_config import DatabaseConfig +from deerflow.persistence import bootstrap as bootstrap_mod +from deerflow.persistence.bootstrap import ( + _FORWARD_COMPATIBLE_REVISION, + _get_alembic_config, + _upgrade, + bootstrap_schema, +) +from deerflow.persistence.engine import close_engine, get_engine, init_engine_from_config +from deerflow.persistence.mcp_tasks import McpTaskRepository +from deerflow.persistence.thread_meta import ThreadMetaRepository + +HEAD = "0018_oauth_identity_pg_partial" +POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI") + + +def _url(tmp_path: Path, name: str) -> str: + return f"sqlite+aiosqlite:///{(tmp_path / name).as_posix()}" + + +def _postgres_url(url: str) -> str: + parts = urlsplit(url) + query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}]) + return urlunsplit(parts._replace(query=query)) + + +async def _database_revision(engine) -> str | None: + async with engine.connect() as conn: + result = await conn.execute(sa.text("SELECT version_num FROM alembic_version")) + return result.scalar_one_or_none() + + +async def _set_database_revision(engine, revision: str) -> None: + async with engine.begin() as conn: + await conn.execute(sa.text("UPDATE alembic_version SET version_num = :revision"), {"revision": revision}) + + +async def _seed_head(engine) -> None: + await bootstrap_schema(engine, backend="sqlite") + assert await _database_revision(engine) == HEAD + + +@pytest.mark.asyncio +async def test_known_older_revision_upgrades_normally(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "known.db")) + try: + cfg = _get_alembic_config(engine) + await asyncio.to_thread(_upgrade, cfg, "0017_personal_access_tokens") + + await bootstrap_schema(engine, backend="sqlite") + + assert await _database_revision(engine) == HEAD + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_exact_forward_revision_skips_upgrade_with_warning( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + engine = create_async_engine(_url(tmp_path, "forward.db")) + try: + await _seed_head(engine) + await _set_database_revision(engine, _FORWARD_COMPATIBLE_REVISION) + + with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"): + await bootstrap_schema(engine, backend="sqlite") + + assert await _database_revision(engine) == _FORWARD_COMPATIBLE_REVISION + assert any(_FORWARD_COMPATIBLE_REVISION in record.getMessage() and "explicitly forward-compatible" in record.getMessage() for record in caplog.records) + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_other_unknown_revision_fails_closed(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "unknown.db")) + try: + await _seed_head(engine) + await _set_database_revision(engine, "9999_unknown") + + with pytest.raises(RuntimeError, match="not known to this build"): + await bootstrap_schema(engine, backend="sqlite") + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_upgrade_race_recovers_when_other_process_applies_forward_revision( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + url = _url(tmp_path, "forward-race.db") + old_gateway = create_async_engine(url) + new_gateway = create_async_engine(url) + upgrade_started = threading.Event() + continue_upgrade = threading.Event() + original_upgrade = bootstrap_mod._upgrade + + def delayed_upgrade(cfg, revision): + upgrade_started.set() + if not continue_upgrade.wait(timeout=5): + raise TimeoutError("timed out waiting for the forward migration") + return original_upgrade(cfg, revision) + + try: + await _seed_head(old_gateway) + monkeypatch.setattr(bootstrap_mod, "_upgrade", delayed_upgrade) + + old_bootstrap = asyncio.create_task(bootstrap_schema(old_gateway, backend="sqlite")) + assert await asyncio.to_thread(upgrade_started.wait, 5) + + await _add_forward_columns(new_gateway) + await _set_database_revision(new_gateway, _FORWARD_COMPATIBLE_REVISION) + + with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"): + continue_upgrade.set() + await old_bootstrap + assert await _database_revision(old_gateway) == _FORWARD_COMPATIBLE_REVISION + assert any("advanced concurrently" in record.getMessage() and _FORWARD_COMPATIBLE_REVISION in record.getMessage() for record in caplog.records) + finally: + continue_upgrade.set() + await old_gateway.dispose() + await new_gateway.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_upgrade_error_stays_fatal_without_forward_revision( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = create_async_engine(_url(tmp_path, "upgrade-error.db")) + try: + await _seed_head(engine) + + def fail_upgrade(_cfg, _revision): + raise CommandError("broken migration") + + monkeypatch.setattr(bootstrap_mod, "_upgrade", fail_upgrade) + with pytest.raises(CommandError, match="broken migration"): + await bootstrap_schema(engine, backend="sqlite") + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_local_forward_migration_error_stays_fatal( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = create_async_engine(_url(tmp_path, "local-forward-error.db")) + try: + await _seed_head(engine) + monkeypatch.setattr( + bootstrap_mod, + "_get_revision_metadata", + lambda: (_FORWARD_COMPATIBLE_REVISION, frozenset({HEAD, _FORWARD_COMPATIBLE_REVISION})), + ) + + def fail_upgrade(_cfg, _revision): + raise CommandError("local 0019 migration failed") + + monkeypatch.setattr(bootstrap_mod, "_upgrade", fail_upgrade) + with pytest.raises(CommandError, match="local 0019 migration failed"): + await bootstrap_schema(engine, backend="sqlite") + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_empty_alembic_version_fails_closed(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "empty-version.db")) + try: + await _seed_head(engine) + async with engine.begin() as conn: + await conn.execute(sa.text("DELETE FROM alembic_version")) + + with pytest.raises(RuntimeError, match="expected exactly one alembic_version row, found 0"): + await bootstrap_schema(engine, backend="sqlite") + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_multiple_alembic_versions_fail_closed(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "multiple-versions.db")) + try: + await _seed_head(engine) + async with engine.begin() as conn: + await conn.execute( + sa.text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), + {"revision": "0017_personal_access_tokens"}, + ) + + with pytest.raises(RuntimeError, match="expected exactly one alembic_version row, found 2"): + await bootstrap_schema(engine, backend="sqlite") + finally: + await engine.dispose() + + +async def _add_forward_columns(engine) -> None: + async with engine.begin() as conn: + await conn.execute(sa.text("ALTER TABLE threads_meta ADD COLUMN incarnation VARCHAR(32)")) + await conn.execute(sa.text("ALTER TABLE mcp_tasks ADD COLUMN thread_incarnation VARCHAR(32)")) + + +@pytest.mark.asyncio +async def test_old_thread_repository_tolerates_forward_nullable_column(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "thread-repository.db")) + try: + await _seed_head(engine) + await _add_forward_columns(engine) + session_factory = async_sessionmaker(engine, expire_on_commit=False) + repository = ThreadMetaRepository(session_factory) + + created = await repository.create("thread-1", user_id=None) + assert created["thread_id"] == "thread-1" + assert "incarnation" not in created + + async with engine.begin() as conn: + await conn.execute( + sa.text("UPDATE threads_meta SET incarnation = :incarnation WHERE thread_id = :thread_id"), + {"incarnation": "a" * 32, "thread_id": "thread-1"}, + ) + + fetched = await repository.get("thread-1", user_id=None) + assert fetched is not None + assert "incarnation" not in fetched + await repository.update_status("thread-1", "busy", user_id=None) + + async with engine.connect() as conn: + incarnation = ( + await conn.execute( + sa.text("SELECT incarnation FROM threads_meta WHERE thread_id = :thread_id"), + {"thread_id": "thread-1"}, + ) + ).scalar_one() + assert incarnation == "a" * 32 + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_old_mcp_task_repository_tolerates_forward_nullable_column(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "mcp-repository.db")) + try: + await _seed_head(engine) + await _add_forward_columns(engine) + session_factory = async_sessionmaker(engine, expire_on_commit=False) + repository = McpTaskRepository(session_factory) + now = datetime.now(UTC) + + created = await repository.create( + task_id="task-1", + user_id="user-1", + thread_id="thread-1", + run_id="run-1", + tool_call_id="call-1", + server_name="reports", + driver_name="fake", + remote_task_id="remote-1", + task_name="Generate report", + status="working", + result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=now - timedelta(seconds=1), + ) + assert created["id"] == "task-1" + assert "thread_incarnation" not in created + + async with engine.begin() as conn: + await conn.execute( + sa.text("UPDATE mcp_tasks SET thread_incarnation = :incarnation WHERE id = :task_id"), + {"incarnation": "b" * 32, "task_id": "task-1"}, + ) + + fetched = await repository.get("task-1", user_id="user-1") + assert fetched is not None + assert "thread_incarnation" not in fetched + claimed = await repository.claim_due_tasks( + now=now, + lease_owner="worker-1", + lease_seconds=60, + limit=1, + ) + assert [task["id"] for task in claimed] == ["task-1"] + + async with engine.connect() as conn: + incarnation = ( + await conn.execute( + sa.text("SELECT thread_incarnation FROM mcp_tasks WHERE id = :task_id"), + {"task_id": "task-1"}, + ) + ).scalar_one() + assert incarnation == "b" * 32 + finally: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not POSTGRES_URL, reason="requires TEST_POSTGRES_URI for a real PostgreSQL restart") +async def test_old_gateway_restarts_against_forward_postgres_revision() -> None: + assert POSTGRES_URL is not None + schema = f"forward_revision_{uuid.uuid4().hex}" + config = DatabaseConfig( + backend="postgres", + postgres_url=_postgres_url(POSTGRES_URL), + postgres_schema=schema, + ) + try: + await init_engine_from_config(config) + engine = get_engine() + assert engine is not None + async with engine.begin() as conn: + await conn.execute(sa.text("ALTER TABLE threads_meta ADD COLUMN incarnation VARCHAR(32)")) + await conn.execute(sa.text("ALTER TABLE mcp_tasks ADD COLUMN thread_incarnation VARCHAR(32)")) + await conn.execute( + sa.text("UPDATE alembic_version SET version_num = :revision"), + {"revision": _FORWARD_COMPATIBLE_REVISION}, + ) + + await close_engine() + await init_engine_from_config(config) + + restarted_engine = get_engine() + assert restarted_engine is not None + assert await _database_revision(restarted_engine) == _FORWARD_COMPATIBLE_REVISION + finally: + engine = get_engine() + if engine is not None: + async with engine.begin() as conn: + await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')) + await close_engine()