mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 13:39:26 +00:00
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 <CorgiBoyG@users.noreply.github.com>
This commit is contained in:
parent
3c7d3303d3
commit
06c827903a
@ -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)
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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))
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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")
|
||||
@ -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",
|
||||
]
|
||||
|
||||
|
||||
@ -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).
|
||||
|
||||
@ -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"),
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
216
backend/tests/test_mcp_task_postgres.py
Normal file
216
backend/tests/test_mcp_task_postgres.py
Normal file
@ -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"]
|
||||
@ -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)
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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"
|
||||
|
||||
356
backend/tests/test_migration_0019_thread_incarnations.py
Normal file
356
backend/tests/test_migration_0019_thread_incarnations.py
Normal file
@ -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()
|
||||
@ -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'}")
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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")
|
||||
|
||||
@ -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.
|
||||
|
||||
|
||||
@ -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.
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user