github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-09 21:49:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(persistence): preserve rollback across the incarnation migration (#5219)

Browse Source 
* fix(persistence): tolerate thread incarnation migration

* docs(persistence): pin forward revision contract

---------

Co-authored-by: CorgiBoyG <CorgiBoyG@users.noreply.github.com>
This commit is contained in:
RongJie G 2026-09-07 15:27:00 +08:00 committed by GitHub
parent d1f1c49dcd
commit 99367100fb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 476 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
backend/packages/harness/deerflow/persistence/bootstrap.py
View File

@ -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}")

29
backend/packages/harness/deerflow/persistence/migrations/AGENTS.md
View File

@ -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)

9
backend/tests/blocking_io/test_persistence_bootstrap.py
View File

@ -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()

357
backend/tests/test_persistence_forward_revision_compat.py Normal file
View File

@ -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()