mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-07 19:20:34 +00:00
fix(runtime): persistently fence stale run event producers (#7619)
* fix(runtime): fence revoked run event producers * fix(runtime): reduce off-mode producer guard locking * docs: keep inherited guidance within size budget * chore(runtime): remove development phase labels * fix(persistence): normalize libpq SSL options for asyncpg connections * docs(runtime): clarify producer fencing behavior in off mode * test(persistence): align migration TLS probes with channel-binding rejection * fix(runtime): authorize fenced Gateway history seeding * fix(scheduler): preserve takeover identity in recovery publication * fix(scheduler): report unconfirmed recovery publication * fix(runtime): preserve history messages before the first human turn * fix(runtime): clarify run creation failure diagnostic --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
01e1ed778b
commit
21175ba46f
63
README.md
63
README.md
@ -144,6 +144,69 @@ Run outcome or prevent stream completion. Recovery with built-in SQL Run and Eve
|
||||
sharing an engine uses one pooled connection for each guarded event write, so it
|
||||
can publish with only one pool slot available. See [Run Event Stream](backend/docs/RUN_EVENT_STREAM.md#authoritative-run-outcome).
|
||||
|
||||
PostgreSQL async ORM connections and the offline command accept `postgres://`
|
||||
and `postgresql://` URLs, including `?sslmode=disable` or `?sslmode=verify-full`.
|
||||
SSL modes are preserved; conflicting `ssl`/`sslmode` settings are rejected.
|
||||
Explicit `channel_binding` options are rejected because asyncpg cannot enforce
|
||||
them. Synchronous psycopg and LangGraph DSNs retain their libpq options.
|
||||
|
||||
With a fencing-capable SQL RunStore and `DbRunEventStore`, host-authorized
|
||||
producer writes validate the persisted Run incarnation, producer epoch and state
|
||||
even with `run_events.producer_fencing: "off"`. Stale Journal flushes after
|
||||
takeover or deletion may therefore be rejected in the default configuration.
|
||||
On PostgreSQL, these writes hold the Thread Fence and RunRow lock through commit.
|
||||
Checkpoint and branch history imports use a separate host grant bound to a
|
||||
durable thread-operation reservation and the captured thread identity. Strict
|
||||
mode permits these imports without treating synthetic turn IDs as admitted Runs;
|
||||
deletion, reservation revocation or lease expiry invalidates their grant.
|
||||
Messages preceding the first human input retain their existing synthetic turn 0.
|
||||
|
||||
On PostgreSQL, `run_events.producer_fencing: "strict"` additionally enables
|
||||
database-level authorization checks against raw SQL and unaware legacy writers
|
||||
after an offline cutover. The `"off"` setting keeps application-level producer
|
||||
validation active. Before activation, the database INSERT trigger skips its
|
||||
Thread/RunRow locks and retains a shared cutover lock until each INSERT
|
||||
transaction ends. This does not remove the application locks above. Off-mode
|
||||
event INSERTs require `READ COMMITTED` (also accepting PostgreSQL's `READ UNCOMMITTED`);
|
||||
`REPEATABLE READ` and `SERIALIZABLE` are rejected to prevent stale policy snapshots
|
||||
from bypassing activation. A write that loses a race with activation rolls back
|
||||
with SQLSTATE `40001`; retry the entire transaction with strict producer authority.
|
||||
Activation requires stopping every Gateway and using the offline cutover command;
|
||||
all restarted Gateways must use the same cutover ID. This does not recover missing
|
||||
terminal events. To activate:
|
||||
|
||||
1. Deploy this version everywhere with `run_events.producer_fencing: "off"`.
|
||||
Complete migrations 0049 and 0050 and recover previously active Runs, then
|
||||
drain admissions and completion hooks. A terminal Run status alone does not
|
||||
prove its producer has drained.
|
||||
2. Stop every Gateway, Worker and other database client, including checkpointer
|
||||
pools and monitoring clients, and keep them stopped during activation.
|
||||
3. From the repository root, using the installed backend environment, run:
|
||||
|
||||
```bash
|
||||
uv run --no-sync --project backend python backend/scripts/run_producer_fencing.py \
|
||||
--config config.yaml --cutover-id deployment-2026-10-27 --all-gateways-stopped
|
||||
```
|
||||
|
||||
The command refuses activation if other database sessions or pending/running
|
||||
Runs remain. Use a deployment ID that every Gateway will share.
|
||||
4. Set `database.backend: postgres`, `run_events.backend: db`,
|
||||
`run_events.producer_fencing: strict` and
|
||||
`run_events.producer_cutover_id: deployment-2026-10-27` everywhere, then restart.
|
||||
Verify a new Run persists its initial event, Journal/Hook, delivery and terminal
|
||||
events and closes normally; unaware old writers must receive a database rejection.
|
||||
|
||||
Activation is durable: missing or mismatched cutover IDs and downgrades to `"off"`
|
||||
or unsupported stores fail startup. This release requires controlled downtime;
|
||||
it does not support a mixed-version rolling cutover. A crashed producer may stay
|
||||
open until explicitly closed or revoked by authorized takeover/deletion; no
|
||||
producer-expiry or terminal-reconciliation job is introduced.
|
||||
|
||||
Scheduled-task orphan recovery carries the captured Run incarnation and retention
|
||||
revision when publishing delivery and terminal evidence, including in strict mode.
|
||||
Unconfirmed publication is logged separately from the committed takeover; it does
|
||||
not undo the recovered Run or queue a new publication attempt.
|
||||
|
||||
Operators can extend lead-agent, subagent, and DeerMem extraction prompts with
|
||||
literal prepend/append configuration without editing source templates. See
|
||||
[prompt overlays](backend/docs/CONFIGURATION.md#prompt-overlays).
|
||||
|
||||
@ -898,6 +898,15 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen
|
||||
run_events_config = getattr(config, "run_events", None)
|
||||
app.state.run_events_config = run_events_config
|
||||
app.state.run_event_store = make_run_event_store(run_events_config)
|
||||
if sf is not None:
|
||||
from deerflow.persistence.run.producer import validate_producer_fencing
|
||||
|
||||
# An activated database also forbids a memory/JSONL downgrade.
|
||||
await validate_producer_fencing(
|
||||
sf,
|
||||
mode=getattr(run_events_config, "producer_fencing", "off"),
|
||||
cutover_id=getattr(run_events_config, "producer_cutover_id", None),
|
||||
)
|
||||
from deerflow.runtime.events.terminal import RecoveryTerminalPublisher
|
||||
|
||||
app.state.recovery_terminal_publisher = RecoveryTerminalPublisher(app.state.run_event_store, app.state.run_store)
|
||||
|
||||
74
backend/app/gateway/history_seed.py
Normal file
74
backend/app/gateway/history_seed.py
Normal file
@ -0,0 +1,74 @@
|
||||
"""Host-authorized history copies under durable thread-operation admission."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from app.gateway.deps import get_run_manager, get_thread_store
|
||||
from app.gateway.internal_auth import get_trusted_internal_owner_user_id
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.persistence.thread_meta.sql import ThreadMetaRepository
|
||||
from deerflow.runtime.events.producer import RunEventProducerRejected, history_seed_authority, producer_scope
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.runs.manager import RunRecord
|
||||
from deerflow.runtime.runs.schemas import ThreadOperationKind
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HistorySeedReservation:
|
||||
record: RunRecord
|
||||
thread: dict | None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def reserve_history_seed(request: Request, thread_id: str) -> AsyncIterator[HistorySeedReservation | None]:
|
||||
"""Keep snapshot reads and seed writes exclusive with delete and admission."""
|
||||
events = request.app.state.run_event_store
|
||||
manager = getattr(request.app.state, "run_manager", None)
|
||||
runs = getattr(manager, "_store", None)
|
||||
threads = getattr(request.app.state, "thread_store", None)
|
||||
if not isinstance(events, DbRunEventStore):
|
||||
yield None
|
||||
return
|
||||
if not isinstance(runs, RunRepository) or not isinstance(threads, ThreadMetaRepository):
|
||||
if events._producer_fencing == "strict":
|
||||
raise RunEventProducerRejected("History seeding requires shared SQL lifecycle stores")
|
||||
yield None
|
||||
return
|
||||
if runs._sf.kw.get("bind") is not events._sf.kw.get("bind") or threads._sf.kw.get("bind") is not events._sf.kw.get("bind"):
|
||||
raise RunEventProducerRejected("History seeding requires shared SQL lifecycle stores")
|
||||
async with get_run_manager(request).reserve_thread_operation(
|
||||
thread_id,
|
||||
kind=ThreadOperationKind.checkpoint_write,
|
||||
user_id=get_trusted_internal_owner_user_id(request),
|
||||
) as reservation:
|
||||
yield HistorySeedReservation(reservation, await threads.get(thread_id, user_id=None))
|
||||
|
||||
|
||||
async def write_history_seed(
|
||||
request: Request,
|
||||
events: list[dict],
|
||||
*,
|
||||
reservation: HistorySeedReservation | None,
|
||||
seed_kind: str,
|
||||
thread: dict | None = None,
|
||||
assistant_id: str | None = None,
|
||||
) -> None:
|
||||
store = request.app.state.run_event_store
|
||||
if reservation is None:
|
||||
await store.put_batch(events)
|
||||
return
|
||||
if thread is None:
|
||||
# Only the checkpoint migration may materialize missing legacy metadata.
|
||||
# The caller read that checkpoint while holding this same reservation.
|
||||
if seed_kind != "checkpoint":
|
||||
raise RunEventProducerRejected("Branch seeding requires its original thread identity")
|
||||
threads = get_thread_store(request)
|
||||
thread = reservation.thread
|
||||
if thread is None:
|
||||
thread = await threads.create(reservation.record.thread_id, assistant_id=assistant_id, user_id=reservation.record.user_id)
|
||||
authority = history_seed_authority(reservation.record, thread, seed_kind=seed_kind)
|
||||
with producer_scope(authority):
|
||||
await store.put_batch(events)
|
||||
@ -38,6 +38,7 @@ from app.gateway.checkpoint_lineage import (
|
||||
resolve_stamp_candidate_versions,
|
||||
)
|
||||
from app.gateway.deps import get_checkpointer, get_run_event_store, get_run_manager, get_run_store
|
||||
from app.gateway.history_seed import reserve_history_seed, write_history_seed
|
||||
from app.gateway.internal_auth import get_trusted_internal_owner_user_id
|
||||
from app.gateway.services import (
|
||||
abuild_checkpoint_state_accessor,
|
||||
@ -1227,9 +1228,9 @@ async def _branch_thread_with_reservation(
|
||||
# into a new thread. Preserve every valid spelling as create does.
|
||||
branch_assistant_id = None
|
||||
|
||||
async def _write_branch_row(project_id: str | None) -> None:
|
||||
async def _write_branch_row(project_id: str | None) -> dict:
|
||||
try:
|
||||
await thread_store.create(
|
||||
return await thread_store.create(
|
||||
new_thread_id,
|
||||
assistant_id=branch_assistant_id,
|
||||
display_name=display_name,
|
||||
@ -1244,13 +1245,13 @@ async def _branch_thread_with_reservation(
|
||||
raise HTTPException(status_code=500, detail="Failed to create branch") from None
|
||||
|
||||
try:
|
||||
await _write_branch_row(source_project_id)
|
||||
branch_record = await _write_branch_row(source_project_id)
|
||||
except ProjectNotAssignableError:
|
||||
# The source project became unassignable (archived, or deleted in a
|
||||
# race that cleared the source row's membership after our read): keep
|
||||
# the branch usable as an unassigned thread — the pre-inheritance
|
||||
# behavior for sources without an active project.
|
||||
await _write_branch_row(None)
|
||||
branch_record = await _write_branch_row(None)
|
||||
|
||||
# The thread feed (GET /messages, /messages/page) reads the run-event
|
||||
# store, not checkpoints, and a fresh branch has no run_events — so the
|
||||
@ -1260,17 +1261,20 @@ async def _branch_thread_with_reservation(
|
||||
# from. Best-effort: on failure the branch stays usable, with history
|
||||
# visible only through the checkpoint overlay until it is re-branched.
|
||||
try:
|
||||
seed_events = build_branch_history_seed_events(
|
||||
_checkpoint_messages(snapshot),
|
||||
thread_id=new_thread_id,
|
||||
run_id_prefix=f"branch-seed-{new_thread_id}",
|
||||
parent_thread_id=thread_id,
|
||||
)
|
||||
if seed_events:
|
||||
await get_run_event_store(request).put_batch(seed_events)
|
||||
history_seed_mode = "seeded"
|
||||
else:
|
||||
history_seed_mode = "skipped_empty"
|
||||
async with reserve_history_seed(request, new_thread_id) as reservation:
|
||||
seed_events = build_branch_history_seed_events(
|
||||
_checkpoint_messages(snapshot),
|
||||
thread_id=new_thread_id,
|
||||
run_id_prefix=f"branch-seed-{new_thread_id}",
|
||||
parent_thread_id=thread_id,
|
||||
)
|
||||
if seed_events:
|
||||
event_store = get_run_event_store(request)
|
||||
if reservation is None or not await event_store.list_messages(new_thread_id, limit=1, user_id=None):
|
||||
await write_history_seed(request, seed_events, reservation=reservation, seed_kind="branch", thread=branch_record)
|
||||
history_seed_mode = "seeded"
|
||||
else:
|
||||
history_seed_mode = "skipped_empty"
|
||||
except Exception:
|
||||
logger.exception("Failed to seed branch history run-events for thread %s", sanitize_log_param(new_thread_id))
|
||||
history_seed_mode = "failed"
|
||||
|
||||
@ -31,6 +31,7 @@ from langgraph.types import Command
|
||||
from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID, AUTH_SOURCE_INTERNAL, get_auth_disabled_user, is_auth_disabled
|
||||
from app.gateway.authz import require_cancel_permission_if, resolve_route_permissions
|
||||
from app.gateway.deps import get_checkpointer, get_local_provider, get_run_context, get_run_manager, get_stream_bridge
|
||||
from app.gateway.history_seed import reserve_history_seed, write_history_seed
|
||||
from app.gateway.internal_auth import (
|
||||
INTERNAL_OWNER_USER_ID_HEADER_NAME,
|
||||
INTERNAL_SYSTEM_ROLE,
|
||||
@ -1925,35 +1926,41 @@ async def ensure_checkpoint_history_seeded(
|
||||
if await event_store.list_messages(thread_id, limit=1, user_id=None):
|
||||
return
|
||||
|
||||
checkpoint_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
checkpoint_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
if await get_checkpointer(request).aget_tuple(checkpoint_config) is None:
|
||||
return
|
||||
|
||||
accessor, config = await abuild_checkpoint_state_accessor(
|
||||
request,
|
||||
thread_id=thread_id,
|
||||
assistant_id=assistant_id,
|
||||
)
|
||||
snapshot = await accessor.aget(config)
|
||||
values = getattr(snapshot, "values", None)
|
||||
messages = values.get("messages") if isinstance(values, dict) else None
|
||||
if not isinstance(messages, list) or not messages:
|
||||
async def checkpoint_messages() -> list:
|
||||
accessor, config = await abuild_checkpoint_state_accessor(request, thread_id=thread_id, assistant_id=assistant_id)
|
||||
snapshot = await accessor.aget(config)
|
||||
values = getattr(snapshot, "values", None)
|
||||
messages = values.get("messages") if isinstance(values, dict) else None
|
||||
return messages if isinstance(messages, list) else []
|
||||
|
||||
# Eligibility is advisory; an empty head needs no migration reservation.
|
||||
# Discard this preview and reread under admission before any publication.
|
||||
if not await checkpoint_messages():
|
||||
return
|
||||
|
||||
events = build_checkpoint_history_seed_events(
|
||||
messages,
|
||||
thread_id=thread_id,
|
||||
run_id_prefix=f"checkpoint-seed-{thread_id}",
|
||||
)
|
||||
if not events:
|
||||
return
|
||||
await event_store.put_batch(events)
|
||||
logger.info("Seeded %d checkpoint-history events for thread %s", len(events), thread_id)
|
||||
async with reserve_history_seed(request, thread_id) as reservation:
|
||||
if reservation is not None and await event_store.list_messages(thread_id, limit=1, user_id=None):
|
||||
return
|
||||
if await get_checkpointer(request).aget_tuple(checkpoint_config) is None:
|
||||
return
|
||||
|
||||
messages = await checkpoint_messages()
|
||||
if not messages:
|
||||
return
|
||||
|
||||
events = build_checkpoint_history_seed_events(
|
||||
messages,
|
||||
thread_id=thread_id,
|
||||
run_id_prefix=f"checkpoint-seed-{thread_id}",
|
||||
)
|
||||
if not events:
|
||||
return
|
||||
await write_history_seed(request, events, reservation=reservation, seed_kind="checkpoint", assistant_id=assistant_id)
|
||||
logger.info("Seeded %d checkpoint-history events for thread %s", len(events), thread_id)
|
||||
|
||||
|
||||
def _message_identifier(message: Any) -> str | None:
|
||||
|
||||
@ -297,13 +297,9 @@ class DatabaseConfig(BaseModel):
|
||||
if self.backend == "sqlite":
|
||||
return f"sqlite+aiosqlite:///{self.sqlite_path}"
|
||||
if self.backend == "postgres":
|
||||
url = self.postgres_url
|
||||
if url.startswith("postgresql://"):
|
||||
url = url.replace("postgresql://", "postgresql+asyncpg://", 1)
|
||||
elif url.startswith("postgres://"):
|
||||
# libpq's short alias: accepted by the psycopg checkpointer, but not a SQLAlchemy dialect.
|
||||
url = url.replace("postgres://", "postgresql+asyncpg://", 1)
|
||||
return url
|
||||
from deerflow.persistence.postgres_url import normalize_asyncpg_url
|
||||
|
||||
return normalize_asyncpg_url(self.postgres_url)
|
||||
raise ValueError(f"No SQLAlchemy URL for backend={self.backend!r}")
|
||||
|
||||
@property
|
||||
|
||||
@ -15,12 +15,21 @@ from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field, ValidationInfo, field_validator
|
||||
from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator
|
||||
|
||||
from deerflow.config._boolean_guards import reject_boolean
|
||||
|
||||
|
||||
class RunEventsConfig(BaseModel):
|
||||
producer_fencing: Literal["off", "strict"] = "off"
|
||||
producer_cutover_id: str | None = Field(default=None, min_length=1, max_length=64)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_producer_fencing(self):
|
||||
if self.producer_fencing == "strict" and (self.backend != "db" or self.producer_cutover_id is None):
|
||||
raise ValueError("Strict producer fencing requires backend=db and an explicit offline cutover ID")
|
||||
return self
|
||||
|
||||
backend: Literal["memory", "db", "jsonl"] = Field(
|
||||
default="memory",
|
||||
description="Storage backend for run events. 'memory' for development (no persistence), 'db' for production (SQL queries), 'jsonl' for lightweight single-node persistence.",
|
||||
|
||||
@ -1,5 +1,11 @@
|
||||
# Persistence lifecycle
|
||||
|
||||
`postgres_url.normalize_asyncpg_url` maps libpq `sslmode` to asyncpg `ssl`,
|
||||
preserving modes and credentials. Reject conflicting/repeated SSL options and
|
||||
unsupported `channel_binding`; never silently drop authentication requirements.
|
||||
`DatabaseConfig.app_sqlalchemy_url` and test fixtures share it; synchronous
|
||||
psycopg/checkpointer DSNs retain their libpq options.
|
||||
|
||||
Postgres bootstrap owns its session-scoped advisory lock until `pg_advisory_unlock` completes. Drain that unlock across host cancellation before leaving the SQLAlchemy connection context; repeated cancellation must not return a pooled session while it still holds the bootstrap mutex. Ordinary database errors remain best-effort and are logged.
|
||||
|
||||
Acquire that lock by polling `pg_try_advisory_lock`, never with a blocking `pg_advisory_lock`: the app engine's asyncpg `command_timeout` bounds every statement, so a blocking acquire fails a second instance's startup whenever a peer's migration outlasts it.
|
||||
@ -41,7 +47,7 @@ Scheduled-task lifecycle rules live in one place each. `ScheduledTaskRepository.
|
||||
|
||||
The per-owner run cap `scheduler.max_concurrent_runs_per_user` is effective `min(per_user, max_concurrent_runs)` (0 = off) and is checked in `claim_queued_run` right after the global count, under the same budget lock. `list_queued_runs` ranks same-thread FIFO heads per owner (`ROW_NUMBER() OVER (PARTITION BY owner)`) and orders the batch by that rank, so the drain is owner-fair even with the cap off; owners at their cap drop out before the batch limit, and orphan rows (task deleted while queued) are never held back by the cap. Scheduled-task repository boundaries coerce serialized timestamps before SQL `DateTime` binding.
|
||||
|
||||
Recovery locks task/run pairs in task-id/run-id order and restores `run_id`, `started_at` and live errors before releasing launch claims. Launch/failure/timeout updates use one parent-first transaction to prevent interleaved claims. Queue timeout fails the occurrence and advances scheduled work to prevent immediate requeue.
|
||||
Recovery locks task/run pairs by task/run ID; restore `run_id`/`started_at`/errors before releasing claims. Bind takeover grants to candidate incarnation/revision; log unconfirmed publication without reauthorizing. Launch/failure/timeout writes lock parent first; queue timeout fails the occurrence and advances scheduled work.
|
||||
|
||||
Single-instance polling also performs a non-destructive terminal-run reconciliation: a `running` occurrence whose durable `runs` row is already terminal is finalized from that row's status, error and goal verdict. Live runs are never taken over or interrupted; multi-instance reconciliation already covers terminal runs through lease ownership. This repairs a completion hook that failed after the run committed its terminal status.
|
||||
|
||||
@ -62,6 +68,16 @@ missing or strictly older leases remain recoverable.
|
||||
|
||||
## Run terminal authority and lease expiry
|
||||
|
||||
Run/thread creation returns committing identity snapshots; Runs mint epochs.
|
||||
Takeover/peer replacement/delete revoke; same-owner finalizers remain open.
|
||||
Lock order: local Thread → optional Change Clock → optional Extension fence →
|
||||
SQL Thread → RunRow (held through Event commit); never acquire Clock afterward.
|
||||
The INSERT trigger takes the shared cutover lock last; offline activation takes
|
||||
only its exclusive counterpart. Off-mode triggers take only the shared gate and
|
||||
recheck policy, aborting on activation; retained-snapshot isolation fails closed.
|
||||
Strict startup cannot downgrade durable policy.
|
||||
[Offline cutover](../../../../../README.md#configuration).
|
||||
|
||||
`RunRepository.update_status(..., owner_worker_id=...)` and
|
||||
`finalize_if_not_cancelled(..., owner_worker_id=...)` require that owner and a
|
||||
live lease. Increment the change clock before taking the run row lock to retain
|
||||
|
||||
@ -6,6 +6,12 @@ DeerFlow's application tables (`runs`, `threads_meta`, `feedback`, `users`, `run
|
||||
|
||||
**Hybrid bootstrap** (`persistence/bootstrap.py::bootstrap_schema`, invoked from `persistence/engine.py::init_engine`):
|
||||
|
||||
0049 adds nullable producer identities (no backfill) and strict policy. Install
|
||||
its PostgreSQL guard in both Alembic and metadata `after_create` for fresh bootstrap.
|
||||
Activated policy forbids downgrade; dispose migration engines even on refusal.
|
||||
0050 refreshes installed guards to admit host-authorized history seeds. 0051 joins
|
||||
that branch with upstream 0049 Run inputs, preserving both revision histories.
|
||||
|
||||
| DB state | Action |
|
||||
|-------------------------------------------|-----------------------------------------|
|
||||
| empty (no DeerFlow tables) | `create_all` + `alembic stamp head` |
|
||||
@ -34,16 +40,13 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a
|
||||
`0039_user_disabled` → `0040_scheduled_group_notice` →
|
||||
`0041_im_notice_preferences` → `0042_channel_leases` →
|
||||
`0043_runtime_im_notices` → `0044_channel_runtime_configs` →
|
||||
`0045_completed_run_evidence` → `0046_skill_mutations` → `0047_mcp_oauth_tokens` → `0048_extension_execution_leases` (current head). The preference
|
||||
revision adds a separate owner/key table with a cascading users foreign key and
|
||||
does not alter users; the project-documents revision adds a new owner-scoped
|
||||
shelf table, and the MCP lease-token revision adds two nullable token columns to
|
||||
`mcp_tasks`. The batch-result revision adds a nullable JSON evidence snapshot
|
||||
to `subagent_batch_items`, without changing existing report text or backfilling
|
||||
historical evidence. The run-idempotency revision adds a nullable private JSON
|
||||
request identity to `runs`; old Gateways ignore the column, so resume digests
|
||||
never enter public `kwargs_json`. Both revisions leave the bootstrap
|
||||
forward-compat floor unchanged.
|
||||
`0045_completed_run_evidence` → `0046_skill_mutations` → `0047_mcp_oauth_tokens` → `0048_extension_execution_leases`, then `0049_run_inputs` and `0049_run_producer_fencing` → `0050_history_seed_authority`, joined by `0051_merge_run_inputs_fencing` (head).
|
||||
The preference revision uses an owner/key table with a cascading users FK;
|
||||
project documents use an owner-scoped shelf. MCP lease tokens, batch evidence,
|
||||
and run-idempotency identity add nullable columns without historical backfill or
|
||||
changing bootstrap's forward-compatibility floor. Resume identity stays private
|
||||
in `idempotency_request_json`, never public `kwargs_json`; legacy rows remain NULL.
|
||||
See the revision entries below for schema and test details.
|
||||
The incarnation revision deliberately retains the exact id audited by the
|
||||
rollback-floor binary; Alembic orders revisions by `down_revision`, not by the
|
||||
numeric prefix.
|
||||
@ -186,7 +189,7 @@ on installs that never enabled it. The convention is:
|
||||
- `migrations/versions/0035_login_throttle.py` — creates `login_throttle` (`ip` PK, `fail_count`, nullable epoch `locked_at`, nullable `lock_duration_seconds`, `updated_at` + `ix_login_throttle_updated_at` for the sweep's stale-counter predicate), the shared failed-login counter behind `POST /api/v1/auth/login/local` selected by `auth.local.throttle_storage` (`persistence/login_throttle/`: `SqlLoginThrottleStore` increments and starts the lock in one dialect-native upsert, applies check decisions with compare-and-set predicates in a separate write transaction, and sweeps served locks plus 24h-idle counters in bounded batches — the sweep's `DELETE` repeats the expiry predicate on its target, not only `ip IN (SELECT … LIMIT n)`, because PostgreSQL READ COMMITTED re-evaluates only the statement's own WHERE on a row it waited for and would otherwise delete a lock a concurrent `check` just extended). Table create/drop is guarded with a fresh inspector in both directions. New table, so the bootstrap forward-compat floor is unchanged; chains after `0034_run_event_seq_watermark`. Test: `tests/test_migration_0035_login_throttle.py`
|
||||
- `migrations/versions/0036_run_idempotency_request.py` — adds nullable `runs.idempotency_request_json` for versioned keyed-resume identities. The field is private persistence state, never part of `kwargs_json` or a run response; this keeps new-writer/old-reader rolling upgrades from exposing resume digests or synthetic input. Legacy rows remain NULL; new readers accept only the exact old-writer identity shape from `kwargs`, while truly identity-less resume fails closed. The bootstrap forward-compat floor is unchanged. Chains after `0035_login_throttle`. Test: `tests/test_migration_0036_run_idempotency_request.py`
|
||||
- `migrations/versions/0037_project_document_summaries.py` — adds nullable `project_documents.summary` (best-effort LLM one-line description, off by default; no backfill, so legacy rows stay NULL and the bootstrap forward-compat floor is unchanged); chains after `0036_run_idempotency_request`. Test: `tests/test_migration_0037_project_document_summaries.py`
|
||||
- `migrations/versions/0038_channel_thread_bindings.py` — creates `channel_thread_bindings` (`key` PK = the legacy `channel_name:chat_id[:topic_id]` string, parsed `channel_name`/`chat_id`/nullable `topic_id`, `thread_id`, `user_id` (`""` when unknown, never NULL), epoch `Double` `created_at`/`updated_at` + `ix_channel_thread_bindings_channel_chat`), the shared successor of the per-process `channels/store.json` that `ChannelManager` used for unbound IM conversations (`persistence/channel_thread_bindings/`: `SqlChannelThreadBindingRepository` upserts with a dialect-native `ON CONFLICT (key) DO UPDATE` that keeps `created_at`, deletes a chat's base + topic rows by key prefix with `startswith(autoescape=True)`, and imports with batched `ON CONFLICT DO NOTHING` so two replicas importing the same file concurrently neither duplicate nor clobber a row; `app/channels/store.py::SqlChannelStore.import_legacy_json` runs that once from `ChannelService.start()` and renames the file to `store.json.migrated`). The shape is 1:1 with the JSON entries so the import is lossless; `split_binding_key` parses legacy keys on the first two colons exactly as the old `list_entries()` did (lookups always match the full `key`). Table create/drop and the index are guarded with a fresh inspector in both directions. New table, so the bootstrap forward-compat floor is unchanged; chains after `0037_project_document_summaries`. Test: `tests/test_migration_0038_channel_thread_bindings.py`
|
||||
- `migrations/versions/0038_channel_thread_bindings.py` — adds shared `channel_thread_bindings`, replacing per-process `channels/store.json`. Keys preserve legacy `channel_name:chat_id[:topic_id]` spelling (`split_binding_key` splits only the first two colons); unknown users remain `""`. `SqlChannelThreadBindingRepository` uses atomic upserts retaining `created_at`, escaped-prefix chat deletion and batched conflict-ignoring imports. `ChannelService.start()` imports once and renames the JSON file `.migrated`. Table/index create and drop are guarded. No bootstrap-floor change; follows 0037. Test: `tests/test_migration_0038_channel_thread_bindings.py`.
|
||||
- `migrations/versions/0042_channel_leases.py` — creates `channel_leases` (`provider` PK, `owner_id`, fencing `epoch` BIGINT, database-clock epoch `Double` `lease_expires_at`/`updated_at`, `transport_running`), the per-provider IM leader lease behind `channels.leader_lease` (`persistence/channel_leases/`: acquire is one dialect-native upsert whose `DO UPDATE ... WHERE lease_expires_at + grace < now OR owner_id = :me` bumps the epoch; renew extends only a live lease at the holder's epoch; release sets the expiry to 0; every `now` is the database clock, never an instance's). Table create/drop is guarded with a fresh inspector in both directions. New table, so the bootstrap forward-compat floor is unchanged; chains after `0041_im_notice_preferences`. Test: `tests/test_migration_0042_channel_leases.py`
|
||||
- `migrations/versions/0044_channel_runtime_configs.py` — creates `channel_runtime_configs` (`provider` PK, nullable `config_ciphertext`, `disconnected`, `version` BIGINT, `updated_at`), the shared successor of `channels/runtime-config.json`. `persistence/channel_runtime_configs/` is a sync repository on `app_sync_sqlalchemy_url` (callers run in threads): `save` upserts and bumps `version`, `insert_missing` uses `ON CONFLICT DO NOTHING`. `app/channels/runtime_config_store.py` encrypts the JSON block with `DEER_FLOW_CREDENTIALS_KEY`; a disconnect has NULL ciphertext. Rows are never deleted. Table-guarded both ways; no data moved (the Gateway imports the JSON file). Test: `tests/test_migration_0044_channel_runtime_configs.py`
|
||||
- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch provisioning decision, locked revision validation, and the narrow 0019 forward-compatibility exception
|
||||
|
||||
@ -120,9 +120,13 @@ async def run_migrations_online() -> None:
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
async with connectable.connect() as connection:
|
||||
await connection.run_sync(do_run_migrations)
|
||||
await connectable.dispose()
|
||||
try:
|
||||
async with connectable.connect() as connection:
|
||||
await connection.run_sync(do_run_migrations)
|
||||
finally:
|
||||
# A refused strict-policy downgrade must not retain an idle client
|
||||
# after the migration worker's event loop exits.
|
||||
await connectable.dispose()
|
||||
|
||||
|
||||
if context.is_offline_mode():
|
||||
|
||||
@ -0,0 +1,31 @@
|
||||
"""Durable Run producer identity, without historical authority backfill."""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from deerflow.persistence.migrations._helpers import safe_add_column, safe_drop_column
|
||||
|
||||
revision = "0049_run_producer_fencing"
|
||||
down_revision = "0048_extension_execution_leases"
|
||||
|
||||
|
||||
def upgrade():
|
||||
for name, length in (("run_incarnation", 64), ("producer_epoch", 64), ("producer_state", 16)):
|
||||
safe_add_column("runs", sa.Column(name, sa.String(length), nullable=True))
|
||||
if not sa.inspect(op.get_bind()).has_table("run_producer_policy"):
|
||||
op.create_table("run_producer_policy", sa.Column("id", sa.Integer(), primary_key=True), sa.Column("strict", sa.Boolean(), nullable=False, server_default=sa.false()), sa.Column("cutover_id", sa.String(64), nullable=True))
|
||||
from deerflow.persistence.run.producer import install_event_guard
|
||||
|
||||
install_event_guard(op.get_bind())
|
||||
|
||||
|
||||
def downgrade():
|
||||
bind = op.get_bind()
|
||||
if bind.dialect.name == "postgresql":
|
||||
if bind.scalar(sa.text("SELECT strict FROM run_producer_policy WHERE id=1")):
|
||||
raise RuntimeError("Cannot downgrade an activated producer fencing database")
|
||||
op.execute("DROP TRIGGER IF EXISTS run_event_producer_guard ON run_events")
|
||||
op.execute("DROP FUNCTION IF EXISTS run_event_producer_guard()")
|
||||
op.drop_table("run_producer_policy")
|
||||
for name in ("producer_state", "producer_epoch", "run_incarnation"):
|
||||
safe_drop_column("runs", name)
|
||||
@ -0,0 +1,20 @@
|
||||
"""Refresh the installed producer guard with host history-seed authorization."""
|
||||
|
||||
from alembic import op
|
||||
|
||||
from deerflow.persistence.run.producer import install_event_guard
|
||||
|
||||
revision = "0050_history_seed_authority"
|
||||
down_revision = "0049_run_producer_fencing"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
install_event_guard(op.get_bind())
|
||||
|
||||
|
||||
def downgrade():
|
||||
# No schema or policy changes to undo. Retain cumulative enforcement until
|
||||
# the producer-fencing migration removes it, which forbids active downgrade.
|
||||
pass
|
||||
@ -0,0 +1,14 @@
|
||||
"""Join Run input and producer-fencing histories without rewriting either."""
|
||||
|
||||
revision = "0051_merge_run_inputs_fencing"
|
||||
down_revision = ("0049_run_inputs", "0050_history_seed_authority")
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
pass
|
||||
|
||||
|
||||
def downgrade():
|
||||
pass
|
||||
@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import JSON, BigInteger, DateTime, Index, String, Text, UniqueConstraint, text
|
||||
from sqlalchemy import JSON, BigInteger, DateTime, Index, String, Text, UniqueConstraint, event, text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.constants import RUN_EVENT_CATEGORY_MAX_LENGTH, RUN_EVENT_TYPE_MAX_LENGTH
|
||||
@ -46,3 +46,11 @@ class RunEventThreadSeqRow(Base):
|
||||
|
||||
thread_id: Mapped[str] = mapped_column(String(64), primary_key=True)
|
||||
seq: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0, server_default=text("0"))
|
||||
|
||||
|
||||
@event.listens_for(Base.metadata, "after_create")
|
||||
def _install_producer_guard(metadata, connection, **kwargs):
|
||||
# Empty-DB bootstrap stamps head without executing migration revisions.
|
||||
from deerflow.persistence.run.producer import install_event_guard
|
||||
|
||||
install_event_guard(connection)
|
||||
|
||||
@ -0,0 +1,28 @@
|
||||
"""Adapt PostgreSQL URLs for the async ORM without weakening security options."""
|
||||
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
from sqlalchemy.engine import make_url
|
||||
|
||||
|
||||
def normalize_asyncpg_url(uri: str) -> str:
|
||||
"""Translate libpq SSL modes and reject unsupported authentication options."""
|
||||
url = make_url(uri)
|
||||
if url.drivername not in {"postgres", "postgresql", "postgresql+asyncpg"}:
|
||||
raise ValueError("Expected a PostgreSQL URL for asyncpg")
|
||||
# SQLAlchemy drops blank query values; validate security options before
|
||||
# an explicit empty setting could be mistaken for a missing setting.
|
||||
options = parse_qs(urlsplit(uri).query, keep_blank_values=True)
|
||||
if "channel_binding" in options:
|
||||
raise ValueError("channel_binding is not supported by asyncpg")
|
||||
for key in ("ssl", "sslmode"):
|
||||
values = options.get(key)
|
||||
if values is not None and (len(values) != 1 or values[0] not in {"disable", "allow", "prefer", "require", "verify-ca", "verify-full"}):
|
||||
raise ValueError(f"Invalid or repeated {key} in PostgreSQL URL")
|
||||
query = dict(url.query)
|
||||
if "sslmode" in query:
|
||||
mode = query.pop("sslmode")
|
||||
if "ssl" in query and query["ssl"] != mode:
|
||||
raise ValueError("Conflicting ssl and sslmode in PostgreSQL URL")
|
||||
query["ssl"] = mode
|
||||
return url.set(drivername="postgresql+asyncpg", query=query).render_as_string(hide_password=False)
|
||||
@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
from datetime import UTC, datetime
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy import JSON, BigInteger, DateTime, Index, Integer, String, Text, UniqueConstraint, text
|
||||
from sqlalchemy import JSON, BigInteger, Boolean, DateTime, Index, Integer, String, Text, UniqueConstraint, false, text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
@ -57,6 +57,10 @@ class RunRow(Base):
|
||||
# Multi-worker run ownership
|
||||
owner_worker_id: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
# Null on migrated rows: historical status is not producer authority.
|
||||
run_incarnation: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
producer_epoch: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
producer_state: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||||
# A non-owning worker records cancellation here; the owner consumes it
|
||||
# while renewing its lease. The first action wins.
|
||||
cancel_action: Mapped[str | None] = mapped_column(String(20), nullable=True)
|
||||
@ -100,6 +104,15 @@ class RunRow(Base):
|
||||
)
|
||||
|
||||
|
||||
class RunProducerPolicyRow(Base):
|
||||
"""Restart-required, database-wide strict event admission policy."""
|
||||
|
||||
__tablename__ = "run_producer_policy"
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True)
|
||||
strict: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, server_default=false())
|
||||
cutover_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
|
||||
|
||||
class RunChangeClockRow(Base):
|
||||
"""Singleton counter allocating backend-owned changed-run positions."""
|
||||
|
||||
|
||||
160
backend/packages/harness/deerflow/persistence/run/producer.py
Normal file
160
backend/packages/harness/deerflow/persistence/run/producer.py
Normal file
@ -0,0 +1,160 @@
|
||||
"""PostgreSQL enforcement and offline activation of Run producer fencing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import func, inspect, select, text
|
||||
|
||||
|
||||
def install_event_guard(connection) -> None:
|
||||
"""Used by both Alembic and metadata bootstrap; never manufacture authority."""
|
||||
if connection.dialect.name != "postgresql":
|
||||
return
|
||||
inspector = inspect(connection)
|
||||
if not all(inspector.has_table(name) for name in ("runs", "run_events", "run_producer_policy")):
|
||||
return
|
||||
if "run_incarnation" not in {column["name"] for column in inspector.get_columns("runs")}:
|
||||
return
|
||||
schema = connection.scalar(text("SELECT current_schema()"))
|
||||
quoted = connection.dialect.identifier_preparer.quote(schema)
|
||||
connection.execute(
|
||||
text(f"""
|
||||
CREATE OR REPLACE FUNCTION {quoted}.run_event_producer_guard() RETURNS trigger
|
||||
LANGUAGE plpgsql VOLATILE SET search_path TO {quoted}, pg_catalog AS $guard$
|
||||
DECLARE r record; t record; a jsonb; strict_mode boolean; seed_prefix text;
|
||||
BEGIN
|
||||
SELECT strict INTO strict_mode FROM run_producer_policy WHERE id=1;
|
||||
IF NOT coalesce(strict_mode, false) THEN
|
||||
-- A retained snapshot cannot observe activation after waiting. Never
|
||||
-- let it bypass strict enforcement using an old off/missing policy.
|
||||
IF current_setting('transaction_isolation') NOT IN ('read committed', 'read uncommitted') THEN
|
||||
RAISE EXCEPTION 'Off-mode run events require READ COMMITTED isolation' USING ERRCODE='0A000';
|
||||
END IF;
|
||||
PERFORM pg_advisory_xact_lock_shared(hashtext('deerflow.event_fence/' || current_schema())::bigint);
|
||||
SELECT strict INTO strict_mode FROM run_producer_policy WHERE id=1;
|
||||
IF coalesce(strict_mode, false) THEN
|
||||
-- Do not acquire earlier Thread/Run locks while holding this gate.
|
||||
RAISE EXCEPTION 'Producer policy changed; retry the transaction' USING ERRCODE='40001';
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
PERFORM pg_advisory_xact_lock(hashtext(NEW.thread_id)::bigint);
|
||||
a := nullif(current_setting('deerflow.event_authority', true), '')::jsonb;
|
||||
IF a->>'source'='history_seed' THEN
|
||||
SELECT * INTO r FROM runs WHERE run_id=a->>'reservation_id' FOR UPDATE;
|
||||
SELECT * INTO t FROM threads_meta WHERE thread_id=NEW.thread_id FOR UPDATE;
|
||||
ELSE
|
||||
SELECT * INTO r FROM runs WHERE run_id=NEW.run_id FOR UPDATE;
|
||||
END IF;
|
||||
-- Cutover takes only this final lock, never a Clock/Thread/Run lock.
|
||||
PERFORM pg_advisory_xact_lock_shared(hashtext('deerflow.event_fence/' || current_schema())::bigint);
|
||||
SELECT strict INTO strict_mode FROM run_producer_policy WHERE id=1;
|
||||
IF NOT coalesce(strict_mode, false) THEN RETURN NEW; END IF;
|
||||
IF a->>'source'='history_seed' THEN
|
||||
seed_prefix := (a->>'seed_kind') || '-seed-' || NEW.thread_id || '-';
|
||||
IF r.run_id IS NULL OR t.thread_id IS NULL
|
||||
OR r.operation_kind <> 'checkpoint_write' OR r.status <> 'pending' OR r.producer_state IS DISTINCT FROM 'open'
|
||||
OR r.thread_id IS DISTINCT FROM NEW.thread_id
|
||||
OR r.user_id IS DISTINCT FROM NEW.user_id
|
||||
OR a->>'thread_id' IS DISTINCT FROM NEW.thread_id
|
||||
OR a->>'user_id' IS DISTINCT FROM NEW.user_id
|
||||
OR r.run_incarnation IS NULL OR a->>'reservation_incarnation' IS DISTINCT FROM r.run_incarnation
|
||||
OR r.producer_epoch IS NULL OR a->>'producer_epoch' IS DISTINCT FROM r.producer_epoch
|
||||
OR r.owner_worker_id IS NULL OR a->>'owner_worker_id' IS DISTINCT FROM r.owner_worker_id
|
||||
OR (a->>'retention_revision')::bigint IS DISTINCT FROM r.evidence_retention_revision
|
||||
OR NOT (a ? 'thread_incarnation') OR a->>'thread_incarnation' IS DISTINCT FROM t.incarnation
|
||||
OR NOT (a ? 'thread_user_id') OR a->>'thread_user_id' IS DISTINCT FROM t.user_id
|
||||
OR (t.user_id IS NOT NULL AND t.user_id IS DISTINCT FROM NEW.user_id)
|
||||
OR (r.lease_expires_at IS NOT NULL AND r.lease_expires_at <= clock_timestamp())
|
||||
OR a->>'seed_kind' IS NULL OR a->>'seed_kind' NOT IN ('checkpoint','branch')
|
||||
OR left(NEW.run_id, length(seed_prefix)) IS DISTINCT FROM seed_prefix
|
||||
OR substring(NEW.run_id FROM length(seed_prefix)+1) !~ '^(0|[1-9][0-9]*)$'
|
||||
OR NEW.category <> 'message' OR NEW.event_type NOT IN ('llm.human.input','llm.ai.response','llm.tool.result') THEN
|
||||
RAISE EXCEPTION 'History seed authority is missing, stale or outside the host grant' USING ERRCODE='42501';
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
IF a IS NULL OR r.run_id IS NULL OR r.operation_kind <> 'run'
|
||||
OR a->>'run_id' IS DISTINCT FROM NEW.run_id
|
||||
OR a->>'thread_id' IS DISTINCT FROM NEW.thread_id
|
||||
OR a->>'user_id' IS DISTINCT FROM NEW.user_id
|
||||
OR r.thread_id IS DISTINCT FROM NEW.thread_id
|
||||
OR r.user_id IS DISTINCT FROM NEW.user_id
|
||||
OR r.run_incarnation IS NULL
|
||||
OR a->>'run_incarnation' IS DISTINCT FROM r.run_incarnation
|
||||
OR (a->>'retention_revision')::bigint IS DISTINCT FROM r.evidence_retention_revision THEN
|
||||
RAISE EXCEPTION 'Run event producer authority is missing or stale' USING ERRCODE='42501';
|
||||
END IF;
|
||||
IF a->>'source' IN ('producer', 'publisher') THEN
|
||||
IF a->>'producer_epoch' IS DISTINCT FROM r.producer_epoch
|
||||
OR a->>'owner_worker_id' IS DISTINCT FROM r.owner_worker_id
|
||||
OR r.producer_state NOT IN ('open', 'closed')
|
||||
OR (a->>'source'='producer' AND (r.producer_state <> 'open' OR NEW.event_type IN ('run.terminal','run.delivery'))) THEN
|
||||
RAISE EXCEPTION 'Run event producer was revoked or closed' USING ERRCODE='42501';
|
||||
END IF;
|
||||
ELSIF a->>'source'='takeover' THEN
|
||||
IF r.status <> 'error' OR r.producer_state <> 'revoked' OR NEW.event_type NOT IN ('run.terminal','run.delivery') THEN
|
||||
RAISE EXCEPTION 'Invalid takeover event authority' USING ERRCODE='42501';
|
||||
END IF;
|
||||
ELSIF a->>'source'='audit' THEN
|
||||
IF NEW.event_type NOT LIKE 'audit.%' THEN
|
||||
RAISE EXCEPTION 'Invalid audit event authority' USING ERRCODE='42501';
|
||||
END IF;
|
||||
ELSE
|
||||
RAISE EXCEPTION 'Invalid event authority source' USING ERRCODE='42501';
|
||||
END IF;
|
||||
IF a->>'source' <> 'producer' AND NOT coalesce((a->'event_types') ? NEW.event_type, false) THEN
|
||||
RAISE EXCEPTION 'Event is outside the host grant' USING ERRCODE='42501';
|
||||
END IF;
|
||||
RETURN NEW;
|
||||
END $guard$
|
||||
""")
|
||||
)
|
||||
connection.execute(text("DROP TRIGGER IF EXISTS run_event_producer_guard ON run_events"))
|
||||
connection.execute(text(f"CREATE TRIGGER run_event_producer_guard BEFORE INSERT ON run_events FOR EACH ROW EXECUTE FUNCTION {quoted}.run_event_producer_guard()"))
|
||||
|
||||
|
||||
async def activate_producer_fencing(session_factory, *, cutover_id: str) -> None:
|
||||
"""Offline-only, monotonic activation; callers must stop every Gateway."""
|
||||
from deerflow.persistence.run.model import RunProducerPolicyRow, RunRow
|
||||
|
||||
if not cutover_id or len(cutover_id) > 64:
|
||||
raise ValueError("A cutover ID of 1 to 64 characters is required")
|
||||
async with session_factory() as session, session.begin():
|
||||
if session.get_bind().dialect.name != "postgresql":
|
||||
raise RuntimeError("Strict producer fencing requires PostgreSQL")
|
||||
others = await session.scalar(text("SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND pid<>pg_backend_pid() AND backend_type='client backend'"))
|
||||
if others:
|
||||
raise RuntimeError("Stop all Gateway/database client sessions before producer cutover")
|
||||
await session.execute(text("SELECT pg_advisory_xact_lock(hashtext('deerflow.event_fence/' || current_schema())::bigint)"))
|
||||
if await session.scalar(select(func.count()).select_from(RunRow).where(RunRow.status.in_(("pending", "running")))):
|
||||
raise RuntimeError("Drain active Run admissions before producer cutover")
|
||||
guard = await session.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard' AND tgenabled='O'"))
|
||||
if not guard:
|
||||
raise RuntimeError("The producer fencing database migration/guard is missing")
|
||||
policy = await session.get(RunProducerPolicyRow, 1)
|
||||
if policy is not None and policy.strict and policy.cutover_id != cutover_id:
|
||||
raise RuntimeError("Producer fencing is already activated under another cutover ID")
|
||||
if policy is None:
|
||||
session.add(RunProducerPolicyRow(id=1, strict=True, cutover_id=cutover_id))
|
||||
else:
|
||||
policy.strict = True
|
||||
policy.cutover_id = cutover_id
|
||||
|
||||
|
||||
async def validate_producer_fencing(session_factory, *, mode: str, cutover_id: str | None) -> None:
|
||||
from deerflow.persistence.run.model import RunProducerPolicyRow
|
||||
|
||||
async with session_factory() as session:
|
||||
if session.get_bind().dialect.name != "postgresql":
|
||||
if mode != "off":
|
||||
raise RuntimeError("Strict producer fencing requires PostgreSQL RunStore + DB EventStore")
|
||||
return
|
||||
policy = await session.get(RunProducerPolicyRow, 1)
|
||||
if mode == "strict":
|
||||
if policy is None or not policy.strict or not cutover_id or policy.cutover_id != cutover_id:
|
||||
raise RuntimeError("Strict producer fencing requires the matching offline cutover")
|
||||
guard = await session.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard' AND tgenabled='O'"))
|
||||
if not guard:
|
||||
raise RuntimeError("Strict producer fencing database guard is unavailable")
|
||||
elif policy is not None and policy.strict:
|
||||
raise RuntimeError("Producer fencing is activated; refusing an unfenced Gateway")
|
||||
@ -46,6 +46,7 @@ def _origin_kind(operation_kind: str, metadata: Any) -> str | None:
|
||||
|
||||
|
||||
class RunRepository(RunStore):
|
||||
supports_event_producer_fencing = True
|
||||
supports_run_inputs = True
|
||||
|
||||
async def mutate_run_inputs(self, run_id, change, *, user_id=None, owner_worker_id=None):
|
||||
@ -232,11 +233,15 @@ class RunRepository(RunStore):
|
||||
values["change_seq"] = await self._next_change_seq(session)
|
||||
row = await session.get(RunRow, run_id)
|
||||
if row is None:
|
||||
session.add(RunRow(run_id=run_id, created_at=created, evidence_seal_state="pending", **values))
|
||||
row = RunRow(run_id=run_id, created_at=created, evidence_seal_state="pending", run_incarnation=uuid4().hex, producer_epoch=uuid4().hex, producer_state="open", **values)
|
||||
session.add(row)
|
||||
else:
|
||||
for key, value in values.items():
|
||||
setattr(row, key, value)
|
||||
await session.flush()
|
||||
admitted = self._row_to_dict(row)
|
||||
await session.commit()
|
||||
return admitted
|
||||
|
||||
async def seal_evidence(self, run_id: str, *, owner_worker_id: str, receipt) -> bool:
|
||||
"""Persist a completed drain and changed-run position in one fenced transaction."""
|
||||
@ -964,6 +969,8 @@ class RunRepository(RunStore):
|
||||
cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds)
|
||||
values: dict[str, Any] = {
|
||||
"status": "error",
|
||||
"producer_state": "revoked",
|
||||
"producer_epoch": uuid4().hex,
|
||||
"error": error,
|
||||
"evidence_seal_state": "partial",
|
||||
"evidence_seal_error": "Worker lease expired without a completed journal seal.",
|
||||
@ -1062,7 +1069,7 @@ class RunRepository(RunStore):
|
||||
operation_kind: str = "run",
|
||||
multitask_strategy: str = "reject",
|
||||
assistant_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
user_id: str | None | _AutoSentinel = AUTO,
|
||||
model_name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
@ -1089,7 +1096,7 @@ class RunRepository(RunStore):
|
||||
"""
|
||||
from deerflow.runtime.runs.manager import ConflictError
|
||||
|
||||
resolved_user_id = resolve_user_id(user_id or AUTO, method_name="RunRepository.create_thread_operation_atomic")
|
||||
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.create_thread_operation_atomic")
|
||||
now = datetime.now(UTC)
|
||||
created = datetime.fromisoformat(created_at) if created_at else now
|
||||
lease_dt = datetime.fromisoformat(lease_expires_at) if lease_expires_at else None
|
||||
@ -1097,6 +1104,9 @@ class RunRepository(RunStore):
|
||||
|
||||
values = {
|
||||
"thread_id": thread_id,
|
||||
"run_incarnation": uuid4().hex,
|
||||
"producer_epoch": uuid4().hex,
|
||||
"producer_state": "open",
|
||||
"evidence_origin": evidence_origin,
|
||||
"evidence_agent_id": evidence_agent_id,
|
||||
"evidence_seal_state": "pending",
|
||||
@ -1169,14 +1179,20 @@ class RunRepository(RunStore):
|
||||
raise ConflictError(f"Thread {thread_id} has an active checkpoint write")
|
||||
row.status = "interrupted"
|
||||
row.error = "Cancelled by newer run"
|
||||
if row.owner_worker_id != owner_worker_id:
|
||||
row.producer_state = "revoked"
|
||||
row.producer_epoch = uuid4().hex
|
||||
row.owner_worker_id = owner_worker_id
|
||||
row.updated_at = now
|
||||
row.change_seq = change_seq
|
||||
claimed.append(self._row_to_dict(row))
|
||||
|
||||
values["change_seq"] = change_seq
|
||||
session.add(RunRow(run_id=run_id, **values))
|
||||
new_row = RunRow(run_id=run_id, **values)
|
||||
session.add(new_row)
|
||||
try:
|
||||
await session.flush()
|
||||
admitted = self._row_to_dict(new_row)
|
||||
await session.commit()
|
||||
except IntegrityError as exc:
|
||||
await session.rollback()
|
||||
@ -1186,5 +1202,4 @@ class RunRepository(RunStore):
|
||||
raise RunIdempotencyConflict(self._row_to_dict(existing)) from exc
|
||||
raise
|
||||
|
||||
new_row = await session.get(RunRow, run_id)
|
||||
return self._row_to_dict(new_row), claimed
|
||||
return admitted, claimed
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import operator
|
||||
import re
|
||||
from datetime import UTC, datetime, timedelta
|
||||
@ -27,6 +28,8 @@ from deerflow.scheduler.schedules import next_run_at as compute_next_run_at
|
||||
from deerflow.scheduler.schedules import should_advance_schedule_cursor
|
||||
from deerflow.utils.time import coerce_iso
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EXECUTING_RUN_STATUSES: tuple[str, ...] = ("launching", "running")
|
||||
_SCHEDULER_BUDGET_LOCK_KEY = 4694001
|
||||
|
||||
@ -1043,7 +1046,14 @@ class ScheduledTaskRunRepository:
|
||||
finally:
|
||||
if self._recovery_terminal_publisher is not None:
|
||||
for target in recovery_targets:
|
||||
await self._recovery_terminal_publisher.publish(target)
|
||||
confirmation = await self._recovery_terminal_publisher.publish(target)
|
||||
if not confirmation.confirmed:
|
||||
logger.warning(
|
||||
"Scheduled recovery terminal publication for run %s returned %s",
|
||||
target.run_id,
|
||||
confirmation.value,
|
||||
extra={"run_id": target.run_id, "recovery_path": "reconcile_active_runs", "confirmation": confirmation.value},
|
||||
)
|
||||
|
||||
async def _reconcile_active_runs(
|
||||
self,
|
||||
@ -1122,6 +1132,8 @@ class ScheduledTaskRunRepository:
|
||||
run_id=candidate.run_id,
|
||||
user_id=candidate.user_id,
|
||||
operation_kind=candidate.operation_kind,
|
||||
run_incarnation=candidate.run_incarnation,
|
||||
retention_revision=candidate.evidence_retention_revision,
|
||||
)
|
||||
if target is not None:
|
||||
recovery_targets.append(target)
|
||||
|
||||
@ -1031,7 +1031,14 @@ class ScheduledTaskRepository:
|
||||
finally:
|
||||
if self._recovery_terminal_publisher is not None:
|
||||
for target in recovery_targets:
|
||||
await self._recovery_terminal_publisher.publish(target)
|
||||
confirmation = await self._recovery_terminal_publisher.publish(target)
|
||||
if not confirmation.confirmed:
|
||||
logger.warning(
|
||||
"Scheduled recovery terminal publication for run %s returned %s",
|
||||
target.run_id,
|
||||
confirmation.value,
|
||||
extra={"run_id": target.run_id, "recovery_path": "reconcile_stuck_once_tasks", "confirmation": confirmation.value},
|
||||
)
|
||||
|
||||
async def _reconcile_stuck_once_tasks(
|
||||
self,
|
||||
@ -1087,6 +1094,8 @@ class ScheduledTaskRepository:
|
||||
run_id=candidate.run_id,
|
||||
user_id=candidate.user_id,
|
||||
operation_kind=candidate.operation_kind,
|
||||
run_incarnation=candidate.run_incarnation,
|
||||
retention_revision=candidate.evidence_retention_revision,
|
||||
)
|
||||
if target is not None:
|
||||
recovery_targets.append(target)
|
||||
|
||||
@ -90,9 +90,11 @@ class ThreadMetaRepository(ThreadMetaStore):
|
||||
updated_at=now,
|
||||
)
|
||||
session.add(row)
|
||||
await session.flush()
|
||||
# A refresh after commit could observe a replacement incarnation.
|
||||
created = self._row_to_dict(row)
|
||||
await session.commit()
|
||||
await session.refresh(row)
|
||||
return self._row_to_dict(row)
|
||||
return created
|
||||
|
||||
async def claim_unowned(self, thread_id: str, owner: str) -> bool:
|
||||
claim_target = table(
|
||||
|
||||
@ -394,6 +394,25 @@ JSONL's single-process deployment constraint. Regression coverage is in
|
||||
|
||||
## Completed-evidence provenance
|
||||
|
||||
SQL Run admissions return their incarnation/epoch from the committing transaction;
|
||||
Manager binds that snapshot without a later identity lookup. `run_agent` installs
|
||||
host authority in a private ContextVar, inherited by owned tasks and offloaded
|
||||
callbacks, never metadata/configurable. Strict DB admission enforces one Run per
|
||||
batch through all original write APIs. Terminal/delivery narrow the grant only
|
||||
after existing final ACK/one-shot takeover checks. Post-terminal audit grants
|
||||
name individual `audit.*` types and retain incarnation/retention fencing.
|
||||
History seeds use a separate host grant from a committed `checkpoint_write`
|
||||
reservation and captured ThreadMeta identity, permitting multiple synthetic turn
|
||||
IDs but only the three message event types. SQL admission and the strict trigger
|
||||
hold Thread fence → reservation row → ThreadMeta row through commit; validate
|
||||
owner, incarnation, epoch, state, retention revision and DB-clock lease after
|
||||
both row-lock waits. Missing legacy metadata is created under that reservation;
|
||||
never recapture a replacement incarnation after reading a checkpoint.
|
||||
Completion hooks remain legal after business terminal commit while the producer
|
||||
is open. The Worker explicitly closes after Journal drain and Subagent buffer
|
||||
completion; uncertain drains conservatively leave it open. Closure is explicit;
|
||||
elapsed time or business terminal status alone grants no recovery publication authority.
|
||||
|
||||
Gateway admission and RunManager validate host-owned evidence origin against the shared `EVIDENCE_ORIGINS` in `runtime/runs/schemas.py`. Preserve the string values and reject client metadata as provenance; only trusted Python entry points set the admission argument.
|
||||
|
||||
Idempotent resume reuse preserves both the private request identity and the original evidence origin/agent across worker and store boundaries. A retry must not relabel existing evidence. Memory/SQLite parity is covered in `tests/test_completed_evidence_admission.py`.
|
||||
|
||||
161
backend/packages/harness/deerflow/runtime/events/producer.py
Normal file
161
backend/packages/harness/deerflow/runtime/events/producer.py
Normal file
@ -0,0 +1,161 @@
|
||||
"""Host-only event authority; never reconstructed from event metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import asdict, dataclass, replace
|
||||
from functools import wraps
|
||||
from typing import Literal
|
||||
|
||||
from deerflow.runtime.events.store.base import RunEventStoreNotCommittedError
|
||||
|
||||
|
||||
class RunEventProducerRejected(RunEventStoreNotCommittedError):
|
||||
"""Admission rejected the whole batch before any event/seq was written."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunEventAuthority:
|
||||
run_id: str
|
||||
thread_id: str
|
||||
user_id: str | None
|
||||
run_incarnation: str
|
||||
producer_epoch: str | None
|
||||
owner_worker_id: str | None
|
||||
retention_revision: int
|
||||
source: Literal["producer", "publisher", "takeover", "audit"] = "producer"
|
||||
event_types: tuple[str, ...] = ()
|
||||
|
||||
def to_storage(self) -> dict:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class HistorySeedAuthority:
|
||||
"""Host grant tied to a durable checkpoint-write reservation, not a Run."""
|
||||
|
||||
thread_id: str
|
||||
user_id: str | None
|
||||
thread_incarnation: str | None
|
||||
thread_user_id: str | None
|
||||
reservation_id: str
|
||||
reservation_incarnation: str
|
||||
producer_epoch: str
|
||||
owner_worker_id: str
|
||||
retention_revision: int
|
||||
seed_kind: Literal["checkpoint", "branch"]
|
||||
source: Literal["history_seed"] = "history_seed"
|
||||
|
||||
def to_storage(self) -> dict:
|
||||
return asdict(self)
|
||||
|
||||
def allows(self, event: dict) -> bool:
|
||||
prefix = f"{self.seed_kind}-seed-{self.thread_id}-"
|
||||
return (
|
||||
self.seed_kind in ("checkpoint", "branch")
|
||||
and event["thread_id"] == self.thread_id
|
||||
and event.get("user_id", self.user_id) == self.user_id
|
||||
and re.fullmatch(re.escape(prefix) + r"(?:0|[1-9][0-9]*)", event["run_id"]) is not None
|
||||
and event.get("category") == "message"
|
||||
and event["event_type"] in ("llm.human.input", "llm.ai.response", "llm.tool.result")
|
||||
)
|
||||
|
||||
|
||||
def history_seed_authority(reservation, thread: dict, *, seed_kind: Literal["checkpoint", "branch"]) -> HistorySeedAuthority:
|
||||
"""Use only the host's committed reservation and captured thread identity."""
|
||||
if (
|
||||
seed_kind not in ("checkpoint", "branch")
|
||||
or getattr(reservation, "operation_kind", None) != "checkpoint_write"
|
||||
or getattr(reservation, "store_only", False)
|
||||
or not getattr(reservation, "run_incarnation", None)
|
||||
or not getattr(reservation, "producer_epoch", None)
|
||||
or not getattr(reservation, "owner_worker_id", None)
|
||||
or thread.get("thread_id") != reservation.thread_id
|
||||
or "incarnation" not in thread
|
||||
or (thread.get("user_id") is not None and thread["user_id"] != reservation.user_id)
|
||||
):
|
||||
raise RunEventProducerRejected("History seeding requires a host reservation and matching thread owner")
|
||||
return HistorySeedAuthority(
|
||||
thread_id=reservation.thread_id,
|
||||
user_id=reservation.user_id,
|
||||
thread_incarnation=thread["incarnation"],
|
||||
thread_user_id=thread.get("user_id"),
|
||||
reservation_id=reservation.run_id,
|
||||
reservation_incarnation=reservation.run_incarnation,
|
||||
producer_epoch=reservation.producer_epoch,
|
||||
owner_worker_id=reservation.owner_worker_id,
|
||||
retention_revision=reservation.evidence_retention_revision,
|
||||
seed_kind=seed_kind,
|
||||
)
|
||||
|
||||
|
||||
_authority: ContextVar[RunEventAuthority | HistorySeedAuthority | None] = ContextVar("deerflow_event_authority", default=None)
|
||||
|
||||
|
||||
def current_producer() -> RunEventAuthority | HistorySeedAuthority | None:
|
||||
return _authority.get()
|
||||
|
||||
|
||||
def producer_authority(record) -> RunEventAuthority | None:
|
||||
"""Bind the actual host admission, not a user-provided context/metadata map."""
|
||||
if getattr(record, "store_only", False):
|
||||
return None
|
||||
incarnation = getattr(record, "run_incarnation", None)
|
||||
epoch = getattr(record, "producer_epoch", None)
|
||||
if not incarnation or not epoch:
|
||||
return None
|
||||
return RunEventAuthority(record.run_id, record.thread_id, record.user_id, incarnation, epoch, getattr(record, "owner_worker_id", None), getattr(record, "evidence_retention_revision", 0))
|
||||
|
||||
|
||||
def run_producer_context(function):
|
||||
"""Keep host authority out of graph/checkpoint/user context carriers."""
|
||||
|
||||
@wraps(function)
|
||||
async def scoped(bridge, run_manager, record, **kwargs):
|
||||
authority = producer_authority(record)
|
||||
if record.store_only or (authority is not None and (run_manager._runs.get(record.run_id) is not record or record.owner_worker_id != run_manager._worker_id)):
|
||||
from deerflow.runtime.runs.manager import RunStartupError
|
||||
|
||||
raise RunStartupError("Only the locally admitted Worker may produce Run events")
|
||||
with producer_scope(authority):
|
||||
return await function(bridge, run_manager, record, **kwargs)
|
||||
|
||||
return scoped
|
||||
|
||||
|
||||
@contextmanager
|
||||
def producer_scope(authority: RunEventAuthority | HistorySeedAuthority | None):
|
||||
token = _authority.set(authority)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_authority.reset(token)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def publisher_scope(event_type: str):
|
||||
"""Called only after the existing host publisher's authority checks."""
|
||||
current = current_producer()
|
||||
if current is not None and current.source == "producer":
|
||||
current = replace(current, source="publisher", event_types=(event_type,))
|
||||
with producer_scope(current):
|
||||
yield
|
||||
|
||||
|
||||
def takeover_authority(*, run_id, thread_id, user_id, run_incarnation, retention_revision=0):
|
||||
"""Issue only from a successful atomic takeover's one-shot target."""
|
||||
if not run_incarnation:
|
||||
return None
|
||||
return RunEventAuthority(run_id, thread_id, user_id, run_incarnation, None, None, retention_revision, "takeover", ("run.terminal", "run.delivery"))
|
||||
|
||||
|
||||
def audit_authority(record, *, event_types: tuple[str, ...]) -> RunEventAuthority:
|
||||
"""Explicit host audit grant; no model/tool-visible bypass flag."""
|
||||
if not event_types or any(not event_type.startswith("audit.") for event_type in event_types):
|
||||
raise ValueError("Audit authority is limited to named audit.* events")
|
||||
authority = producer_authority(record)
|
||||
if authority is None:
|
||||
raise ValueError("Audit authority requires an admitted Run incarnation")
|
||||
return replace(authority, source="audit", producer_epoch=None, event_types=event_types)
|
||||
@ -11,11 +11,13 @@ def make_run_event_store(config=None) -> RunEventStore:
|
||||
|
||||
sf = get_session_factory()
|
||||
if sf is None:
|
||||
if getattr(config, "producer_fencing", "off") == "strict":
|
||||
raise RuntimeError("Strict producer fencing requires a persistent PostgreSQL RunStore")
|
||||
# database.backend=memory but run_events.backend=db -> fallback
|
||||
return MemoryRunEventStore()
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
return DbRunEventStore(sf, max_trace_content=config.max_trace_content)
|
||||
return DbRunEventStore(sf, max_trace_content=config.max_trace_content, producer_fencing=getattr(config, "producer_fencing", "off"), producer_cutover_id=getattr(config, "producer_cutover_id", None))
|
||||
if config.backend == "jsonl":
|
||||
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
|
||||
|
||||
|
||||
@ -15,13 +15,15 @@ import weakref
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import delete, func, or_, select, text, update
|
||||
from sqlalchemy import and_, delete, func, or_, select, text, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunRow
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
||||
from deerflow.runtime.events.message_identity import message_identity
|
||||
from deerflow.runtime.events.producer import HistorySeedAuthority, RunEventProducerRejected, current_producer
|
||||
from deerflow.runtime.events.store.base import RunEventLifecycleRejected, RunEventStore, SessionAwareRunEventLifecycleGuard
|
||||
from deerflow.runtime.user_context import AUTO, _AutoSentinel, get_current_user, resolve_user_id
|
||||
from deerflow.utils.time import coerce_iso
|
||||
@ -37,9 +39,11 @@ def _parse_created_at(value: str | None) -> datetime:
|
||||
|
||||
|
||||
class DbRunEventStore(RunEventStore):
|
||||
def __init__(self, session_factory: async_sessionmaker[AsyncSession], *, max_trace_content: int = 10240):
|
||||
def __init__(self, session_factory: async_sessionmaker[AsyncSession], *, max_trace_content: int = 10240, producer_fencing: str = "off", producer_cutover_id: str | None = None):
|
||||
self._sf = session_factory
|
||||
self._max_trace_content = max_trace_content
|
||||
self._producer_fencing = producer_fencing
|
||||
self._producer_cutover_id = producer_cutover_id
|
||||
# Per-thread asyncio locks serialize seq assignment for concurrent
|
||||
# in-process writers on the same thread. The DB-level FOR UPDATE /
|
||||
# advisory lock guards cross-process races; this guards the common
|
||||
@ -65,6 +69,11 @@ class DbRunEventStore(RunEventStore):
|
||||
self._write_lock_pins[thread_id] = lock
|
||||
return lock
|
||||
|
||||
async def validate_producer_fencing(self) -> None:
|
||||
from deerflow.persistence.run.producer import validate_producer_fencing
|
||||
|
||||
await validate_producer_fencing(self._sf, mode=self._producer_fencing, cutover_id=self._producer_cutover_id)
|
||||
|
||||
@staticmethod
|
||||
def _row_to_dict(row: RunEventRow) -> dict:
|
||||
d = row.to_dict()
|
||||
@ -211,6 +220,82 @@ class DbRunEventStore(RunEventStore):
|
||||
elif watermark.seq < seq:
|
||||
watermark.seq = seq
|
||||
|
||||
async def _admit_events(self, session: AsyncSession, events: list[dict]) -> None:
|
||||
"""Hold T then R until COMMIT; no status-derived producer authority."""
|
||||
authority = current_producer()
|
||||
if authority is None:
|
||||
if self._producer_fencing == "strict":
|
||||
raise RunEventProducerRejected("Run event write requires host producer authority")
|
||||
return
|
||||
if isinstance(authority, HistorySeedAuthority):
|
||||
await self._admit_history_seed(session, events, authority)
|
||||
return
|
||||
for event in events:
|
||||
if (event["thread_id"], event["run_id"], event.get("user_id", authority.user_id)) != (authority.thread_id, authority.run_id, authority.user_id):
|
||||
raise RunEventProducerRejected("Event batch must match one admitted Run and owner")
|
||||
reserved = event["event_type"] in {"run.terminal", "run.delivery"}
|
||||
if authority.source == "producer":
|
||||
if reserved:
|
||||
raise RunEventProducerRejected("Terminal evidence requires publisher authority")
|
||||
elif event["event_type"] not in authority.event_types:
|
||||
raise RunEventProducerRejected("Event type is outside the host grant")
|
||||
await self._acquire_thread_mutation_fence(session, authority.thread_id)
|
||||
row = await session.scalar(select(RunRow).where(RunRow.run_id == authority.run_id).with_for_update().execution_options(populate_existing=True))
|
||||
if row is None or (row.thread_id, row.user_id, row.run_incarnation, row.evidence_retention_revision, row.operation_kind) != (authority.thread_id, authority.user_id, authority.run_incarnation, authority.retention_revision, "run"):
|
||||
raise RunEventProducerRejected("Run producer identity is stale or deleted")
|
||||
if authority.source in {"producer", "publisher"}:
|
||||
states = ("open",) if authority.source == "producer" else ("open", "closed")
|
||||
if row.producer_epoch != authority.producer_epoch or row.producer_state not in states or row.owner_worker_id != authority.owner_worker_id:
|
||||
raise RunEventProducerRejected("Run producer was closed or revoked")
|
||||
elif authority.source == "takeover":
|
||||
if row.status != "error" or row.producer_state != "revoked":
|
||||
raise RunEventProducerRejected("Takeover publisher no longer matches the committed Run")
|
||||
elif authority.source != "audit":
|
||||
raise RunEventProducerRejected("Unsupported host event authority")
|
||||
if session.get_bind().dialect.name == "postgresql":
|
||||
await session.execute(text("SELECT set_config('deerflow.event_authority', :authority, true)"), {"authority": json.dumps(authority.to_storage())})
|
||||
|
||||
async def _admit_history_seed(self, session: AsyncSession, events: list[dict], authority: HistorySeedAuthority) -> None:
|
||||
if not all(authority.allows(event) for event in events):
|
||||
raise RunEventProducerRejected("History seed batch is outside the host grant")
|
||||
await self._acquire_thread_mutation_fence(session, authority.thread_id)
|
||||
reservation = await session.scalar(select(RunRow).where(RunRow.run_id == authority.reservation_id).with_for_update().execution_options(populate_existing=True))
|
||||
thread = await session.scalar(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == authority.thread_id).with_for_update().execution_options(populate_existing=True))
|
||||
if (
|
||||
reservation is None
|
||||
or thread is None
|
||||
or (reservation.thread_id, reservation.user_id, reservation.run_incarnation, reservation.producer_epoch, reservation.owner_worker_id, reservation.evidence_retention_revision)
|
||||
!= (authority.thread_id, authority.user_id, authority.reservation_incarnation, authority.producer_epoch, authority.owner_worker_id, authority.retention_revision)
|
||||
or (reservation.operation_kind, reservation.status, reservation.producer_state) != ("checkpoint_write", "pending", "open")
|
||||
or (thread.incarnation, thread.user_id) != (authority.thread_incarnation, authority.thread_user_id)
|
||||
or (thread.user_id is not None and thread.user_id != authority.user_id)
|
||||
):
|
||||
raise RunEventProducerRejected("History seed reservation or thread was revoked or replaced")
|
||||
# Evaluate after every row-lock wait. NULL is the existing single-worker
|
||||
# reservation contract; leased reservations never receive expiry grace.
|
||||
if reservation.lease_expires_at is not None and not await session.scalar(select(RunRepository._live_lease_predicate(session)).where(RunRow.run_id == authority.reservation_id)):
|
||||
raise RunEventProducerRejected("History seed reservation lease expired")
|
||||
if session.get_bind().dialect.name == "postgresql":
|
||||
await session.execute(text("SELECT set_config('deerflow.event_authority', :authority, true)"), {"authority": json.dumps(authority.to_storage())})
|
||||
|
||||
async def close_producer(self, authority) -> bool:
|
||||
"""Persist explicit drained closure; never infer closure from Run.status."""
|
||||
if authority is None or authority.source != "producer":
|
||||
return False
|
||||
async with self._get_write_lock(authority.thread_id):
|
||||
async with self._sf() as session, session.begin():
|
||||
await self._acquire_thread_mutation_fence(session, authority.thread_id)
|
||||
row = await session.scalar(select(RunRow).where(RunRow.run_id == authority.run_id).with_for_update())
|
||||
if (
|
||||
row is None
|
||||
or (row.thread_id, row.user_id, row.run_incarnation, row.producer_epoch, row.evidence_retention_revision)
|
||||
!= (authority.thread_id, authority.user_id, authority.run_incarnation, authority.producer_epoch, authority.retention_revision)
|
||||
or row.producer_state not in ("open", "closed")
|
||||
):
|
||||
return False
|
||||
row.producer_state = "closed"
|
||||
return True
|
||||
|
||||
async def put(self, *, thread_id, run_id, event_type, category, content="", metadata=None, created_at=None): # noqa: D401
|
||||
"""Write a single event — low-frequency path only.
|
||||
|
||||
@ -222,10 +307,12 @@ class DbRunEventStore(RunEventStore):
|
||||
"""
|
||||
content, metadata = self._truncate_trace(category, content, metadata)
|
||||
db_content, metadata = self._content_to_db(content, metadata)
|
||||
user_id = self._user_id_from_context()
|
||||
authority = current_producer()
|
||||
user_id = authority.user_id if authority is not None else self._user_id_from_context()
|
||||
async with self._get_write_lock(thread_id):
|
||||
async with self._sf() as session:
|
||||
async with session.begin():
|
||||
await self._admit_events(session, [{"thread_id": thread_id, "run_id": run_id, "user_id": user_id, "event_type": event_type}])
|
||||
max_seq = await self._max_seq_for_thread(session, thread_id)
|
||||
seq = (max_seq or 0) + 1
|
||||
row = RunEventRow(
|
||||
@ -250,12 +337,14 @@ class DbRunEventStore(RunEventStore):
|
||||
thread_ids = {e["thread_id"] for e in events}
|
||||
if len(thread_ids) > 1:
|
||||
raise ValueError(f"put_batch requires all events to belong to the same thread; got {thread_ids!r}")
|
||||
user_id = self._user_id_from_context()
|
||||
authority = current_producer()
|
||||
user_id = authority.user_id if authority is not None else self._user_id_from_context()
|
||||
# All events belong to the same thread (validated above).
|
||||
thread_id = events[0]["thread_id"]
|
||||
async with self._get_write_lock(thread_id):
|
||||
async with self._sf() as session:
|
||||
async with session.begin():
|
||||
await self._admit_events(session, events)
|
||||
max_seq = await self._max_seq_for_thread(session, thread_id)
|
||||
seq = max_seq or 0
|
||||
rows = []
|
||||
@ -306,10 +395,12 @@ class DbRunEventStore(RunEventStore):
|
||||
"""
|
||||
content, metadata = self._truncate_trace(category, content, metadata)
|
||||
db_content, metadata = self._content_to_db(content, metadata)
|
||||
user_id = self._user_id_from_context() if user_id is AUTO else user_id
|
||||
authority = current_producer()
|
||||
user_id = (authority.user_id if authority is not None else self._user_id_from_context()) if user_id is AUTO else user_id
|
||||
async with self._get_write_lock(thread_id):
|
||||
async with self._sf() as session:
|
||||
async with session.begin():
|
||||
await self._admit_events(session, [{"thread_id": thread_id, "run_id": run_id, "user_id": user_id, "event_type": event_type}])
|
||||
max_seq = await self._max_seq_for_thread(session, thread_id)
|
||||
if lifecycle_guard is not None:
|
||||
admitted = await lifecycle_guard.check_in_session(session) if isinstance(lifecycle_guard, SessionAwareRunEventLifecycleGuard) else await lifecycle_guard()
|
||||
@ -403,11 +494,14 @@ class DbRunEventStore(RunEventStore):
|
||||
if user_id is not None:
|
||||
conditions.append(RunRow.user_id == user_id)
|
||||
if run_id is not None:
|
||||
conditions.append(RunRow.run_id == run_id)
|
||||
# Synthetic turns have no RunRow. Revoke outstanding history grants
|
||||
# on any per-run deletion so they cannot reinsert the deleted turn.
|
||||
conditions.append(or_(RunRow.run_id == run_id, and_(RunRow.operation_kind == "checkpoint_write", RunRow.status.in_(("pending", "running")))))
|
||||
await session.execute(
|
||||
update(RunRow)
|
||||
.where(*conditions)
|
||||
.values(
|
||||
producer_state="revoked",
|
||||
evidence_retention_revision=RunRow.evidence_retention_revision + 1,
|
||||
evidence_seal_state="partial",
|
||||
change_seq=change_seq,
|
||||
|
||||
@ -9,6 +9,7 @@ from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
|
||||
from deerflow.runtime.events.catalog import RUN_TERMINAL_EVENT
|
||||
from deerflow.runtime.events.producer import producer_scope, publisher_scope, takeover_authority
|
||||
from deerflow.runtime.events.store.base import RunEventLifecycleRejected, RunEventStore, SessionAwareRunEventLifecycleGuard
|
||||
from deerflow.runtime.runs.store.base import RunStore
|
||||
|
||||
@ -48,6 +49,8 @@ class _RecoveryTargetIdentity:
|
||||
user_id: str | None
|
||||
operation_kind: str
|
||||
committed_status: str
|
||||
run_incarnation: str | None = None
|
||||
retention_revision: int = 0
|
||||
|
||||
|
||||
class RecoveryTerminalTarget:
|
||||
@ -64,6 +67,8 @@ class RecoveryTerminalTarget:
|
||||
operation_kind: str,
|
||||
committed_status: str,
|
||||
_issuer: object,
|
||||
run_incarnation: str | None = None,
|
||||
retention_revision: int = 0,
|
||||
) -> None:
|
||||
if _issuer is not _RECOVERY_TARGET_ISSUER:
|
||||
raise TypeError("RecoveryTerminalTarget must be issued by RecoveryTerminalPublisher")
|
||||
@ -75,7 +80,7 @@ class RecoveryTerminalTarget:
|
||||
raise ValueError("Recovery terminal target requires thread_id and run_id")
|
||||
if user_id is not None and not isinstance(user_id, str):
|
||||
raise TypeError("Recovery terminal target user_id must be a string or None")
|
||||
self._identity = _RecoveryTargetIdentity(thread_id, run_id, user_id, operation_kind, committed_status)
|
||||
self._identity = _RecoveryTargetIdentity(thread_id, run_id, user_id, operation_kind, committed_status, run_incarnation, retention_revision)
|
||||
self._claim_state = _RecoveryTargetClaimState()
|
||||
|
||||
def __setattr__(self, name: str, value: object) -> None:
|
||||
@ -158,6 +163,8 @@ class RecoveryTerminalPublisher:
|
||||
run_id: str,
|
||||
user_id: str | None,
|
||||
operation_kind: str,
|
||||
run_incarnation: str | None = None,
|
||||
retention_revision: int = 0,
|
||||
) -> RecoveryTerminalTarget | None:
|
||||
"""Issue one-shot publication authority only for a committed run takeover."""
|
||||
if claimed is not True or operation_kind != "run":
|
||||
@ -169,6 +176,8 @@ class RecoveryTerminalPublisher:
|
||||
operation_kind=operation_kind,
|
||||
committed_status="error",
|
||||
_issuer=_RECOVERY_TARGET_ISSUER,
|
||||
run_incarnation=run_incarnation,
|
||||
retention_revision=retention_revision,
|
||||
)
|
||||
|
||||
async def publish(self, target: RecoveryTerminalTarget) -> TerminalEventConfirmation:
|
||||
@ -176,6 +185,11 @@ class RecoveryTerminalPublisher:
|
||||
if not isinstance(target, RecoveryTerminalTarget):
|
||||
raise TypeError("RecoveryTerminalPublisher.publish requires a RecoveryTerminalTarget")
|
||||
identity = target._consume()
|
||||
authority = takeover_authority(run_id=identity.run_id, thread_id=identity.thread_id, user_id=identity.user_id, run_incarnation=identity.run_incarnation, retention_revision=identity.retention_revision)
|
||||
with producer_scope(authority):
|
||||
return await self._publish_identity(identity)
|
||||
|
||||
async def _publish_identity(self, identity: _RecoveryTargetIdentity) -> TerminalEventConfirmation:
|
||||
guard = self.lifecycle_guard(thread_id=identity.thread_id, run_id=identity.run_id, user_id=identity.user_id)
|
||||
await self._ensure_delivery_receipt(identity, guard)
|
||||
try:
|
||||
@ -298,7 +312,8 @@ async def confirm_terminal_event(
|
||||
}
|
||||
if lifecycle_guard is not None:
|
||||
put_kwargs["lifecycle_guard"] = lifecycle_guard
|
||||
result = await store.put_if_absent(**put_kwargs)
|
||||
with publisher_scope("run.terminal"):
|
||||
result = await store.put_if_absent(**put_kwargs)
|
||||
except RunEventLifecycleRejected:
|
||||
return TerminalEventConfirmation.CONFLICT
|
||||
except Exception:
|
||||
|
||||
@ -329,6 +329,11 @@ class RunJournal(BaseCallbackHandler):
|
||||
self.thread_id = thread_id
|
||||
self._evidence_owner_id = user_id if isinstance(event_store, DbRunEventStore) else AUTO
|
||||
self._store: RunEventStore | None = event_store
|
||||
from deerflow.runtime.events.producer import current_producer
|
||||
|
||||
# Callback threads may have an empty Context. Keep the host admission
|
||||
# captured at construction; never infer a replacement from event data.
|
||||
self._event_producer = current_producer()
|
||||
self._closed = False
|
||||
self._accepting_evidence = True
|
||||
self._producer_lock = threading.Lock()
|
||||
@ -1216,7 +1221,10 @@ class RunJournal(BaseCallbackHandler):
|
||||
store = self._store
|
||||
if store is None:
|
||||
return
|
||||
store_task: asyncio.Task[None] = asyncio.ensure_future(store.put_batch(batch))
|
||||
from deerflow.runtime.events.producer import producer_scope
|
||||
|
||||
with producer_scope(self._event_producer):
|
||||
store_task: asyncio.Task[None] = asyncio.ensure_future(store.put_batch(batch))
|
||||
try:
|
||||
await asyncio.shield(store_task)
|
||||
except asyncio.CancelledError:
|
||||
|
||||
@ -286,6 +286,9 @@ class RunRecord:
|
||||
first_human_message: str | None = None
|
||||
finalizing: bool = False
|
||||
owner_worker_id: str | None = None
|
||||
run_incarnation: str | None = None
|
||||
producer_epoch: str | None = None
|
||||
evidence_retention_revision: int = 0
|
||||
lease_expires_at: str | None = None
|
||||
# Process-local fencing signal. Once set, this worker must not perform
|
||||
# further durable run/thread finalization because its lease ownership is
|
||||
@ -546,11 +549,23 @@ class RunManager:
|
||||
"""
|
||||
if self._store is None:
|
||||
return
|
||||
await self._call_store_with_retry(
|
||||
admitted = await self._call_store_with_retry(
|
||||
"put",
|
||||
record.run_id,
|
||||
lambda: self._store.put(record.run_id, **self._store_put_payload(record)),
|
||||
)
|
||||
self._bind_admitted_producer(record, admitted)
|
||||
|
||||
def _bind_admitted_producer(self, record: RunRecord, row: dict | None) -> None:
|
||||
# Bind only a successful local admission, never snapshot recreation or a
|
||||
# same-status historical row read. Unsupported stores remain legacy.
|
||||
if self._store is None or getattr(self._store, "supports_event_producer_fencing", False) is not True:
|
||||
return
|
||||
if row is None or row.get("run_id") != record.run_id or row.get("owner_worker_id") != self._worker_id or row.get("thread_id") != record.thread_id or row.get("user_id") != record.user_id:
|
||||
raise RunStartupError("Run producer admission identity could not be bound")
|
||||
record.run_incarnation = row.get("run_incarnation")
|
||||
record.producer_epoch = row.get("producer_epoch")
|
||||
record.evidence_retention_revision = row.get("evidence_retention_revision", 0)
|
||||
|
||||
async def _persist_to_store(self, record: RunRecord, *, error: str | None = None) -> bool:
|
||||
"""Best-effort persist run record to backing store."""
|
||||
@ -821,6 +836,9 @@ class RunManager:
|
||||
last_ai_message=row.get("last_ai_message"),
|
||||
first_human_message=row.get("first_human_message"),
|
||||
owner_worker_id=row.get("owner_worker_id"),
|
||||
run_incarnation=row.get("run_incarnation"),
|
||||
producer_epoch=row.get("producer_epoch"),
|
||||
evidence_retention_revision=row.get("evidence_retention_revision", 0),
|
||||
evidence_origin=row.get("evidence_origin") or "unknown",
|
||||
evidence_agent_id=row.get("evidence_agent_id"),
|
||||
lease_expires_at=row.get("lease_expires_at"),
|
||||
@ -1043,7 +1061,7 @@ class RunManager:
|
||||
try:
|
||||
await self._persist_new_run_to_store(record)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist run %s; record was not registered", run_id, exc_info=True)
|
||||
logger.warning("Failed to create run %s; record was not registered", run_id, exc_info=True)
|
||||
raise
|
||||
async with self._lock:
|
||||
self._runs[run_id] = record
|
||||
@ -1921,6 +1939,8 @@ class RunManager:
|
||||
run_id=run_id,
|
||||
user_id=row.get("user_id"),
|
||||
operation_kind=row.get("operation_kind", "run"),
|
||||
run_incarnation=row.get("run_incarnation"),
|
||||
retention_revision=row.get("evidence_retention_revision", 0),
|
||||
)
|
||||
if target is not None:
|
||||
await self._recovery_terminal_publisher.publish(target)
|
||||
@ -2209,7 +2229,7 @@ class RunManager:
|
||||
if idempotency_request is not None:
|
||||
create_kwargs["idempotency_request"] = idempotency_request
|
||||
try:
|
||||
await self._call_store_with_retry(
|
||||
admitted, _ = await self._call_store_with_retry(
|
||||
"create_thread_operation_atomic",
|
||||
run_id,
|
||||
lambda: self._store.create_thread_operation_atomic(**create_kwargs),
|
||||
@ -2252,7 +2272,7 @@ class RunManager:
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
await self._atomic_replacement(
|
||||
admitted, _ = await self._atomic_replacement(
|
||||
local_inflight,
|
||||
multitask_strategy,
|
||||
lambda: self._call_store_with_retry(
|
||||
@ -2281,6 +2301,8 @@ class RunManager:
|
||||
|
||||
async with self._lock:
|
||||
# 3) Only now safe to register locally — store insert succeeded.
|
||||
if self._store is not None:
|
||||
self._bind_admitted_producer(record, admitted)
|
||||
self._runs[run_id] = record
|
||||
self._index_run_locked(record)
|
||||
self._arm_lease_deadline(record)
|
||||
@ -2346,7 +2368,7 @@ class RunManager:
|
||||
*,
|
||||
kind: ThreadOperationKind,
|
||||
user_id: str | None = None,
|
||||
) -> AsyncIterator[None]:
|
||||
) -> AsyncIterator[RunRecord]:
|
||||
"""Hold exclusive durable admission for a non-run thread operation.
|
||||
|
||||
The reservation is a short-lived pending row, so the same durable
|
||||
@ -2372,7 +2394,7 @@ class RunManager:
|
||||
lease_lost = record.abort_event.is_set()
|
||||
if lease_lost:
|
||||
raise asyncio.CancelledError()
|
||||
yield
|
||||
yield record
|
||||
except asyncio.CancelledError:
|
||||
if record.abort_event.is_set():
|
||||
raise ConflictError(f"Thread {thread_id} reservation lease was lost") from None
|
||||
@ -2499,6 +2521,8 @@ class RunManager:
|
||||
run_id=record.run_id,
|
||||
user_id=record.user_id,
|
||||
operation_kind=record.operation_kind.value,
|
||||
run_incarnation=record.run_incarnation,
|
||||
retention_revision=record.evidence_retention_revision,
|
||||
)
|
||||
if target is not None:
|
||||
await publisher.publish(target)
|
||||
|
||||
@ -69,6 +69,7 @@ from deerflow.runtime.context_keys import (
|
||||
checkpoint_agent_binding_metadata,
|
||||
)
|
||||
from deerflow.runtime.events.message_identity import attach_message_seq, message_identity
|
||||
from deerflow.runtime.events.producer import current_producer, publisher_scope, run_producer_context
|
||||
from deerflow.runtime.goal import (
|
||||
DEFAULT_MAX_GOAL_CONTINUATIONS,
|
||||
DEFAULT_MAX_NO_PROGRESS_CONTINUATIONS,
|
||||
@ -320,13 +321,14 @@ async def _persist_delivery_receipt(
|
||||
attempts = len(_DELIVERY_RECEIPT_RETRY_DELAYS_SECONDS) + 1
|
||||
for attempt in range(attempts):
|
||||
try:
|
||||
await event_store.put_if_absent(
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
event_type="run.delivery",
|
||||
category="outputs",
|
||||
content=content,
|
||||
)
|
||||
with publisher_scope("run.delivery"):
|
||||
await event_store.put_if_absent(
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
event_type="run.delivery",
|
||||
category="outputs",
|
||||
content=content,
|
||||
)
|
||||
return True
|
||||
except Exception:
|
||||
if attempt == attempts - 1:
|
||||
@ -990,6 +992,7 @@ async def _drain_run_finalization(task: asyncio.Task[None], *, on_host_cancel: C
|
||||
raise deferred
|
||||
|
||||
|
||||
@run_producer_context
|
||||
async def run_agent(
|
||||
bridge: StreamBridge,
|
||||
run_manager: RunManager,
|
||||
@ -2229,6 +2232,12 @@ async def run_agent(
|
||||
raise RuntimeError(evidence_failure)
|
||||
|
||||
evidence_receipt = await journal.finish_evidence(before_receipt=confirm_terminal)
|
||||
if evidence_receipt is not None and (subagent_events is None or not subagent_events._pending):
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
if isinstance(event_store, DbRunEventStore) and current_producer() is not None:
|
||||
if not await event_store.close_producer(current_producer()):
|
||||
raise RuntimeError("Run evidence producer could not be closed")
|
||||
if evidence_receipt is not None:
|
||||
if not evidence_execution_started:
|
||||
evidence_failure = "Lead graph execution did not begin; producer coverage is partial."
|
||||
|
||||
46
backend/scripts/run_producer_fencing.py
Normal file
46
backend/scripts/run_producer_fencing.py
Normal file
@ -0,0 +1,46 @@
|
||||
"""Offline, monotonic activation of PostgreSQL Run producer enforcement."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
|
||||
async def activate(*, config_path: str | None, cutover_id: str) -> None:
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
|
||||
config = AppConfig.from_file(config_path)
|
||||
if config.database.backend != "postgres":
|
||||
raise RuntimeError("Producer cutover requires PostgreSQL")
|
||||
try:
|
||||
# Bootstrap the migrated schema, but never echo credential-bearing SQL.
|
||||
await init_engine_from_config(config.database.model_copy(update={"echo_sql": False}))
|
||||
# Bootstrap may have opened several pooled connections. The offline
|
||||
# session check must count external clients, not our own idle pool.
|
||||
await get_engine().dispose()
|
||||
await activate_producer_fencing(get_session_factory(), cutover_id=cutover_id)
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--config", help="Operator config.yaml (normal config resolution when omitted)")
|
||||
parser.add_argument("--cutover-id", required=True, help="Deployment ID; must match every Gateway's startup config")
|
||||
parser.add_argument("--all-gateways-stopped", required=True, action="store_true", help="Acknowledge that all Gateways and producers have stopped")
|
||||
args = parser.parse_args(argv)
|
||||
try:
|
||||
asyncio.run(activate(config_path=args.config, cutover_id=args.cutover_id))
|
||||
except Exception as exc:
|
||||
# Driver/config exceptions can contain credentials and config contents.
|
||||
print(f"Producer cutover refused ({type(exc).__name__}); verify migration, stopped clients, active Runs and cutover ID.", file=sys.stderr)
|
||||
return 1
|
||||
print("Producer fencing activated. Restart every Gateway with strict mode and the matching cutover ID.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@ -210,14 +210,13 @@ dangling links and removal before retry. Keep filesystem work off the event loop
|
||||
|
||||
## PostgreSQL batch fixtures
|
||||
|
||||
Batch fixtures use `support.postgres.asyncpg_test_url` to map libpq `sslmode`
|
||||
to asyncpg `ssl` and remove unsupported `channel_binding` before constructing
|
||||
database config. Preserve TLS modes, credentials and other query options;
|
||||
reject conflicting `ssl`/`sslmode` values. CI uses `?sslmode=disable`; validate
|
||||
Batch fixtures use `support.postgres.asyncpg_test_url`, delegating to production
|
||||
`normalize_asyncpg_url` to map libpq `sslmode` to asyncpg `ssl`. Preserve TLS
|
||||
modes, credentials and other query options; reject conflicting `ssl`/`sslmode`
|
||||
and unsupported `channel_binding`. CI uses `?sslmode=disable`; validate
|
||||
that URI shape against a real test database, not only a parameter-free local
|
||||
URI. Reuse the normalized config for reopening and teardown, and drop only the
|
||||
fixture's UUID schema. This is test-only handling; production connection and
|
||||
TLS policy are unchanged.
|
||||
fixture's UUID schema. Keep the raw URI for production connection and CLI tests.
|
||||
|
||||
The 0033 batch-evidence migration fixture uses the same adapter. Its connection
|
||||
contract probes execute the actual migration test setup through SQLAlchemy's
|
||||
|
||||
97
backend/tests/support/history_seed_process.py
Normal file
97
backend/tests/support/history_seed_process.py
Normal file
@ -0,0 +1,97 @@
|
||||
"""Independent-process actors for history-seed commit/deletion ordering."""
|
||||
|
||||
import asyncio
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
|
||||
from support.producer_process import CONTEXT, ProducerActors, proceed, receive
|
||||
|
||||
|
||||
async def actor(uri, schema, job, identity, pipe):
|
||||
from deerflow.runtime.events.producer import HistorySeedAuthority, RunEventProducerRejected, producer_scope
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
engine = create_async_engine(uri, connect_args={"server_settings": {"search_path": schema, "statement_timeout": "10000", "lock_timeout": "8000"}})
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
try:
|
||||
async with engine.connect() as conn:
|
||||
pid = await conn.scalar(text("SELECT pg_backend_pid()"))
|
||||
pipe.send({"stage": "started", "pid": pid})
|
||||
await proceed(pipe)
|
||||
events = DbRunEventStore(sf, producer_fencing="strict")
|
||||
if job == "delete_locked":
|
||||
original = events._invalidate_evidence
|
||||
|
||||
async def held(*args):
|
||||
await original(*args)
|
||||
pipe.send({"stage": "revoked_uncommitted"})
|
||||
await proceed(pipe)
|
||||
|
||||
events._invalidate_evidence = held
|
||||
if identity.get("delete_run"):
|
||||
await events.delete_by_run("thread", "checkpoint-seed-thread-1", user_id="alice")
|
||||
else:
|
||||
await events.delete_by_thread("thread", user_id="alice")
|
||||
pipe.send({"stage": "done", "deleted": True})
|
||||
elif job == "metadata_locked":
|
||||
from sqlalchemy import select
|
||||
|
||||
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
||||
|
||||
async with sf() as session, session.begin():
|
||||
await session.scalar(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == "thread").with_for_update())
|
||||
pipe.send({"stage": "metadata_held"})
|
||||
await proceed(pipe)
|
||||
pipe.send({"stage": "done", "released": True})
|
||||
elif job == "admission":
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime import ConflictError, RunManager
|
||||
|
||||
manager = RunManager(store=RunRepository(sf), worker_id="other-host")
|
||||
try:
|
||||
await manager.create_or_reject("thread", user_id="alice")
|
||||
except ConflictError:
|
||||
pipe.send({"stage": "done", "admitted": False})
|
||||
else:
|
||||
pipe.send({"stage": "done", "admitted": True})
|
||||
else:
|
||||
authority = HistorySeedAuthority(**identity["authority"])
|
||||
original = events._admit_events
|
||||
|
||||
async def held(session, batch):
|
||||
await original(session, batch)
|
||||
if job == "seed_locked":
|
||||
pipe.send({"stage": "admitted"})
|
||||
await proceed(pipe)
|
||||
|
||||
events._admit_events = held
|
||||
try:
|
||||
with producer_scope(authority):
|
||||
await events.put_batch(identity["events"])
|
||||
except RunEventProducerRejected:
|
||||
pipe.send({"stage": "done", "written": False})
|
||||
else:
|
||||
pipe.send({"stage": "done", "written": True})
|
||||
except BaseException as exc:
|
||||
pipe.send({"stage": "error", "error": f"{type(exc).__name__}: {exc}"})
|
||||
raise
|
||||
finally:
|
||||
await engine.dispose()
|
||||
pipe.close()
|
||||
|
||||
|
||||
def process_entry(*args):
|
||||
asyncio.run(actor(*args))
|
||||
|
||||
|
||||
class HistorySeedActors(ProducerActors):
|
||||
async def start(self, job):
|
||||
parent, child = CONTEXT.Pipe()
|
||||
process = CONTEXT.Process(target=process_entry, args=(self.uri, self.schema, job, self.identity, child))
|
||||
process.start()
|
||||
child.close()
|
||||
self.processes.append((process, parent))
|
||||
message = await receive(parent)
|
||||
assert message["stage"] == "started", message
|
||||
return parent, message["pid"]
|
||||
@ -1,15 +1,7 @@
|
||||
"""Adapt libpq test URIs to SQLAlchemy's asyncpg connection options."""
|
||||
|
||||
from sqlalchemy.engine import make_url
|
||||
from deerflow.persistence.postgres_url import normalize_asyncpg_url
|
||||
|
||||
|
||||
def asyncpg_test_url(uri: str) -> str:
|
||||
url = make_url(uri)
|
||||
query = dict(url.query)
|
||||
if "sslmode" in query:
|
||||
sslmode = query.pop("sslmode")
|
||||
if "ssl" in query and query["ssl"] != sslmode:
|
||||
raise ValueError("Conflicting ssl and sslmode in the test PostgreSQL URI")
|
||||
query["ssl"] = sslmode
|
||||
query.pop("channel_binding", None)
|
||||
return url.set(drivername="postgresql+asyncpg", query=query).render_as_string(hide_password=False)
|
||||
return normalize_asyncpg_url(uri)
|
||||
|
||||
168
backend/tests/support/producer_process.py
Normal file
168
backend/tests/support/producer_process.py
Normal file
@ -0,0 +1,168 @@
|
||||
"""OS-process actors for the production producer-fencing protocol."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import multiprocessing
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
|
||||
CONTEXT = multiprocessing.get_context("spawn")
|
||||
|
||||
|
||||
async def receive(pipe, timeout=12):
|
||||
def bounded():
|
||||
assert pipe.poll(timeout), "Producer actor failed to reach its barrier"
|
||||
return pipe.recv()
|
||||
|
||||
return await asyncio.to_thread(bounded)
|
||||
|
||||
|
||||
async def proceed(pipe):
|
||||
assert await receive(pipe) == "go"
|
||||
|
||||
|
||||
async def actor(uri, schema, job, identity, pipe):
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
from deerflow.runtime.events.store.base import RunEventStoreNotCommittedError
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.events.terminal import RecoveryTerminalPublisher
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
engine = create_async_engine(uri, pool_size=1, max_overflow=0, pool_timeout=1, connect_args={"server_settings": {"search_path": schema, "statement_timeout": "10000", "lock_timeout": "8000"}})
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
try:
|
||||
async with engine.connect() as connection:
|
||||
pid = await connection.scalar(text("SELECT pg_backend_pid()"))
|
||||
pipe.send({"stage": "started", "pid": pid})
|
||||
await proceed(pipe)
|
||||
runs = RunRepository(sf)
|
||||
events = DbRunEventStore(sf, producer_fencing="strict")
|
||||
authority = producer_authority(SimpleNamespace(**identity))
|
||||
if job in ("writer_locked", "writer_late", "writer"):
|
||||
original = events._admit_events
|
||||
|
||||
async def barrier(session, batch):
|
||||
if job == "writer_late":
|
||||
pipe.send({"stage": "before_admission"})
|
||||
await proceed(pipe)
|
||||
await original(session, batch)
|
||||
if job == "writer_locked":
|
||||
pipe.send({"stage": "admitted"})
|
||||
await proceed(pipe)
|
||||
|
||||
events._admit_events = barrier
|
||||
with producer_scope(authority):
|
||||
journal = RunJournal(identity["run_id"], identity["thread_id"], events, user_id=identity["user_id"], flush_threshold=100)
|
||||
journal.record_middleware("race", name="race", hook="process", action="accepted", changes={})
|
||||
try:
|
||||
await journal.flush()
|
||||
pipe.send({"stage": "done", "written": True})
|
||||
except RunEventStoreNotCommittedError:
|
||||
pipe.send({"stage": "done", "written": False})
|
||||
finally:
|
||||
await journal.close(flush=False)
|
||||
elif job in ("takeover", "takeover_locked"):
|
||||
if job == "takeover_locked":
|
||||
original_commit = AsyncSession.commit
|
||||
|
||||
async def held(session):
|
||||
pipe.send({"stage": "revoked_uncommitted"})
|
||||
await proceed(pipe)
|
||||
return await original_commit(session)
|
||||
|
||||
AsyncSession.commit = held
|
||||
claimed = await runs.claim_for_takeover(identity["run_id"], grace_seconds=0, error="crashed worker")
|
||||
if job == "takeover_locked":
|
||||
AsyncSession.commit = original_commit
|
||||
assert claimed
|
||||
publisher = RecoveryTerminalPublisher(events, runs)
|
||||
target = publisher.authorize_takeover(
|
||||
claimed=claimed,
|
||||
thread_id=identity["thread_id"],
|
||||
run_id=identity["run_id"],
|
||||
user_id=identity["user_id"],
|
||||
operation_kind="run",
|
||||
run_incarnation=identity["run_incarnation"],
|
||||
retention_revision=identity["evidence_retention_revision"],
|
||||
)
|
||||
assert (await publisher.publish(target)).confirmed
|
||||
pipe.send({"stage": "done", "taken": True})
|
||||
elif job in ("delete_locked", "delete_waiting"):
|
||||
if job == "delete_waiting":
|
||||
original_thread_fence = events._acquire_thread_mutation_fence
|
||||
|
||||
async def after_clock(session, thread_id):
|
||||
pipe.send({"stage": "clock_held"})
|
||||
await proceed(pipe)
|
||||
await original_thread_fence(session, thread_id)
|
||||
|
||||
DbRunEventStore._acquire_thread_mutation_fence = staticmethod(after_clock)
|
||||
original_invalidate = events._invalidate_evidence
|
||||
|
||||
async def held(*args, **kwargs):
|
||||
await original_invalidate(*args, **kwargs)
|
||||
if job == "delete_locked":
|
||||
pipe.send({"stage": "deleted_uncommitted"})
|
||||
await proceed(pipe)
|
||||
|
||||
events._invalidate_evidence = held
|
||||
await events.delete_by_run(identity["thread_id"], identity["run_id"], user_id=identity["user_id"])
|
||||
pipe.send({"stage": "done", "deleted": True})
|
||||
else:
|
||||
raise AssertionError(job)
|
||||
except BaseException as exc:
|
||||
pipe.send({"stage": "error", "error": f"{type(exc).__name__}: {exc}"})
|
||||
raise
|
||||
finally:
|
||||
await engine.dispose()
|
||||
pipe.close()
|
||||
|
||||
|
||||
def process_entry(*args):
|
||||
asyncio.run(actor(*args))
|
||||
|
||||
|
||||
class ProducerActors:
|
||||
def __init__(self, uri, schema, identity, observer):
|
||||
self.uri, self.schema, self.identity, self.observer = uri, schema, identity, observer
|
||||
self.processes = []
|
||||
|
||||
async def start(self, job):
|
||||
parent, child = CONTEXT.Pipe()
|
||||
process = CONTEXT.Process(target=process_entry, args=(self.uri, self.schema, job, self.identity, child))
|
||||
process.start()
|
||||
child.close()
|
||||
self.processes.append((process, parent))
|
||||
message = await receive(parent)
|
||||
assert message["stage"] == "started", message
|
||||
return parent, message["pid"]
|
||||
|
||||
async def blocked(self, waiter, holder):
|
||||
deadline = time.monotonic() + 5
|
||||
while time.monotonic() < deadline:
|
||||
async with self.observer.connect() as connection:
|
||||
blockers = await connection.scalar(text("SELECT pg_blocking_pids(:pid)"), {"pid": waiter})
|
||||
if holder in blockers:
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError("Expected PostgreSQL lock wait was not observed")
|
||||
|
||||
async def close(self):
|
||||
for process, pipe in self.processes:
|
||||
if process.is_alive():
|
||||
try:
|
||||
for _ in range(3):
|
||||
pipe.send("go")
|
||||
except (BrokenPipeError, EOFError):
|
||||
pass
|
||||
for process, pipe in self.processes:
|
||||
await asyncio.to_thread(process.join, 3)
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
await asyncio.to_thread(process.join, 3)
|
||||
pipe.close()
|
||||
@ -349,6 +349,7 @@ async def test_sql_runtime_shares_run_repository_with_scheduler(monkeypatch):
|
||||
# This test asserts repository identity, not SQL startup/recovery. Its
|
||||
# session factory is intentionally an opaque sentinel, not a usable DB.
|
||||
monkeypatch.setattr("deerflow.extensions.host_capabilities.HostCapabilities.start", noop)
|
||||
monkeypatch.setattr("deerflow.persistence.run.producer.validate_producer_fencing", noop)
|
||||
monkeypatch.setattr(engine_module, "init_engine_from_config", noop)
|
||||
monkeypatch.setattr(engine_module, "get_session_factory", lambda: session_factory)
|
||||
monkeypatch.setattr(engine_module, "close_engine", noop)
|
||||
|
||||
353
backend/tests/test_history_seed_fencing.py
Normal file
353
backend/tests/test_history_seed_fencing.py
Normal file
@ -0,0 +1,353 @@
|
||||
"""History seed grants retain thread lifecycle and reservation fencing."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import replace
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
from sqlalchemy import select, text, update
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunProducerPolicyRow, RunRow
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
||||
from deerflow.persistence.thread_meta.sql import ThreadMetaRepository
|
||||
from deerflow.runtime import RunManager
|
||||
from deerflow.runtime.events import producer
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.journal import build_branch_history_seed_events, build_checkpoint_history_seed_events
|
||||
from deerflow.runtime.runs.schemas import ThreadOperationKind
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def seed_db():
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "history_seed_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), pool_size=1, max_overflow=0, pool_timeout=2, connect_args={"server_settings": {"search_path": schema}})
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await conn.run_sync(
|
||||
lambda c: RunRow.metadata.create_all(
|
||||
c, tables=[RunRow.__table__, RunChangeClockRow.__table__, CompletedRunSnapshotRow.__table__, RunProducerPolicyRow.__table__, RunEventRow.__table__, RunEventThreadSeqRow.__table__, ThreadMetaRow.__table__]
|
||||
)
|
||||
)
|
||||
await conn.execute(text("INSERT INTO run_producer_policy(id,strict) VALUES (1,true)"))
|
||||
runs = RunRepository(sf)
|
||||
threads = ThreadMetaRepository(sf)
|
||||
thread = await threads.create("thread", user_id="alice")
|
||||
yield SimpleNamespace(
|
||||
uri=asyncpg_test_url(uri), schema=schema, engine=engine, sf=sf, runs=runs, threads=threads, thread=thread, manager=RunManager(store=runs, worker_id="host"), events=DbRunEventStore(sf, producer_fencing="strict")
|
||||
)
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def history():
|
||||
return build_checkpoint_history_seed_events(
|
||||
[HumanMessage(id="h1", content="one"), AIMessage(id="a1", content="answer"), HumanMessage(id="h2", content="two"), AIMessage(id="a2", content="answer two")], thread_id="thread", run_id_prefix="checkpoint-seed-thread"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("seed_kind", ["checkpoint", "branch"])
|
||||
@pytest.mark.parametrize("mode", ["off", "strict"])
|
||||
async def test_seed_preserves_messages_before_first_human_turn(seed_db, seed_kind, mode):
|
||||
db = seed_db
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(update(RunProducerPolicyRow).where(RunProducerPolicyRow.id == 1).values(strict=mode == "strict"))
|
||||
events = DbRunEventStore(db.sf, producer_fencing=mode)
|
||||
builder = build_checkpoint_history_seed_events if seed_kind == "checkpoint" else build_branch_history_seed_events
|
||||
batch = builder(
|
||||
[AIMessage(id="a0", content="inherited"), ToolMessage(id="t0", content="result", tool_call_id="call"), HumanMessage(id="h1", content="next"), AIMessage(id="a1", content="answer")],
|
||||
thread_id="thread",
|
||||
run_id_prefix=f"{seed_kind}-seed-thread",
|
||||
**({"parent_thread_id": "parent"} if seed_kind == "branch" else {}),
|
||||
)
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind=seed_kind)
|
||||
with producer.producer_scope(authority):
|
||||
rows = await events.put_batch(batch)
|
||||
assert [row["run_id"] for row in rows] == [f"{seed_kind}-seed-thread-0", f"{seed_kind}-seed-thread-0", f"{seed_kind}-seed-thread-1", f"{seed_kind}-seed-thread-1"]
|
||||
assert [row["event_type"] for row in rows] == ["llm.ai.response", "llm.tool.result", "llm.human.input", "llm.ai.response"]
|
||||
assert [row["seq"] for row in rows] == [1, 2, 3, 4]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_seed_commits_multiple_synthetic_turns(seed_db):
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
with producer.producer_scope(authority):
|
||||
rows = await db.events.put_batch(history())
|
||||
assert [row["run_id"] for row in rows] == ["checkpoint-seed-thread-1", "checkpoint-seed-thread-1", "checkpoint-seed-thread-2", "checkpoint-seed-thread-2"]
|
||||
assert [row["seq"] for row in rows] == [1, 2, 3, 4]
|
||||
assert all(row["user_id"] == "alice" for row in rows)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mutation", ["thread", "owner", "incarnation", "reservation", "worker", "revision", "run_id", "category", "type"])
|
||||
async def test_invalid_seed_grant_rejects_whole_batch(seed_db, mutation):
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
batch = history()
|
||||
if mutation == "thread":
|
||||
batch[-1]["thread_id"] = "other"
|
||||
elif mutation == "owner":
|
||||
batch[-1]["user_id"] = "mallory"
|
||||
elif mutation == "incarnation":
|
||||
authority = replace(authority, thread_incarnation="stale")
|
||||
elif mutation == "reservation":
|
||||
authority = replace(authority, reservation_incarnation="stale")
|
||||
elif mutation == "worker":
|
||||
authority = replace(authority, owner_worker_id="other")
|
||||
elif mutation == "revision":
|
||||
authority = replace(authority, retention_revision=99)
|
||||
elif mutation == "run_id":
|
||||
batch[-1]["run_id"] = "checkpoint-seed-thread-arbitrary"
|
||||
elif mutation == "category":
|
||||
batch[-1]["category"] = "trace"
|
||||
else:
|
||||
batch[-1]["event_type"] = "run.terminal"
|
||||
with producer.producer_scope(authority), pytest.raises((producer.RunEventProducerRejected, ValueError)):
|
||||
await db.events.put_batch(batch)
|
||||
async with db.sf() as session:
|
||||
assert await session.scalar(select(RunEventRow.id).limit(1)) is None
|
||||
assert await session.scalar(select(RunEventThreadSeqRow.seq).limit(1)) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_and_prefix_do_not_authorize_seeding(seed_db):
|
||||
with pytest.raises(producer.RunEventProducerRejected):
|
||||
await seed_db.events.put_batch(history())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("change", ["release", "recreate", "delete_events", "expire", "takeover"])
|
||||
async def test_seed_grant_cannot_outlive_its_lifecycle(seed_db, change):
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
if change == "release":
|
||||
await db.runs.delete_thread_operation(reservation.run_id, user_id="alice")
|
||||
elif change == "recreate":
|
||||
await db.threads.delete("thread", user_id="alice")
|
||||
await db.threads.create("thread", user_id="alice")
|
||||
elif change == "delete_events":
|
||||
await db.events.delete_by_thread("thread", user_id="alice")
|
||||
assert await db.threads.get("thread", user_id="alice") is not None
|
||||
else:
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(update(RunRow).where(RunRow.run_id == reservation.run_id).values(lease_expires_at=datetime.now(UTC) - timedelta(seconds=60)))
|
||||
if change == "takeover":
|
||||
assert await db.runs.claim_for_takeover(reservation.run_id, grace_seconds=0, error="expired")
|
||||
with producer.producer_scope(authority), pytest.raises(producer.RunEventProducerRejected):
|
||||
await db.events.put_batch(history())
|
||||
assert await db.events.list_messages("thread", user_id=None) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("turn", ["0", "1"])
|
||||
async def test_trigger_independently_checks_seed_grant(seed_db, turn):
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(text("SELECT set_config('deerflow.event_authority', :value, true)"), {"value": json.dumps(authority.to_storage())})
|
||||
session.add(RunEventRow(thread_id="thread", run_id=f"checkpoint-seed-thread-{turn}", user_id="alice", event_type="llm.human.input", category="message", seq=1))
|
||||
await db.events.delete_by_thread("thread", user_id="alice")
|
||||
with pytest.raises(DBAPIError):
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(text("SELECT set_config('deerflow.event_authority', :value, true)"), {"value": json.dumps(authority.to_storage())})
|
||||
session.add(RunEventRow(thread_id="thread", run_id="checkpoint-seed-thread-1", user_id="alice", event_type="llm.human.input", category="message", seq=2))
|
||||
assert await db.events.list_messages("thread", user_id=None) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("seed_first", [False, True])
|
||||
@pytest.mark.parametrize("delete_run", [False, True])
|
||||
async def test_cross_process_seed_and_delete_serialize_at_event_commit(seed_db, seed_first, delete_run):
|
||||
from support.history_seed_process import HistorySeedActors
|
||||
from support.producer_process import receive
|
||||
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
actors = HistorySeedActors(db.uri, db.schema, {"authority": authority.to_storage(), "events": history(), "delete_run": delete_run}, db.engine)
|
||||
try:
|
||||
first, first_pid = await actors.start("seed_locked" if seed_first else "delete_locked")
|
||||
first.send("go")
|
||||
assert (await receive(first))["stage"] == ("admitted" if seed_first else "revoked_uncommitted")
|
||||
second, second_pid = await actors.start("delete_locked" if seed_first else "seed")
|
||||
second.send("go")
|
||||
await actors.blocked(second_pid, first_pid)
|
||||
first.send("go")
|
||||
result = await receive(first)
|
||||
assert result.get("written" if seed_first else "deleted") is True, result
|
||||
if seed_first:
|
||||
assert (await receive(second))["stage"] == "revoked_uncommitted"
|
||||
second.send("go")
|
||||
assert (await receive(second))["deleted"] is True
|
||||
else:
|
||||
assert (await receive(second))["written"] is False
|
||||
# Metadata deliberately remains during the cleanup gap; existence
|
||||
# alone must never authorize resurrection after event deletion.
|
||||
assert await db.threads.get("thread", user_id="alice") is not None
|
||||
remaining = await db.events.list_messages("thread", user_id=None)
|
||||
assert [row["content"]["id"] for row in remaining] == (["h2", "a2"] if seed_first and delete_run else [])
|
||||
finally:
|
||||
await actors.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cross_process_first_run_cannot_overtake_seed_reservation(seed_db):
|
||||
from support.history_seed_process import HistorySeedActors
|
||||
from support.producer_process import receive
|
||||
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
actors = HistorySeedActors(db.uri, db.schema, {"authority": authority.to_storage(), "events": history()}, db.engine)
|
||||
try:
|
||||
admission, _ = await actors.start("admission")
|
||||
admission.send("go")
|
||||
assert (await receive(admission))["admitted"] is False
|
||||
with producer.producer_scope(authority):
|
||||
await db.events.put_batch(history())
|
||||
finally:
|
||||
await actors.close()
|
||||
record = await db.manager.create_or_reject("thread", user_id="alice")
|
||||
with producer.producer_scope(producer.producer_authority(record)):
|
||||
result = await db.events.put(thread_id="thread", run_id=record.run_id, event_type="llm.human.input", category="message", content={"type": "human", "id": "new"})
|
||||
assert result["seq"] == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mutation", ["owner", "incarnation", "epoch", "kind", "prefix", "type", "category", "missing_incarnation"])
|
||||
async def test_strict_trigger_rejects_forged_or_overbroad_seed_grants(seed_db, mutation):
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
grant = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint").to_storage()
|
||||
event = {"thread_id": "thread", "run_id": "checkpoint-seed-thread-1", "user_id": "alice", "event_type": "llm.human.input", "category": "message"}
|
||||
if mutation == "owner":
|
||||
grant["user_id"] = "mallory"
|
||||
elif mutation == "incarnation":
|
||||
grant["thread_incarnation"] = "stale"
|
||||
elif mutation == "epoch":
|
||||
grant["producer_epoch"] = "stale"
|
||||
elif mutation == "kind":
|
||||
grant["seed_kind"] = "arbitrary"
|
||||
elif mutation == "prefix":
|
||||
event["run_id"] = "checkpoint-seed-thread-arbitrary"
|
||||
elif mutation == "type":
|
||||
event["event_type"] = "run.terminal"
|
||||
elif mutation == "category":
|
||||
event["category"] = "trace"
|
||||
else:
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == "thread").values(incarnation=None))
|
||||
del grant["thread_incarnation"]
|
||||
with pytest.raises(DBAPIError):
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(text("SELECT set_config('deerflow.event_authority', :value, true)"), {"value": json.dumps(grant)})
|
||||
session.add(RunEventRow(**event, seq=1))
|
||||
assert await db.events.list_messages("thread", user_id=None) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("trigger_active", [False, True])
|
||||
async def test_seed_lease_is_rechecked_after_metadata_lock_wait(seed_db, trigger_active):
|
||||
import asyncio
|
||||
|
||||
from support.history_seed_process import HistorySeedActors
|
||||
from support.producer_process import receive
|
||||
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(update(RunProducerPolicyRow).values(strict=trigger_active))
|
||||
await session.execute(update(RunRow).where(RunRow.run_id == reservation.run_id).values(lease_expires_at=datetime.now(UTC) + timedelta(seconds=3)))
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
actors = HistorySeedActors(db.uri, db.schema, {"authority": authority.to_storage(), "events": history()}, db.engine)
|
||||
try:
|
||||
holder, holder_pid = await actors.start("metadata_locked")
|
||||
holder.send("go")
|
||||
assert (await receive(holder))["stage"] == "metadata_held"
|
||||
seed, seed_pid = await actors.start("seed")
|
||||
seed.send("go")
|
||||
await actors.blocked(seed_pid, holder_pid)
|
||||
async with asyncio.timeout(8):
|
||||
while True:
|
||||
async with db.sf() as session:
|
||||
expired = await session.scalar(select(RunRow.lease_expires_at <= text("clock_timestamp()")).where(RunRow.run_id == reservation.run_id))
|
||||
if expired:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
holder.send("go")
|
||||
assert (await receive(holder))["released"] is True
|
||||
assert (await receive(seed))["written"] is False
|
||||
assert await db.events.list_messages("thread", user_id=None) == []
|
||||
finally:
|
||||
await actors.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_synthetic_turn_revokes_outstanding_seed_grant(seed_db):
|
||||
db = seed_db
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = producer.history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
with producer.producer_scope(authority):
|
||||
await db.events.put_batch(history())
|
||||
assert await db.events.delete_by_run("thread", "checkpoint-seed-thread-1", user_id="alice") == 2
|
||||
with producer.producer_scope(authority), pytest.raises(producer.RunEventProducerRejected):
|
||||
await db.events.put_batch(history()[:2])
|
||||
assert [row["content"]["id"] for row in await db.events.list_messages("thread", user_id=None)] == ["h2", "a2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thread_create_returns_its_committed_incarnation_after_recreation(seed_db, monkeypatch):
|
||||
import asyncio
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
db = seed_db
|
||||
committed = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
original_commit = AsyncSession.commit
|
||||
creator = None
|
||||
|
||||
async def pause_after_commit(session):
|
||||
await original_commit(session)
|
||||
if asyncio.current_task() is creator:
|
||||
committed.set()
|
||||
await release.wait()
|
||||
|
||||
monkeypatch.setattr(AsyncSession, "commit", pause_after_commit)
|
||||
creator = asyncio.create_task(db.threads.create("branch", user_id="alice"))
|
||||
try:
|
||||
async with asyncio.timeout(5):
|
||||
await committed.wait()
|
||||
original = await db.threads.get("branch", user_id="alice")
|
||||
await db.threads.delete("branch", user_id="alice")
|
||||
replacement = await db.threads.create("branch", user_id="alice")
|
||||
assert replacement["incarnation"] != original["incarnation"]
|
||||
release.set()
|
||||
created = await asyncio.wait_for(asyncio.shield(creator), timeout=5)
|
||||
assert created["incarnation"] == original["incarnation"]
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(asyncio.gather(creator, return_exceptions=True), timeout=5)
|
||||
383
backend/tests/test_history_seed_gateway.py
Normal file
383
backend/tests/test_history_seed_gateway.py
Normal file
@ -0,0 +1,383 @@
|
||||
"""Gateway history backfills must remain usable with strict producer fencing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from _agent_e2e_helpers import FakeToolCallingModel
|
||||
from _router_auth_helpers import make_authed_test_app
|
||||
from fastapi import Request
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
from sqlalchemy import select, text, update
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
from test_threads_router import _patch_checkpoint_state_builder as _patch_checkpoint_state_builder
|
||||
from test_threads_router import _RawStateAccessor, _write_checkpoint
|
||||
|
||||
from app.gateway import services
|
||||
from app.gateway.auth.models import User
|
||||
from app.gateway.routers import threads
|
||||
from app.gateway.routers.thread_runs import RunCreateRequest
|
||||
from deerflow.config import paths as paths_module
|
||||
from deerflow.config.app_config import AppConfig, pop_current_app_config, push_current_app_config
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.projects import ProjectRepository
|
||||
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunProducerPolicyRow, RunRow
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.persistence.thread_meta import ThreadMetaRepository
|
||||
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
||||
from deerflow.runtime import ConflictError, RunManager, RunStatus
|
||||
from deerflow.runtime.events.producer import RunEventProducerRejected
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.stream_bridge.memory import MemoryStreamBridge
|
||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def history_gateway(monkeypatch, tmp_path):
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "history_seed_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema}})
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
tables = [RunRow.__table__, RunChangeClockRow.__table__, RunEventRow.__table__, RunEventThreadSeqRow.__table__, CompletedRunSnapshotRow.__table__, RunProducerPolicyRow.__table__, ThreadMetaRow.__table__]
|
||||
manager = None
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await conn.run_sync(lambda c: RunRow.metadata.create_all(c, tables=tables))
|
||||
await conn.execute(text("INSERT INTO run_producer_policy(id,strict,cutover_id) VALUES (1,true,'history-seed-test')"))
|
||||
events = DbRunEventStore(sf, producer_fencing="strict", producer_cutover_id="history-seed-test")
|
||||
await events.validate_producer_fencing()
|
||||
runs = RunRepository(sf)
|
||||
manager = RunManager(store=runs, event_store=events, worker_id="history-seed-worker")
|
||||
user = User(id=uuid4(), email="history-seed@example.com", password_hash="unused", system_role="user")
|
||||
app = make_authed_test_app(user_factory=lambda: user, bind_current_user=True)
|
||||
app.state.thread_store = ThreadMetaRepository(sf)
|
||||
app.state.run_manager = manager
|
||||
app.state.run_event_store = events
|
||||
app.state.checkpointer = InMemorySaver()
|
||||
app.state.store = InMemoryStore()
|
||||
app.state.project_repo = ProjectRepository(sf)
|
||||
app.state.stream_bridge = MemoryStreamBridge()
|
||||
app.include_router(threads.router)
|
||||
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path / "home"))
|
||||
monkeypatch.setattr(paths_module, "_paths", Paths(tmp_path / "home"))
|
||||
extensions_path = tmp_path / "extensions_config.json"
|
||||
extensions_path.write_text('{"mcpServers": {}, "skills": {}}', encoding="utf-8")
|
||||
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(extensions_path))
|
||||
config = AppConfig.model_validate(
|
||||
{
|
||||
"models": [{"name": "history-test-model", "use": "langchain_openai:ChatOpenAI", "model": "unused", "api_key": "unused"}],
|
||||
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
||||
"skills": {"path": str(tmp_path / "skills")},
|
||||
"title": {"enabled": False},
|
||||
"memory": {"enabled": False},
|
||||
"run_events": {"backend": "db", "producer_fencing": "strict", "producer_cutover_id": "history-seed-test"},
|
||||
}
|
||||
)
|
||||
app.state.run_events_config = config.run_events
|
||||
|
||||
def accessor(request, *, thread_id, assistant_id=None, checkpoint_id=None):
|
||||
config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
if checkpoint_id is not None:
|
||||
config["configurable"]["checkpoint_id"] = checkpoint_id
|
||||
return _RawStateAccessor(request.app.state.checkpointer), config
|
||||
|
||||
monkeypatch.setattr(services, "build_checkpoint_state_accessor", accessor)
|
||||
token = set_current_user(user)
|
||||
push_current_app_config(config)
|
||||
try:
|
||||
yield SimpleNamespace(app=app, sf=sf, runs=runs, events=events, manager=manager, user=user, user_id=str(user.id))
|
||||
finally:
|
||||
try:
|
||||
await manager.shutdown()
|
||||
finally:
|
||||
manager = None
|
||||
await app.state.stream_bridge.close()
|
||||
pop_current_app_config()
|
||||
reset_current_user(token)
|
||||
finally:
|
||||
if manager is not None:
|
||||
await manager.shutdown()
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def _request(db):
|
||||
request = Request({"type": "http", "app": db.app, "method": "POST", "path": "/", "headers": []})
|
||||
request.state.user = db.user
|
||||
return request
|
||||
|
||||
|
||||
async def _wait_for_gate(event, task):
|
||||
waiter = asyncio.create_task(event.wait())
|
||||
try:
|
||||
done, _ = await asyncio.wait((waiter, task), timeout=5, return_when=asyncio.FIRST_COMPLETED)
|
||||
if task in done:
|
||||
await task
|
||||
assert waiter in done, "history seed did not reach the observation gate"
|
||||
finally:
|
||||
waiter.cancel()
|
||||
await asyncio.gather(waiter, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("metadata_state", ["current", "absent", "null_incarnation"])
|
||||
@pytest.mark.parametrize("turns", [1, 2])
|
||||
@pytest.mark.parametrize("ownerless", [False, True], ids=["owned", "ownerless"])
|
||||
async def test_strict_checkpoint_only_history_is_seeded(history_gateway, metadata_state, turns, ownerless):
|
||||
db = history_gateway
|
||||
thread_id = "checkpoint-only"
|
||||
owner = None if ownerless else db.user_id
|
||||
if metadata_state != "absent":
|
||||
await db.app.state.thread_store.create(thread_id, user_id=owner)
|
||||
if metadata_state == "null_incarnation":
|
||||
async with db.sf() as session:
|
||||
await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(incarnation=None))
|
||||
await session.commit()
|
||||
messages = [HumanMessage(id="h1", content="old question"), AIMessage(id="a1", content="old answer")]
|
||||
if turns == 2:
|
||||
messages.extend([HumanMessage(id="h2", content="another question"), AIMessage(id="a2", content="another answer")])
|
||||
await _write_checkpoint(
|
||||
db.app.state.checkpointer,
|
||||
thread_id,
|
||||
"legacy-head",
|
||||
messages,
|
||||
step=1,
|
||||
)
|
||||
request = _request(db)
|
||||
|
||||
token = set_current_user(None) if ownerless else None
|
||||
try:
|
||||
await services.ensure_checkpoint_history_seeded(request, thread_id=thread_id, assistant_id=None)
|
||||
rows = await db.events.list_messages(thread_id, user_id=None)
|
||||
await services.ensure_checkpoint_history_seeded(request, thread_id=thread_id, assistant_id=None)
|
||||
assert await db.events.list_messages(thread_id, user_id=None) == rows
|
||||
finally:
|
||||
if token is not None:
|
||||
reset_current_user(token)
|
||||
|
||||
assert [row["content"]["id"] for row in rows] == ["h1", "a1", "h2", "a2"][: 2 * turns]
|
||||
assert all(row["metadata"]["checkpoint_history_seed"] is True for row in rows)
|
||||
assert [row["run_id"] for row in rows] == [f"checkpoint-seed-{thread_id}-{turn}" for turn in range(1, turns + 1) for _ in range(2)]
|
||||
assert [row["seq"] for row in rows] == list(range(1, 2 * turns + 1))
|
||||
assert all(row["user_id"] == owner for row in rows)
|
||||
metadata = await db.app.state.thread_store.get(thread_id, user_id=None)
|
||||
assert metadata is not None
|
||||
assert metadata["user_id"] == owner
|
||||
if metadata_state == "null_incarnation":
|
||||
assert metadata["incarnation"] is None
|
||||
async with db.sf() as session:
|
||||
assert await session.scalar(select(RunRow.run_id).where(RunRow.thread_id == thread_id)) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("turns", [1, 2])
|
||||
async def test_strict_branch_history_is_seeded(history_gateway, turns):
|
||||
db = history_gateway
|
||||
source_id = "branch-source"
|
||||
await db.app.state.thread_store.create(source_id)
|
||||
messages = []
|
||||
parent_config = None
|
||||
for turn in range(1, turns + 2):
|
||||
messages.extend([HumanMessage(id=f"h{turn}", content=f"question {turn}"), AIMessage(id=f"a{turn}", content=f"answer {turn}")])
|
||||
parent_config = await _write_checkpoint(db.app.state.checkpointer, source_id, f"turn-{turn}", list(messages), step=turn, parent_config=parent_config)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=db.app), base_url="http://test") as client:
|
||||
response = await client.post(f"/api/threads/{source_id}/branches", json={"message_id": f"a{turns}"})
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
branch = response.json()
|
||||
assert branch["history_seed_mode"] == "seeded", branch
|
||||
rows = await db.events.list_messages(branch["thread_id"], user_id=db.user_id)
|
||||
assert [row["content"]["id"] for row in rows] == [f"{role}{turn}" for turn in range(1, turns + 1) for role in ("h", "a")]
|
||||
assert [row["run_id"] for row in rows] == [f"branch-seed-{branch['thread_id']}-{turn}" for turn in range(1, turns + 1) for _ in range(2)]
|
||||
assert all(row["metadata"]["branch_parent_thread_id"] == source_id for row in rows)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("turns", [1, 2])
|
||||
async def test_strict_first_branch_run_is_admitted_with_inherited_history(history_gateway, monkeypatch, turns):
|
||||
from deerflow.agents.lead_agent import agent as lead_agent
|
||||
|
||||
db = history_gateway
|
||||
# The raw router helper uses integer channel versions. Let the real
|
||||
# mutation graph write the branch consumed by the real agent graph.
|
||||
monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", services.build_checkpoint_state_mutation_accessor)
|
||||
source_id = "first-run-source"
|
||||
await db.app.state.thread_store.create(source_id)
|
||||
messages = [HumanMessage(id="h1", content="old question"), AIMessage(id="a1", content="old answer")]
|
||||
if turns == 2:
|
||||
messages.extend([HumanMessage(id="h2", content="second question"), AIMessage(id="a2", content="second answer")])
|
||||
parent_config = await _write_checkpoint(db.app.state.checkpointer, source_id, "inherited-turn", messages, step=turns)
|
||||
await _write_checkpoint(
|
||||
db.app.state.checkpointer,
|
||||
source_id,
|
||||
"later-turn",
|
||||
[*messages, HumanMessage(id="later-human", content="later question"), AIMessage(id="later-ai", content="later answer")],
|
||||
step=turns + 1,
|
||||
parent_config=parent_config,
|
||||
)
|
||||
async with AsyncClient(transport=ASGITransport(app=db.app), base_url="http://test") as client:
|
||||
response = await client.post(f"/api/threads/{source_id}/branches", json={"message_id": f"a{turns}"})
|
||||
assert response.status_code == 200, response.text
|
||||
branch = response.json()
|
||||
model = FakeToolCallingModel(responses=[AIMessage(id="new-ai", content="new answer")])
|
||||
monkeypatch.setattr(lead_agent, "create_chat_model", lambda **kwargs: model)
|
||||
|
||||
record = await services.start_run(
|
||||
RunCreateRequest(assistant_id="lead_agent", input={"messages": [{"role": "user", "id": "new-human", "content": "new question"}]}, context={"thinking_enabled": False}),
|
||||
branch["thread_id"],
|
||||
_request(db),
|
||||
)
|
||||
assert record.task is not None
|
||||
await asyncio.wait_for(asyncio.shield(record.task), timeout=15)
|
||||
|
||||
durable = await db.runs.get(record.run_id, user_id=db.user_id)
|
||||
assert durable["status"] == RunStatus.success.value, durable
|
||||
assert branch["history_seed_mode"] == "seeded", branch
|
||||
rows = await db.events.list_messages(branch["thread_id"], user_id=db.user_id)
|
||||
inherited = ["h1", "a1", "h2", "a2"][: 2 * turns]
|
||||
assert [row["content"]["id"] for row in rows] == [*inherited, "new-human__user", "new-ai"]
|
||||
assert [row["run_id"] for row in rows[2 * turns :]] == [record.run_id, record.run_id]
|
||||
assert all(row["metadata"]["branch_seed"] is True for row in rows[: 2 * turns])
|
||||
assert all(row["user_id"] == db.user_id for row in rows)
|
||||
async with db.sf() as session:
|
||||
assert list(await session.scalars(select(RunRow.run_id).where(RunRow.thread_id == branch["thread_id"]))) == [record.run_id]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_seed_rechecks_stale_empty_feed_across_managers(history_gateway, monkeypatch):
|
||||
db = history_gateway
|
||||
thread_id = "concurrent-checkpoint-only"
|
||||
await _write_checkpoint(
|
||||
db.app.state.checkpointer,
|
||||
thread_id,
|
||||
"legacy-head",
|
||||
[HumanMessage(id="h1", content="first"), AIMessage(id="a1", content="answer one"), HumanMessage(id="h2", content="second"), AIMessage(id="a2", content="answer two")],
|
||||
step=1,
|
||||
)
|
||||
peer_events = DbRunEventStore(db.sf, producer_fencing="strict", producer_cutover_id="history-seed-test")
|
||||
peer_manager = RunManager(store=RunRepository(db.sf), event_store=peer_events, worker_id="history-seed-peer")
|
||||
peer = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(**db.app.state._state)), user=db.user)
|
||||
peer.app.state.run_manager = peer_manager
|
||||
peer.app.state.run_event_store = peer_events
|
||||
empty_observed = asyncio.Event()
|
||||
release_probe = asyncio.Event()
|
||||
list_messages = peer_events.list_messages
|
||||
|
||||
async def pause_first_empty_probe(*args, **kwargs):
|
||||
rows = await list_messages(*args, **kwargs)
|
||||
if not empty_observed.is_set():
|
||||
assert rows == []
|
||||
empty_observed.set()
|
||||
await release_probe.wait()
|
||||
return rows
|
||||
|
||||
monkeypatch.setattr(peer_events, "list_messages", pause_first_empty_probe)
|
||||
pending = asyncio.create_task(services.ensure_checkpoint_history_seeded(_request(peer), thread_id=thread_id, assistant_id=None))
|
||||
try:
|
||||
await _wait_for_gate(empty_observed, pending)
|
||||
assert await db.app.state.thread_store.get(thread_id, user_id=None) is None
|
||||
await services.ensure_checkpoint_history_seeded(_request(db), thread_id=thread_id, assistant_id=None)
|
||||
metadata = await db.app.state.thread_store.get(thread_id, user_id=None)
|
||||
assert metadata is not None
|
||||
rows = await db.events.list_messages(thread_id, user_id=None)
|
||||
assert [row["content"]["id"] for row in rows] == ["h1", "a1", "h2", "a2"]
|
||||
|
||||
release_probe.set()
|
||||
await asyncio.wait_for(asyncio.shield(pending), timeout=5)
|
||||
|
||||
assert await peer_events.list_messages(thread_id, user_id=None) == rows
|
||||
assert await db.app.state.thread_store.get(thread_id, user_id=None) == metadata
|
||||
async with db.sf() as session:
|
||||
assert await session.scalar(select(RunRow.run_id).where(RunRow.thread_id == thread_id)) is None
|
||||
finally:
|
||||
release_probe.set()
|
||||
await asyncio.wait_for(asyncio.gather(pending, return_exceptions=True), timeout=5)
|
||||
await peer_manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_seed_rejects_metadata_replaced_after_snapshot_read(history_gateway, monkeypatch):
|
||||
db = history_gateway
|
||||
thread_id = "recreated-checkpoint-thread"
|
||||
original = await db.app.state.thread_store.create(thread_id)
|
||||
await _write_checkpoint(db.app.state.checkpointer, thread_id, "legacy-head", [HumanMessage(id="h1", content="old question"), AIMessage(id="a1", content="old answer")], step=1)
|
||||
snapshot_read = asyncio.Event()
|
||||
release_snapshot = asyncio.Event()
|
||||
build_accessor = services.abuild_checkpoint_state_accessor
|
||||
|
||||
async def pause_after_snapshot(request, **kwargs):
|
||||
accessor, config = await build_accessor(request, **kwargs)
|
||||
if request.app is not db.app:
|
||||
return accessor, config
|
||||
read = accessor.aget
|
||||
|
||||
async def gated_read(config):
|
||||
snapshot = await read(config)
|
||||
async with db.sf() as session:
|
||||
reservation = await session.scalar(select(RunRow.run_id).where(RunRow.thread_id == thread_id, RunRow.operation_kind == "checkpoint_write", RunRow.status == "pending"))
|
||||
if reservation is None:
|
||||
return snapshot
|
||||
snapshot_read.set()
|
||||
await release_snapshot.wait()
|
||||
return snapshot
|
||||
|
||||
monkeypatch.setattr(accessor, "aget", gated_read)
|
||||
return accessor, config
|
||||
|
||||
monkeypatch.setattr(services, "abuild_checkpoint_state_accessor", pause_after_snapshot)
|
||||
pending = asyncio.create_task(services.ensure_checkpoint_history_seeded(_request(db), thread_id=thread_id, assistant_id=None))
|
||||
peer_manager = RunManager(store=RunRepository(db.sf), event_store=DbRunEventStore(db.sf, producer_fencing="strict"), worker_id="history-seed-peer")
|
||||
try:
|
||||
await _wait_for_gate(snapshot_read, pending)
|
||||
# A second Gateway must respect the reservation around the real read.
|
||||
peer = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(**db.app.state._state)), user=db.user)
|
||||
peer.app.state.run_manager = peer_manager
|
||||
with pytest.raises(ConflictError):
|
||||
await services.ensure_checkpoint_history_seeded(_request(peer), thread_id=thread_id, assistant_id=None)
|
||||
await db.app.state.thread_store.delete(thread_id)
|
||||
replacement = await db.app.state.thread_store.create(thread_id)
|
||||
assert replacement["incarnation"] != original["incarnation"]
|
||||
|
||||
release_snapshot.set()
|
||||
with pytest.raises(RunEventProducerRejected):
|
||||
await asyncio.wait_for(asyncio.shield(pending), timeout=5)
|
||||
|
||||
assert await db.events.list_messages(thread_id, user_id=None) == []
|
||||
assert await db.app.state.thread_store.get(thread_id, user_id=None) == replacement
|
||||
async with db.sf() as session:
|
||||
assert await session.scalar(select(RunEventThreadSeqRow.seq).where(RunEventThreadSeqRow.thread_id == thread_id)) is None
|
||||
assert await session.scalar(select(RunRow.run_id).where(RunRow.thread_id == thread_id)) is None
|
||||
finally:
|
||||
release_snapshot.set()
|
||||
await asyncio.wait_for(asyncio.gather(pending, return_exceptions=True), timeout=5)
|
||||
await peer_manager.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("empty_checkpoint", [False, True])
|
||||
async def test_no_history_does_not_reserve_an_active_thread(history_gateway, empty_checkpoint):
|
||||
db = history_gateway
|
||||
record = await db.manager.create_or_reject("fresh-active-thread", user_id=db.user_id)
|
||||
if empty_checkpoint:
|
||||
await _write_checkpoint(db.app.state.checkpointer, record.thread_id, "empty-head", [], step=-1)
|
||||
await services.ensure_checkpoint_history_seeded(_request(db), thread_id=record.thread_id, assistant_id=None)
|
||||
assert await db.events.list_messages(record.thread_id, user_id=None) == []
|
||||
async with db.sf() as session:
|
||||
assert list(await session.scalars(select(RunRow.run_id).where(RunRow.thread_id == record.thread_id))) == [record.run_id]
|
||||
@ -78,7 +78,7 @@ class _StopBeforeDatabase(Exception):
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sslmode", [None, "disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
|
||||
async def test_postgres_migration_setup_preserves_tls_policy(monkeypatch, tmp_path, sslmode):
|
||||
uri = "postgresql://user:password@localhost:5432/test?target_session_attrs=read-write&channel_binding=prefer"
|
||||
uri = "postgresql://user:password@localhost:5432/test?target_session_attrs=read-write"
|
||||
if sslmode is not None:
|
||||
uri += f"&sslmode={sslmode}"
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", uri)
|
||||
@ -109,3 +109,16 @@ async def test_postgres_migration_setup_rejects_conflicting_tls_policy(monkeypat
|
||||
monkeypatch.setitem(globals(), "create_async_engine", reject_engine_creation)
|
||||
with pytest.raises(ValueError, match="Conflicting ssl and sslmode"):
|
||||
await test_upgrade_downgrade_and_reupgrade_preserve_report(tmp_path, "postgres")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("channel_binding", ["disable", "prefer", "require"])
|
||||
async def test_postgres_migration_setup_rejects_channel_binding_before_engine_creation(monkeypatch, tmp_path, channel_binding):
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", f"postgresql://user:password@localhost:5432/test?sslmode=verify-full&channel_binding={channel_binding}")
|
||||
|
||||
def reject_engine_creation(*args, **kwargs):
|
||||
pytest.fail("Unsupported Channel Binding must be rejected before engine creation")
|
||||
|
||||
monkeypatch.setitem(globals(), "create_async_engine", reject_engine_creation)
|
||||
with pytest.raises(ValueError, match="channel_binding"):
|
||||
await test_upgrade_downgrade_and_reupgrade_preserve_report(tmp_path, "postgres")
|
||||
|
||||
@ -103,7 +103,7 @@ async def test_0037_upgrade_keeps_rows_and_downgrade_drops_summary(tmp_path, bac
|
||||
|
||||
def test_postgres_uri_tls_policy_is_preserved(tmp_path, monkeypatch):
|
||||
"""The fixture must translate, never drop, an explicit sslmode (tests/AGENTS.md)."""
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", "postgresql://u:p@db:5432/deerflow?sslmode=verify-full&channel_binding=require")
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", "postgresql://u:p@db:5432/deerflow?sslmode=verify-full")
|
||||
captured = {}
|
||||
|
||||
def capture_engine(url, **kwargs):
|
||||
@ -123,3 +123,15 @@ def test_postgres_uri_conflicting_ssl_settings_rejected(tmp_path, monkeypatch):
|
||||
monkeypatch.setitem(globals(), "create_async_engine", lambda *a, **k: None)
|
||||
with pytest.raises(ValueError, match="Conflicting ssl and sslmode"):
|
||||
_engine(tmp_path, "postgres")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channel_binding", ["disable", "prefer", "require"])
|
||||
async def test_postgres_uri_channel_binding_rejected_before_engine_creation(tmp_path, monkeypatch, channel_binding):
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", f"postgresql://u:p@db:5432/deerflow?sslmode=verify-full&channel_binding={channel_binding}")
|
||||
|
||||
def reject_engine_creation(*args, **kwargs):
|
||||
pytest.fail("Unsupported Channel Binding must be rejected before engine creation")
|
||||
|
||||
monkeypatch.setitem(globals(), "create_async_engine", reject_engine_creation)
|
||||
with pytest.raises(ValueError, match="channel_binding"):
|
||||
_engine(tmp_path, "postgres")
|
||||
|
||||
@ -44,9 +44,12 @@ pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def test_0047_remains_in_the_single_migration_chain():
|
||||
assert bootstrap._get_head_revision() == "0049_run_inputs"
|
||||
assert bootstrap._get_head_revision() == "0051_merge_run_inputs_fencing"
|
||||
script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR))
|
||||
assert script.get_heads() == ["0049_run_inputs"]
|
||||
assert script.get_heads() == ["0051_merge_run_inputs_fencing"]
|
||||
assert set(script.get_revision("0051_merge_run_inputs_fencing").down_revision) == {"0049_run_inputs", "0050_history_seed_authority"}
|
||||
assert script.get_revision("0050_history_seed_authority").down_revision == "0049_run_producer_fencing"
|
||||
assert script.get_revision("0049_run_producer_fencing").down_revision == "0048_extension_execution_leases"
|
||||
assert script.get_revision("0049_run_inputs").down_revision == "0048_extension_execution_leases"
|
||||
assert script.get_revision("0048_extension_execution_leases").down_revision == REVISION
|
||||
assert script.get_revision(REVISION).down_revision == PREVIOUS
|
||||
|
||||
85
backend/tests/test_migration_0049_run_producer_fencing.py
Normal file
85
backend/tests/test_migration_0049_run_producer_fencing.py
Normal file
@ -0,0 +1,85 @@
|
||||
"""Historical rows gain no authority; migration and fresh bootstrap agree."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from alembic import command
|
||||
from sqlalchemy import inspect, text
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence import bootstrap
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
|
||||
REVISION = "0049_run_producer_fencing"
|
||||
PREVIOUS = "0048_extension_execution_leases"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
||||
@pytest.mark.parametrize("fresh", [False, True])
|
||||
async def test_migration_and_bootstrap(tmp_path, backend, fresh):
|
||||
schema = None
|
||||
if backend == "postgres":
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "fence_migration_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema, "application_name": schema}})
|
||||
else:
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}")
|
||||
cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
|
||||
try:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
if fresh:
|
||||
await bootstrap.bootstrap_schema(engine, backend=backend, postgres_schema=schema or "")
|
||||
else:
|
||||
await asyncio.to_thread(command.upgrade, cfg, PREVIOUS)
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text("""INSERT INTO runs(
|
||||
run_id,thread_id,user_id,status,multitask_strategy,metadata_json,kwargs_json,
|
||||
message_count,total_input_tokens,total_output_tokens,total_tokens,llm_call_count,
|
||||
lead_agent_tokens,subagent_tokens,middleware_tokens,created_at,updated_at
|
||||
) VALUES ('old','thread','alice','success','reject','{}','{}',0,0,0,0,0,0,0,0,CURRENT_TIMESTAMP,CURRENT_TIMESTAMP)""")
|
||||
)
|
||||
await asyncio.to_thread(command.upgrade, cfg, REVISION)
|
||||
async with engine.connect() as conn:
|
||||
columns = await conn.run_sync(lambda sync: {c["name"]: c for c in inspect(sync).get_columns("runs")})
|
||||
for name in ("run_incarnation", "producer_epoch", "producer_state"):
|
||||
assert columns[name]["nullable"]
|
||||
assert columns[name]["default"] is None
|
||||
if not fresh:
|
||||
assert (await conn.execute(text("SELECT run_incarnation,producer_epoch,producer_state FROM runs WHERE run_id='old'"))).one() == (None, None, None)
|
||||
assert await conn.scalar(text("SELECT count(*) FROM run_producer_policy")) == 0
|
||||
if schema:
|
||||
assert await conn.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard'")) == 1
|
||||
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
|
||||
async with engine.connect() as conn:
|
||||
assert "producer_epoch" not in await conn.run_sync(lambda sync: {c["name"] for c in inspect(sync).get_columns("runs")})
|
||||
await asyncio.to_thread(command.upgrade, cfg, REVISION)
|
||||
if schema:
|
||||
# Dispose bootstrap's idle pooled connections before offline cutover.
|
||||
await engine.dispose()
|
||||
try:
|
||||
await activate_producer_fencing(async_sessionmaker(engine), cutover_id="migration-check")
|
||||
except RuntimeError:
|
||||
async with engine.connect() as conn:
|
||||
clients = (
|
||||
await conn.execute(
|
||||
text("SELECT pid,state,application_name,backend_start,query_start,left(query,250) FROM pg_stat_activity WHERE datname=current_database() AND pid<>pg_backend_pid() AND backend_type='client backend'")
|
||||
)
|
||||
).all()
|
||||
pytest.fail(f"Unexpected live clients at isolated offline cutover: {clients}")
|
||||
await engine.dispose()
|
||||
with pytest.raises(RuntimeError, match="Cannot downgrade"):
|
||||
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
|
||||
finally:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
43
backend/tests/test_migration_0050_history_seed_authority.py
Normal file
43
backend/tests/test_migration_0050_history_seed_authority.py
Normal file
@ -0,0 +1,43 @@
|
||||
"""Existing strict databases must install history-seed grants during upgrade."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from alembic import command
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
from test_history_seed_fencing import history
|
||||
from test_history_seed_fencing import seed_db as seed_db
|
||||
|
||||
from deerflow.persistence import bootstrap
|
||||
from deerflow.runtime.events.producer import history_seed_authority, producer_scope
|
||||
from deerflow.runtime.journal import build_checkpoint_history_seed_events
|
||||
from deerflow.runtime.runs.schemas import ThreadOperationKind
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("leading_ai", [False, True])
|
||||
async def test_existing_strict_database_refreshes_persisted_guard(seed_db, leading_ai):
|
||||
db = seed_db
|
||||
batch = (
|
||||
build_checkpoint_history_seed_events([AIMessage(id="a0", content="inherited"), HumanMessage(id="h1", content="next"), AIMessage(id="a1", content="answer")], thread_id="thread", run_id_prefix="checkpoint-seed-thread")
|
||||
if leading_ai
|
||||
else history()
|
||||
)
|
||||
cfg = bootstrap._get_alembic_config(db.engine, postgres_schema=db.schema)
|
||||
await asyncio.to_thread(command.stamp, cfg, "0049_run_producer_fencing")
|
||||
async with db.engine.begin() as conn:
|
||||
definition = await conn.scalar(text("SELECT pg_get_functiondef('run_event_producer_guard()'::regprocedure)"))
|
||||
# Disable the new grant in the persisted function to reproduce a
|
||||
# pre-upgrade database; the application already knows the new grant.
|
||||
await conn.execute(text(definition.replace("a->>'source'='history_seed'", "false")))
|
||||
async with db.manager.reserve_thread_operation("thread", kind=ThreadOperationKind.checkpoint_write, user_id="alice") as reservation:
|
||||
authority = history_seed_authority(reservation, db.thread, seed_kind="checkpoint")
|
||||
with producer_scope(authority), pytest.raises(DBAPIError):
|
||||
await db.events.put_batch(batch)
|
||||
await asyncio.to_thread(command.upgrade, cfg, "head")
|
||||
with producer_scope(authority):
|
||||
rows = await db.events.put_batch(batch)
|
||||
assert [row["seq"] for row in rows] == ([1, 2, 3] if leading_ai else [1, 2, 3, 4])
|
||||
assert rows[0]["run_id"] == ("checkpoint-seed-thread-0" if leading_ai else "checkpoint-seed-thread-1")
|
||||
69
backend/tests/test_migration_0051_run_inputs_fencing.py
Normal file
69
backend/tests/test_migration_0051_run_inputs_fencing.py
Normal file
@ -0,0 +1,69 @@
|
||||
"""Both deployed migration branches converge without losing Run data or policy."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from alembic import command
|
||||
from sqlalchemy import inspect, text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence import bootstrap
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
||||
@pytest.mark.parametrize("previous", ["0049_run_inputs", "0049_run_producer_fencing", "0050_history_seed_authority"])
|
||||
async def test_both_migration_branches_upgrade_to_one_head(tmp_path, backend, previous):
|
||||
schema = None
|
||||
if backend == "postgres":
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "merged_runs_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema}})
|
||||
else:
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}")
|
||||
cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
|
||||
try:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await asyncio.to_thread(command.upgrade, cfg, previous)
|
||||
async with engine.begin() as conn:
|
||||
columns = await conn.run_sync(lambda sync: {c["name"] for c in inspect(sync).get_columns("runs")})
|
||||
if previous == "0049_run_inputs":
|
||||
assert "run_inputs_json" in columns
|
||||
assert "run_incarnation" not in columns
|
||||
assert not await conn.run_sync(lambda sync: inspect(sync).has_table("run_producer_policy"))
|
||||
else:
|
||||
assert "run_inputs_json" not in columns
|
||||
assert {"run_incarnation", "producer_epoch", "producer_state"} <= columns
|
||||
if schema:
|
||||
await conn.execute(text("INSERT INTO run_producer_policy(id,strict,cutover_id) VALUES (1,true,'preserved-policy')"))
|
||||
await conn.execute(
|
||||
text("""INSERT INTO runs(
|
||||
run_id,thread_id,user_id,status,multitask_strategy,metadata_json,kwargs_json,
|
||||
message_count,total_input_tokens,total_output_tokens,total_tokens,llm_call_count,
|
||||
lead_agent_tokens,subagent_tokens,middleware_tokens,created_at,updated_at
|
||||
) VALUES ('old','thread','alice','success','reject','{}','{}',0,0,0,0,0,0,0,0,CURRENT_TIMESTAMP,CURRENT_TIMESTAMP)""")
|
||||
)
|
||||
await asyncio.to_thread(command.upgrade, cfg, "head")
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(text("SELECT version_num FROM alembic_version"))).scalars().all() == ["0051_merge_run_inputs_fencing"]
|
||||
columns = await conn.run_sync(lambda sync: {c["name"] for c in inspect(sync).get_columns("runs")})
|
||||
assert {"run_inputs_json", "run_incarnation", "producer_epoch", "producer_state"} <= columns
|
||||
assert (await conn.execute(text("SELECT status,run_inputs_json,run_incarnation,producer_epoch,producer_state FROM runs WHERE run_id='old'"))).one() == ("success", None, None, None, None)
|
||||
if schema:
|
||||
assert await conn.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard'")) == 1
|
||||
if previous != "0049_run_inputs":
|
||||
assert (await conn.execute(text("SELECT strict,cutover_id FROM run_producer_policy WHERE id=1"))).one() == (True, "preserved-policy")
|
||||
# Bootstrap must also recognize the merged stamp on a later restart.
|
||||
await bootstrap.bootstrap_schema(engine, backend=backend, postgres_schema=schema or "")
|
||||
finally:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
104
backend/tests/test_postgres_asyncpg_url.py
Normal file
104
backend/tests/test_postgres_asyncpg_url.py
Normal file
@ -0,0 +1,104 @@
|
||||
"""Async ORM URLs preserve TLS requirements without changing libpq consumers."""
|
||||
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import make_url
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence import engine as engine_mod
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheme", ["postgres", "postgresql", "postgresql+asyncpg"])
|
||||
@pytest.mark.parametrize("mode", [None, "disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
|
||||
def test_async_orm_url_preserves_tls_credentials_and_query_options(scheme, mode):
|
||||
uri = f"{scheme}://us%40er%3Aname:p%40ss%2Fword%3A%25@localhost:5432/deerflow_test?target_session_attrs=read-write"
|
||||
if mode is not None:
|
||||
uri += f"&sslmode={mode}"
|
||||
config = DatabaseConfig(backend="postgres", postgres_url=uri)
|
||||
url = make_url(config.app_sqlalchemy_url)
|
||||
_, kwargs = url.get_dialect()().create_connect_args(url)
|
||||
assert url.drivername == "postgresql+asyncpg"
|
||||
assert kwargs["user"] == "us@er:name"
|
||||
assert kwargs["password"] == "p@ss/word:%"
|
||||
assert kwargs["target_session_attrs"] == "read-write"
|
||||
assert "sslmode" not in kwargs
|
||||
if mode is None:
|
||||
assert "ssl" not in kwargs
|
||||
else:
|
||||
assert kwargs["ssl"] == mode
|
||||
assert config.postgres_url == uri
|
||||
assert make_url(config.app_sync_sqlalchemy_url).query == make_url(uri).query
|
||||
assert asyncpg_test_url(uri) == config.app_sqlalchemy_url
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", ["ssl=require", "ssl=require&sslmode=require"])
|
||||
def test_explicit_matching_ssl_options_are_preserved(query):
|
||||
config = DatabaseConfig(backend="postgres", postgres_url=f"postgresql://user:password@localhost/test?{query}")
|
||||
assert make_url(config.app_sqlalchemy_url).query == {"ssl": "require"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"query",
|
||||
[
|
||||
"ssl=require&sslmode=disable",
|
||||
"sslmode=invalid",
|
||||
"sslmode=",
|
||||
"ssl=invalid",
|
||||
"ssl=",
|
||||
"sslmode=require&sslmode=disable",
|
||||
"ssl=require&ssl=disable",
|
||||
"channel_binding=require",
|
||||
"channel_binding=prefer",
|
||||
"channel_binding=disable",
|
||||
"channel_binding=invalid",
|
||||
"channel_binding=",
|
||||
],
|
||||
)
|
||||
def test_unsupported_or_ambiguous_security_options_are_rejected_without_secrets(query):
|
||||
config = DatabaseConfig(backend="postgres", postgres_url=f"postgresql://user:private-password@localhost/test?{query}")
|
||||
with pytest.raises(ValueError) as caught:
|
||||
_ = config.app_sqlalchemy_url
|
||||
assert "private-password" not in str(caught.value)
|
||||
assert make_url(config.app_sync_sqlalchemy_url).query == make_url(config.postgres_url).query
|
||||
|
||||
|
||||
def test_asyncpg_url_normalization_is_idempotent():
|
||||
from deerflow.persistence.postgres_url import normalize_asyncpg_url
|
||||
|
||||
normalized = normalize_asyncpg_url("postgres://user:p%40ss@localhost/test?sslmode=verify-full")
|
||||
assert normalize_asyncpg_url(normalized) == normalized
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheme", ["sqlite", "mysql", "postgresql+psycopg"])
|
||||
def test_asyncpg_url_rejects_other_drivers(scheme):
|
||||
from deerflow.persistence.postgres_url import normalize_asyncpg_url
|
||||
|
||||
with pytest.raises(ValueError, match="PostgreSQL"):
|
||||
normalize_asyncpg_url(f"{scheme}://user:password@localhost/test")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_libpq_url_bootstraps_through_database_config():
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "asyncpg_url_" + uuid4().hex
|
||||
config = DatabaseConfig(backend="postgres", postgres_url=uri, postgres_schema=schema)
|
||||
cleanup = create_async_engine(asyncpg_test_url(uri))
|
||||
try:
|
||||
await engine_mod.init_engine_from_config(config)
|
||||
async with engine_mod.get_engine().connect() as conn:
|
||||
assert await conn.scalar(text("SELECT current_schema()")) == schema
|
||||
assert await conn.scalar(text("SELECT version_num FROM alembic_version"))
|
||||
assert await conn.scalar(text("SELECT to_regclass('runs')")) == "runs"
|
||||
assert config.postgres_url == uri
|
||||
finally:
|
||||
await engine_mod.close_engine()
|
||||
async with cleanup.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await cleanup.dispose()
|
||||
176
backend/tests/test_run_producer_cutover.py
Normal file
176
backend/tests/test_run_producer_cutover.py
Normal file
@ -0,0 +1,176 @@
|
||||
"""Strict fencing is a durable, executable cutover, not a hot config toggle."""
|
||||
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.config.run_events_config import RunEventsConfig
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def cutover_db():
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
import deerflow.persistence.models # noqa: F401
|
||||
|
||||
schema = "cutover_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), pool_size=1, max_overflow=0, connect_args={"server_settings": {"search_path": schema, "application_name": schema}})
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
try:
|
||||
yield engine, async_sessionmaker(engine, expire_on_commit=False), schema
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def test_strict_requires_db_backend_and_explicit_cutover():
|
||||
with pytest.raises(ValueError):
|
||||
RunEventsConfig(backend="memory", producer_fencing="strict", producer_cutover_id="upgrade-1")
|
||||
with pytest.raises(ValueError):
|
||||
RunEventsConfig(backend="db", producer_fencing="strict")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unaware_sql_writer_is_rejected_after_cutover(cutover_db):
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
|
||||
engine, sf, _ = cutover_db
|
||||
runs = RunRepository(sf)
|
||||
await runs.put("run", thread_id="thread", user_id="alice", status="success")
|
||||
await activate_producer_fencing(sf, cutover_id="upgrade-1")
|
||||
with pytest.raises(DBAPIError):
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text("INSERT INTO run_events(thread_id,run_id,user_id,event_type,category,content,event_metadata,seq,created_at) VALUES ('thread','run','alice','middleware:old','trace','','{}',1,clock_timestamp())"))
|
||||
async with engine.connect() as conn:
|
||||
assert await conn.scalar(text("SELECT count(*) FROM run_events")) == 0
|
||||
with pytest.raises(RuntimeError):
|
||||
await DbRunEventStore(sf).validate_producer_fencing()
|
||||
with pytest.raises(RuntimeError):
|
||||
await DbRunEventStore(sf, producer_fencing="strict", producer_cutover_id="wrong").validate_producer_fencing()
|
||||
await DbRunEventStore(sf, producer_fencing="strict", producer_cutover_id="upgrade-1").validate_producer_fencing()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cutover_refuses_live_run(cutover_db):
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="running")
|
||||
with pytest.raises(RuntimeError, match="active"):
|
||||
await activate_producer_fencing(sf, cutover_id="upgrade-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cutover_refuses_other_database_sessions(cutover_db):
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
|
||||
_, sf, schema = cutover_db
|
||||
other = create_async_engine(asyncpg_test_url(os.environ["TEST_POSTGRES_URI"]), connect_args={"server_settings": {"search_path": schema}})
|
||||
try:
|
||||
async with other.connect():
|
||||
with pytest.raises(RuntimeError, match="sessions"):
|
||||
await activate_producer_fencing(sf, cutover_id="upgrade-1")
|
||||
finally:
|
||||
await other.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["put", "put_batch", "put_if_absent"])
|
||||
async def test_original_apis_reject_missing_authority_and_metadata_forgery(cutover_db, method):
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
from deerflow.runtime.events.producer import RunEventProducerRejected
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
runs = RunRepository(sf)
|
||||
await runs.put("run", thread_id="thread", user_id="alice", status="success")
|
||||
row = await runs.get("run", user_id="alice")
|
||||
await activate_producer_fencing(sf, cutover_id="upgrade-1")
|
||||
store = DbRunEventStore(sf, producer_fencing="strict", producer_cutover_id="upgrade-1")
|
||||
event = dict(thread_id="thread", run_id="run", event_type="middleware:old", category="trace", metadata={"producer_epoch": row["producer_epoch"], "run_incarnation": row["run_incarnation"], "trusted": True})
|
||||
with pytest.raises(RunEventProducerRejected):
|
||||
if method == "put_batch":
|
||||
await store.put_batch([event])
|
||||
else:
|
||||
await getattr(store, method)(**event)
|
||||
assert await store.list_events("thread", "run", user_id="alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_grant_is_named_and_independent_of_closed_producer(cutover_db):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
from deerflow.runtime.events.producer import RunEventProducerRejected, audit_authority, producer_authority, producer_scope
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
runs = RunRepository(sf)
|
||||
await runs.put("run", thread_id="thread", user_id="alice", status="success")
|
||||
record = SimpleNamespace(**await runs.get("run", user_id="alice"))
|
||||
await activate_producer_fencing(sf, cutover_id="upgrade-1")
|
||||
store = DbRunEventStore(sf, producer_fencing="strict", producer_cutover_id="upgrade-1")
|
||||
assert await store.close_producer(producer_authority(record))
|
||||
with pytest.raises(ValueError):
|
||||
audit_authority(record, event_types=("run.terminal",))
|
||||
with producer_scope(audit_authority(record, event_types=("audit.review",))):
|
||||
await store.put(thread_id="thread", run_id="run", event_type="audit.review", category="trace")
|
||||
with pytest.raises(RunEventProducerRejected):
|
||||
await store.put(thread_id="thread", run_id="run", event_type="audit.other", category="trace")
|
||||
await store.delete_by_run("thread", "run", user_id="alice")
|
||||
with pytest.raises(RunEventProducerRejected):
|
||||
await store.put(thread_id="thread", run_id="run", event_type="audit.review", category="trace")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sqlite_strict_startup_fails(tmp_path):
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'strict.db'}")
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="PostgreSQL"):
|
||||
await DbRunEventStore(async_sessionmaker(engine), producer_fencing="strict", producer_cutover_id="upgrade-1").validate_producer_fencing()
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["memory", "jsonl"])
|
||||
async def test_gateway_rejects_activated_database_with_downgraded_event_backend(cutover_db, monkeypatch, backend):
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.gateway.deps import langgraph_runtime
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
await activate_producer_fencing(sf, cutover_id="upgrade-1")
|
||||
|
||||
@asynccontextmanager
|
||||
async def resource(_config):
|
||||
yield object()
|
||||
|
||||
async def noop(*args):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr("deerflow.runtime.checkpointer.async_provider.make_checkpointer", resource)
|
||||
monkeypatch.setattr("deerflow.runtime.make_stream_bridge", resource)
|
||||
monkeypatch.setattr("deerflow.runtime.make_store", resource)
|
||||
monkeypatch.setattr("deerflow.persistence.engine.init_engine_from_config", noop)
|
||||
monkeypatch.setattr("deerflow.persistence.engine.close_engine", noop)
|
||||
monkeypatch.setattr("deerflow.persistence.engine.get_session_factory", lambda: sf)
|
||||
config = SimpleNamespace(database=SimpleNamespace(backend="postgres", checkpoint_channel_mode="full", checkpoint_delta=SimpleNamespace(snapshot_frequency=10)), run_events=RunEventsConfig(backend=backend))
|
||||
with pytest.raises(RuntimeError, match="unfenced Gateway"):
|
||||
async with langgraph_runtime(FastAPI(), config):
|
||||
pytest.fail("An activated SQL database allowed an unfenced event backend")
|
||||
94
backend/tests/test_run_producer_cutover_cli.py
Normal file
94
backend/tests/test_run_producer_cutover_cli.py
Normal file
@ -0,0 +1,94 @@
|
||||
"""Offline activation must be explicit and must not disclose database secrets."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
|
||||
def test_cli_requires_stopped_cluster_acknowledgement():
|
||||
from scripts.run_producer_fencing import main
|
||||
|
||||
with pytest.raises(SystemExit) as stopped:
|
||||
main(["--cutover-id", "upgrade-1"])
|
||||
assert stopped.value.code == 2
|
||||
|
||||
|
||||
def test_cli_redacts_connection_failure(monkeypatch, capsys):
|
||||
from scripts import run_producer_fencing as cli
|
||||
|
||||
monkeypatch.setattr(cli, "activate", AsyncMock(side_effect=RuntimeError("postgresql://user:secret@host/db")))
|
||||
assert cli.main(["--cutover-id", "upgrade-1", "--all-gateways-stopped"]) == 1
|
||||
output = capsys.readouterr()
|
||||
assert "secret" not in output.out + output.err
|
||||
assert "RuntimeError" in output.err
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_activation_uses_configured_schema_and_closes_engine(monkeypatch):
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.persistence import engine
|
||||
from deerflow.persistence.run import producer
|
||||
from scripts import run_producer_fencing as cli
|
||||
|
||||
config = SimpleNamespace(database=SimpleNamespace(backend="postgres", model_copy=lambda **kw: "database-config"))
|
||||
monkeypatch.setattr(AppConfig, "from_file", lambda path: config)
|
||||
initialize, close, dispose, activate = AsyncMock(), AsyncMock(), AsyncMock(), AsyncMock()
|
||||
monkeypatch.setattr(engine, "init_engine_from_config", initialize)
|
||||
monkeypatch.setattr(engine, "get_engine", lambda: SimpleNamespace(dispose=dispose))
|
||||
monkeypatch.setattr(engine, "get_session_factory", lambda: "session-factory")
|
||||
monkeypatch.setattr(engine, "close_engine", close)
|
||||
monkeypatch.setattr(producer, "activate_producer_fencing", activate)
|
||||
await cli.activate(config_path="operator.yaml", cutover_id="upgrade-1")
|
||||
initialize.assert_awaited_once_with("database-config")
|
||||
dispose.assert_awaited_once()
|
||||
activate.assert_awaited_once_with("session-factory", cutover_id="upgrade-1")
|
||||
close.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_offline_command_bootstraps_and_activates_postgres(tmp_path):
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import yaml
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "cutover_cli_" + uuid4().hex
|
||||
config = tmp_path / "operator.yaml"
|
||||
config.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"database": {"backend": "postgres", "postgres_url": "$TEST_POSTGRES_URI", "postgres_schema": schema},
|
||||
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"},
|
||||
"run_events": {"backend": "db", "producer_fencing": "strict", "producer_cutover_id": "cli-test"},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema}})
|
||||
try:
|
||||
script = Path(__file__).resolve().parents[1] / "scripts" / "run_producer_fencing.py"
|
||||
process = await asyncio.create_subprocess_exec(sys.executable, str(script), "--config", str(config), "--cutover-id", "cli-test", "--all-gateways-stopped", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), 30)
|
||||
finally:
|
||||
if process.returncode is None:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
assert process.returncode == 0, (stdout, stderr)
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(text("SELECT strict,cutover_id FROM run_producer_policy WHERE id=1"))).one() == (True, "cli-test")
|
||||
assert await conn.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard'")) == 1
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
214
backend/tests/test_run_producer_fencing.py
Normal file
214
backend/tests/test_run_producer_fencing.py
Normal file
@ -0,0 +1,214 @@
|
||||
"""Producer revocation must fence all event entry points, not only a new API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import select, text
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunRow
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.store.base import RunEventStoreNotCommittedError
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(params=["sqlite", "postgresql"])
|
||||
async def producer_db(request, tmp_path):
|
||||
schema = None
|
||||
if request.param == "postgresql":
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "producer_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema, "application_name": schema}})
|
||||
else:
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'producer.db'}")
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
if schema:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await conn.run_sync(lambda c: RunRow.metadata.create_all(c, tables=[RunRow.__table__, RunChangeClockRow.__table__, RunEventRow.__table__, RunEventThreadSeqRow.__table__, CompletedRunSnapshotRow.__table__]))
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
yield SimpleNamespace(engine=engine, sf=sf, runs=RunRepository(sf), events=DbRunEventStore(sf), backend=request.param)
|
||||
finally:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def admit(db, *, run_id="run", thread_id="thread", expired=False, status="running", owner="owner"):
|
||||
await db.runs.put(run_id, thread_id=thread_id, user_id="alice", status=status, owner_worker_id=owner, lease_expires_at=(datetime.now(UTC) + timedelta(seconds=-60 if expired else 60)).isoformat())
|
||||
return await db.runs.get(run_id, user_id="alice")
|
||||
|
||||
|
||||
def batch(run_id="run", thread_id="thread"):
|
||||
return [{"thread_id": thread_id, "run_id": run_id, "user_id": "alice", "event_type": "middleware:test", "category": "trace"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_admission_persists_nonreusable_producer_identity(producer_db):
|
||||
db = producer_db
|
||||
old = await admit(db)
|
||||
assert old.get("run_incarnation"), "new Run has no durable incarnation"
|
||||
assert old.get("producer_epoch"), "new Run has no durable producer identity"
|
||||
assert old.get("producer_state") == "open"
|
||||
await db.runs.delete("run", user_id="alice")
|
||||
new = await admit(db)
|
||||
assert new["run_incarnation"] != old["run_incarnation"]
|
||||
assert new["producer_epoch"] != old["producer_epoch"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_takeover_atomically_revokes_producer(producer_db):
|
||||
db = producer_db
|
||||
row = await admit(db, expired=True)
|
||||
assert await db.runs.claim_for_takeover("run", grace_seconds=0, error="expired")
|
||||
current = await db.runs.get("run", user_id="alice")
|
||||
assert current.get("producer_state") == "revoked", "takeover left old event producer admitted"
|
||||
assert current["producer_epoch"] != row["producer_epoch"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_old_authority_cannot_write_after_takeover(producer_db):
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
|
||||
db = producer_db
|
||||
row = await admit(db, expired=True)
|
||||
authority = producer_authority(SimpleNamespace(**row))
|
||||
assert await db.runs.claim_for_takeover("run", grace_seconds=0, error="expired")
|
||||
with producer_scope(authority):
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await db.events.put_batch(batch())
|
||||
assert await db.events.list_events("thread", "run", user_id="alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mutation", ["epoch", "incarnation", "mixed_run", "mixed_owner", "reserved_type"])
|
||||
async def test_invalid_batch_commits_no_rows_or_watermark(producer_db, mutation):
|
||||
from dataclasses import replace
|
||||
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
|
||||
db = producer_db
|
||||
authority = producer_authority(SimpleNamespace(**await admit(db)))
|
||||
events = batch() + batch()
|
||||
if mutation == "epoch":
|
||||
authority = replace(authority, producer_epoch="obsolete")
|
||||
elif mutation == "incarnation":
|
||||
authority = replace(authority, run_incarnation="obsolete")
|
||||
elif mutation == "mixed_run":
|
||||
events[1]["run_id"] = "other"
|
||||
elif mutation == "mixed_owner":
|
||||
events[1]["user_id"] = "mallory"
|
||||
else:
|
||||
events[1]["event_type"] = "run.terminal"
|
||||
with producer_scope(authority):
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await db.events.put_batch(events)
|
||||
async with db.sf() as session:
|
||||
assert await session.scalar(select(RunEventRow.id).limit(1)) is None
|
||||
assert await session.scalar(select(RunEventThreadSeqRow.seq).limit(1)) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_status_does_not_close_live_producer(producer_db):
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
|
||||
db = producer_db
|
||||
authority = producer_authority(SimpleNamespace(**await admit(db)))
|
||||
assert (await db.runs.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")).finalized
|
||||
with producer_scope(authority):
|
||||
written = await db.events.put_batch(batch())
|
||||
assert written[0]["seq"] == 1
|
||||
assert (await db.runs.get("run", user_id="alice"))["producer_state"] == "open"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("delete", ["events_run", "events_thread", "run"])
|
||||
async def test_delete_fences_old_writer_including_same_id_recreation(producer_db, delete):
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
|
||||
db = producer_db
|
||||
authority = producer_authority(SimpleNamespace(**await admit(db)))
|
||||
if delete == "events_run":
|
||||
await db.events.delete_by_run("thread", "run", user_id="alice")
|
||||
elif delete == "events_thread":
|
||||
await db.events.delete_by_thread("thread", user_id="alice")
|
||||
else:
|
||||
await db.runs.delete("run", user_id="alice")
|
||||
await admit(db)
|
||||
with producer_scope(authority):
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await db.events.put_batch(batch())
|
||||
assert await db.events.list_events("thread", "run", user_id="alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_explicit_close_is_idempotent_and_never_reopens(producer_db):
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
|
||||
db = producer_db
|
||||
authority = producer_authority(SimpleNamespace(**await admit(db)))
|
||||
assert await db.events.close_producer(authority)
|
||||
assert await db.events.close_producer(authority)
|
||||
with producer_scope(authority):
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await db.events.put_batch(batch())
|
||||
await db.runs.put("run", thread_id="thread", user_id="alice", status="success")
|
||||
assert (await db.runs.get("run", user_id="alice"))["producer_state"] == "closed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("owner,expected", [("owner", "open"), ("peer", "revoked")])
|
||||
async def test_atomic_admission_preserves_only_local_finalizer(producer_db, owner, expected):
|
||||
db = producer_db
|
||||
await admit(db, expired=True)
|
||||
await db.runs.create_thread_operation_atomic("replacement", thread_id="thread", owner_worker_id=owner, lease_expires_at=None, multitask_strategy="rollback", user_id="alice", grace_seconds=0)
|
||||
assert (await db.runs.get("run", user_id="alice")).get("producer_state") == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("admission", ["put", "atomic"])
|
||||
async def test_manager_binds_committed_identity_without_post_commit_read(producer_db, admission):
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
from deerflow.runtime.runs.manager import RunManager
|
||||
|
||||
db = producer_db
|
||||
|
||||
class RecreatedAdmission(RunRepository):
|
||||
original = None
|
||||
|
||||
async def recreate(self, result):
|
||||
self.original = result.copy()
|
||||
await super().delete(result["run_id"], user_id="alice")
|
||||
await super().put(result["run_id"], thread_id="thread", user_id="alice", status="success", owner_worker_id="owner")
|
||||
|
||||
async def put(self, *args, **kwargs):
|
||||
result = await super().put(*args, **kwargs)
|
||||
await self.recreate(result)
|
||||
return result
|
||||
|
||||
async def create_thread_operation_atomic(self, *args, **kwargs):
|
||||
result = await super().create_thread_operation_atomic(*args, **kwargs)
|
||||
await self.recreate(result[0])
|
||||
return result
|
||||
|
||||
async def get(self, *args, **kwargs):
|
||||
raise AssertionError("Admission must not await a later Row read to manufacture authority")
|
||||
|
||||
store = RecreatedAdmission(db.sf)
|
||||
manager = RunManager(store=store, worker_id="owner")
|
||||
record = await (manager.create("thread", user_id="alice") if admission == "put" else manager.create_or_reject("thread", user_id="alice"))
|
||||
assert record.run_incarnation == store.original["run_incarnation"]
|
||||
with producer_scope(producer_authority(record)):
|
||||
with pytest.raises(RunEventStoreNotCommittedError):
|
||||
await db.events.put_batch(batch(run_id=record.run_id))
|
||||
196
backend/tests/test_run_producer_off.py
Normal file
196
backend/tests/test_run_producer_off.py
Normal file
@ -0,0 +1,196 @@
|
||||
"""Off-mode INSERT contention and atomic cutover on independent PostgreSQL sessions."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
from test_run_producer_cutover import cutover_db as cutover_db
|
||||
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
INSERT = text("INSERT INTO run_events(thread_id,run_id,user_id,event_type,category,content,event_metadata,seq,created_at) VALUES ('thread','run','alice','middleware:off','trace','','{}',:seq,clock_timestamp())")
|
||||
GATE = "SELECT pg_advisory_xact_lock(hashtext('deerflow.event_fence/' || current_schema())::bigint)"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(params=["missing", "false"])
|
||||
async def off_db(cutover_db, request):
|
||||
if request.param == "false":
|
||||
async with cutover_db[0].begin() as conn:
|
||||
await conn.execute(text("INSERT INTO run_producer_policy(id,strict) VALUES (1,false)"))
|
||||
return cutover_db
|
||||
|
||||
|
||||
def peer_engine(schema, **kwargs):
|
||||
return create_async_engine(asyncpg_test_url(os.environ["TEST_POSTGRES_URI"]), pool_size=2, max_overflow=0, connect_args={"server_settings": {"search_path": schema, "statement_timeout": "6000"}}, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("lock", "api"), [("thread", False), ("run", False), ("run", True)])
|
||||
async def test_off_insert_ignores_trigger_added_thread_and_run_locks(off_db, lock, api):
|
||||
engine, sf, schema = off_db
|
||||
await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="success")
|
||||
holder = peer_engine(schema)
|
||||
try:
|
||||
async with holder.begin() as conn:
|
||||
await conn.execute(text("SELECT pg_advisory_xact_lock(hashtext('thread')::bigint)" if lock == "thread" else "SELECT run_id FROM runs WHERE run_id='run' FOR UPDATE"))
|
||||
async with engine.begin() as writer:
|
||||
await writer.execute(text("SET LOCAL lock_timeout='250ms'"))
|
||||
if not api:
|
||||
await writer.execute(INSERT, {"seq": 1})
|
||||
locks = (await writer.execute(text("SELECT mode FROM pg_locks WHERE pid=pg_backend_pid() AND locktype='advisory'"))).scalars().all()
|
||||
assert locks == ["ShareLock"]
|
||||
if api:
|
||||
await asyncio.wait_for(DbRunEventStore(sf).put_batch([dict(thread_id="thread", run_id="run", user_id="alice", event_type="middleware:off", category="trace")]), 1)
|
||||
assert len(await DbRunEventStore(sf).list_events("thread", "run", user_id="alice")) == 1
|
||||
finally:
|
||||
await holder.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("isolation", ["REPEATABLE READ", "SERIALIZABLE"])
|
||||
async def test_off_snapshot_isolation_is_fail_closed(off_db, isolation):
|
||||
_, sf, schema = off_db
|
||||
await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="success")
|
||||
writer = peer_engine(schema, isolation_level=isolation)
|
||||
try:
|
||||
with pytest.raises(DBAPIError) as caught:
|
||||
async with writer.begin() as conn:
|
||||
await conn.execute(INSERT, {"seq": 1})
|
||||
assert caught.value.orig.sqlstate == "0A000"
|
||||
finally:
|
||||
await writer.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("writer_first", [False, True])
|
||||
async def test_off_cutover_serializes_without_acquiring_prior_locks(off_db, writer_first):
|
||||
engine, sf, schema = off_db
|
||||
await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="success")
|
||||
reached, release = asyncio.Event(), asyncio.Event()
|
||||
pid = {}
|
||||
|
||||
class GatedCutover(AsyncSession):
|
||||
async def execute(self, statement, *args, **kwargs):
|
||||
if str(statement) == GATE:
|
||||
pid["cutover"] = await self.scalar(text("SELECT pg_backend_pid()"))
|
||||
if writer_first:
|
||||
reached.set()
|
||||
await release.wait()
|
||||
return await super().execute(statement, *args, **kwargs)
|
||||
result = await super().execute(statement, *args, **kwargs)
|
||||
reached.set()
|
||||
await release.wait()
|
||||
return result
|
||||
return await super().execute(statement, *args, **kwargs)
|
||||
|
||||
activation = asyncio.create_task(activate_producer_fencing(async_sessionmaker(engine, class_=GatedCutover, expire_on_commit=False), cutover_id="off-race"))
|
||||
peer = peer_engine(schema)
|
||||
writer_task = None
|
||||
try:
|
||||
# Production's stopped-client check runs before these racing peers
|
||||
# connect. Exercise the lock even against an unexpected reconnect.
|
||||
await asyncio.wait_for(reached.wait(), 3)
|
||||
async with peer.connect() as writer, peer.connect() as observer:
|
||||
writer_pid = await writer.scalar(text("SELECT pg_backend_pid()"))
|
||||
if writer_first:
|
||||
await writer.execute(INSERT, [{"seq": 1}, {"seq": 2}])
|
||||
release.set()
|
||||
await blocked_peer(observer, pid["cutover"], writer_pid)
|
||||
await writer.commit()
|
||||
await asyncio.wait_for(activation, 3)
|
||||
else:
|
||||
# Holding earlier locks makes reverse lock acquisition fail
|
||||
# visibly: the off writer must abort at the gate after cutover.
|
||||
await observer.execute(text("SELECT pg_advisory_xact_lock(hashtext('thread')::bigint)"))
|
||||
await observer.execute(text("SELECT run_id FROM runs WHERE run_id='run' FOR UPDATE"))
|
||||
writer_task = asyncio.create_task(writer.execute(INSERT, [{"seq": 1}, {"seq": 2}]))
|
||||
try:
|
||||
await blocked_peer(observer, writer_pid, pid["cutover"])
|
||||
release.set()
|
||||
await asyncio.wait_for(activation, 3)
|
||||
with pytest.raises(DBAPIError) as caught:
|
||||
await asyncio.wait_for(writer_task, 3)
|
||||
assert caught.value.orig.sqlstate == "40001"
|
||||
finally:
|
||||
release.set()
|
||||
await observer.rollback()
|
||||
await asyncio.wait_for(asyncio.gather(writer_task, return_exceptions=True), 3)
|
||||
await writer.rollback()
|
||||
async with engine.connect() as conn:
|
||||
assert await conn.scalar(text("SELECT count(*) FROM run_events")) == (2 if writer_first else 0)
|
||||
assert await conn.scalar(text("SELECT strict FROM run_producer_policy WHERE id=1")) is True
|
||||
finally:
|
||||
release.set()
|
||||
for task in (activation, writer_task):
|
||||
if task is not None:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
await peer.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("isolation", ["REPEATABLE READ", "SERIALIZABLE"])
|
||||
async def test_retained_off_snapshot_cannot_bypass_committed_cutover(off_db, isolation):
|
||||
engine, sf, schema = off_db
|
||||
await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="success")
|
||||
reached, release = asyncio.Event(), asyncio.Event()
|
||||
|
||||
class GatedCutover(AsyncSession):
|
||||
async def execute(self, statement, *args, **kwargs):
|
||||
if str(statement) == GATE:
|
||||
reached.set()
|
||||
await release.wait()
|
||||
return await super().execute(statement, *args, **kwargs)
|
||||
|
||||
activation = asyncio.create_task(activate_producer_fencing(async_sessionmaker(engine, class_=GatedCutover), cutover_id="snapshot-race"))
|
||||
peer = peer_engine(schema, isolation_level=isolation)
|
||||
try:
|
||||
await asyncio.wait_for(reached.wait(), 3)
|
||||
async with peer.connect() as writer:
|
||||
assert not await writer.scalar(text("SELECT strict FROM run_producer_policy WHERE id=1"))
|
||||
release.set()
|
||||
await asyncio.wait_for(activation, 3)
|
||||
# The old snapshot still says off even after activation commits.
|
||||
assert not await writer.scalar(text("SELECT strict FROM run_producer_policy WHERE id=1"))
|
||||
with pytest.raises(DBAPIError) as caught:
|
||||
await writer.execute(INSERT, [{"seq": 1}, {"seq": 2}])
|
||||
await writer.commit()
|
||||
assert caught.value.orig.sqlstate == "0A000"
|
||||
await writer.rollback()
|
||||
async with engine.connect() as conn:
|
||||
assert await conn.scalar(text("SELECT count(*) FROM run_events")) == 0
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(asyncio.gather(activation, return_exceptions=True), 3)
|
||||
await peer.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_off_read_uncommitted_remains_accepted(off_db):
|
||||
_, sf, schema = off_db
|
||||
await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="success")
|
||||
peer = peer_engine(schema)
|
||||
try:
|
||||
async with peer.begin() as writer:
|
||||
await writer.execute(text("SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED"))
|
||||
await writer.execute(INSERT, {"seq": 1})
|
||||
finally:
|
||||
await peer.dispose()
|
||||
|
||||
|
||||
async def blocked_peer(observer, waiter, holder):
|
||||
deadline = time.monotonic() + 3
|
||||
while time.monotonic() < deadline:
|
||||
if holder in await observer.scalar(text("SELECT pg_blocking_pids(:pid)"), {"pid": waiter}):
|
||||
return
|
||||
await asyncio.sleep(0.01)
|
||||
raise AssertionError("Expected PostgreSQL blocking relationship was not observed")
|
||||
154
backend/tests/test_run_producer_postgres.py
Normal file
154
backend/tests/test_run_producer_postgres.py
Normal file
@ -0,0 +1,154 @@
|
||||
"""Real PostgreSQL, actual Journal/Store APIs, independent OS-process races."""
|
||||
|
||||
import os
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
from support.producer_process import ProducerActors, receive
|
||||
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunProducerPolicyRow, RunRow
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def producer_pg():
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
uri = asyncpg_test_url(uri)
|
||||
schema = "producer_race_" + uuid4().hex
|
||||
engine = create_async_engine(uri, pool_size=1, max_overflow=0, pool_timeout=1, connect_args={"server_settings": {"search_path": schema}})
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await conn.run_sync(
|
||||
lambda c: RunRow.metadata.create_all(c, tables=[RunRow.__table__, RunChangeClockRow.__table__, RunEventRow.__table__, RunEventThreadSeqRow.__table__, CompletedRunSnapshotRow.__table__, RunProducerPolicyRow.__table__])
|
||||
)
|
||||
# This fixture tests contention after activation; cutover itself has
|
||||
# separate tests requiring all other DB clients to be stopped.
|
||||
await conn.execute(text("INSERT INTO run_producer_policy(id,strict,cutover_id) VALUES (1,true,'race-fixture')"))
|
||||
runs = RunRepository(sf)
|
||||
await runs.put("run", thread_id="thread", user_id="alice", status="running", owner_worker_id="old", lease_expires_at=(datetime.now(UTC) - timedelta(seconds=60)).isoformat())
|
||||
identity = await runs.get("run", user_id="alice")
|
||||
actors = ProducerActors(uri, schema, identity, engine)
|
||||
try:
|
||||
yield SimpleNamespace(engine=engine, sf=sf, runs=runs, events=DbRunEventStore(sf, producer_fencing="strict"), identity=identity, actors=actors)
|
||||
finally:
|
||||
await actors.close()
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_writer_first_takeover_waits_for_real_event_commit(producer_pg):
|
||||
db = producer_pg
|
||||
writer, writer_pid = await db.actors.start("writer_locked")
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["stage"] == "admitted"
|
||||
taker, taker_pid = await db.actors.start("takeover")
|
||||
taker.send("go")
|
||||
await db.actors.blocked(taker_pid, writer_pid)
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["written"] is True
|
||||
assert (await receive(taker))["taken"] is True
|
||||
events = await db.events.list_events("thread", "run", user_id="alice")
|
||||
assert [event["event_type"] for event in events] == ["middleware:race", "run.delivery", "run.terminal"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_takeover_first_writer_waits_then_rejects_old_epoch(producer_pg):
|
||||
db = producer_pg
|
||||
taker, taker_pid = await db.actors.start("takeover_locked")
|
||||
taker.send("go")
|
||||
assert (await receive(taker))["stage"] == "revoked_uncommitted"
|
||||
writer, writer_pid = await db.actors.start("writer")
|
||||
writer.send("go")
|
||||
await db.actors.blocked(writer_pid, taker_pid)
|
||||
taker.send("go")
|
||||
assert (await receive(taker))["taken"] is True
|
||||
assert (await receive(writer))["written"] is False
|
||||
assert [event["event_type"] for event in await db.events.list_events("thread", "run", user_id="alice")] == ["run.delivery", "run.terminal"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("delete", [False, True])
|
||||
async def test_delayed_journal_cannot_commit_after_recovery_publication(producer_pg, delete):
|
||||
db = producer_pg
|
||||
writer, _ = await db.actors.start("writer_late")
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["stage"] == "before_admission"
|
||||
taker, _ = await db.actors.start("takeover")
|
||||
taker.send("go")
|
||||
assert (await receive(taker))["taken"] is True
|
||||
if delete:
|
||||
await db.events.delete_by_run("thread", "run", user_id="alice")
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["written"] is False
|
||||
types = [event["event_type"] for event in await db.events.list_events("thread", "run", user_id="alice")]
|
||||
assert types == ([] if delete else ["run.delivery", "run.terminal"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_first_waiting_journal_cannot_resurrect_events(producer_pg):
|
||||
db = producer_pg
|
||||
deleter, delete_pid = await db.actors.start("delete_locked")
|
||||
deleter.send("go")
|
||||
assert (await receive(deleter))["stage"] == "deleted_uncommitted"
|
||||
writer, writer_pid = await db.actors.start("writer")
|
||||
writer.send("go")
|
||||
await db.actors.blocked(writer_pid, delete_pid)
|
||||
deleter.send("go")
|
||||
assert (await receive(deleter))["deleted"] is True
|
||||
assert (await receive(writer))["written"] is False
|
||||
assert await db.events.list_events("thread", "run", user_id="alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fenced_write_uses_one_connection_and_never_needs_change_clock(producer_pg):
|
||||
db = producer_pg
|
||||
# pool_size=1, max_overflow=0: a nested guard Session would self-block.
|
||||
with producer_scope(producer_authority(SimpleNamespace(**db.identity))):
|
||||
result = await db.events.put_batch([{"thread_id": "thread", "run_id": "run", "event_type": "middleware:pool", "category": "trace"}])
|
||||
assert result[0]["seq"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_writer_can_commit_while_deletion_holds_clock_waiting_for_thread(producer_pg):
|
||||
db = producer_pg
|
||||
writer, writer_pid = await db.actors.start("writer_locked")
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["stage"] == "admitted"
|
||||
deleter, delete_pid = await db.actors.start("delete_waiting")
|
||||
deleter.send("go")
|
||||
assert (await receive(deleter))["stage"] == "clock_held"
|
||||
deleter.send("go")
|
||||
await db.actors.blocked(delete_pid, writer_pid)
|
||||
# A writer acquiring C after T/R would deadlock against this deleter.
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["written"] is True
|
||||
assert (await receive(deleter))["deleted"] is True
|
||||
assert await db.events.list_events("thread", "run", user_id="alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cross_process_same_id_recreation_rejects_old_incarnation(producer_pg):
|
||||
db = producer_pg
|
||||
writer, _ = await db.actors.start("writer_late")
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["stage"] == "before_admission"
|
||||
await db.runs.delete("run", user_id="alice")
|
||||
new = await db.runs.put("run", thread_id="thread", user_id="alice", status="success", owner_worker_id="old")
|
||||
assert new["run_incarnation"] != db.identity["run_incarnation"]
|
||||
writer.send("go")
|
||||
assert (await receive(writer))["written"] is False
|
||||
assert await db.events.list_events("thread", "run", user_id="alice") == []
|
||||
160
backend/tests/test_run_producer_worker.py
Normal file
160
backend/tests/test_run_producer_worker.py
Normal file
@ -0,0 +1,160 @@
|
||||
"""Strict producer permissions survive the actual Worker and publisher paths."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from test_run_producer_cutover import cutover_db as cutover_db
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_worker_keeps_terminal_hook_open_then_closes_after_drain(cutover_db):
|
||||
from deerflow.config.run_ownership_config import RunOwnershipConfig
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.producer import current_producer
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.runs.manager import RunManager
|
||||
from deerflow.runtime.runs.worker import RunContext, _SubagentEventBuffer, run_agent
|
||||
from deerflow.runtime.stream_bridge import MemoryStreamBridge
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
await activate_producer_fencing(sf, cutover_id="worker-test")
|
||||
events = DbRunEventStore(sf, producer_fencing="strict", producer_cutover_id="worker-test")
|
||||
runs = RunRepository(sf)
|
||||
manager = RunManager(store=runs, event_store=events, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=300))
|
||||
record = await manager.create_or_reject("thread", user_id="alice")
|
||||
holder = {}
|
||||
|
||||
class RootAgent:
|
||||
async def astream(self, graph_input, config=None, **kwargs):
|
||||
holder["journal"] = config["context"]["__run_journal"]
|
||||
yield {"messages": []}
|
||||
|
||||
async def hook(completed):
|
||||
row = await runs.get(completed.run_id, user_id="alice")
|
||||
assert row["status"] == "success"
|
||||
assert row["producer_state"] == "open"
|
||||
assert current_producer().run_incarnation == record.run_incarnation
|
||||
assert await asyncio.to_thread(lambda: current_producer().producer_epoch) == record.producer_epoch
|
||||
holder["journal"].record_middleware("completion", name="hook", hook="completion", action="accepted", changes={})
|
||||
# Same production buffer used by the Worker for delegated task events.
|
||||
buffer = _SubagentEventBuffer(events, record.thread_id, record.run_id)
|
||||
await buffer.add({"type": "task_completed", "task_id": "subtask", "result": "done"})
|
||||
await buffer.flush()
|
||||
|
||||
task = asyncio.create_task(
|
||||
run_agent(
|
||||
MemoryStreamBridge(), manager, record, ctx=RunContext(checkpointer=None, event_store=events, on_run_completed=hook), agent_factory=lambda **kwargs: RootAgent(), graph_input={}, config={"configurable": {"thread_id": "thread"}}
|
||||
)
|
||||
)
|
||||
record.task = task
|
||||
await asyncio.wait_for(task, 15)
|
||||
row = await runs.get(record.run_id, user_id="alice")
|
||||
written = await events.list_events("thread", record.run_id, user_id="alice")
|
||||
types = [event["event_type"] for event in written]
|
||||
assert "middleware:completion" in types
|
||||
assert "subagent.end" in types
|
||||
assert types.count("run.delivery") == types.count("run.terminal") == 1
|
||||
terminal_seq = next(event["seq"] for event in written if event["event_type"] == "run.terminal")
|
||||
assert all(event["seq"] < terminal_seq for event in written if event["event_type"] in {"middleware:completion", "subagent.end"})
|
||||
assert row["producer_state"] == "closed"
|
||||
assert row["evidence_seal_state"] == "sealed"
|
||||
assert current_producer() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_takeover_publisher_has_separate_retention_bound_authority(cutover_db):
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from deerflow.persistence.run.producer import activate_producer_fencing
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.events.terminal import RecoveryTerminalPublisher
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
await activate_producer_fencing(sf, cutover_id="takeover-test")
|
||||
runs = RunRepository(sf)
|
||||
events = DbRunEventStore(sf, producer_fencing="strict")
|
||||
await runs.put("run", thread_id="thread", user_id="alice", status="running", owner_worker_id="old", lease_expires_at=(datetime.now(UTC) - timedelta(seconds=60)).isoformat())
|
||||
old = await runs.get("run", user_id="alice")
|
||||
claimed = await runs.claim_for_takeover("run", grace_seconds=0, error="orphan")
|
||||
publisher = RecoveryTerminalPublisher(events, runs)
|
||||
target = publisher.authorize_takeover(claimed=claimed, thread_id="thread", run_id="run", user_id="alice", operation_kind="run", run_incarnation=old["run_incarnation"], retention_revision=old["evidence_retention_revision"])
|
||||
assert (await publisher.publish(target)).confirmed
|
||||
assert [event["event_type"] for event in await events.list_events("thread", "run", user_id="alice")] == ["run.delivery", "run.terminal"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_journal_keeps_host_identity_across_contextless_callback(cutover_db):
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
|
||||
row = await RunRepository(sf).put("run", thread_id="thread", user_id="alice", owner_worker_id="owner")
|
||||
events = DbRunEventStore(sf, producer_fencing="strict")
|
||||
with producer_scope(producer_authority(SimpleNamespace(**row))):
|
||||
journal = RunJournal("run", "thread", events, user_id="alice", flush_threshold=1)
|
||||
|
||||
def callback():
|
||||
journal.record_middleware("thread", name="raw-thread", hook="tool", action="accepted", changes={})
|
||||
|
||||
thread = threading.Thread(target=callback)
|
||||
thread.start()
|
||||
thread.join(timeout=1)
|
||||
assert not thread.is_alive()
|
||||
try:
|
||||
await journal.flush()
|
||||
assert [event["event_type"] for event in await events.list_events("thread", "run", user_id="alice")] == ["middleware:thread"]
|
||||
finally:
|
||||
await journal.close(flush=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_detached_idempotent_handle_cannot_publish_owner_receipt(cutover_db):
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.producer import producer_authority
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.runs.manager import RunManager, RunStartupError
|
||||
from deerflow.runtime.runs.worker import RunContext, run_agent
|
||||
from deerflow.runtime.stream_bridge import MemoryStreamBridge
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
runs = RunRepository(sf)
|
||||
owner = RunManager(store=runs, worker_id="owner")
|
||||
admitted = await owner.create_or_reject("thread", user_id="alice", idempotency_key="same-request")
|
||||
peer = RunManager(store=runs, worker_id="peer")
|
||||
handle = await peer.create_or_reject("thread", user_id="alice", idempotency_key="same-request")
|
||||
assert handle.store_only and handle.idempotency_reused
|
||||
assert producer_authority(handle) is None
|
||||
events = DbRunEventStore(sf, producer_fencing="strict")
|
||||
with pytest.raises(RunStartupError):
|
||||
await run_agent(MemoryStreamBridge(), peer, handle, ctx=RunContext(checkpointer=None, event_store=events), agent_factory=lambda **kwargs: pytest.fail("Detached handle executed the agent"), graph_input={}, config={})
|
||||
assert (await runs.get(admitted.run_id, user_id="alice"))["status"] == "pending"
|
||||
assert await events.list_events("thread", admitted.run_id, user_id="alice") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workspace_recorder_preserves_producer_grant(cutover_db, tmp_path, monkeypatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.runtime.events.producer import producer_authority, producer_scope
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.workspace_changes import WorkspaceRoot, record_workspace_changes, scan_workspace_roots
|
||||
|
||||
_, sf, _ = cutover_db
|
||||
row = await RunRepository(sf).put("run", thread_id="thread", user_id="alice", status="success", owner_worker_id="owner")
|
||||
roots = [WorkspaceRoot(name="workspace", host_path=tmp_path, virtual_prefix="/mnt/user-data/workspace")]
|
||||
monkeypatch.setattr("deerflow.workspace_changes.recorder.build_thread_workspace_roots", lambda *args, **kwargs: roots)
|
||||
before = scan_workspace_roots(roots)
|
||||
(tmp_path / "report.txt").write_text("completed", encoding="utf-8")
|
||||
with producer_scope(producer_authority(SimpleNamespace(**row))):
|
||||
event = await record_workspace_changes(DbRunEventStore(sf, producer_fencing="strict"), "thread", "run", before, user_id="alice")
|
||||
assert event["event_type"] == "workspace_changes"
|
||||
assert event["user_id"] == "alice"
|
||||
79
backend/tests/test_scheduled_recovery_confirmation.py
Normal file
79
backend/tests/test_scheduled_recovery_confirmation.py
Normal file
@ -0,0 +1,79 @@
|
||||
"""Unconfirmed scheduled recovery publication must remain observable after commit."""
|
||||
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
from test_scheduled_recovery_fencing import prepare_recovery, reconcile
|
||||
from test_scheduled_recovery_fencing import recovery_db as recovery_db
|
||||
|
||||
from deerflow.runtime.events import terminal
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "recovery_path", "repository_logger"),
|
||||
[
|
||||
("once", "reconcile_stuck_once_tasks", "deerflow.persistence.scheduled_tasks.sql"),
|
||||
("active", "reconcile_active_runs", "deerflow.persistence.scheduled_task_runs.sql"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("failure", "expected_confirmation"),
|
||||
[("UNKNOWN", "UNKNOWN"), ("CONFLICT", "CONFLICT"), ("exception", "UNKNOWN")],
|
||||
)
|
||||
async def test_unconfirmed_recovery_warns_without_undoing_or_repeating_takeover(recovery_db, path, recovery_path, repository_logger, failure, expected_confirmation, monkeypatch, caplog):
|
||||
db = recovery_db
|
||||
now, before = await prepare_recovery(db, path)
|
||||
committed_at_confirmation = []
|
||||
targets = []
|
||||
publish = db.publisher.publish
|
||||
|
||||
async def confirm_with_failure(store, **kwargs):
|
||||
assert store is db.events
|
||||
assert (kwargs["thread_id"], kwargs["run_id"], kwargs["user_id"], kwargs["status"]) == ("thread", "run", "alice", "error")
|
||||
assert kwargs["recovered"] is True
|
||||
assert kwargs["authorized"]()
|
||||
# Read through a separate real session at the confirmation boundary.
|
||||
# An uncommitted takeover must not satisfy these assertions.
|
||||
row = await db.runs.get("run", user_id="alice")
|
||||
assert row["status"] == "error"
|
||||
assert row["producer_state"] == "revoked"
|
||||
assert row["producer_epoch"] != before["producer_epoch"]
|
||||
assert row["run_incarnation"] == before["run_incarnation"]
|
||||
assert row["evidence_retention_revision"] == 7
|
||||
committed_at_confirmation.append(row)
|
||||
if failure == "exception":
|
||||
raise RuntimeError("injected confirmation failure")
|
||||
return terminal.TerminalEventConfirmation(failure)
|
||||
|
||||
async def observe_publish(target):
|
||||
targets.append(target)
|
||||
return await publish(target)
|
||||
|
||||
monkeypatch.setattr(terminal, "confirm_terminal_event", confirm_with_failure)
|
||||
monkeypatch.setattr(db.publisher, "publish", observe_publish)
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
assert await reconcile(db, path, now) == 1
|
||||
assert len(committed_at_confirmation) == 1
|
||||
committed = await db.runs.get("run", user_id="alice")
|
||||
assert committed == committed_at_confirmation[0]
|
||||
assert committed["error"] == "owner expired"
|
||||
assert committed["evidence_seal_state"] == "partial"
|
||||
|
||||
# Neither a later scan nor reuse of the one-shot target can publish again.
|
||||
assert await reconcile(db, path, now) == 0
|
||||
assert await db.runs.get("run", user_id="alice") == committed
|
||||
assert len(targets) == 1
|
||||
with pytest.raises(terminal.RecoveryTerminalTargetConsumed):
|
||||
await publish(targets[0])
|
||||
assert len(committed_at_confirmation) == 1
|
||||
events = await db.events.list_events("thread", "run", user_id="alice")
|
||||
assert [event["event_type"] for event in events] == ["run.delivery"]
|
||||
|
||||
warnings = [record for record in caplog.records if record.name == repository_logger and record.levelno == logging.WARNING]
|
||||
assert len(warnings) == 1
|
||||
warning = warnings[0]
|
||||
assert warning.run_id == "run"
|
||||
assert warning.recovery_path == recovery_path
|
||||
assert warning.confirmation == expected_confirmation
|
||||
273
backend/tests/test_scheduled_recovery_fencing.py
Normal file
273
backend/tests/test_scheduled_recovery_fencing.py
Normal file
@ -0,0 +1,273 @@
|
||||
"""Scheduled recovery must publish using the identity of its committed takeover."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import text, update
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
|
||||
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunProducerPolicyRow, RunRow
|
||||
from deerflow.persistence.run.sql import RunRepository
|
||||
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
|
||||
from deerflow.persistence.scheduled_task_runs.sql import ScheduledTaskRunRepository
|
||||
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
|
||||
from deerflow.persistence.scheduled_tasks.sql import ScheduledTaskRepository
|
||||
from deerflow.runtime.events.store.db import DbRunEventStore
|
||||
from deerflow.runtime.events.terminal import RecoveryTerminalPublisher, RecoveryTerminalTargetConsumed, TerminalEventConfirmation
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(params=["off", "strict"])
|
||||
async def recovery_db(request):
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI")
|
||||
schema = "scheduled_recovery_" + uuid4().hex
|
||||
engine = create_async_engine(asyncpg_test_url(uri), pool_size=2, max_overflow=0, pool_timeout=2, connect_args={"server_settings": {"search_path": schema, "statement_timeout": "8000"}})
|
||||
sf = async_sessionmaker(engine, expire_on_commit=False)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
||||
await conn.run_sync(
|
||||
lambda c: RunRow.metadata.create_all(
|
||||
c,
|
||||
tables=[
|
||||
RunRow.__table__,
|
||||
RunChangeClockRow.__table__,
|
||||
CompletedRunSnapshotRow.__table__,
|
||||
RunProducerPolicyRow.__table__,
|
||||
RunEventRow.__table__,
|
||||
RunEventThreadSeqRow.__table__,
|
||||
ScheduledTaskRow.__table__,
|
||||
ScheduledTaskRunRow.__table__,
|
||||
],
|
||||
)
|
||||
)
|
||||
await conn.execute(text("INSERT INTO run_producer_policy(id,strict) VALUES (1,:strict)"), {"strict": request.param == "strict"})
|
||||
assert await conn.scalar(text("SELECT count(*) FROM pg_trigger WHERE tgrelid='run_events'::regclass AND tgname='run_event_producer_guard'")) == 1
|
||||
runs = RunRepository(sf)
|
||||
events = DbRunEventStore(sf, producer_fencing=request.param)
|
||||
publisher = RecoveryTerminalPublisher(events, runs)
|
||||
yield SimpleNamespace(
|
||||
url=asyncpg_test_url(uri),
|
||||
schema=schema,
|
||||
mode=request.param,
|
||||
sf=sf,
|
||||
runs=runs,
|
||||
events=events,
|
||||
publisher=publisher,
|
||||
tasks=ScheduledTaskRepository(sf, run_repository=runs, recovery_terminal_publisher=publisher),
|
||||
occurrences=ScheduledTaskRunRepository(sf, run_repository=runs, recovery_terminal_publisher=publisher),
|
||||
)
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def prepare_recovery(db, path):
|
||||
now = datetime.now(UTC)
|
||||
await db.tasks.create(
|
||||
task_id="task",
|
||||
user_id="alice",
|
||||
thread_id=None,
|
||||
context_mode="fresh_thread_per_run",
|
||||
assistant_id="lead_agent",
|
||||
title="recover",
|
||||
prompt="p",
|
||||
schedule_type="once" if path == "once" else "cron",
|
||||
schedule_spec={"run_at": now.isoformat()} if path == "once" else {"cron": "* * * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
if path == "once":
|
||||
await db.tasks.update("task", user_id="alice", updates={"status": "running", "last_run_id": "run"})
|
||||
else:
|
||||
await db.occurrences.create(run_record_id="occurrence", task_id="task", thread_id="thread", scheduled_for=now, trigger="scheduled", status="running")
|
||||
await db.occurrences.update_status("occurrence", status="running", run_id="run")
|
||||
await db.runs.put(
|
||||
"run",
|
||||
thread_id="thread",
|
||||
user_id="alice",
|
||||
status="running",
|
||||
owner_worker_id="dead-owner",
|
||||
lease_expires_at=(now - timedelta(seconds=60)).isoformat(),
|
||||
metadata={"scheduled_task_id": "task", "scheduled_task_run_id": "occurrence"} if path == "active" else {},
|
||||
)
|
||||
# Use a non-default revision so omitting it is independently detectable.
|
||||
async with db.sf() as session, session.begin():
|
||||
row = await session.get(RunRow, "run")
|
||||
row.evidence_retention_revision = 7
|
||||
before = await db.runs.get("run", user_id="alice")
|
||||
assert before["run_incarnation"] and before["producer_epoch"]
|
||||
return now, before
|
||||
|
||||
|
||||
async def reconcile(db, path, now):
|
||||
method = db.tasks.reconcile_stuck_once_tasks if path == "once" else db.occurrences.reconcile_active_runs
|
||||
return await method(error="owner expired", now=now, lease_grace_seconds=0)
|
||||
|
||||
|
||||
async def assert_recovered_events(db):
|
||||
rows = await db.events.list_events("thread", "run", user_id=None)
|
||||
assert [row["event_type"] for row in rows] == ["run.delivery", "run.terminal"]
|
||||
assert all((row["thread_id"], row["run_id"], row["user_id"]) == ("thread", "run", "alice") for row in rows)
|
||||
assert rows[0]["seq"] < rows[1]["seq"]
|
||||
assert rows[0]["content"] == {"presented": 0, "paths": [], "by_tool": {}}
|
||||
assert rows[1]["content"] == {"status": "error"}
|
||||
assert rows[1]["metadata"]["protocol_version"] == 1
|
||||
assert rows[1]["metadata"]["recovered"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["once", "active"])
|
||||
async def test_scheduled_takeover_publishes_matching_terminal_and_delivery(recovery_db, path):
|
||||
db = recovery_db
|
||||
now, before = await prepare_recovery(db, path)
|
||||
assert await reconcile(db, path, now) == 1
|
||||
row = await db.runs.get("run", user_id="alice")
|
||||
assert row["status"] == "error"
|
||||
assert row["producer_state"] == "revoked"
|
||||
assert row["producer_epoch"] != before["producer_epoch"]
|
||||
assert row["run_incarnation"] == before["run_incarnation"]
|
||||
assert row["evidence_retention_revision"] == 7
|
||||
assert row["evidence_seal_state"] == "partial"
|
||||
await assert_recovered_events(db)
|
||||
assert await reconcile(db, path, now) == 0
|
||||
await assert_recovered_events(db)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["once", "active"])
|
||||
async def test_competing_reconcilers_have_one_takeover_and_publisher(recovery_db, path, monkeypatch):
|
||||
db = recovery_db
|
||||
now, _ = await prepare_recovery(db, path)
|
||||
peer_engine = create_async_engine(db.url, pool_size=2, max_overflow=0, pool_timeout=2, connect_args={"server_settings": {"search_path": db.schema, "statement_timeout": "8000"}})
|
||||
peer_sf = async_sessionmaker(peer_engine, expire_on_commit=False)
|
||||
peer_runs = RunRepository(peer_sf)
|
||||
peer_publisher = RecoveryTerminalPublisher(DbRunEventStore(peer_sf, producer_fencing=db.mode), peer_runs)
|
||||
peer = SimpleNamespace(
|
||||
tasks=ScheduledTaskRepository(peer_sf, run_repository=peer_runs, recovery_terminal_publisher=peer_publisher),
|
||||
occurrences=ScheduledTaskRunRepository(peer_sf, run_repository=peer_runs, recovery_terminal_publisher=peer_publisher),
|
||||
)
|
||||
claims = []
|
||||
for repo in (db.runs, peer_runs):
|
||||
original = repo.claim_for_takeover
|
||||
|
||||
async def record_claim(*args, _original=original, **kwargs):
|
||||
claimed = await _original(*args, **kwargs)
|
||||
claims.append(claimed)
|
||||
return claimed
|
||||
|
||||
monkeypatch.setattr(repo, "claim_for_takeover", record_claim)
|
||||
try:
|
||||
outcomes = await asyncio.wait_for(asyncio.gather(reconcile(db, path, now), reconcile(peer, path, now)), timeout=10)
|
||||
assert sorted(outcomes) == [0, 1]
|
||||
assert claims.count(True) == 1
|
||||
await assert_recovered_events(db)
|
||||
finally:
|
||||
await peer_engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["once", "active"])
|
||||
@pytest.mark.parametrize("change", ["delete_run", "delete_events", "recreate", "revision"])
|
||||
async def test_scheduled_publication_rejects_changed_lifecycle(recovery_db, path, change, monkeypatch):
|
||||
db = recovery_db
|
||||
now, before = await prepare_recovery(db, path)
|
||||
original = db.publisher.publish
|
||||
confirmations = []
|
||||
|
||||
async def mutate_then_publish(target):
|
||||
row = await db.runs.get("run", user_id="alice")
|
||||
assert row["status"] == "error" and row["producer_state"] == "revoked"
|
||||
if change in ("delete_run", "recreate"):
|
||||
await db.runs.delete("run", user_id="alice")
|
||||
if change == "recreate":
|
||||
await db.runs.put("run", thread_id="thread", user_id="alice", status="error")
|
||||
# Match every field except incarnation: a reused ID, status and
|
||||
# retention revision must not revive the old publication grant.
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(update(RunRow).where(RunRow.run_id == "run").values(producer_state="revoked", evidence_retention_revision=7))
|
||||
assert (await db.runs.get("run", user_id="alice"))["run_incarnation"] != before["run_incarnation"]
|
||||
elif change == "delete_events":
|
||||
await db.events.delete_by_run("thread", "run", user_id="alice")
|
||||
else:
|
||||
async with db.sf() as session, session.begin():
|
||||
await session.execute(update(RunRow).where(RunRow.run_id == "run").values(evidence_retention_revision=8))
|
||||
confirmation = await original(target)
|
||||
confirmations.append(confirmation)
|
||||
with pytest.raises(RecoveryTerminalTargetConsumed):
|
||||
await original(target)
|
||||
return confirmation
|
||||
|
||||
monkeypatch.setattr(db.publisher, "publish", mutate_then_publish)
|
||||
assert await reconcile(db, path, now) == 1
|
||||
assert len(confirmations) == 1 and not confirmations[0].confirmed
|
||||
assert await db.events.list_events("thread", "run", user_id=None) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("missing", ["incarnation", "revision", "unclaimed"])
|
||||
async def test_strict_takeover_identity_contract(recovery_db, missing):
|
||||
db = recovery_db
|
||||
if db.mode != "strict":
|
||||
pytest.skip("missing identity is only forbidden by the strict protocol")
|
||||
_, before = await prepare_recovery(db, "active")
|
||||
claimed = await db.runs.claim_for_takeover("run", grace_seconds=0, error="expired")
|
||||
assert claimed
|
||||
target = db.publisher.authorize_takeover(
|
||||
claimed=claimed if missing != "unclaimed" else False,
|
||||
thread_id="thread",
|
||||
run_id="run",
|
||||
user_id="alice",
|
||||
operation_kind="run",
|
||||
run_incarnation=before["run_incarnation"] if missing != "incarnation" else None,
|
||||
retention_revision=7 if missing != "revision" else 0,
|
||||
)
|
||||
if missing == "unclaimed":
|
||||
assert target is None
|
||||
else:
|
||||
assert target is not None
|
||||
assert await db.publisher.publish(target) == TerminalEventConfirmation.UNKNOWN
|
||||
with pytest.raises(RecoveryTerminalTargetConsumed):
|
||||
await db.publisher.publish(target)
|
||||
assert await db.events.list_events("thread", "run", user_id=None) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["once", "active"])
|
||||
async def test_scheduled_recovery_lost_ack_confirms_existing_singleton(recovery_db, path, monkeypatch):
|
||||
db = recovery_db
|
||||
now, _ = await prepare_recovery(db, path)
|
||||
original_put = db.events.put_if_absent
|
||||
attempts = []
|
||||
|
||||
async def lose_terminal_ack(**kwargs):
|
||||
result = await original_put(**kwargs)
|
||||
if kwargs["event_type"] == "run.terminal":
|
||||
attempts.append(result[1])
|
||||
if len(attempts) == 1:
|
||||
raise OSError("terminal ACK lost after commit")
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(db.events, "put_if_absent", lose_terminal_ack)
|
||||
original_publish = db.publisher.publish
|
||||
confirmations = []
|
||||
|
||||
async def record_confirmation(target):
|
||||
result = await original_publish(target)
|
||||
confirmations.append(result)
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(db.publisher, "publish", record_confirmation)
|
||||
assert await reconcile(db, path, now) == 1
|
||||
assert confirmations == [TerminalEventConfirmation.CONFIRMED_EXISTING]
|
||||
assert attempts == [True, False]
|
||||
await assert_recovered_events(db)
|
||||
@ -8,7 +8,7 @@ from support.postgres import asyncpg_test_url
|
||||
@pytest.mark.parametrize("scheme", ["postgres", "postgresql", "postgresql+asyncpg"])
|
||||
@pytest.mark.parametrize("sslmode", [None, "disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
|
||||
def test_asyncpg_test_url_preserves_tls_credentials_and_other_options(scheme, sslmode):
|
||||
uri = f"{scheme}://user:p%40ss%2Fword@localhost:5432/deerflow_test?target_session_attrs=read-write&channel_binding=prefer"
|
||||
uri = f"{scheme}://user:p%40ss%2Fword@localhost:5432/deerflow_test?target_session_attrs=read-write"
|
||||
if sslmode is not None:
|
||||
uri += f"&sslmode={sslmode}"
|
||||
url = make_url(asyncpg_test_url(uri))
|
||||
@ -33,3 +33,9 @@ def test_asyncpg_test_url_preserves_explicit_ssl(query):
|
||||
def test_asyncpg_test_url_rejects_conflicting_tls_options():
|
||||
with pytest.raises(ValueError, match="Conflicting ssl and sslmode"):
|
||||
asyncpg_test_url("postgresql://user:password@localhost/test?ssl=require&sslmode=disable")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["require", "prefer", "disable"])
|
||||
def test_asyncpg_test_url_rejects_unsupported_channel_binding(mode):
|
||||
with pytest.raises(ValueError, match="channel_binding"):
|
||||
asyncpg_test_url(f"postgresql://user:password@localhost/test?channel_binding={mode}")
|
||||
|
||||
@ -2790,6 +2790,16 @@ run_events:
|
||||
backend: memory
|
||||
max_trace_content: 10240
|
||||
track_token_usage: true
|
||||
# With a fencing-capable SQL RunStore and DB event store:
|
||||
# off: host-authorized writes still check Run incarnation, epoch and state;
|
||||
# stale Journal flushes after takeover or deletion may be rejected.
|
||||
# strict: additionally checks raw SQL and unaware legacy writers via the
|
||||
# database trigger after an offline cutover.
|
||||
# Restart-required; strict requires database.backend: postgres and backend: db.
|
||||
# Stop every Gateway, activate the database policy, then restart all with the
|
||||
# matching cutover ID. See README.md#configuration for the offline cutover.
|
||||
producer_fencing: "off"
|
||||
# producer_cutover_id: deployment-2026-10-27
|
||||
|
||||
# ============================================================================
|
||||
# Agent Storage Configuration
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user