From 06c827903a8c9e515b306508a9c91d734074e7f7 Mon Sep 17 00:00:00 2001 From: RongJie G <111257566+CorgiBoyG@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:52:09 +0800 Subject: [PATCH] feat(persistence): add expand-phase thread incarnation storage (#5216) * feat(persistence): expand thread incarnation storage Add nullable thread and MCP task incarnation columns while preserving mixed-version writes. New thread records receive stable incarnation IDs, and new task rows copy the matching owned or shared thread incarnation without changing any read, claim, session, or deletion behavior. * test(persistence): pin incarnation rollback compatibility * test(api): pin internal thread response boundary * fix(persistence): rebase incarnation rollout after projects --------- Co-authored-by: CorgiBoyG --- backend/app/gateway/routers/threads.py | 23 +- .../harness/deerflow/persistence/bootstrap.py | 102 ++-- .../deerflow/persistence/mcp_tasks/model.py | 1 + .../deerflow/persistence/mcp_tasks/sql.py | 15 + .../deerflow/persistence/migrations/AGENTS.md | 89 ++-- .../versions/0019_thread_incarnations.py | 100 ++++ .../persistence/thread_meta/__init__.py | 3 +- .../deerflow/persistence/thread_meta/base.py | 12 + .../persistence/thread_meta/memory.py | 124 +++-- .../deerflow/persistence/thread_meta/model.py | 1 + .../deerflow/persistence/thread_meta/sql.py | 42 +- backend/tests/test_mcp_task_postgres.py | 216 ++++++++ backend/tests/test_mcp_task_repository.py | 170 ++++++- .../test_memory_thread_meta_isolation.py | 203 ++++++++ ...est_migration_0004_run_ownership_dedupe.py | 3 +- ...ration_0007_scheduled_run_active_dedupe.py | 3 +- ...t_migration_0015_scheduled_task_enqueue.py | 4 +- ...test_migration_0019_thread_incarnations.py | 356 ++++++++++++++ .../test_migration_0021_batch_acceptance.py | 14 + backend/tests/test_persistence_bootstrap.py | 15 +- .../test_persistence_bootstrap_concurrency.py | 4 +- .../test_persistence_bootstrap_regression.py | 5 +- ...est_persistence_forward_revision_compat.py | 461 ++++++++++++++---- backend/tests/test_thread_meta_repo.py | 44 ++ backend/tests/test_threads_router.py | 131 +++++ docs/database-forward-revision-recovery.md | 35 +- 26 files changed, 1917 insertions(+), 259 deletions(-) create mode 100644 backend/packages/harness/deerflow/persistence/migrations/versions/0019_thread_incarnations.py create mode 100644 backend/tests/test_mcp_task_postgres.py create mode 100644 backend/tests/test_migration_0019_thread_incarnations.py diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 39cf463bd..6441a3739 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -21,7 +21,7 @@ from typing import Any from fastapi import APIRouter, BackgroundTasks, HTTPException, Request from langgraph.checkpoint.base import empty_checkpoint from langgraph.types import Overwrite -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator from sqlalchemy.exc import IntegrityError from app.gateway.authz import require_permission @@ -46,7 +46,7 @@ from app.gateway.utils import sanitize_log_param from deerflow.agents.thread_state import THREAD_STATE_REDUCER_FIELDS from deerflow.config.paths import Paths, get_paths from deerflow.config.summarization_config import ContextSize -from deerflow.persistence.thread_meta import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, THREAD_PROJECT_METADATA_KEY +from deerflow.persistence.thread_meta import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, THREAD_PROJECT_METADATA_KEY, ThreadOwnershipConflictError from deerflow.runtime import ThreadOperationKind, serialize_channel_values_for_api from deerflow.runtime.checkpoint_mode import CheckpointModeMismatchError, CheckpointModeReconfigurationError from deerflow.runtime.checkpoint_state import graph_reducer_channels, graph_state_schema, graph_writable_channels @@ -429,6 +429,10 @@ class _MetadataRedactingResponse(BaseModel): class ThreadResponse(_MetadataRedactingResponse): """Response model for a single thread.""" + # ThreadMetaStore records include internal lifecycle fields such as + # ``incarnation``. Keep the HTTP response as an explicit public projection. + model_config = ConfigDict(extra="ignore") + thread_id: str = Field(description="Unique thread identifier") status: str = Field(default="idle", description="Thread status: idle, busy, interrupted, error") created_at: str = Field(default="", description="ISO timestamp") @@ -784,11 +788,8 @@ async def _resolve_existing_thread( """ existing_record = await thread_store.get(thread_id, **thread_owner_kwargs) if existing_record is None and thread_owner_user_id: - unscoped_record = await thread_store.get(thread_id, user_id=None) - if unscoped_record is not None: - if unscoped_record.get("user_id") != thread_owner_user_id: - await thread_store.update_owner(thread_id, thread_owner_user_id, user_id=None) - existing_record = await thread_store.get(thread_id, **thread_owner_kwargs) + await thread_store.claim_unowned(thread_id, thread_owner_user_id) + existing_record = await thread_store.get(thread_id, **thread_owner_kwargs) return existing_record @@ -842,14 +843,18 @@ async def create_thread(body: ThreadCreateRequest, request: Request) -> ThreadRe # Fail closed: missing, foreign, or archived projects are # indistinguishable at the API surface. raise HTTPException(status_code=404, detail="Project not found") from None + except ThreadOwnershipConflictError: + # Do not reveal that a caller-chosen id belongs to another user. + raise HTTPException(status_code=404, detail="Thread not found") from None except IntegrityError: # The idempotency read above and this insert are not atomic: a # concurrent request for the same thread_id can commit in between, so # the SQL-backed store rejects ours on the duplicate primary key. # Honour the documented idempotency contract by resolving the # now-existing record — running the same owner reconciliation the fast - # path does — instead of surfacing the conflict as a 500. (The memory - # store overwrites rather than raising, so it never reaches here.) + # path does — instead of surfacing the conflict as a 500. The memory + # store serializes same-id creates under its per-thread lock and keeps + # its historical overwrite behavior, so it does not reach this branch. existing_record = await _resolve_existing_thread(thread_store, thread_id, thread_owner_user_id, thread_owner_kwargs) if existing_record is not None: return _existing_thread_response(thread_id, existing_record) diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py index 7a8356d0c..de4862a50 100644 --- a/backend/packages/harness/deerflow/persistence/bootstrap.py +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -72,11 +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 + During the compatibility window, a rollback-floor 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. + only when another process advanced the file to the explicitly reviewed + incarnation revision and that revision is absent from the rollback binary's + 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. @@ -109,23 +109,67 @@ _MIGRATIONS_DIR = Path(__file__).resolve().parent / "migrations" _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 +# One additive revision may be present when the 0020 rollback-floor 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. The exception also requires every current ORM table -# and column: the original 0018 + incarnation columns shape lacks projects and -# is no longer compatible with this build. Both skip paths validate that floor. -# Note: this tree's own chain already carries ``0019_projects`` / -# ``0020_threads_meta_project_id`` off ``0018_oauth_identity_pg_partial``; the -# ``0019_`` numeric prefix is intentionally reused. When the owning rollout -# revision merges it must re-parent onto the current head (see -# ``migrations/AGENTS.md``) so alembic never sees two heads off 0018. +# server default, table, index, constraint, or data backfill. Tests cross-pin +# this revision id and schema shape; amending that DDL requires re-auditing old +# repository reads and writes before this exception remains valid. The +# exception also requires every rollback-floor ORM table and column. +# +# The revision id was deployed briefly as an out-of-tree child of 0018 before +# becoming the in-tree successor to ``0020_threads_meta_project_id``. Current +# and future binaries that know this reused stamp validate the fixed +# canonical-0019 floor before migrating onward. This code also uses that fixed +# floor when tests remove 0019 from the local revision set to exercise the +# unknown-revision path. The published 0020 rollback binary instead validates +# its own ORM floor before skipping the revision. _FORWARD_COMPATIBLE_REVISION = "0019_thread_incarnations" +# Fixed table/column floor produced by the canonical in-tree 0019 revision. +# Never derive this from Base.metadata: future ORM fields belong to later +# migrations and must not be required before those migrations run. +_CANONICAL_0019_SCHEMA_FLOOR: dict[str, frozenset[str]] = { + "agents": frozenset("config created_at id name soul updated_at user_id".split()), + "channel_connections": frozenset( + "bot_user_id capabilities_json created_at external_account_id external_account_name id last_error_at last_seen_at metadata_json owner_user_id provider scopes_json status updated_at workspace_id workspace_name".split() + ), + "channel_conversations": frozenset("connection_id created_at external_conversation_id external_topic_id id owner_user_id provider thread_id updated_at".split()), + "channel_credentials": frozenset("connection_id encrypted_access_token encrypted_extra_json encrypted_refresh_token expires_at refresh_expires_at token_type updated_at version".split()), + "channel_oauth_states": frozenset("code_verifier_encrypted consumed_at created_at expires_at metadata_json nonce_hash owner_user_id provider redirect_after requested_scopes_json state_hash".split()), + "feedback": frozenset("comment created_at feedback_id message_id rating run_id thread_id user_id".split()), + "managed_subagents": frozenset("created_at definition id name updated_at".split()), + "mcp_tasks": frozenset( + """cancel_attempt_count cancel_requested_at completed_at consecutive_poll_error_count created_at dispatch_attempt dispatch_event dispatch_version driver_data driver_name error event_fingerprint event_version id input_required + last_cancel_error last_poll_error last_polled_at lease_expires_at lease_owner next_cancel_at next_notification_at next_poll_at notification_attempt_count notification_error notification_lease_expires_at + notification_lease_owner notification_run_id notification_status notified_version poll_attempt_count remote_task_id result result_artifact result_preview result_truncated run_id server_name status task_name + thread_id thread_incarnation tool_call_id updated_at user_id""".split() + ), + "personal_access_tokens": frozenset("created_at expires_at id last_used_at name revoked_at scopes token_digest user_id".split()), + "projects": frozenset("created_at id instructions name presentation status updated_at user_id".split()), + "run_events": frozenset("category content created_at event_metadata event_type id run_id seq thread_id user_id".split()), + "runs": frozenset( + """assistant_id cancel_action cancel_requested_at created_at error first_human_message follow_up_to_run_id idempotency_key kwargs_json last_ai_message lead_agent_tokens lease_expires_at llm_call_count + message_count metadata_json middleware_tokens model_name multitask_strategy operation_kind owner_worker_id run_id status stop_reason subagent_tokens thread_id token_usage_by_model total_input_tokens + total_output_tokens total_tokens updated_at user_id""".split() + ), + "scheduled_task_runs": frozenset("attempt_count created_at error finished_at id lease_expires_at lease_owner run_id scheduled_for started_at status task_id thread_id trigger".split()), + "scheduled_tasks": frozenset( + """assistant_id context_mode created_at id last_error last_run_at last_run_id last_thread_id lease_expires_at lease_owner next_run_at overlap_policy prompt run_count schedule_spec schedule_type + status thread_id timezone title updated_at user_id""".split() + ), + "subagent_batch_items": frozenset( + """acceptance_criteria acceptance_verdict attempt batch_id cancel_requested_at completed_at created_at error id item_key lease_expires_at lease_owner model_name position prompt result result_preview + result_truncated started_at status stop_reason token_usage updated_at""".split() + ), + "subagent_batches": frozenset("completed_at created_at execution_spec id max_attempts max_live_items max_running_items run_id status subagent_type submission_key thread_id title tool_call_id total_items updated_at user_id".split()), + "threads_meta": frozenset("assistant_id created_at display_name incarnation metadata_json project_id status thread_id updated_at user_id".split()), + "users": frozenset("created_at email id needs_setup oauth_id oauth_provider password_hash system_role token_version".split()), + "webhook_deliveries": frozenset("channel chat_id first_seen message_id workspace_id".split()), +} + # Baseline (stamp target for legacy DBs). Pinned here so the bootstrap layer # fails loudly if the baseline revision is ever renamed without updating the # stamp call. ``tests/test_persistence_bootstrap.py`` asserts this string is a @@ -325,25 +369,22 @@ async def _read_database_revision(conn: Any) -> str: def _validate_forward_schema(sync_conn: Any) -> None: - """Require the local repository schema before skipping unknown migrations. + """Require the fixed canonical-0019 schema before acting on its stamp. This is a presence check, not a general schema compatibility proof. The - allowlisted additive DDL still needs its separate read/write audit. Derive - the local floor from ORM metadata so a new mapped column cannot silently - invalidate the existing exception again. + allowlisted additive DDL still needs its separate read/write audit. The + floor remains independent of current ORM metadata so future revisions can + add their own schema after this validation succeeds. """ - import deerflow.persistence.models # noqa: F401 - from deerflow.persistence.base import Base - inspector = sa_inspect(sync_conn) tables = set(inspector.get_table_names()) missing = [] - for name, table in sorted(Base.metadata.tables.items()): + for name, required_columns in sorted(_CANONICAL_0019_SCHEMA_FLOOR.items()): if name not in tables: missing.append(name) continue columns = {column["name"] for column in inspector.get_columns(name)} - missing.extend(f"{name}.{column.name}" for column in table.columns if column.name not in columns) + missing.extend(f"{name}.{column_name}" for column_name in sorted(required_columns - columns)) if missing: raise RuntimeError( f"bootstrap: revision {_FORWARD_COMPATIBLE_REVISION!r} is missing required local schema: {', '.join(missing)}; refusing to start. See docs/database-forward-revision-recovery.md for the audited offline migration path." @@ -612,6 +653,13 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema await asyncio.to_thread(_upgrade, cfg, "head") elif decision == "versioned": + # The same revision id once named a different out-of-tree schema. + # Validate canonical 0019 before upgrading a known revision or + # accepting it through the unknown-revision path that current + # tests use to simulate old 0020. + if database_revision == _FORWARD_COMPATIBLE_REVISION: + async with engine.connect() as conn: + await conn.run_sync(_validate_forward_schema) if database_revision in known_revisions: logger.info( "bootstrap: branch=versioned revision=%s -> upgrade head (%s)", @@ -638,8 +686,6 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema current_revision, ) elif database_revision == _FORWARD_COMPATIBLE_REVISION: - async with engine.connect() as conn: - await conn.run_sync(_validate_forward_schema) logger.warning( "bootstrap: database revision %s is explicitly forward-compatible with local head %s and has its required tables and columns; skipping migration", database_revision, diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py index 4ccf080d8..73c851e9b 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/model.py @@ -20,6 +20,7 @@ class McpTaskRow(Base): id: Mapped[str] = mapped_column(String(64), primary_key=True) user_id: Mapped[str] = mapped_column(String(64), index=True) thread_id: Mapped[str] = mapped_column(String(64), index=True) + thread_incarnation: Mapped[str | None] = mapped_column(String(32), nullable=True) run_id: Mapped[str | None] = mapped_column(String(64), nullable=True) tool_call_id: Mapped[str | None] = mapped_column(String(128), nullable=True) server_name: Mapped[str] = mapped_column(String(MCP_TASK_SERVER_NAME_MAX_LENGTH)) diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index db9a9b2bf..8766f57ac 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.mcp.tasks import ATTENTION_TASK_STATUSES, POLLABLE_TASK_STATUSES, TERMINAL_TASK_STATUSES from deerflow.persistence.mcp_tasks.model import McpTaskRow +from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.utils.time import coerce_iso _POLLABLE_STATUS_VALUES = tuple(status.value for status in POLLABLE_TASK_STATUSES) @@ -98,6 +99,7 @@ class McpTaskRepository: @staticmethod def _row_to_dict(row: McpTaskRow) -> dict[str, Any]: data = row.to_dict() + data.pop("thread_incarnation", None) for key in _TIMESTAMP_FIELDS: if data.get(key) is not None: data[key] = coerce_iso(data[key]) @@ -152,6 +154,19 @@ class McpTaskRepository: ) _record_event_if_changed(row, tracking_degraded=False, now=now) async with self._sf() as session: + matching_thread = select(ThreadMetaRow.incarnation).where( + ThreadMetaRow.thread_id == thread_id, + or_(ThreadMetaRow.user_id == user_id, ThreadMetaRow.user_id.is_(None)), + ) + if session.get_bind().dialect.name == "sqlite": + # Keep lookup and write in one SQLite statement. A preliminary + # read would leave a delete/recreate window before the INSERT. + row.thread_incarnation = matching_thread.scalar_subquery() + else: + # FOR SHARE also conflicts with the FOR NO KEY UPDATE lock taken + # by an older writer's plain owner UPDATE. KEY SHARE would not, + # leaving a mixed-version ownership race before this INSERT. + row.thread_incarnation = (await session.execute(matching_thread.with_for_update(read=True))).scalar_one_or_none() session.add(row) try: await session.commit() diff --git a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md index dbb1d605c..3a491b4e3 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md @@ -19,47 +19,58 @@ The legacy branch handles pre-alembic databases that already have at least one D 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**: the local chain head is `0021_batch_acceptance` -(`0018_oauth_identity_pg_partial` → `0019_projects` → `0020_threads_meta_project_id` → `0021_batch_acceptance`). -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 conditionally allowlisted: -bootstrap first requires every current ORM table and column, then logs a -warning and leaves the schema untouched. The original rollout shape (0018 plus -the two incarnation columns) is now rejected: it lacks `projects` and -`threads_meta.project_id`, and the batch acceptance columns. Seeding current head is only a positive compatibility -fixture; tests must also construct the original 0018-based schema and assert -rejection on both the direct startup and SQLite race-recovery paths. The check -uses `conn.run_sync` reflection and derives its local floor from `Base.metadata` -so future ORM additions cannot silently bypass it. This checks presence only; -the additive DDL audit below still owns type/constraint compatibility. 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 owning -`0019_thread_incarnations` migration must cross-pin its revision id and schema -shape against the bootstrap contract; amending that DDL requires a fresh -old-repository compatibility audit. The `0019_` numeric prefix is intentionally -reused: `0019_projects` is this tree's in-chain revision, while -`0019_thread_incarnations` is the reserved, out-of-tree rollout id allowlisted -above — revision ids only need to be unique, not numerically ordered, but the -owning rollout revision must re-parent from `0018_oauth_identity_pg_partial` -onto this tree's head when it merges so `alembic` never sees two heads off -0018. 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 with all current ORM tables and columns 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. +**Rolling forward compatibility**: the local chain is +`0018_oauth_identity_pg_partial` → `0019_projects` → +`0020_threads_meta_project_id` → `0021_batch_acceptance` → +`0019_thread_incarnations`. The final revision deliberately retains the exact +id audited by the rollback-floor binary; Alembic orders revisions by +`down_revision`, not by the numeric prefix. + +The deployed `0020_threads_meta_project_id` rollback-floor binary knows neither +`0021_batch_acceptance` nor `0019_thread_incarnations`. It treats only the final +incarnation revision as forward-compatible, after reflection confirms every +table and column in its own ORM schema. The intervening acceptance columns and +the incarnation columns are nullable and have no server default, so old +repositories may omit them. Tests must prove old reads and writes across both +additive revisions; do not model the rollback binary with `0021` in its local +revision set. + +The same incarnation revision id existed briefly as an out-of-tree child of +0018. Current and future binaries that know the reused +`0019_thread_incarnations` id validate a fixed table/column snapshot of the +canonical in-tree 0019 schema whenever they see that stamp. The original +0018-plus-incarnation shape is rejected because it lacks Projects and +batch-acceptance schema; matching the revision string alone is not proof of the +new ancestry. The fixed floor is deliberately not derived from `Base.metadata`: +a future binary may add mapped columns after 0019 and must validate this floor +before Alembic adds them. A binary that knows 0019 upgrades normally after the +check. + +Tests that model the rollback binary by removing `0021_batch_acceptance` and +`0019_thread_incarnations` from the mocked local revision set exercise the +unknown-revision branch of the current implementation. They therefore use the +new fixed canonical-0019 floor; they do not execute the ORM-derived check from +the published 0020 binary. Reflection checks presence only; each migration +preflight owns type, nullability, and default compatibility. + +Any other unknown revision, an empty version table, or multiple version rows +fails closed. Do not broaden the allowlist without proving that the rollback +repositories can read, insert, and update through the newer schema. The +exception is reviewed only for the additive `0021` JSON columns plus nullable +VARCHAR(32) `threads_meta.incarnation` and `mcp_tasks.thread_incarnation` +columns, all without non-NULL defaults, constraints, or data backfills. The +incarnation migration and tests cross-pin its revision id and DDL shape; +changing either requires a fresh old-repository compatibility audit. SQLite's +stale-upgrade recovery remains available only to a rollback binary whose +migration tree owns neither post-`0020` revision. A current binary owns both, +so migration failures remain fatal. For an existing database with the original 0018-plus-incarnation shape, use the [audited offline recovery procedure](../../../../../../docs/database-forward-revision-recovery.md). Bootstrap never re-stamps an unknown revision automatically. After stopping all writers, backing up, and verifying the exact additive schema, the operator may -purge-stamp the known 0018 parent and apply this tree's 0019/0020 migrations; +purge-stamp the known 0018 parent and apply the Projects, acceptance, and +incarnation migrations through the current head; the extra nullable columns and their data remain intact. A regression exercises that procedure from the original schema and verifies repository reads/inserts and preservation of incarnation data. @@ -137,9 +148,9 @@ on installs that never enabled it. The convention is: - `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 - `migrations/versions/0019_projects.py` — creates the `projects` table (id/user_id/name/instructions/presentation/status + timestamps) for the Projects Phase-1 organization feature; chains after `0018_oauth_identity_pg_partial` -- `migrations/versions/0020_threads_meta_project_id.py` — adds nullable `threads_meta.project_id` plus `ix_threads_meta_project_id` (no FK by design: project delete clears membership first, and the reserved `deerflow_project_id` metadata key stays in sync); chains after `0019_projects`. The `0019_` numeric prefix is reused by the reserved out-of-tree `0019_thread_incarnations` — see the rolling-forward section above +- `migrations/versions/0020_threads_meta_project_id.py` — adds nullable `threads_meta.project_id` plus `ix_threads_meta_project_id` (no FK by design: project delete clears membership first, and the reserved `deerflow_project_id` metadata key stays in sync); chains after `0019_projects` +- `migrations/versions/0021_batch_acceptance.py` — adds nullable per-item acceptance criteria and verdict JSON columns after `0020_threads_meta_project_id`; legacy rows remain unchecked +- `migrations/versions/0019_thread_incarnations.py` — current head; chains after `0021_batch_acceptance` while retaining the exact revision id audited by the rollback-floor binary. Adds nullable `threads_meta.incarnation` / `mcp_tasks.thread_incarnation` columns. New thread rows get a random 32-character incarnation. Memory mutations serialize per thread; an overwrite inherits the existing incarnation, while a delete/recreate gets a new one. SQLite MCP task INSERTs copy the owner-or-shared incarnation with a scalar subquery in the same statement. PostgreSQL task creation holds `FOR SHARE`, which conflicts with both current `FOR UPDATE` mutations and an older writer's plain owner update (`FOR NO KEY UPDATE`). Missing or differently owned threads store NULL, old writers may omit both columns, and current API/task serialization hides them. The migration preflights both tables before DDL and its SQLite downgrade cleans only safe remnants from its own interrupted batch-copy attempt - `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) - -- `migrations/versions/0021_batch_acceptance.py` — adds nullable per-item acceptance criteria and verdict JSON columns after `0020_threads_meta_project_id`; legacy rows remain unchecked. diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0019_thread_incarnations.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0019_thread_incarnations.py new file mode 100644 index 000000000..094ca3721 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0019_thread_incarnations.py @@ -0,0 +1,100 @@ +"""add nullable thread incarnation columns. + +Revision ID: 0019_thread_incarnations +Revises: 0021_batch_acceptance +Create Date: 2026-09-05 + +This is the expand-only schema step. Existing rows remain nullable and no +runtime behavior consumes either column in this phase. +""" + +from __future__ import annotations + +import re +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +from deerflow.persistence.migrations._helpers import safe_add_column, safe_drop_column + +revision: str = "0019_thread_incarnations" +down_revision: str | Sequence[str] | None = "0021_batch_acceptance" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_SQLITE_BATCH_TABLES = ( + ("mcp_tasks", "_alembic_tmp_mcp_tasks"), + ("threads_meta", "_alembic_tmp_threads_meta"), +) +_NULL_VARCHAR_CAST = re.compile(r"(?:character\s+varying|varchar|text)(?:\s*\(\s*\d+\s*\))?", re.IGNORECASE) + + +def _cleanup_sqlite_downgrade_retry() -> None: + """Remove batch tables left by an interrupted prior downgrade attempt.""" + bind = op.get_bind() + if bind.dialect.name != "sqlite": + return + table_names = set(sa.inspect(bind).get_table_names()) + for source_name, temp_name in _SQLITE_BATCH_TABLES: + if temp_name not in table_names: + continue + if source_name not in table_names: + raise RuntimeError(f"Refusing to drop {temp_name}: {source_name} is missing and the batch table may be the only remaining copy") + op.drop_table(temp_name) + + +def _is_null_server_default(value: object) -> bool: + """Return whether reflection reports no default or a SQL NULL default.""" + if value is None: + return True + text = str(value).strip() + while text.startswith("(") and text.endswith(")"): + text = text[1:-1].strip() + base, separator, cast = text.partition("::") + while base.startswith("(") and base.endswith(")"): + base = base[1:-1].strip() + return base.casefold() == "null" and (not separator or _NULL_VARCHAR_CAST.fullmatch(cast.strip()) is not None) + + +def _assert_existing_column_compatible(table: str, column_name: str) -> None: + bind = op.get_bind() + inspector = sa.inspect(bind) + if table not in inspector.get_table_names(): + return + existing = {column["name"]: column for column in inspector.get_columns(table)} + actual = existing.get(column_name) + if actual is None: + return + + actual_type = actual.get("type") + length = getattr(actual_type, "length", None) + varchar_compatible = isinstance(actual_type, sa.VARCHAR) and (length is None or length >= 32) + actual_nullable = actual.get("nullable", True) + nullable_compatible = bool(actual_nullable) + actual_default = actual.get("default") + default_compatible = _is_null_server_default(actual_default) + if not varchar_compatible or not nullable_compatible or not default_compatible: + raise RuntimeError( + f"Incompatible pre-existing column {table}.{column_name}: expected nullable VARCHAR(32) or wider with no default or DEFAULT NULL, got type={actual_type!r}, nullable={actual_nullable!r}, server_default={actual_default!r}" + ) + + +def upgrade() -> None: + columns = ( + ("threads_meta", sa.Column("incarnation", sa.VARCHAR(length=32), nullable=True)), + ("mcp_tasks", sa.Column("thread_incarnation", sa.VARCHAR(length=32), nullable=True)), + ) + # Preflight every pre-existing column before changing either table. Expand + # must not silently accept a narrow, non-VARCHAR, NOT NULL, or value-defaulted + # manual column. In particular, finish both tables' preflight before DDL. + for table, column in columns: + _assert_existing_column_compatible(table, str(column.name)) + for table, column in columns: + safe_add_column(table, column) + + +def downgrade() -> None: + _cleanup_sqlite_downgrade_retry() + safe_drop_column("mcp_tasks", "thread_incarnation") + safe_drop_column("threads_meta", "incarnation") diff --git a/backend/packages/harness/deerflow/persistence/thread_meta/__init__.py b/backend/packages/harness/deerflow/persistence/thread_meta/__init__.py index 7908fc10e..ef4517525 100644 --- a/backend/packages/harness/deerflow/persistence/thread_meta/__init__.py +++ b/backend/packages/harness/deerflow/persistence/thread_meta/__init__.py @@ -4,7 +4,7 @@ from __future__ import annotations from typing import TYPE_CHECKING -from deerflow.persistence.thread_meta.base import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, THREAD_PROJECT_METADATA_KEY, InvalidMetadataFilterError, ThreadMetaStore +from deerflow.persistence.thread_meta.base import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, THREAD_PROJECT_METADATA_KEY, InvalidMetadataFilterError, ThreadMetaStore, ThreadOwnershipConflictError from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.persistence.thread_meta.sql import ThreadMetaRepository @@ -24,6 +24,7 @@ __all__ = [ "ThreadMetaRepository", "ThreadMetaRow", "ThreadMetaStore", + "ThreadOwnershipConflictError", "make_thread_store", ] diff --git a/backend/packages/harness/deerflow/persistence/thread_meta/base.py b/backend/packages/harness/deerflow/persistence/thread_meta/base.py index 8b5205f7e..621604355 100644 --- a/backend/packages/harness/deerflow/persistence/thread_meta/base.py +++ b/backend/packages/harness/deerflow/persistence/thread_meta/base.py @@ -52,6 +52,10 @@ class InvalidMetadataFilterError(ValueError): """Raised when all client-supplied metadata filter keys are rejected.""" +class ThreadOwnershipConflictError(Exception): + """Raised when create would overwrite a thread owned by another user.""" + + class ThreadMetaStore(abc.ABC): @abc.abstractmethod async def create( @@ -68,6 +72,14 @@ class ThreadMetaStore(abc.ABC): project inside the insert transaction and raise ``ProjectNotAssignableError`` on failure (no partial row).""" + @abc.abstractmethod + async def claim_unowned(self, thread_id: str, owner: str) -> bool: + """Atomically claim a legacy row whose owner is ``None``. + + Returns ``True`` only when this call changed ``user_id`` from ``None`` + to ``owner``. Missing and already-owned rows return ``False``. + """ + @abc.abstractmethod async def set_project(self, thread_id: str, project_id: str | None, *, user_id: str | None | _AutoSentinel = AUTO) -> bool: """Atomically move a thread into/out of a project (RFC v2 §5.2). diff --git a/backend/packages/harness/deerflow/persistence/thread_meta/memory.py b/backend/packages/harness/deerflow/persistence/thread_meta/memory.py index 8b4bb9a87..c1913932f 100644 --- a/backend/packages/harness/deerflow/persistence/thread_meta/memory.py +++ b/backend/packages/harness/deerflow/persistence/thread_meta/memory.py @@ -7,12 +7,14 @@ router for thread records. from __future__ import annotations +import uuid from typing import Any from langgraph.store.base import BaseStore from deerflow.persistence.json_compat import json_value_matches -from deerflow.persistence.thread_meta.base import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, ThreadMetaStore, _ProjectFilterUnset +from deerflow.persistence.thread_meta.base import PROJECT_FILTER_UNSET, THREAD_ARCHIVED_METADATA_KEY, THREAD_PINNED_METADATA_KEY, ThreadMetaStore, ThreadOwnershipConflictError, _ProjectFilterUnset +from deerflow.runtime.keyed_lock import AsyncKeyedLockTable from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id from deerflow.utils.time import coerce_iso, now_iso @@ -23,6 +25,7 @@ SEARCH_PAGE_SIZE = 500 class MemoryThreadMetaStore(ThreadMetaStore): def __init__(self, store: BaseStore) -> None: self._store = store + self._thread_locks = AsyncKeyedLockTable[str]() async def _get_owned_record( self, @@ -61,20 +64,35 @@ class MemoryThreadMetaStore(ThreadMetaStore): raise ProjectNotAssignableError(project_id) resolved_user_id = resolve_user_id(user_id, method_name="MemoryThreadMetaStore.create") - now = now_iso() - record: dict[str, Any] = { - "thread_id": thread_id, - "assistant_id": assistant_id, - "user_id": resolved_user_id, - "display_name": display_name, - "status": "idle", - "metadata": metadata or {}, - "values": {}, - "created_at": now, - "updated_at": now, - } - await self._store.aput(THREADS_NS, thread_id, record) - return record + async with self._thread_locks.hold(thread_id): + existing = await self._store.aget(THREADS_NS, thread_id) + if existing is not None and resolved_user_id is not None and existing.value.get("user_id") != resolved_user_id: + raise ThreadOwnershipConflictError(thread_id) + now = now_iso() + record: dict[str, Any] = { + "thread_id": thread_id, + "incarnation": (existing.value.get("incarnation") if existing is not None else None) or uuid.uuid4().hex, + "assistant_id": assistant_id, + "user_id": resolved_user_id, + "display_name": display_name, + "status": "idle", + "metadata": metadata or {}, + "values": {}, + "created_at": now, + "updated_at": now, + } + await self._store.aput(THREADS_NS, thread_id, record) + return record + + async def claim_unowned(self, thread_id: str, owner: str) -> bool: + async with self._thread_locks.hold(thread_id): + item = await self._store.aget(THREADS_NS, thread_id) + if item is None or item.value.get("user_id") is not None: + return False + record = dict(item.value) + record["user_id"] = owner + await self._store.aput(THREADS_NS, thread_id, record) + return True async def set_project(self, thread_id: str, project_id: str | None, *, user_id: str | None | _AutoSentinel = AUTO) -> bool: # Memory mode has no projects backend in Phase 1: membership moves @@ -152,49 +170,54 @@ class MemoryThreadMetaStore(ThreadMetaStore): remove_metadata_keys: tuple[str, ...] = (), user_id: str | None | _AutoSentinel = AUTO, ) -> None: - record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_display_name") - if record is None: - return - record["display_name"] = display_name - metadata = dict(record.get("metadata") or {}) - for key in remove_metadata_keys: - metadata.pop(key, None) - record["metadata"] = metadata - record["updated_at"] = now_iso() - await self._store.aput(THREADS_NS, thread_id, record) + async with self._thread_locks.hold(thread_id): + record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_display_name") + if record is None: + return + record["display_name"] = display_name + metadata = dict(record.get("metadata") or {}) + for key in remove_metadata_keys: + metadata.pop(key, None) + record["metadata"] = metadata + record["updated_at"] = now_iso() + await self._store.aput(THREADS_NS, thread_id, record) async def update_status(self, thread_id: str, status: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None: - record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_status") - if record is None: - return - record["status"] = status - record["updated_at"] = now_iso() - await self._store.aput(THREADS_NS, thread_id, record) + async with self._thread_locks.hold(thread_id): + record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_status") + if record is None: + return + record["status"] = status + record["updated_at"] = now_iso() + await self._store.aput(THREADS_NS, thread_id, record) async def update_metadata(self, thread_id: str, metadata: dict, *, touch: bool = True, user_id: str | None | _AutoSentinel = AUTO) -> None: - record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_metadata") - if record is None: - return - merged = dict(record.get("metadata") or {}) - merged.update(metadata) - record["metadata"] = merged - if touch: - record["updated_at"] = now_iso() - await self._store.aput(THREADS_NS, thread_id, record) + async with self._thread_locks.hold(thread_id): + record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_metadata") + if record is None: + return + merged = dict(record.get("metadata") or {}) + merged.update(metadata) + record["metadata"] = merged + if touch: + record["updated_at"] = now_iso() + await self._store.aput(THREADS_NS, thread_id, record) async def update_owner(self, thread_id: str, owner_user_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None: - record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_owner") - if record is None: - return - record["user_id"] = owner_user_id - record["updated_at"] = now_iso() - await self._store.aput(THREADS_NS, thread_id, record) + async with self._thread_locks.hold(thread_id): + record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_owner") + if record is None: + return + record["user_id"] = owner_user_id + record["updated_at"] = now_iso() + await self._store.aput(THREADS_NS, thread_id, record) async def delete(self, thread_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None: - record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.delete") - if record is None: - return - await self._store.adelete(THREADS_NS, thread_id) + async with self._thread_locks.hold(thread_id): + record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.delete") + if record is None: + return + await self._store.adelete(THREADS_NS, thread_id) @staticmethod def _item_to_dict(item) -> dict[str, Any]: @@ -202,6 +225,7 @@ class MemoryThreadMetaStore(ThreadMetaStore): val = item.value return { "thread_id": item.key, + "incarnation": val.get("incarnation"), "assistant_id": val.get("assistant_id"), "user_id": val.get("user_id"), "display_name": val.get("display_name"), diff --git a/backend/packages/harness/deerflow/persistence/thread_meta/model.py b/backend/packages/harness/deerflow/persistence/thread_meta/model.py index 939b12f44..5f95aea3d 100644 --- a/backend/packages/harness/deerflow/persistence/thread_meta/model.py +++ b/backend/packages/harness/deerflow/persistence/thread_meta/model.py @@ -14,6 +14,7 @@ class ThreadMetaRow(Base): __tablename__ = "threads_meta" thread_id: Mapped[str] = mapped_column(String(64), primary_key=True) + incarnation: Mapped[str | None] = mapped_column(String(32), nullable=True) assistant_id: Mapped[str | None] = mapped_column(String(128), index=True) user_id: Mapped[str | None] = mapped_column(String(64), index=True) project_id: Mapped[str | None] = mapped_column(String(64), index=True, nullable=True) diff --git a/backend/packages/harness/deerflow/persistence/thread_meta/sql.py b/backend/packages/harness/deerflow/persistence/thread_meta/sql.py index 28ac16d2e..b1fbd0858 100644 --- a/backend/packages/harness/deerflow/persistence/thread_meta/sql.py +++ b/backend/packages/harness/deerflow/persistence/thread_meta/sql.py @@ -3,10 +3,11 @@ from __future__ import annotations import logging +import uuid from datetime import UTC, datetime from typing import Any -from sqlalchemy import case, select, text, update +from sqlalchemy import case, column, select, table, text, update from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm.attributes import flag_modified @@ -78,6 +79,7 @@ class ThreadMetaRepository(ThreadMetaStore): raise ProjectNotAssignableError(project_id) row = ThreadMetaRow( thread_id=thread_id, + incarnation=uuid.uuid4().hex, assistant_id=assistant_id, user_id=resolved_user_id, display_name=display_name, @@ -92,6 +94,24 @@ class ThreadMetaRepository(ThreadMetaStore): await session.refresh(row) return self._row_to_dict(row) + async def claim_unowned(self, thread_id: str, owner: str) -> bool: + claim_target = table( + ThreadMetaRow.__tablename__, + column(ThreadMetaRow.thread_id.key), + column(ThreadMetaRow.user_id.key), + ) + async with self._sf() as session: + result = await session.execute( + update(claim_target) + .where( + claim_target.c.thread_id == thread_id, + claim_target.c.user_id.is_(None), + ) + .values(user_id=owner) + ) + await session.commit() + return result.rowcount > 0 + async def set_project( self, thread_id: str, @@ -357,9 +377,15 @@ class ThreadMetaRepository(ThreadMetaStore): """Move a thread metadata row to ``owner_user_id``.""" resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_owner") async with self._sf() as session: - if not await self._check_ownership(session, thread_id, resolved_user_id): + if session.get_bind().dialect.name == "sqlite": + await session.execute(text("BEGIN IMMEDIATE")) + row = await session.get(ThreadMetaRow, thread_id) + else: + row = (await session.execute(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).with_for_update())).scalar_one_or_none() + if row is None or (resolved_user_id is not None and row.user_id != resolved_user_id): return - await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(user_id=owner_user_id, updated_at=datetime.now(UTC))) + row.user_id = owner_user_id + row.updated_at = datetime.now(UTC) await session.commit() async def delete( @@ -370,10 +396,12 @@ class ThreadMetaRepository(ThreadMetaStore): ) -> None: resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.delete") async with self._sf() as session: - row = await session.get(ThreadMetaRow, thread_id) - if row is None: - return - if resolved_user_id is not None and row.user_id != resolved_user_id: + if session.get_bind().dialect.name == "sqlite": + await session.execute(text("BEGIN IMMEDIATE")) + row = await session.get(ThreadMetaRow, thread_id) + else: + row = (await session.execute(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).with_for_update())).scalar_one_or_none() + if row is None or (resolved_user_id is not None and row.user_id != resolved_user_id): return await session.delete(row) await session.commit() diff --git a/backend/tests/test_mcp_task_postgres.py b/backend/tests/test_mcp_task_postgres.py new file mode 100644 index 000000000..3375f1711 --- /dev/null +++ b/backend/tests/test_mcp_task_postgres.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import asyncio +import os +import uuid +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import pytest +import pytest_asyncio +from sqlalchemy import event, select, text, update +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from deerflow.config.database_config import DatabaseConfig +from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.mcp_tasks import McpTaskRepository +from deerflow.persistence.mcp_tasks.model import McpTaskRow +from deerflow.persistence.thread_meta import ThreadMetaRepository +from deerflow.persistence.thread_meta.model import ThreadMetaRow + +POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI") + +pytestmark = pytest.mark.skipif( + not POSTGRES_URL, + reason="requires TEST_POSTGRES_URI (real Postgres for row-lock interleaving)", +) + + +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)) + + +@pytest_asyncio.fixture() +async def postgres_repositories(): + assert POSTGRES_URL is not None + schema = f"mcp_incarnation_{uuid.uuid4().hex}" + await init_engine_from_config( + DatabaseConfig( + backend="postgres", + postgres_url=_postgres_url(POSTGRES_URL), + postgres_schema=schema, + ) + ) + session_factory = get_session_factory() + assert session_factory is not None + try: + yield ThreadMetaRepository(session_factory), McpTaskRepository(session_factory), session_factory + finally: + engine = get_engine() + assert engine is not None + async with engine.begin() as connection: + await connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')) + await close_engine() + + +async def _create_task(repo: McpTaskRepository, task_id: str) -> None: + await repo.create( + task_id=task_id, + user_id="user-1", + thread_id="thread-1", + run_id=None, + tool_call_id=None, + server_name="reports", + driver_name="fake", + remote_task_id=f"remote-{task_id}", + 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=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mutation", ["delete", "update_owner"]) +async def test_postgres_task_create_serializes_with_thread_mutation(postgres_repositories, mutation: str) -> None: + thread_repo, task_repo, session_factory = postgres_repositories + created = await thread_repo.create("thread-1", user_id="user-1") + + async with session_factory() as blocker: + locked = (await blocker.execute(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == "thread-1").with_for_update())).scalar_one() + assert locked.incarnation == created["incarnation"] + + if mutation == "delete": + mutation_task = asyncio.create_task(thread_repo.delete("thread-1", user_id=None)) + else: + mutation_task = asyncio.create_task(thread_repo.update_owner("thread-1", "user-2", user_id=None)) + await asyncio.sleep(0.1) + assert not mutation_task.done() + + # PostgreSQL grants the mutation's earlier queued row-lock request + # before this later FOR SHARE request. The task therefore observes the + # committed delete/owner change rather than the pre-mutation row. + create_task = asyncio.create_task(_create_task(task_repo, f"task-{mutation}")) + await asyncio.sleep(0.1) + assert not create_task.done() + await blocker.commit() + + await asyncio.wait_for(mutation_task, timeout=5) + await asyncio.wait_for(create_task, timeout=5) + + async with session_factory() as session: + task = await session.get(McpTaskRow, f"task-{mutation}") + assert task is not None + assert task.thread_incarnation is None + + +@pytest.mark.asyncio +async def test_postgres_task_create_uses_share_lock(postgres_repositories) -> None: + thread_repo, task_repo, _session_factory = postgres_repositories + await thread_repo.create("thread-1", user_id="user-1") + engine = get_engine() + assert engine is not None + statements: list[str] = [] + + def capture_statement(_conn, _cursor, statement, _parameters, _context, _executemany): + statements.append(" ".join(statement.upper().split())) + + event.listen(engine.sync_engine, "before_cursor_execute", capture_statement) + try: + await _create_task(task_repo, "task-share-lock") + finally: + event.remove(engine.sync_engine, "before_cursor_execute", capture_statement) + + assert any(statement.endswith("FOR SHARE") for statement in statements) + assert not any(statement.endswith("FOR KEY SHARE") for statement in statements) + + +@pytest.mark.asyncio +async def test_postgres_repository_holds_share_lock_until_task_commit(postgres_repositories) -> None: + thread_repo, _task_repo, session_factory = postgres_repositories + created = await thread_repo.create("thread-1", user_id="user-1") + engine = get_engine() + assert engine is not None + task_commit_entered = asyncio.Event() + allow_task_commit = asyncio.Event() + owner_update_started = asyncio.Event() + task_backend_pid: int | None = None + owner_backend_pid: int | None = None + + class PausingTaskCommitSession(AsyncSession): + async def commit(self) -> None: + nonlocal task_backend_pid + contains_target_task = any(isinstance(instance, McpTaskRow) and instance.id == "task-lock-lifetime" for instance in self.new) + if contains_target_task: + task_backend_pid = await self.scalar(text("SELECT pg_backend_pid()")) + task_commit_entered.set() + await allow_task_commit.wait() + await super().commit() + + task_session_factory = async_sessionmaker( + engine, + expire_on_commit=False, + class_=PausingTaskCommitSession, + ) + task_repo = McpTaskRepository(task_session_factory) + + def observe_owner_update(_conn, _cursor, statement, _parameters, _context, _executemany): + normalized = " ".join(statement.upper().split()) + if normalized.startswith("UPDATE THREADS_META SET USER_ID"): + owner_update_started.set() + + async def legacy_update_owner() -> None: + nonlocal owner_backend_pid + async with session_factory() as legacy_writer: + owner_backend_pid = await legacy_writer.scalar(text("SELECT pg_backend_pid()")) + await legacy_writer.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == "thread-1").values(user_id="user-2")) + await legacy_writer.commit() + + event.listen(engine.sync_engine, "before_cursor_execute", observe_owner_update) + task_create = asyncio.create_task(_create_task(task_repo, "task-lock-lifetime")) + owner_update = None + try: + await asyncio.wait_for(task_commit_entered.wait(), timeout=5) + owner_update = asyncio.create_task(legacy_update_owner()) + await asyncio.wait_for(owner_update_started.wait(), timeout=5) + assert task_backend_pid is not None + assert owner_backend_pid is not None + async with session_factory() as observer: + async with asyncio.timeout(5): + while True: + if owner_update.done(): + await owner_update + pytest.fail("owner update completed before the task transaction committed") + blockers = await observer.scalar( + text("SELECT pg_blocking_pids(:pid)"), + {"pid": owner_backend_pid}, + ) + if task_backend_pid in blockers: + break + await asyncio.sleep(0.01) + allow_task_commit.set() + await asyncio.wait_for(task_create, timeout=5) + await asyncio.wait_for(owner_update, timeout=5) + finally: + allow_task_commit.set() + event.remove(engine.sync_engine, "before_cursor_execute", observe_owner_update) + if not task_create.done(): + task_create.cancel() + await asyncio.gather(task_create, return_exceptions=True) + if owner_update is not None and not owner_update.done(): + owner_update.cancel() + await asyncio.gather(owner_update, return_exceptions=True) + + async with session_factory() as session: + thread = await session.get(ThreadMetaRow, "thread-1") + task = await session.get(McpTaskRow, "task-lock-lifetime") + assert thread is not None + assert thread.user_id == "user-2" + assert task is not None + assert task.thread_incarnation == created["incarnation"] diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index 79e5af32f..87db91e66 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -1,12 +1,17 @@ +import contextlib +import sqlite3 from datetime import UTC, datetime, timedelta import pytest import pytest_asyncio +from sqlalchemy import event from sqlalchemy.exc import IntegrityError from deerflow.config.database_config import DatabaseConfig -from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError, McpTaskRepository +from deerflow.persistence.mcp_tasks.model import McpTaskRow +from deerflow.persistence.thread_meta.model import ThreadMetaRow @pytest_asyncio.fixture(autouse=True) @@ -52,6 +57,169 @@ async def _create_working_task( ) +@pytest.mark.asyncio +async def test_legacy_task_writer_leaves_thread_incarnation_null(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + async with repo._sf() as session: + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation="owned-incarnation", + user_id="user-1", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + session.add( + McpTaskRow( + id="legacy-writer", + user_id="user-1", + thread_id="thread-1", + server_name="reports", + driver_name="fake", + remote_task_id="remote-legacy-writer", + task_name="Generate report", + status="working", + driver_data={}, + next_poll_at=now, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + + async with repo._sf() as session: + row = await session.get(McpTaskRow, "legacy-writer") + assert row is not None + assert row.thread_incarnation is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("thread_owner", "expected_incarnation"), + [ + ("user-1", "matching-owner"), + (None, "shared-thread"), + ("user-2", None), + ], +) +async def test_create_atomically_copies_accessible_thread_incarnation( + tmp_path, + thread_owner, + expected_incarnation, +): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + incarnation = expected_incarnation or "different-owner" + async with repo._sf() as session: + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation=incarnation, + user_id=thread_owner, + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + + task = await _create_working_task(repo, task_id="new-writer", now=now) + + assert "thread_incarnation" not in task + async with repo._sf() as session: + row = await session.get(McpTaskRow, "new-writer") + assert row is not None + assert row.thread_incarnation == expected_incarnation + + +@pytest.mark.asyncio +async def test_create_leaves_incarnation_null_without_matching_thread(tmp_path): + repo = await _make_repo(tmp_path) + + await _create_working_task(repo, task_id="missing-thread", now=datetime.now(UTC)) + + async with repo._sf() as session: + row = await session.get(McpTaskRow, "missing-thread") + assert row is not None + assert row.thread_incarnation is None + + +@pytest.mark.asyncio +async def test_create_observes_delete_and_recreate_at_insert_boundary(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + async with repo._sf() as session: + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation="old-incarnation", + user_id="user-1", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + + engine = get_engine() + assert engine is not None + replaced = False + insert_statement = None + + def replace_thread_before_task_insert( + _conn, + _cursor, + statement, + _parameters, + _context, + _executemany, + ): + nonlocal insert_statement, replaced + if replaced or not statement.lstrip().upper().startswith("INSERT INTO MCP_TASKS"): + return + replaced = True + insert_statement = statement + with contextlib.closing(sqlite3.connect(tmp_path / "deerflow.db")) as connection: + with connection: + connection.execute("DELETE FROM threads_meta WHERE thread_id = ?", ("thread-1",)) + connection.execute( + """ + INSERT INTO threads_meta ( + thread_id, incarnation, user_id, status, metadata_json, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + "thread-1", + "replacement-incarnation", + "user-1", + "idle", + "{}", + now.isoformat(), + now.isoformat(), + ), + ) + + event.listen(engine.sync_engine, "before_cursor_execute", replace_thread_before_task_insert) + try: + await _create_working_task(repo, task_id="racing-task", now=now) + finally: + event.remove(engine.sync_engine, "before_cursor_execute", replace_thread_before_task_insert) + + assert replaced is True + assert insert_statement is not None + normalized_insert = " ".join(insert_statement.upper().split()) + assert "SELECT THREADS_META.INCARNATION" in normalized_insert + assert "INSERT INTO MCP_TASKS" in normalized_insert + async with repo._sf() as session: + row = await session.get(McpTaskRow, "racing-task") + assert row is not None + assert row.thread_incarnation == "replacement-incarnation" + + @pytest.mark.asyncio async def test_remote_task_id_is_unique_per_user_and_server(tmp_path): repo = await _make_repo(tmp_path) diff --git a/backend/tests/test_memory_thread_meta_isolation.py b/backend/tests/test_memory_thread_meta_isolation.py index fdc54e6bc..b5a2b5b72 100644 --- a/backend/tests/test_memory_thread_meta_isolation.py +++ b/backend/tests/test_memory_thread_meta_isolation.py @@ -6,12 +6,14 @@ the in-memory LangGraph Store backend used when database.backend=memory. from __future__ import annotations +import asyncio from types import SimpleNamespace import pytest from langgraph.store.memory import InMemoryStore from deerflow.persistence.projects import ProjectNotAssignableError +from deerflow.persistence.thread_meta import ThreadOwnershipConflictError from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore from deerflow.runtime.user_context import reset_current_user, set_current_user @@ -36,6 +38,207 @@ def store(): return MemoryThreadMetaStore(InMemoryStore()) +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_create_generates_stable_incarnation(store): + with _as_user(USER_A): + created = await store.create("incarnation-thread") + fetched = await store.get("incarnation-thread") + + assert len(created["incarnation"]) == 32 + assert fetched is not None + assert fetched["incarnation"] == created["incarnation"] + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_concurrent_create_preserves_overwrite_semantics_and_incarnation(store): + with _as_user(USER_A): + outcomes = await asyncio.gather( + store.create("same-thread", display_name="first"), + store.create("same-thread", display_name="second"), + ) + fetched = await store.get("same-thread") + + assert {outcome["display_name"] for outcome in outcomes} == {"first", "second"} + assert len({outcome["incarnation"] for outcome in outcomes}) == 1 + assert fetched is not None + assert fetched["display_name"] in {"first", "second"} + assert fetched["incarnation"] == outcomes[0]["incarnation"] + assert not store._thread_locks._entries_by_loop + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_duplicate_create_overwrites_but_inherits_incarnation(store): + with _as_user(USER_A): + original = await store.create("duplicate", display_name="original") + replacement = await store.create("duplicate", display_name="replacement") + fetched = await store.get("duplicate") + + assert replacement["incarnation"] == original["incarnation"] + assert replacement["display_name"] == "replacement" + assert fetched == replacement + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_duplicate_create_rejects_different_owner_without_overwrite(store): + with _as_user(USER_A): + original = await store.create("foreign-duplicate", display_name="original") + + with _as_user(USER_B): + with pytest.raises(ThreadOwnershipConflictError): + await store.create("foreign-duplicate", display_name="replacement") + + with _as_user(USER_A): + fetched = await store.get("foreign-duplicate") + assert fetched == original + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_concurrent_create_allows_one_owner_and_rejects_the_other(store): + async def create_for(user, display_name): + with _as_user(user): + return await store.create("owner-race", display_name=display_name) + + outcomes = await asyncio.gather(create_for(USER_A, "A"), create_for(USER_B, "B"), return_exceptions=True) + + records = [outcome for outcome in outcomes if isinstance(outcome, dict)] + conflicts = [outcome for outcome in outcomes if isinstance(outcome, ThreadOwnershipConflictError)] + assert len(records) == 1 + assert len(conflicts) == 1 + assert await store.get("owner-race", user_id=None) == records[0] + assert not store._thread_locks._entries_by_loop + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_duplicate_create_with_explicit_none_preserves_unscoped_overwrite(store): + with _as_user(USER_A): + original = await store.create("admin-overwrite", display_name="original") + + replacement = await store.create("admin-overwrite", display_name="replacement", user_id=None) + + assert replacement["incarnation"] == original["incarnation"] + assert replacement["display_name"] == "replacement" + assert replacement["user_id"] is None + assert await store.get("admin-overwrite", user_id=None) == replacement + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_claim_unowned_only_changes_none_owner(store): + await store.create("legacy", user_id=None) + await store.create("owned", user_id="original-owner") + + assert await store.claim_unowned("missing", "owner-a") is False + assert await store.claim_unowned("owned", "owner-a") is False + assert await store.claim_unowned("legacy", "owner-a") is True + assert await store.claim_unowned("legacy", "owner-b") is False + + assert (await store.get("owned", user_id=None))["user_id"] == "original-owner" + assert (await store.get("legacy", user_id=None))["user_id"] == "owner-a" + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_concurrent_claim_unowned_has_exactly_one_winner(store): + await store.create("legacy-race", user_id=None) + + outcomes = await asyncio.gather( + store.claim_unowned("legacy-race", "owner-a"), + store.claim_unowned("legacy-race", "owner-b"), + ) + + assert sorted(outcomes) == [False, True] + record = await store.get("legacy-race", user_id=None) + assert record["user_id"] in {"owner-a", "owner-b"} + assert not store._thread_locks._entries_by_loop + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +@pytest.mark.parametrize( + "contender", + ["create", "claim_unowned", "update_display_name", "update_status", "update_metadata", "update_owner", "delete"], +) +async def test_all_memory_mutations_share_the_per_thread_lock(contender): + class PausingGetStore(InMemoryStore): + def __init__(self): + super().__init__() + self.pause_next_get = False + self.get_entered = asyncio.Event() + self.allow_get = asyncio.Event() + + async def aget(self, namespace, key): + item = await super().aget(namespace, key) + if self.pause_next_get: + self.pause_next_get = False + self.get_entered.set() + await self.allow_get.wait() + return item + + backend = PausingGetStore() + store = MemoryThreadMetaStore(backend) + await store.create("locked-thread", user_id=None) + backend.pause_next_get = True + holder = asyncio.create_task(store.update_metadata("locked-thread", {"holder": True}, user_id=None)) + await backend.get_entered.wait() + + operations = { + "create": lambda: store.create("locked-thread", display_name="replacement", user_id=None), + "claim_unowned": lambda: store.claim_unowned("locked-thread", "new-owner"), + "update_display_name": lambda: store.update_display_name("locked-thread", "renamed", user_id=None), + "update_status": lambda: store.update_status("locked-thread", "busy", user_id=None), + "update_metadata": lambda: store.update_metadata("locked-thread", {"contender": True}, user_id=None), + "update_owner": lambda: store.update_owner("locked-thread", "new-owner", user_id=None), + "delete": lambda: store.delete("locked-thread", user_id=None), + } + waiting = asyncio.create_task(operations[contender]()) + await asyncio.sleep(0) + assert not waiting.done(), f"{contender} bypassed the per-thread lock" + + backend.allow_get.set() + await asyncio.gather(holder, waiting) + assert not store._thread_locks._entries_by_loop + + +@pytest.mark.anyio +@pytest.mark.no_auto_user +async def test_delete_and_recreate_are_serialized_per_thread(): + class PausingDeleteStore(InMemoryStore): + def __init__(self): + super().__init__() + self.delete_entered = asyncio.Event() + self.allow_delete = asyncio.Event() + + async def adelete(self, namespace, key): + self.delete_entered.set() + await self.allow_delete.wait() + await super().adelete(namespace, key) + + backend = PausingDeleteStore() + store = MemoryThreadMetaStore(backend) + with _as_user(USER_A): + original = await store.create("replace-me") + delete_task = asyncio.create_task(store.delete("replace-me")) + await backend.delete_entered.wait() + create_task = asyncio.create_task(store.create("replace-me")) + await asyncio.sleep(0) + assert not create_task.done() + + backend.allow_delete.set() + await delete_task + replacement = await create_task + fetched = await store.get("replace-me") + + assert replacement["incarnation"] != original["incarnation"] + assert fetched == replacement + assert not store._thread_locks._entries_by_loop + + @pytest.mark.anyio @pytest.mark.no_auto_user async def test_search_isolation(store): diff --git a/backend/tests/test_migration_0004_run_ownership_dedupe.py b/backend/tests/test_migration_0004_run_ownership_dedupe.py index 16c3ee605..2e36f4d65 100644 --- a/backend/tests/test_migration_0004_run_ownership_dedupe.py +++ b/backend/tests/test_migration_0004_run_ownership_dedupe.py @@ -32,6 +32,7 @@ from sqlalchemy.orm import Session import deerflow.persistence.models # noqa: F401 -- registers ORM models from deerflow.persistence.base import Base +from deerflow.persistence.bootstrap import _get_head_revision from deerflow.persistence.engine import close_engine, init_engine from deerflow.persistence.run.model import RunRow @@ -157,7 +158,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p with sqlite3.connect(db_path) as raw: version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() # Bootstrap upgrades through the later revisions after 0004. - assert version_row[0] == "0021_batch_acceptance" + assert version_row[0] == _get_head_revision() # Sanity: the invariant the index enforces is now true — at most one # active row per thread. diff --git a/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py b/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py index 65c9527c9..64e4af011 100644 --- a/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py +++ b/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py @@ -35,6 +35,7 @@ from sqlalchemy.orm import Session import deerflow.persistence.models # noqa: F401 -- registers ORM models from deerflow.persistence.base import Base +from deerflow.persistence.bootstrap import _get_head_revision from deerflow.persistence.engine import close_engine, init_engine from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow @@ -173,7 +174,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm with sqlite3.connect(db_path) as raw: version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0021_batch_acceptance" + assert version_row[0] == _get_head_revision() # Sanity: the invariant the index enforces now holds — at most one # active row per task_id. diff --git a/backend/tests/test_migration_0015_scheduled_task_enqueue.py b/backend/tests/test_migration_0015_scheduled_task_enqueue.py index a9f998d8b..312d306f6 100644 --- a/backend/tests/test_migration_0015_scheduled_task_enqueue.py +++ b/backend/tests/test_migration_0015_scheduled_task_enqueue.py @@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import create_async_engine import deerflow.persistence.models # noqa: F401 from deerflow.persistence.base import Base -from deerflow.persistence.bootstrap import bootstrap_schema +from deerflow.persistence.bootstrap import _get_head_revision, bootstrap_schema pytestmark = pytest.mark.asyncio @@ -57,7 +57,7 @@ async def test_migration_interrupts_legacy_queue_and_adds_claim_fields(tmp_path: # Bootstrap always advances to the repository head after exercising # the 0015 migration behavior below. - assert version == "0021_batch_acceptance" + assert version == _get_head_revision() assert {"lease_owner", "lease_expires_at", "attempt_count"} <= columns.keys() assert columns["attempt_count"]["nullable"] is False assert overlap_policy == "enqueue" diff --git a/backend/tests/test_migration_0019_thread_incarnations.py b/backend/tests/test_migration_0019_thread_incarnations.py new file mode 100644 index 000000000..7b4f2ad98 --- /dev/null +++ b/backend/tests/test_migration_0019_thread_incarnations.py @@ -0,0 +1,356 @@ +from __future__ import annotations + +import asyncio +import importlib +import os +import uuid +from pathlib import Path +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import pytest +import sqlalchemy as sa +from alembic import command as alembic_command +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.dialects import postgresql +from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy.schema import CreateTable + +from deerflow.persistence.bootstrap import _FORWARD_COMPATIBLE_REVISION, _get_alembic_config + +_PREVIOUS = "0021_batch_acceptance" +_REVISION = "0019_thread_incarnations" +_MIGRATION_MODULE = "deerflow.persistence.migrations.versions.0019_thread_incarnations" + + +def _asyncpg_url(url: str | None) -> str | None: + if not url: + return url + parts = urlsplit(url) + scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme + 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(scheme=scheme, query=query)) + + +_POSTGRES_URL = _asyncpg_url(os.getenv("DEERFLOW_TEST_POSTGRES_URL") or os.getenv("TEST_POSTGRES_URI")) + + +def test_0019_matches_reviewed_rollback_contract(monkeypatch: pytest.MonkeyPatch) -> None: + migration = importlib.import_module(_MIGRATION_MODULE) + events: list[tuple[str, str, str, int | None, bool, object]] = [] + + def capture_preflight(table: str, column_name: str) -> None: + events.append(("preflight", table, column_name, None, True, None)) + + def capture_add(table: str, column: sa.Column) -> None: + assert isinstance(column.type, sa.VARCHAR) + events.append(("add", table, str(column.name), column.type.length, bool(column.nullable), column.server_default)) + + class NoAdditionalOperations: + def __getattr__(self, name: str): + raise AssertionError(f"0019 rollback contract does not allow direct Alembic operation: {name}") + + monkeypatch.setattr(migration, "_assert_existing_column_compatible", capture_preflight) + monkeypatch.setattr(migration, "safe_add_column", capture_add) + monkeypatch.setattr(migration, "op", NoAdditionalOperations()) + + migration.upgrade() + + assert migration.revision == _FORWARD_COMPATIBLE_REVISION == _REVISION + assert migration.down_revision == _PREVIOUS + assert events == [ + ("preflight", "threads_meta", "incarnation", None, True, None), + ("preflight", "mcp_tasks", "thread_incarnation", None, True, None), + ("add", "threads_meta", "incarnation", 32, True, None), + ("add", "mcp_tasks", "thread_incarnation", 32, True, None), + ] + + +@pytest.mark.asyncio +async def test_sqlite_0019_adds_and_drops_nullable_columns(tmp_path: Path) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}") + cfg = _get_alembic_config(engine) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS) + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + async with engine.connect() as conn: + thread_columns = {column["name"]: column for column in await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("threads_meta"))} + task_columns = {column["name"]: column for column in await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks"))} + version = await conn.scalar(sa.text("SELECT version_num FROM alembic_version")) + + assert version == _REVISION + assert thread_columns["incarnation"]["nullable"] is True + assert task_columns["thread_incarnation"]["nullable"] is True + + await asyncio.to_thread(alembic_command.downgrade, cfg, _PREVIOUS) + async with engine.connect() as conn: + thread_columns = {column["name"] for column in await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("threads_meta"))} + task_columns = {column["name"] for column in await conn.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks"))} + assert "incarnation" not in thread_columns + assert "thread_incarnation" not in task_columns + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_0019_round_trip_preserves_parent_schema_data(tmp_path: Path) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'projects.db'}") + cfg = _get_alembic_config(engine) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS) + async with engine.begin() as connection: + await connection.execute( + sa.text( + "INSERT INTO projects " + "(id, user_id, name, instructions, presentation, status, created_at, updated_at) " + "VALUES ('project-1', 'user-1', 'Project', 'Keep me', " + "'{\"theme\":\"dark\"}', 'active', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" + ) + ) + await connection.execute( + sa.text("INSERT INTO threads_meta (thread_id, user_id, status, metadata_json, project_id, created_at, updated_at) VALUES ('thread-1', 'user-1', 'idle', '{}', 'project-1', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)") + ) + await connection.execute( + sa.text( + "INSERT INTO subagent_batches " + "(id, user_id, thread_id, submission_key, title, subagent_type, " + "status, total_items, max_live_items, max_running_items, " + "max_attempts, execution_spec, created_at, updated_at) " + "VALUES ('batch-1', 'user-1', 'thread-1', 'submission-1', " + "'Batch', 'general-purpose', 'completed', 1, 1, 1, 2, '{}', " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" + ) + ) + await connection.execute( + sa.text( + "INSERT INTO subagent_batch_items " + "(id, batch_id, item_key, position, prompt, acceptance_criteria, " + "acceptance_verdict, status, attempt, result_truncated, created_at, updated_at) " + "VALUES ('item-1', 'batch-1', 'item', 0, 'Prompt', " + ":criteria, :verdict, 'succeeded', 1, 0, " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" + ), + { + "criteria": '["must pass"]', + "verdict": '{"passed":true}', + }, + ) + + async def assert_preserved(*, incarnation_columns: bool) -> None: + async with engine.connect() as connection: + project = (await connection.execute(sa.text("SELECT id, user_id, name, instructions, presentation, status FROM projects WHERE id = 'project-1'"))).one() + membership = await connection.scalar(sa.text("SELECT project_id FROM threads_meta WHERE thread_id = 'thread-1'")) + indexes = {index["name"] for index in await connection.run_sync(lambda sync: sa.inspect(sync).get_indexes("threads_meta"))} + thread_columns = {column["name"] for column in await connection.run_sync(lambda sync: sa.inspect(sync).get_columns("threads_meta"))} + task_columns = {column["name"] for column in await connection.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks"))} + acceptance = (await connection.execute(sa.text("SELECT json_extract(acceptance_criteria, '$[0]'), json_extract(acceptance_verdict, '$.passed') FROM subagent_batch_items WHERE id = 'item-1'"))).one() + + assert tuple(project[:4]) == ("project-1", "user-1", "Project", "Keep me") + assert project.status == "active" + assert '"theme"' in str(project.presentation) and '"dark"' in str(project.presentation) + assert membership == "project-1" + assert tuple(acceptance) == ("must pass", 1) + assert "ix_threads_meta_project_id" in indexes + assert ("incarnation" in thread_columns) is incarnation_columns + assert ("thread_incarnation" in task_columns) is incarnation_columns + + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + await assert_preserved(incarnation_columns=True) + + await asyncio.to_thread(alembic_command.downgrade, cfg, _PREVIOUS) + await assert_preserved(incarnation_columns=False) + + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + await assert_preserved(incarnation_columns=True) + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_0019_reapply_does_not_report_varchar_drift( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}") + cfg = _get_alembic_config(engine) + migration = importlib.import_module(_MIGRATION_MODULE) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + def reapply(sync_connection) -> None: + context = MigrationContext.configure(sync_connection) + with Operations.context(context): + migration.upgrade() + + with caplog.at_level("WARNING", logger="deerflow.persistence.migrations._helpers"): + async with engine.begin() as connection: + await connection.run_sync(reapply) + + assert "drifts from the model definition" not in caplog.text + finally: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "column_sql", + [ + "VARCHAR(16) NULL", + "TEXT NULL", + 'VARCHAR(32) NOT NULL DEFAULT "legacy"', + "VARCHAR(32) NULL DEFAULT 'legacy'", + ], +) +async def test_sqlite_0019_fails_fast_on_incompatible_existing_column(tmp_path: Path, column_sql: str) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}") + cfg = _get_alembic_config(engine) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS) + async with engine.begin() as connection: + await connection.execute(sa.text(f"ALTER TABLE mcp_tasks ADD COLUMN thread_incarnation {column_sql}")) + + with pytest.raises(RuntimeError, match="with no default or DEFAULT NULL"): + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + async with engine.connect() as connection: + version = await connection.scalar(sa.text("SELECT version_num FROM alembic_version")) + thread_columns = {column["name"] for column in await connection.run_sync(lambda sync: sa.inspect(sync).get_columns("threads_meta"))} + assert version == _PREVIOUS + assert "incarnation" not in thread_columns + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_0019_accepts_wider_nullable_varchar(tmp_path: Path) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}") + cfg = _get_alembic_config(engine) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS) + async with engine.begin() as connection: + await connection.execute(sa.text("ALTER TABLE mcp_tasks ADD COLUMN thread_incarnation VARCHAR(64) NULL")) + + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + async with engine.connect() as connection: + columns = {column["name"]: column for column in await connection.run_sync(lambda sync: sa.inspect(sync).get_columns("mcp_tasks"))} + assert columns["thread_incarnation"]["type"].length == 64 + assert columns["thread_incarnation"]["nullable"] is True + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_0019_accepts_default_null_and_legacy_writer_omission(tmp_path: Path) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}") + cfg = _get_alembic_config(engine) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS) + async with engine.begin() as connection: + await connection.execute(sa.text("ALTER TABLE threads_meta ADD COLUMN incarnation VARCHAR(32) NULL DEFAULT NULL")) + + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + async with engine.begin() as connection: + await connection.execute(sa.text("INSERT INTO threads_meta (thread_id, status, metadata_json, created_at, updated_at) VALUES ('legacy-writer', 'idle', '{}', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")) + incarnation = await connection.scalar(sa.text("SELECT incarnation FROM threads_meta WHERE thread_id = 'legacy-writer'")) + assert incarnation is None + finally: + await engine.dispose() + + +def test_postgresql_0019_column_ddl_compiles_without_value_default() -> None: + table = sa.Table( + "incarnation_compile_check", + sa.MetaData(), + sa.Column("incarnation", sa.VARCHAR(length=32), nullable=True), + ) + + ddl = str(CreateTable(table).compile(dialect=postgresql.dialect())) + + assert "incarnation VARCHAR(32)" in ddl + assert "DEFAULT" not in ddl + + +def test_postgresql_reflected_default_null_cast_is_semantically_null() -> None: + migration = importlib.import_module(_MIGRATION_MODULE) + + assert migration._is_null_server_default("NULL::character varying") is True + assert migration._is_null_server_default("(NULL)::character varying") is True + assert migration._is_null_server_default("'legacy'::character varying") is False + assert migration._is_null_server_default("NULL::integer IS NULL") is False + + +def test_postgresql_ci_url_is_normalized_for_asyncpg() -> None: + assert _asyncpg_url("postgresql://user:pass@localhost/db?sslmode=disable") == "postgresql+asyncpg://user:pass@localhost/db" + + +@pytest.mark.asyncio +@pytest.mark.skipif(not _POSTGRES_URL, reason="set TEST_POSTGRES_URI or DEERFLOW_TEST_POSTGRES_URL to run live PostgreSQL tests") +async def test_postgresql_0019_accepts_default_null_and_legacy_writer_omission() -> None: + schema = f"deerflow_0019_{uuid.uuid4().hex[:12]}" + engine = create_async_engine(_POSTGRES_URL or "") + cfg = _get_alembic_config(engine, postgres_schema=schema) + try: + async with engine.begin() as connection: + await connection.execute(sa.text(f'CREATE SCHEMA "{schema}"')) + await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS) + async with engine.begin() as connection: + await connection.execute(sa.text(f'ALTER TABLE "{schema}".threads_meta ADD COLUMN incarnation VARCHAR(32) NULL DEFAULT NULL')) + + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + async with engine.begin() as connection: + await connection.execute(sa.text(f"INSERT INTO \"{schema}\".threads_meta (thread_id, status, metadata_json, created_at, updated_at) VALUES ('legacy-writer', 'idle', '{{}}', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")) + incarnation = await connection.scalar(sa.text(f"SELECT incarnation FROM \"{schema}\".threads_meta WHERE thread_id = 'legacy-writer'")) + assert incarnation is None + finally: + async with engine.begin() as connection: + await connection.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')) + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_0019_downgrade_retry_cleans_failed_batch_table( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}") + cfg = _get_alembic_config(engine) + migration = importlib.import_module(_MIGRATION_MODULE) + try: + await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION) + + def fail_then_retry(sync_connection) -> None: + context = MigrationContext.configure(sync_connection) + original_safe_drop = migration.safe_drop_column + injected = False + + def fail_after_batch_temp_create(table: str, column_name: str) -> None: + nonlocal injected + if table == "mcp_tasks" and not injected: + injected = True + sync_connection.exec_driver_sql("CREATE TABLE _alembic_tmp_mcp_tasks AS SELECT * FROM mcp_tasks") + raise RuntimeError("injected batch downgrade failure") + original_safe_drop(table, column_name) + + with Operations.context(context): + monkeypatch.setattr(migration, "safe_drop_column", fail_after_batch_temp_create) + with pytest.raises(RuntimeError, match="injected"): + migration.downgrade() + sync_connection.commit() + + monkeypatch.setattr(migration, "safe_drop_column", original_safe_drop) + migration.downgrade() + sync_connection.commit() + + tables = set(sa.inspect(sync_connection).get_table_names()) + assert "_alembic_tmp_mcp_tasks" not in tables + assert "thread_incarnation" not in {column["name"] for column in sa.inspect(sync_connection).get_columns("mcp_tasks")} + assert "incarnation" not in {column["name"] for column in sa.inspect(sync_connection).get_columns("threads_meta")} + + async with engine.connect() as connection: + await connection.run_sync(fail_then_retry) + finally: + await engine.dispose() diff --git a/backend/tests/test_migration_0021_batch_acceptance.py b/backend/tests/test_migration_0021_batch_acceptance.py index 54db0bccd..c4adfbb34 100644 --- a/backend/tests/test_migration_0021_batch_acceptance.py +++ b/backend/tests/test_migration_0021_batch_acceptance.py @@ -60,7 +60,21 @@ async def test_forward_revision_cannot_skip_required_batch_columns(tmp_path, mon # Keep the project schema present so only the batch-column guard can # reject this database, on both direct and concurrent-startup paths. await asyncio.to_thread(bootstrap._upgrade, cfg, "0020_threads_meta_project_id") + 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)")) if race: + current_head, current_revisions = bootstrap._get_revision_metadata() + assert current_head == "0019_thread_incarnations" + assert {"0020_threads_meta_project_id", "0021_batch_acceptance", current_head} <= current_revisions + monkeypatch.setattr( + bootstrap, + "_get_revision_metadata", + lambda: ( + "0020_threads_meta_project_id", + current_revisions - {"0021_batch_acceptance", current_head}, + ), + ) def raced_upgrade(*args): sync = sa.create_engine(f"sqlite:///{tmp_path / 'forward.db'}") diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py index cefeb2076..c12b083e9 100644 --- a/backend/tests/test_persistence_bootstrap.py +++ b/backend/tests/test_persistence_bootstrap.py @@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default asyncio_test = pytest.mark.asyncio -HEAD = "0021_batch_acceptance" +HEAD = _get_head_revision() BASELINE = "0001_baseline" @@ -61,9 +61,14 @@ async def _table_names(engine) -> set[str]: return await conn.run_sync(lambda c: set(sa.inspect(c).get_table_names())) -async def _runs_columns(engine) -> set[str]: +async def _table_columns(engine, table_name: str) -> dict[str, dict]: async with engine.connect() as conn: - return await conn.run_sync(lambda c: {col["name"] for col in sa.inspect(c).get_columns("runs")}) + columns = await conn.run_sync(lambda c: sa.inspect(c).get_columns(table_name)) + return {column["name"]: column for column in columns} + + +async def _runs_columns(engine) -> set[str]: + return set(await _table_columns(engine, "runs")) async def _runs_column_meta(engine, column_name: str) -> dict: @@ -149,6 +154,10 @@ async def test_empty_branch_creates_all_and_stamps_head(tmp_path: Path) -> None: assert "token_usage_by_model" in await _runs_columns(engine) assert "cancel_action" in await _runs_columns(engine) assert "cancel_requested_at" in await _runs_columns(engine) + thread_columns = await _table_columns(engine, "threads_meta") + task_columns = await _table_columns(engine, "mcp_tasks") + assert thread_columns["incarnation"]["nullable"] is True + assert task_columns["thread_incarnation"]["nullable"] is True operation_kind = await _runs_column_meta(engine, "operation_kind") assert operation_kind["nullable"] is False assert await _alembic_version(engine) == HEAD diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py index 4fb604e39..e74510990 100644 --- a/backend/tests/test_persistence_bootstrap_concurrency.py +++ b/backend/tests/test_persistence_bootstrap_concurrency.py @@ -23,12 +23,12 @@ from sqlalchemy.ext.asyncio import create_async_engine import deerflow.persistence.models # noqa: F401 from deerflow.persistence import bootstrap as bootstrap_mod -from deerflow.persistence.bootstrap import bootstrap_schema +from deerflow.persistence.bootstrap import _get_head_revision, bootstrap_schema pytestmark = pytest.mark.asyncio -HEAD = "0021_batch_acceptance" +HEAD = _get_head_revision() def _url(tmp_path: Path) -> str: diff --git a/backend/tests/test_persistence_bootstrap_regression.py b/backend/tests/test_persistence_bootstrap_regression.py index 7d261d9ca..92c855d22 100644 --- a/backend/tests/test_persistence_bootstrap_regression.py +++ b/backend/tests/test_persistence_bootstrap_regression.py @@ -27,6 +27,7 @@ import sqlalchemy as sa import deerflow.persistence.models # noqa: F401 -- registers ORM models from deerflow.persistence.base import Base +from deerflow.persistence.bootstrap import _get_head_revision from deerflow.persistence.engine import close_engine, get_session_factory, init_engine from deerflow.persistence.run import RunRepository @@ -76,7 +77,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()} assert "token_usage_by_model" in cols version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0021_batch_acceptance" + assert version_row[0] == _get_head_revision() # And the read path that originally 500'd must now succeed. sf = get_session_factory() @@ -116,6 +117,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path # No duplicate column -- list, not set, to catch dupes. assert cols.count("token_usage_by_model") == 1 version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0021_batch_acceptance" + assert version_row[0] == _get_head_revision() finally: await close_engine() diff --git a/backend/tests/test_persistence_forward_revision_compat.py b/backend/tests/test_persistence_forward_revision_compat.py index 56806d9f6..789cd35d1 100644 --- a/backend/tests/test_persistence_forward_revision_compat.py +++ b/backend/tests/test_persistence_forward_revision_compat.py @@ -20,16 +20,19 @@ 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 ( + _CANONICAL_0019_SCHEMA_FLOOR, _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 +from deerflow.persistence.thread_meta.sql import ThreadMetaRepository -HEAD = "0021_batch_acceptance" +CURRENT_HEAD = "0019_thread_incarnations" +ROLLBACK_HEAD = "0020_threads_meta_project_id" +INCARNATION_PARENT = "0021_batch_acceptance" +ORIGINAL_INCARNATION_PARENT = "0018_oauth_identity_pg_partial" POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI") @@ -54,18 +57,66 @@ async def _set_database_revision(engine, revision: str) -> None: await conn.execute(sa.text("UPDATE alembic_version SET version_num = :revision"), {"revision": revision}) -async def _seed_head(engine) -> None: +async def _seed_current_head(engine) -> None: await bootstrap_schema(engine, backend="sqlite") - assert await _database_revision(engine) == HEAD + assert await _database_revision(engine) == CURRENT_HEAD + + +async def _seed_rollback_head(engine) -> None: + cfg = _get_alembic_config(engine) + await asyncio.to_thread(_upgrade, cfg, ROLLBACK_HEAD) + assert await _database_revision(engine) == ROLLBACK_HEAD + + +async def _seed_incarnation_parent(engine) -> None: + cfg = _get_alembic_config(engine) + await asyncio.to_thread(_upgrade, cfg, INCARNATION_PARENT) + assert await _database_revision(engine) == INCARNATION_PARENT + + +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)")) + + +def _simulate_rollback_binary(monkeypatch: pytest.MonkeyPatch) -> None: + current_head, current_revisions = bootstrap_mod._get_revision_metadata() + assert current_head == CURRENT_HEAD == _FORWARD_COMPATIBLE_REVISION + assert ROLLBACK_HEAD in current_revisions + assert INCARNATION_PARENT in current_revisions + assert CURRENT_HEAD in current_revisions + monkeypatch.setattr( + bootstrap_mod, + "_get_revision_metadata", + lambda: (ROLLBACK_HEAD, current_revisions - {INCARNATION_PARENT, CURRENT_HEAD}), + ) async def _seed_original_forward_schema(engine) -> None: # The rollout predates projects: seeding today's head masks missing columns. - await asyncio.to_thread(_upgrade, _get_alembic_config(engine), "0018_oauth_identity_pg_partial") + await asyncio.to_thread(_upgrade, _get_alembic_config(engine), ORIGINAL_INCARNATION_PARENT) await _add_forward_columns(engine) await _set_database_revision(engine, _FORWARD_COMPATIBLE_REVISION) +@pytest.mark.asyncio +async def test_canonical_0019_floor_matches_migration_schema(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "canonical-floor.db")) + try: + await asyncio.to_thread(_upgrade, _get_alembic_config(engine), CURRENT_HEAD) + async with engine.connect() as conn: + + def reflect(sync_conn): + inspector = sa.inspect(sync_conn) + return {table: frozenset(column["name"] for column in inspector.get_columns(table)) for table in inspector.get_table_names() if table != "alembic_version"} + + reflected = await conn.run_sync(reflect) + assert reflected == _CANONICAL_0019_SCHEMA_FLOOR + finally: + await engine.dispose() + + @pytest.mark.asyncio @pytest.mark.parametrize("concurrent", [False, True]) async def test_original_forward_schema_fails_closed(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, concurrent: bool) -> None: @@ -73,7 +124,8 @@ async def test_original_forward_schema_fails_closed(tmp_path: Path, monkeypatch: try: await _seed_original_forward_schema(engine) if concurrent: - await _set_database_revision(engine, "0018_oauth_identity_pg_partial") + await _set_database_revision(engine, ORIGINAL_INCARNATION_PARENT) + _simulate_rollback_binary(monkeypatch) def concurrent_upgrade(_cfg, _revision): asyncio.run(_set_database_revision(engine, _FORWARD_COMPATIBLE_REVISION)) @@ -99,12 +151,16 @@ async def test_original_forward_schema_fails_closed(tmp_path: Path, monkeypatch: ("DROP TABLE projects", "projects"), ("ALTER TABLE projects DROP COLUMN instructions", "projects.instructions"), ("ALTER TABLE threads_meta DROP COLUMN project_id", "threads_meta.project_id"), + ("ALTER TABLE subagent_batch_items DROP COLUMN acceptance_criteria", "subagent_batch_items.acceptance_criteria"), + ("ALTER TABLE subagent_batch_items DROP COLUMN acceptance_verdict", "subagent_batch_items.acceptance_verdict"), + ("ALTER TABLE threads_meta DROP COLUMN incarnation", "threads_meta.incarnation"), + ("ALTER TABLE mcp_tasks DROP COLUMN thread_incarnation", "mcp_tasks.thread_incarnation"), ], ) -async def test_forward_revision_rejects_partial_project_schema(tmp_path: Path, ddl: str, missing: str) -> None: +async def test_current_incarnation_revision_rejects_incomplete_schema(tmp_path: Path, ddl: str, missing: str) -> None: engine = create_async_engine(_url(tmp_path, "partial-projects.db")) try: - await _seed_head(engine) + await _seed_incarnation_parent(engine) await _add_forward_columns(engine) async with engine.begin() as conn: if "DROP COLUMN project_id" in ddl: @@ -132,10 +188,10 @@ async def test_audited_original_forward_schema_can_upgrade_preserving_incarnatio # Documented offline operator recovery, only after verifying the exact # 0018 + two nullable columns shape. Bootstrap never re-stamps an unknown DB. - await asyncio.to_thread(alembic_command.stamp, _get_alembic_config(engine), "0018_oauth_identity_pg_partial", purge=True) + await asyncio.to_thread(alembic_command.stamp, _get_alembic_config(engine), ORIGINAL_INCARNATION_PARENT, purge=True) await bootstrap_schema(engine, backend="sqlite") - assert await _database_revision(engine) == HEAD + assert await _database_revision(engine) == CURRENT_HEAD repository = ThreadMetaRepository(async_sessionmaker(engine, expire_on_commit=False)) assert [row["thread_id"] for row in await repository.search(user_id=None)] == ["existing"] assert (await repository.create("new", user_id=None))["thread_id"] == "new" @@ -147,6 +203,67 @@ async def test_audited_original_forward_schema_can_upgrade_preserving_incarnatio await engine.dispose() +@pytest.mark.asyncio +async def test_known_canonical_0019_validates_fixed_floor_then_upgrades_to_future_head( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from deerflow.persistence.base import Base + + engine = create_async_engine(_url(tmp_path, "future-head.db")) + future_table = None + calls: list[str] = [] + try: + await _seed_current_head(engine) + # Model a future binary whose ORM includes schema that only its next + # migration can add. Canonical 0019 must not be rejected for lacking it. + future_table = sa.Table("future_after_0019", Base.metadata, sa.Column("id", sa.String(), primary_key=True)) + current_revisions = bootstrap_mod._get_known_revisions() + monkeypatch.setattr( + bootstrap_mod, + "_get_revision_metadata", + lambda: ("0022_future", current_revisions | {"0022_future"}), + ) + monkeypatch.setattr(bootstrap_mod, "_upgrade", lambda _cfg, revision: calls.append(revision)) + + await bootstrap_schema(engine, backend="sqlite") + + assert calls == ["head"] + finally: + if future_table is not None: + Base.metadata.remove(future_table) + await engine.dispose() + + +@pytest.mark.asyncio +async def test_known_canonical_0019_rejects_missing_floor_before_future_upgrade( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = create_async_engine(_url(tmp_path, "future-head-missing-floor.db")) + upgrade_called = False + try: + await _seed_original_forward_schema(engine) + current_revisions = bootstrap_mod._get_known_revisions() + monkeypatch.setattr( + bootstrap_mod, + "_get_revision_metadata", + lambda: ("0022_future", current_revisions | {"0022_future"}), + ) + + def future_upgrade(_cfg, _revision): + nonlocal upgrade_called + upgrade_called = True + + monkeypatch.setattr(bootstrap_mod, "_upgrade", future_upgrade) + + with pytest.raises(RuntimeError, match="missing required local schema: projects"): + await bootstrap_schema(engine, backend="sqlite") + assert upgrade_called is False + finally: + await engine.dispose() + + @pytest.mark.asyncio async def test_known_older_revision_upgrades_normally(tmp_path: Path) -> None: engine = create_async_engine(_url(tmp_path, "known.db")) @@ -156,7 +273,7 @@ async def test_known_older_revision_upgrades_normally(tmp_path: Path) -> None: await bootstrap_schema(engine, backend="sqlite") - assert await _database_revision(engine) == HEAD + assert await _database_revision(engine) == CURRENT_HEAD finally: await engine.dispose() @@ -164,12 +281,13 @@ async def test_known_older_revision_upgrades_normally(tmp_path: Path) -> None: @pytest.mark.asyncio async def test_exact_forward_revision_skips_upgrade_with_warning( tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, 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) + await _seed_current_head(engine) + _simulate_rollback_binary(monkeypatch) with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"): await bootstrap_schema(engine, backend="sqlite") @@ -184,7 +302,7 @@ async def test_exact_forward_revision_skips_upgrade_with_warning( 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 _seed_current_head(engine) await _set_database_revision(engine, "9999_unknown") with pytest.raises(RuntimeError, match="not known to this build"): @@ -204,23 +322,24 @@ async def test_sqlite_upgrade_race_recovers_when_other_process_applies_forward_r new_gateway = create_async_engine(url) upgrade_started = threading.Event() continue_upgrade = threading.Event() - original_upgrade = bootstrap_mod._upgrade - def delayed_upgrade(cfg, revision): + def delayed_old_upgrade(_cfg, revision): + assert revision == "head" upgrade_started.set() if not continue_upgrade.wait(timeout=5): raise TimeoutError("timed out waiting for the forward migration") - return original_upgrade(cfg, revision) + raise CommandError(f"Can't locate revision identified by '{CURRENT_HEAD}'") try: - await _seed_head(old_gateway) - monkeypatch.setattr(bootstrap_mod, "_upgrade", delayed_upgrade) + await _seed_rollback_head(old_gateway) + _simulate_rollback_binary(monkeypatch) + monkeypatch.setattr(bootstrap_mod, "_upgrade", delayed_old_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) + new_cfg = _get_alembic_config(new_gateway) + await asyncio.to_thread(_upgrade, new_cfg, CURRENT_HEAD) with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"): continue_upgrade.set() @@ -240,7 +359,8 @@ async def test_sqlite_upgrade_error_stays_fatal_without_forward_revision( ) -> None: engine = create_async_engine(_url(tmp_path, "upgrade-error.db")) try: - await _seed_head(engine) + await _seed_rollback_head(engine) + _simulate_rollback_binary(monkeypatch) def fail_upgrade(_cfg, _revision): raise CommandError("broken migration") @@ -259,12 +379,7 @@ async def test_local_forward_migration_error_stays_fatal( ) -> 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})), - ) + await _seed_rollback_head(engine) def fail_upgrade(_cfg, _revision): raise CommandError("local 0019 migration failed") @@ -280,7 +395,7 @@ async def test_local_forward_migration_error_stays_fatal( 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) + await _seed_current_head(engine) async with engine.begin() as conn: await conn.execute(sa.text("DELETE FROM alembic_version")) @@ -294,7 +409,7 @@ async def test_empty_alembic_version_fails_closed(tmp_path: Path) -> None: 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) + await _seed_current_head(engine) async with engine.begin() as conn: await conn.execute( sa.text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), @@ -307,24 +422,99 @@ async def test_multiple_alembic_versions_fail_closed(tmp_path: Path) -> None: 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_rollback_batch_writer_tolerates_acceptance_columns(tmp_path: Path) -> None: + engine = create_async_engine(_url(tmp_path, "batch-repository.db")) + try: + await _seed_current_head(engine) + old_items = sa.table( + "subagent_batch_items", + sa.column("id"), + sa.column("batch_id"), + sa.column("item_key"), + sa.column("position", sa.Integer()), + sa.column("prompt"), + sa.column("status"), + sa.column("attempt", sa.Integer()), + sa.column("result"), + sa.column("result_truncated", sa.Boolean()), + sa.column("created_at", sa.DateTime(timezone=True)), + sa.column("updated_at", sa.DateTime(timezone=True)), + ) + now = datetime.now(UTC) + async with engine.begin() as conn: + await conn.execute( + sa.text( + "INSERT INTO subagent_batches " + "(id, user_id, thread_id, submission_key, title, subagent_type, " + "status, total_items, max_live_items, max_running_items, " + "max_attempts, execution_spec, created_at, updated_at) " + "VALUES ('batch-1', 'user-1', 'thread-1', 'submission-1', " + "'Batch', 'general-purpose', 'queued', 1, 1, 1, 2, '{}', " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" + ) + ) + await conn.execute( + old_items.insert().values( + id="item-1", + batch_id="batch-1", + item_key="item", + position=0, + prompt="Prompt", + status="queued", + attempt=0, + result=None, + result_truncated=False, + created_at=now, + updated_at=now, + ) + ) + fetched = (await conn.execute(sa.select(*old_items.c).where(old_items.c.id == "item-1"))).mappings().one() + assert "acceptance_criteria" not in fetched + assert "acceptance_verdict" not in fetched + await conn.execute(old_items.update().where(old_items.c.id == "item-1").values(status="succeeded", result="legacy result", updated_at=now)) + + async with engine.connect() as conn: + row = (await conn.execute(sa.text("SELECT status, result, acceptance_criteria, acceptance_verdict FROM subagent_batch_items WHERE id = 'item-1'"))).one() + assert row == ("succeeded", "legacy result", None, None) + finally: + await engine.dispose() @pytest.mark.asyncio -async def test_old_thread_repository_tolerates_forward_nullable_column(tmp_path: Path) -> None: +async def test_rollback_thread_writer_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 + await _seed_current_head(engine) + # This is the complete 0020 table shape. Keeping it independent from + # the current ORM prevents a future model change from silently making + # this rollback-writer test aware of the forward column. + old_threads = sa.table( + "threads_meta", + sa.column("thread_id"), + sa.column("assistant_id"), + sa.column("user_id"), + sa.column("display_name"), + sa.column("status"), + sa.column("metadata_json", sa.JSON()), + sa.column("project_id"), + sa.column("created_at", sa.DateTime(timezone=True)), + sa.column("updated_at", sa.DateTime(timezone=True)), + ) + now = datetime.now(UTC) + async with engine.begin() as conn: + await conn.execute( + old_threads.insert().values( + thread_id="thread-1", + assistant_id=None, + user_id=None, + display_name=None, + status="idle", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) async with engine.begin() as conn: await conn.execute( @@ -332,54 +522,110 @@ async def test_old_thread_repository_tolerates_forward_nullable_column(tmp_path: {"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) + fetched = (await conn.execute(sa.select(*old_threads.c).where(old_threads.c.thread_id == "thread-1"))).mappings().one() + assert fetched["thread_id"] == "thread-1" + assert "incarnation" not in fetched + await conn.execute(old_threads.update().where(old_threads.c.thread_id == "thread-1").values(status="busy", updated_at=datetime.now(UTC))) async with engine.connect() as conn: - incarnation = ( + row = ( await conn.execute( - sa.text("SELECT incarnation FROM threads_meta WHERE thread_id = :thread_id"), + sa.text("SELECT status, incarnation FROM threads_meta WHERE thread_id = :thread_id"), {"thread_id": "thread-1"}, ) - ).scalar_one() - assert incarnation == "a" * 32 + ).one() + assert row == ("busy", "a" * 32) finally: await engine.dispose() @pytest.mark.asyncio -async def test_old_mcp_task_repository_tolerates_forward_nullable_column(tmp_path: Path) -> None: +async def test_rollback_shaped_mcp_task_sql_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) + await _seed_current_head(engine) + # This is the complete 0020 task table shape, deliberately excluding + # only the forward thread_incarnation column. + old_tasks = sa.table( + "mcp_tasks", + sa.column("id"), + sa.column("user_id"), + sa.column("thread_id"), + sa.column("run_id"), + sa.column("tool_call_id"), + sa.column("server_name"), + sa.column("driver_name"), + sa.column("remote_task_id"), + sa.column("task_name"), + sa.column("status"), + sa.column("result", sa.JSON()), + sa.column("result_preview"), + sa.column("result_truncated", sa.Boolean()), + sa.column("result_artifact", sa.JSON()), + sa.column("error"), + sa.column("input_required", sa.JSON()), + sa.column("driver_data", sa.JSON()), + sa.column("notification_status"), + sa.column("event_fingerprint"), + sa.column("event_version", sa.Integer()), + sa.column("notified_version", sa.Integer()), + sa.column("dispatch_version", sa.Integer()), + sa.column("dispatch_attempt", sa.Integer()), + sa.column("dispatch_event", sa.JSON()), + sa.column("notification_run_id"), + sa.column("notification_error"), + sa.column("notification_attempt_count", sa.Integer()), + sa.column("next_notification_at", sa.DateTime(timezone=True)), + sa.column("notification_lease_owner"), + sa.column("notification_lease_expires_at", sa.DateTime(timezone=True)), + sa.column("next_poll_at", sa.DateTime(timezone=True)), + sa.column("last_polled_at", sa.DateTime(timezone=True)), + sa.column("last_poll_error"), + sa.column("poll_attempt_count", sa.Integer()), + sa.column("consecutive_poll_error_count", sa.Integer()), + sa.column("lease_owner"), + sa.column("lease_expires_at", sa.DateTime(timezone=True)), + sa.column("cancel_requested_at", sa.DateTime(timezone=True)), + sa.column("cancel_attempt_count", sa.Integer()), + sa.column("next_cancel_at", sa.DateTime(timezone=True)), + sa.column("last_cancel_error"), + sa.column("completed_at", sa.DateTime(timezone=True)), + sa.column("created_at", sa.DateTime(timezone=True)), + sa.column("updated_at", sa.DateTime(timezone=True)), + ) 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( + old_tasks.insert().values( + 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, + error=None, + input_required=None, + driver_data={}, + notification_status="none", + next_poll_at=now - timedelta(seconds=1), + last_polled_at=None, + last_poll_error=None, + poll_attempt_count=0, + consecutive_poll_error_count=0, + lease_owner=None, + lease_expires_at=None, + cancel_requested_at=None, + completed_at=None, + created_at=now, + updated_at=now, + ) + ) async with engine.begin() as conn: await conn.execute( @@ -387,32 +633,52 @@ async def test_old_mcp_task_repository_tolerates_forward_nullable_column(tmp_pat {"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"] + fetched = ( + ( + await conn.execute( + sa.select(*old_tasks.c).where( + old_tasks.c.id == "task-1", + old_tasks.c.user_id == "user-1", + ) + ) + ) + .mappings() + .one() + ) + assert fetched["id"] == "task-1" + assert "thread_incarnation" not in fetched + await conn.execute( + old_tasks.update() + .where( + old_tasks.c.id == "task-1", + old_tasks.c.status == "working", + old_tasks.c.next_poll_at <= now, + ) + .values( + lease_owner="worker-1", + lease_expires_at=now + timedelta(seconds=60), + poll_attempt_count=old_tasks.c.poll_attempt_count + 1, + updated_at=now, + ) + ) async with engine.connect() as conn: - incarnation = ( + row = ( await conn.execute( - sa.text("SELECT thread_incarnation FROM mcp_tasks WHERE id = :task_id"), + sa.text("SELECT lease_owner, poll_attempt_count, thread_incarnation FROM mcp_tasks WHERE id = :task_id"), {"task_id": "task-1"}, ) - ).scalar_one() - assert incarnation == "b" * 32 + ).one() + assert row == ("worker-1", 1, "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: +async def test_old_gateway_restarts_against_forward_postgres_revision( + monkeypatch: pytest.MonkeyPatch, +) -> None: assert POSTGRES_URL is not None schema = f"forward_revision_{uuid.uuid4().hex}" config = DatabaseConfig( @@ -424,15 +690,10 @@ async def test_old_gateway_restarts_against_forward_postgres_revision() -> None: 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}, - ) + assert await _database_revision(engine) == CURRENT_HEAD await close_engine() + _simulate_rollback_binary(monkeypatch) await init_engine_from_config(config) restarted_engine = get_engine() diff --git a/backend/tests/test_thread_meta_repo.py b/backend/tests/test_thread_meta_repo.py index e8f64cc94..e284fe65d 100644 --- a/backend/tests/test_thread_meta_repo.py +++ b/backend/tests/test_thread_meta_repo.py @@ -4,6 +4,7 @@ import asyncio import logging import pytest +from sqlalchemy.exc import IntegrityError from deerflow.persistence.thread_meta import THREAD_PINNED_METADATA_KEY, InvalidMetadataFilterError, ThreadMetaRepository @@ -24,11 +25,13 @@ class TestThreadMetaRepository: record = await repo.create("t1") assert record["thread_id"] == "t1" assert record["status"] == "idle" + assert len(record["incarnation"]) == 32 assert "created_at" in record fetched = await repo.get("t1") assert fetched is not None assert fetched["thread_id"] == "t1" + assert fetched["incarnation"] == record["incarnation"] @pytest.mark.anyio async def test_create_with_assistant_id(self, repo): @@ -46,6 +49,46 @@ class TestThreadMetaRepository: record = await repo.create("t1", metadata={"key": "value"}) assert record["metadata"] == {"key": "value"} + @pytest.mark.anyio + async def test_duplicate_create_raises_integrity_error(self, repo): + await repo.create("t1", display_name="original") + + with pytest.raises(IntegrityError): + await repo.create("t1", display_name="replacement") + + record = await repo.get("t1") + assert record is not None + assert record["display_name"] == "original" + + @pytest.mark.anyio + async def test_claim_unowned_only_updates_null_owner(self, repo): + legacy = await repo.create("legacy", user_id=None) + await repo.create("owned", user_id="original-owner") + + assert await repo.claim_unowned("missing", "owner-a") is False + assert await repo.claim_unowned("owned", "owner-a") is False + assert await repo.claim_unowned("legacy", "owner-a") is True + assert await repo.claim_unowned("legacy", "owner-b") is False + + owned = await repo.get("owned", user_id=None) + claimed = await repo.get("legacy", user_id=None) + assert owned["user_id"] == "original-owner" + assert claimed["user_id"] == "owner-a" + assert claimed["updated_at"] == legacy["updated_at"] + + @pytest.mark.anyio + async def test_concurrent_claim_unowned_has_exactly_one_winner(self, repo): + await repo.create("legacy-race", user_id=None) + + outcomes = await asyncio.gather( + repo.claim_unowned("legacy-race", "owner-a"), + repo.claim_unowned("legacy-race", "owner-b"), + ) + + assert sorted(outcomes) == [False, True] + record = await repo.get("legacy-race", user_id=None) + assert record["user_id"] in {"owner-a", "owner-b"} + @pytest.mark.anyio async def test_update_display_name_can_remove_stale_metadata_atomically(self, repo): await repo.create("t1", display_name="Original (2)", metadata={"branch_title_sequence": 2, "keep": True}) @@ -587,6 +630,7 @@ class TestThreadMetaRepository: p = await projects.create(name="P", user_id="u1") record = await repo.create("t1", user_id="u1", project_id=p["id"]) assert record["metadata"]["deerflow_project_id"] == p["id"] + assert len(record["incarnation"]) == 32 with pytest.raises(ProjectNotAssignableError): await repo.create("t2", user_id="u1", project_id="missing") diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index 840f00d73..b221d1b11 100644 --- a/backend/tests/test_threads_router.py +++ b/backend/tests/test_threads_router.py @@ -26,6 +26,7 @@ from deerflow.persistence.thread_meta import ( THREAD_PROJECT_METADATA_KEY, InvalidMetadataFilterError, ThreadMetaRepository, + ThreadOwnershipConflictError, ) from deerflow.persistence.thread_meta.memory import THREADS_NS, MemoryThreadMetaStore from deerflow.runtime import ConflictError, ThreadOperationKind @@ -97,6 +98,17 @@ def _build_thread_app() -> tuple[FastAPI, InMemoryStore, InMemorySaver]: return app, store, checkpointer +def test_thread_response_excludes_internal_incarnation() -> None: + response = threads.ThreadResponse.model_validate( + { + "thread_id": "thread-with-incarnation", + "incarnation": "a" * 32, + } + ) + + assert "incarnation" not in response.model_dump() + + def test_compact_rejects_run_owned_by_another_worker(monkeypatch) -> None: """The HTTP guard must consult the shared store, not only local run memory.""" from deerflow.runtime import RunManager, RunStatus @@ -799,6 +811,125 @@ def test_insert_race_recovery_claims_unscoped_row_for_trusted_owner() -> None: assert unscoped_lookup["user_id"] == "owner-1" +def test_fast_path_concurrent_trusted_claims_have_one_winner() -> None: + thread_store = MemoryThreadMetaStore(InMemoryStore()) + + async def _scenario(): + await thread_store.create("legacy-fast-race", user_id=None) + owners = ("owner-a", "owner-b") + outcomes = await asyncio.gather( + *( + threads._resolve_existing_thread( + thread_store, + "legacy-fast-race", + owner, + {"user_id": owner}, + ) + for owner in owners + ) + ) + return owners, outcomes, await thread_store.get("legacy-fast-race", user_id=None) + + owners, outcomes, final_record = asyncio.run(_scenario()) + + winners = [owner for owner, outcome in zip(owners, outcomes, strict=True) if outcome is not None] + assert winners == [final_record["user_id"]] + assert final_record["user_id"] in owners + + +def test_fast_path_trusted_claim_does_not_take_over_owned_row() -> None: + thread_store = MemoryThreadMetaStore(InMemoryStore()) + + async def _scenario(): + await thread_store.create("already-owned", user_id="owner-a") + outcome = await threads._resolve_existing_thread( + thread_store, + "already-owned", + "owner-b", + {"user_id": "owner-b"}, + ) + return outcome, await thread_store.get("already-owned", user_id=None) + + outcome, final_record = asyncio.run(_scenario()) + + assert outcome is None + assert final_record["user_id"] == "owner-a" + + +def test_insert_race_concurrent_trusted_claims_have_one_winner() -> None: + from sqlalchemy.exc import IntegrityError + + from app.gateway.auth_disabled import AUTH_SOURCE_INTERNAL + from app.gateway.internal_auth import INTERNAL_OWNER_USER_ID_HEADER_NAME, INTERNAL_SYSTEM_ROLE + + class _ConcurrentInsertRaceStore(MemoryThreadMetaStore): + def __init__(self): + super().__init__(InMemoryStore()) + self._create_arrivals = 0 + self._create_lock = asyncio.Lock() + self._legacy_row_committed = asyncio.Event() + + async def create(self, thread_id, *, assistant_id=None, user_id=None, display_name=None, metadata=None, project_id=None): # type: ignore[override] + async with self._create_lock: + self._create_arrivals += 1 + if self._create_arrivals == 2: + await super().create(thread_id, user_id=None, metadata=metadata) + self._legacy_row_committed.set() + await self._legacy_row_committed.wait() + raise IntegrityError( + "INSERT INTO threads_meta", + {}, + Exception("UNIQUE constraint failed: threads_meta.thread_id"), + ) + + thread_store = _ConcurrentInsertRaceStore() + checkpointer = InMemorySaver() + + def _request(owner): + return SimpleNamespace( + headers={INTERNAL_OWNER_USER_ID_HEADER_NAME: owner}, + state=SimpleNamespace(user=SimpleNamespace(id="default", system_role=INTERNAL_SYSTEM_ROLE), auth_source=AUTH_SOURCE_INTERNAL), + app=SimpleNamespace(state=SimpleNamespace(checkpointer=checkpointer, thread_store=thread_store)), + ) + + async def _scenario(): + owners = ("owner-a", "owner-b") + outcomes = await asyncio.gather( + *( + threads.create_thread( + threads.ThreadCreateRequest(thread_id="legacy-insert-race"), + _request(owner), + ) + for owner in owners + ), + return_exceptions=True, + ) + return owners, outcomes, await thread_store.get("legacy-insert-race", user_id=None) + + owners, outcomes, final_record = asyncio.run(_scenario()) + + winners = [owner for owner, outcome in zip(owners, outcomes, strict=True) if isinstance(outcome, threads.ThreadResponse)] + failures = [outcome for outcome in outcomes if isinstance(outcome, HTTPException)] + assert winners == [final_record["user_id"]] + assert final_record["user_id"] in owners + assert len(failures) == 1 + assert failures[0].status_code == 500 + + +def test_create_thread_maps_memory_owner_conflict_to_404() -> None: + app, _store, _checkpointer = _build_thread_app() + app.state.thread_store = SimpleNamespace( + get=AsyncMock(return_value=None), + create=AsyncMock(side_effect=ThreadOwnershipConflictError("foreign-thread")), + ) + + with TestClient(app) as client: + response = client.post("/api/threads", json={"thread_id": "foreign-thread"}) + + assert response.status_code == 404 + assert response.json()["detail"] == "Thread not found" + + def test_create_thread_does_not_swallow_non_integrity_errors() -> None: """A non-race insert failure must surface as 500, even when a row now exists. diff --git a/docs/database-forward-revision-recovery.md b/docs/database-forward-revision-recovery.md index fce9ab22e..fe53b1446 100644 --- a/docs/database-forward-revision-recovery.md +++ b/docs/database-forward-revision-recovery.md @@ -7,10 +7,16 @@ only `0018_oauth_identity_pg_partial` plus two nullable `VARCHAR(32)` columns: serve this build's repositories. Startup now rejects it without changing the schema or revision, and reports the missing tables/columns. -Normal databases on this tree's known migration chain upgrade automatically. +Normal databases on the known migration chain in this tree upgrade automatically. The procedure below is only for the exact original incarnation rollout shape. -An incarnation-stamped database that already has all current ORM tables and -columns can still use the audited compatibility exception without re-stamping. +Current and future builds that know the reused `0019_thread_incarnations` id +first check the fixed canonical-0019 table/column floor, then upgrade to their +local head. The published 0020 rollback binary does not use this fixed snapshot: +it validates its own ORM floor before skipping the unknown revision. Current +tests remove 0019 from the mocked local revision set to exercise the same +unknown-revision path, but they still run the current fixed-floor validator. +Because the fixed floor is not derived from current ORM metadata, future columns +are not required before their migration runs. ## Offline migration @@ -29,8 +35,9 @@ columns can still use the audited compatibility exception without re-stamping. `DEERFLOW_RECOVERY_POSTGRES_SCHEMA` to the configured application schema, if one is used. Keep credentials out of shell history. 4. Rebase the version marker to the verified common parent and run the normal - migrations. `purge=True` is necessary because this tree does not contain the - out-of-tree revision; it replaces the version row, not application data. + migrations. `purge=True` replaces the version row, not application data; + it is necessary because the same revision id was previously deployed with a + different parent. ```bash uv run python - <<'PY' @@ -53,18 +60,20 @@ columns can still use the audited compatibility exception without re-stamping. PY ``` - This applies `0019_projects` and `0020_threads_meta_project_id`, preserving - the two incarnation columns and their existing values. Do not stamp directly - to head: that would skip the DDL and reproduce the missing-column failure. -5. Confirm the version is `0020_threads_meta_project_id`, the project table and + This applies `0019_projects`, `0020_threads_meta_project_id`, + `0021_batch_acceptance`, and the idempotent + `0019_thread_incarnations` head, preserving the two incarnation columns and + their existing values. Do not stamp directly to head: that would skip the + Projects and batch-acceptance DDL and reproduce the missing-schema failure. +5. Confirm the version is `0019_thread_incarnations`, the project table and membership column/index exist, and existing incarnation values are retained. Start this build, verify existing conversations load and a new conversation - can be created, then resume service. Do not restart older binaries that - cannot read this tree's head revision. + can be created, then resume service. Rolling back is supported only to the + audited `0020_threads_meta_project_id` compatibility build. Bootstrap never performs this re-stamp itself. The regression in `backend/tests/test_persistence_forward_revision_compat.py` constructs the original schema, verifies startup rejection, and exercises the recovery while -checking thread reads/inserts and preservation of incarnation data. The future -incarnation migration must chain from the current local head and handle these +checking thread reads/inserts and preservation of incarnation data. The +incarnation migration chains from `0021_batch_acceptance` and handles these already-present nullable columns idempotently.