mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-31 07:30:43 +00:00
33f7b43400
* fix(runtime): backfill run.terminal on confirmed recovery takeover * fix(runtime): isolate recovery terminal publication failures * fix(runtime): fence recovery terminal publication lifecycle * fix(runtime): unify recovery receipt publication * docs(runtime): document lifecycle guard lock constraints * fix(runtime): reuse event transaction for recovery guards Read Run lifecycle through the active event session when SQL stores share an engine, preserving deletion fencing without a second pool checkout. Keep custom lookups and different engines on their ordinary read path. Scope memory recovery guards and deletions to per-thread locks. Add SQLite/PostgreSQL pool pressure, stale identity, fallback, and independent-thread regressions, and document the contracts. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
1152 lines
50 KiB
Python
1152 lines
50 KiB
Python
"""SQLAlchemy-backed RunStore implementation.
|
|
|
|
Each method acquires and releases its own short-lived session.
|
|
Run status updates happen from background workers that may live
|
|
minutes -- we don't hold connections across long execution.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import and_, case, delete, func, literal, or_, select, text, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.run.model import CompletedRunSnapshotRow, RunChangeClockRow, RunRow
|
|
from deerflow.runtime.run_origin import origin_kind_of
|
|
from deerflow.runtime.runs.store.base import (
|
|
LeaseRenewal,
|
|
RunIdempotencyConflict,
|
|
RunStore,
|
|
StatusFinalization,
|
|
canonical_run_created_at,
|
|
normalize_run_created_at_iso,
|
|
)
|
|
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id
|
|
from deerflow.utils.time import coerce_iso
|
|
|
|
|
|
def _lease_expired_or_null(lease_col, cutoff: datetime):
|
|
"""SQLAlchemy filter: True when the lease is NULL or has expired past *cutoff*."""
|
|
return or_(lease_col.is_(None), lease_col < cutoff)
|
|
|
|
|
|
def _origin_kind(operation_kind: str, metadata: Any) -> str | None:
|
|
"""Denormalized ``runs.origin_kind``: the server-owned origin of an agent run.
|
|
|
|
Thread operations (checkpoint writes, deletes) never carry one, and a run
|
|
without a well-formed ``deerflow_origin`` stays NULL (interactive/legacy).
|
|
"""
|
|
return origin_kind_of(metadata) if operation_kind == "run" else None
|
|
|
|
|
|
class RunRepository(RunStore):
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
@staticmethod
|
|
async def _next_change_seq(session: AsyncSession) -> int:
|
|
dialect = session.bind.dialect.name
|
|
if dialect == "postgresql":
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
elif dialect == "sqlite":
|
|
from sqlalchemy.dialects.sqlite import insert
|
|
else: # pragma: no cover - configured databases are SQLite/Postgres
|
|
raise RuntimeError(f"unsupported run-change clock dialect: {dialect}")
|
|
await session.execute(insert(RunChangeClockRow).values(id=1, value=0).on_conflict_do_nothing(index_elements=[RunChangeClockRow.id]))
|
|
value = await session.scalar(update(RunChangeClockRow).where(RunChangeClockRow.id == 1).values(value=RunChangeClockRow.value + 1).returning(RunChangeClockRow.value))
|
|
if value is None:
|
|
raise RuntimeError("run-change clock did not return a position")
|
|
return int(value)
|
|
|
|
@staticmethod
|
|
def _live_lease_predicate(session: AsyncSession, lease=RunRow.lease_expires_at):
|
|
# PostgreSQL clock_timestamp is evaluated after the row lock below,
|
|
# unlike transaction-start now(). SQLite evaluates its clock at UPDATE.
|
|
if session.bind.dialect.name == "postgresql":
|
|
return lease > func.clock_timestamp()
|
|
return func.julianday(lease) > func.julianday("now")
|
|
|
|
@staticmethod
|
|
async def _lock_owned_row(session: AsyncSession, run_id: str) -> None:
|
|
if session.bind.dialect.name == "postgresql":
|
|
await session.execute(select(RunRow.run_id).where(RunRow.run_id == run_id).with_for_update())
|
|
|
|
@staticmethod
|
|
def _normalize_model_name(model_name: str | None) -> str | None:
|
|
"""Normalize model_name for storage: strip whitespace, truncate to 128 chars."""
|
|
if model_name is None:
|
|
return None
|
|
if not isinstance(model_name, str):
|
|
model_name = str(model_name)
|
|
normalized = model_name.strip()
|
|
if len(normalized) > 128:
|
|
normalized = normalized[:128]
|
|
return normalized
|
|
|
|
@staticmethod
|
|
def _safe_json(obj: Any) -> Any:
|
|
"""Ensure obj is JSON-serializable. Falls back to model_dump() or str()."""
|
|
if obj is None:
|
|
return None
|
|
if isinstance(obj, (str, int, float, bool)):
|
|
return obj
|
|
if isinstance(obj, dict):
|
|
return {k: RunRepository._safe_json(v) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [RunRepository._safe_json(v) for v in obj]
|
|
if hasattr(obj, "model_dump"):
|
|
try:
|
|
return obj.model_dump()
|
|
except Exception:
|
|
pass
|
|
if hasattr(obj, "dict"):
|
|
try:
|
|
return obj.dict()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
json.dumps(obj)
|
|
return obj
|
|
except (TypeError, ValueError):
|
|
return str(obj)
|
|
|
|
@staticmethod
|
|
def _row_to_dict(row: RunRow) -> dict[str, Any]:
|
|
d = row.to_dict()
|
|
# Remap JSON columns to match RunStore interface
|
|
d["metadata"] = d.pop("metadata_json", {})
|
|
d["kwargs"] = d.pop("kwargs_json", {})
|
|
d["idempotency_request"] = d.pop("idempotency_request_json", None)
|
|
# Convert datetime to ISO string for consistency with MemoryRunStore.
|
|
# SQLite drops tzinfo on read despite ``DateTime(timezone=True)`` —
|
|
# ``coerce_iso`` normalizes naive datetimes as UTC.
|
|
for key in ("created_at", "updated_at", "lease_expires_at", "cancel_requested_at"):
|
|
val = d.get(key)
|
|
if isinstance(val, datetime):
|
|
d[key] = coerce_iso(val)
|
|
return d
|
|
|
|
async def put(
|
|
self,
|
|
run_id,
|
|
*,
|
|
thread_id,
|
|
assistant_id=None,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
model_name: str | None = None,
|
|
status="pending",
|
|
operation_kind: str = "run",
|
|
multitask_strategy="reject",
|
|
metadata=None,
|
|
kwargs=None,
|
|
error=None,
|
|
stop_reason: str | None = None,
|
|
goal_verdict: dict[str, Any] | None = None,
|
|
created_at=None,
|
|
follow_up_to_run_id=None,
|
|
owner_worker_id: str | None = None,
|
|
lease_expires_at: str | None = None,
|
|
idempotency_key: str | None = None,
|
|
evidence_origin: str = "unknown",
|
|
evidence_agent_id: str | None = None,
|
|
idempotency_request: dict[str, Any] | None = None,
|
|
):
|
|
"""Insert or update a run row.
|
|
|
|
``RunManager`` retries ``put`` after transient SQLite failures. Making
|
|
this operation idempotent prevents a successful-but-unacknowledged first
|
|
commit from turning the retry into a primary-key failure.
|
|
"""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.put")
|
|
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
|
|
values = {
|
|
"thread_id": thread_id,
|
|
"evidence_origin": evidence_origin,
|
|
"evidence_agent_id": evidence_agent_id,
|
|
"assistant_id": assistant_id,
|
|
"user_id": resolved_user_id,
|
|
"model_name": self._normalize_model_name(model_name),
|
|
"status": status,
|
|
"operation_kind": operation_kind,
|
|
"multitask_strategy": multitask_strategy,
|
|
"metadata_json": self._safe_json(metadata) or {},
|
|
"kwargs_json": self._safe_json(kwargs) or {},
|
|
"error": error,
|
|
"stop_reason": stop_reason,
|
|
"goal_verdict": self._safe_json(goal_verdict),
|
|
"follow_up_to_run_id": follow_up_to_run_id,
|
|
"owner_worker_id": owner_worker_id,
|
|
"lease_expires_at": lease_dt,
|
|
"idempotency_key": idempotency_key,
|
|
"idempotency_request_json": self._safe_json(idempotency_request),
|
|
"origin_kind": _origin_kind(operation_kind, metadata),
|
|
"updated_at": now,
|
|
}
|
|
async with self._sf() as session:
|
|
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))
|
|
else:
|
|
for key, value in values.items():
|
|
setattr(row, key, value)
|
|
await session.commit()
|
|
|
|
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."""
|
|
from deerflow.runtime.events.evidence import JournalSealReceipt
|
|
|
|
if not isinstance(receipt, JournalSealReceipt) or receipt.session_factory is not self._sf or receipt.run_id != run_id:
|
|
return False
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.thread_id == receipt.thread_id,
|
|
RunRow.owner_worker_id == owner_worker_id,
|
|
or_(RunRow.lease_expires_at.is_(None), RunRow.lease_expires_at > datetime.now(UTC)),
|
|
RunRow.evidence_retention_revision == receipt.retention_revision,
|
|
RunRow.status.in_(("success", "error", "timeout", "interrupted")),
|
|
RunRow.evidence_seal_state == "pending",
|
|
)
|
|
.values(evidence_seal_state="sealed", evidence_revision=uuid4().hex, evidence_upper_seq=receipt.upper_event_seq, evidence_event_count=receipt.event_count, change_seq=change_seq)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def mark_evidence_partial(self, run_id: str, *, owner_worker_id: str, error: str) -> bool:
|
|
"""Make a failed seal discoverable without changing business outcome."""
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.owner_worker_id == owner_worker_id,
|
|
or_(RunRow.lease_expires_at.is_(None), RunRow.lease_expires_at > datetime.now(UTC)),
|
|
RunRow.evidence_seal_state == "pending",
|
|
)
|
|
.values(evidence_seal_state="partial", evidence_seal_error=error[:512], evidence_revision=uuid4().hex, change_seq=change_seq)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def get(
|
|
self,
|
|
run_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.get")
|
|
async with self._sf() as session:
|
|
row = await session.get(RunRow, run_id)
|
|
if row is None:
|
|
return None
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return None
|
|
return self._row_to_dict(row)
|
|
|
|
async def get_in_session(
|
|
self,
|
|
session: AsyncSession,
|
|
run_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
"""Read through a compatible caller transaction without another checkout.
|
|
|
|
Keep custom ``get`` implementations and differently bound databases
|
|
on their normal read path. The caller owns this session's lifetime.
|
|
"""
|
|
if type(self).get is not RunRepository.get or session.bind is not self._sf.kw.get("bind"):
|
|
return await self.get(run_id, user_id=user_id)
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.get")
|
|
row = await session.get(RunRow, run_id)
|
|
if row is None or (resolved_user_id is not None and row.user_id != resolved_user_id):
|
|
return None
|
|
return self._row_to_dict(row)
|
|
|
|
async def list_changed(
|
|
self,
|
|
*,
|
|
after_change_seq: int,
|
|
after_run_id: str,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
limit: int = 100,
|
|
) -> list[dict[str, Any]]:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_changed")
|
|
stmt = select(RunRow).where(
|
|
RunRow.operation_kind == "run",
|
|
or_(
|
|
RunRow.change_seq > after_change_seq,
|
|
and_(RunRow.change_seq == after_change_seq, RunRow.run_id > after_run_id),
|
|
),
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(RunRow.change_seq.asc(), RunRow.run_id.asc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(row) for row in result.scalars()]
|
|
|
|
async def latest_change(self, *, user_id: str) -> tuple[int, str] | None:
|
|
"""The caller's newest run-change position ``(change_seq, run_id)``, or ``None``.
|
|
|
|
Seeds the activity feed cursor: one backward seek on
|
|
``ix_runs_user_change_seq``. ``user_id`` is always an explicit owner id.
|
|
"""
|
|
stmt = select(RunRow.change_seq, RunRow.run_id).where(RunRow.user_id == user_id).order_by(RunRow.change_seq.desc(), RunRow.run_id.desc()).limit(1)
|
|
async with self._sf() as session:
|
|
row = (await session.execute(stmt)).first()
|
|
return (int(row.change_seq), str(row.run_id)) if row is not None else None
|
|
|
|
async def list_by_thread(
|
|
self,
|
|
thread_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
limit=100,
|
|
before_created_at: str | None = None,
|
|
before_run_id: str | None = None,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_by_thread")
|
|
stmt = select(RunRow).where(RunRow.thread_id == thread_id, RunRow.operation_kind == "run")
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
if before_created_at and before_run_id:
|
|
cursor_dt = datetime.fromisoformat(normalize_run_created_at_iso(before_created_at))
|
|
if cursor_dt.tzinfo is None:
|
|
cursor_dt = cursor_dt.replace(tzinfo=UTC)
|
|
else:
|
|
cursor_dt = cursor_dt.astimezone(UTC)
|
|
stmt = stmt.where(
|
|
or_(
|
|
RunRow.created_at < cursor_dt,
|
|
and_(RunRow.created_at == cursor_dt, RunRow.run_id < before_run_id),
|
|
)
|
|
)
|
|
# Keyset pages filter on (created_at, run_id) after thread_id. Existing
|
|
# indexes are (thread_id) and (thread_id, status), so each page still
|
|
# sorts matching rows. A covering (thread_id, created_at, run_id) index
|
|
# is a follow-up if deep paging shows up in profiles.
|
|
stmt = stmt.order_by(RunRow.created_at.desc(), RunRow.run_id.desc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def list_by_thread_created_at(self, thread_id, *, user_id, created_at):
|
|
timestamp = datetime.fromisoformat(canonical_run_created_at(created_at))
|
|
async with self._sf() as session:
|
|
result = await session.execute(select(RunRow).where(RunRow.thread_id == thread_id, RunRow.user_id == user_id, RunRow.created_at == timestamp))
|
|
return [self._row_to_dict(row) for row in result.scalars()]
|
|
|
|
async def list_successful_regenerate_sources(
|
|
self,
|
|
thread_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_successful_regenerate_sources")
|
|
source = RunRow.metadata_json["regenerate_from_run_id"].as_string()
|
|
stmt = select(source).where(
|
|
RunRow.thread_id == thread_id,
|
|
RunRow.operation_kind == "run",
|
|
RunRow.status == "success",
|
|
source.is_not(None),
|
|
source != "",
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return {value for value in result.scalars() if isinstance(value, str) and value}
|
|
|
|
async def list_edit_regenerate_runs(
|
|
self,
|
|
thread_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_edit_regenerate_runs")
|
|
replay_kind = RunRow.metadata_json["replay_kind"].as_string()
|
|
source = RunRow.metadata_json["regenerate_from_run_id"].as_string()
|
|
stmt = select(RunRow).where(
|
|
RunRow.thread_id == thread_id,
|
|
replay_kind == "edit",
|
|
source.is_not(None),
|
|
source != "",
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(RunRow.created_at.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(row) for row in result.scalars()]
|
|
|
|
async def get_many_by_thread(
|
|
self,
|
|
thread_id,
|
|
run_ids,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
if not run_ids:
|
|
return {}
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.get_many_by_thread")
|
|
stmt = select(RunRow).where(RunRow.thread_id == thread_id, RunRow.operation_kind == "run", RunRow.run_id.in_(run_ids))
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return {row.run_id: self._row_to_dict(row) for row in result.scalars()}
|
|
|
|
async def update_status(self, run_id, status, *, error=None, stop_reason=None, goal_verdict=None, owner_worker_id=None) -> bool:
|
|
values: dict[str, Any] = {"status": status, "updated_at": datetime.now(UTC)}
|
|
if error is not None:
|
|
values["error"] = error
|
|
if stop_reason is not None:
|
|
values["stop_reason"] = stop_reason
|
|
if goal_verdict is not None:
|
|
values["goal_verdict"] = self._safe_json(goal_verdict)
|
|
# Guard: only transition rows that are still active. ``interrupted`` is
|
|
# included because the rollback path goes ``running → interrupted``
|
|
# (cancel acknowledged) then ``interrupted → error`` (task finalize).
|
|
# ``error`` and ``success`` remain locked so a peer's takeover (or a
|
|
# completed run) cannot be overwritten by a late writer.
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
conditions = [RunRow.run_id == run_id, RunRow.status.in_(("pending", "running", "interrupted"))]
|
|
if owner_worker_id is not None:
|
|
await self._lock_owned_row(session, run_id)
|
|
conditions.extend([RunRow.owner_worker_id == owner_worker_id, self._live_lease_predicate(session)])
|
|
if status != "error":
|
|
conditions.append(RunRow.status.in_(("pending", "running")))
|
|
if status == "success":
|
|
conditions.append(RunRow.cancel_action.is_(None))
|
|
result = await session.execute(update(RunRow).where(*conditions).values(**values))
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def start_run(self, run_id: str) -> bool:
|
|
"""Start only a still-pending run; cancelled rows must not be resurrected."""
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status == "pending",
|
|
)
|
|
.values(status="running", updated_at=datetime.now(UTC), change_seq=change_seq)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def update_model_name(self, run_id, model_name):
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
await session.execute(
|
|
update(RunRow)
|
|
.where(RunRow.run_id == run_id)
|
|
.values(
|
|
model_name=self._normalize_model_name(model_name),
|
|
updated_at=datetime.now(UTC),
|
|
change_seq=change_seq,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
async def delete(
|
|
self,
|
|
run_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.delete")
|
|
async with self._sf() as session:
|
|
if session.bind.dialect.name == "sqlite":
|
|
await session.execute(text("BEGIN IMMEDIATE"))
|
|
row = await session.get(RunRow, run_id, with_for_update=True)
|
|
if row is None:
|
|
return
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return
|
|
await session.execute(delete(CompletedRunSnapshotRow).where(CompletedRunSnapshotRow.run_id == run_id))
|
|
await session.delete(row)
|
|
await session.commit()
|
|
|
|
async def delete_by_thread(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> int:
|
|
"""Delete a thread's historical runs, keeping internal operation rows.
|
|
|
|
Only ``operation_kind == "run"`` rows are removed. The same table holds
|
|
durable thread-operation reservations (checkpoint writes, artifact
|
|
writes, branches and thread deletions) that their own release path owns
|
|
through ``delete_thread_operation``; deleting the reservation currently
|
|
protecting a thread deletion would drop the cross-worker exclusion in
|
|
the middle of that request.
|
|
|
|
The bulk delete deliberately does not bump the run-change clock, matching
|
|
the single-row :meth:`delete` and keeping thread cleanup away from
|
|
``run_change_clock`` (#5516). ``user_id`` follows the same three-state
|
|
convention as the rest of this repository: ``AUTO`` reads the request
|
|
context (raising when there is none), an explicit id scopes the delete,
|
|
and ``None`` skips the owner filter for migration/CLI callers.
|
|
"""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.delete_by_thread")
|
|
|
|
conditions = [
|
|
RunRow.thread_id == thread_id,
|
|
RunRow.operation_kind == "run",
|
|
]
|
|
if resolved_user_id is not None:
|
|
conditions.append(RunRow.user_id == resolved_user_id)
|
|
|
|
async with self._sf() as session:
|
|
if session.bind.dialect.name == "sqlite":
|
|
await session.execute(text("BEGIN IMMEDIATE"))
|
|
# The run-row lock serializes snapshot creation and cleanup on PG;
|
|
# SQLite's immediate transaction supplies the equivalent fence.
|
|
# Acquire locks in a separate statement: READ COMMITTED must take a
|
|
# fresh snapshot after waiting for an in-flight snapshot creator.
|
|
await session.execute(select(RunRow.run_id).where(*conditions).with_for_update())
|
|
await session.execute(delete(CompletedRunSnapshotRow).where(CompletedRunSnapshotRow.run_id.in_(select(RunRow.run_id).where(*conditions))))
|
|
count = await session.scalar(select(func.count()).select_from(RunRow).where(*conditions)) or 0
|
|
if count:
|
|
await session.execute(delete(RunRow).where(*conditions))
|
|
await session.commit()
|
|
return count
|
|
|
|
async def delete_thread_operation(self, run_id: str, *, user_id: str | None) -> None:
|
|
"""Release a reservation using its captured owner, not request context."""
|
|
await self.delete(run_id, user_id=user_id)
|
|
|
|
async def list_pending(self, *, before=None):
|
|
if before is None:
|
|
before_dt = datetime.now(UTC)
|
|
elif isinstance(before, datetime):
|
|
before_dt = before
|
|
else:
|
|
before_dt = datetime.fromisoformat(before)
|
|
stmt = select(RunRow).where(RunRow.operation_kind == "run", RunRow.status == "pending", RunRow.created_at <= before_dt).order_by(RunRow.created_at.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def list_inflight(self, *, before=None):
|
|
"""Return persisted active runs for startup recovery."""
|
|
if before is None:
|
|
before_dt = datetime.now(UTC)
|
|
elif isinstance(before, datetime):
|
|
before_dt = before
|
|
else:
|
|
before_dt = datetime.fromisoformat(before)
|
|
stmt = (
|
|
select(RunRow)
|
|
.where(
|
|
RunRow.status.in_(("pending", "running")),
|
|
RunRow.created_at <= before_dt,
|
|
)
|
|
.order_by(RunRow.created_at.asc())
|
|
)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def update_run_completion(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
status: str,
|
|
require_terminal: bool = False,
|
|
owner_worker_id: str | None = None,
|
|
total_input_tokens: int = 0,
|
|
total_output_tokens: int = 0,
|
|
total_tokens: int = 0,
|
|
llm_call_count: int = 0,
|
|
lead_agent_tokens: int = 0,
|
|
subagent_tokens: int = 0,
|
|
middleware_tokens: int = 0,
|
|
token_usage_by_model: dict[str, dict[str, int]] | None = None,
|
|
message_count: int = 0,
|
|
last_ai_message: str | None = None,
|
|
first_human_message: str | None = None,
|
|
error: str | None = None,
|
|
goal_verdict: dict[str, Any] | None = None,
|
|
) -> bool:
|
|
"""Update status + token usage + convenience fields on run completion.
|
|
|
|
Returns ``False`` when the row is missing or already has a conflicting
|
|
terminal outcome.
|
|
"""
|
|
values: dict[str, Any] = {
|
|
"status": status,
|
|
"total_input_tokens": total_input_tokens,
|
|
"total_output_tokens": total_output_tokens,
|
|
"total_tokens": total_tokens,
|
|
"llm_call_count": llm_call_count,
|
|
"lead_agent_tokens": lead_agent_tokens,
|
|
"subagent_tokens": subagent_tokens,
|
|
"middleware_tokens": middleware_tokens,
|
|
"token_usage_by_model": self._safe_json(token_usage_by_model) or {},
|
|
"message_count": message_count,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
if last_ai_message is not None:
|
|
values["last_ai_message"] = last_ai_message[:2000]
|
|
if first_human_message is not None:
|
|
values["first_human_message"] = first_human_message[:2000]
|
|
if not require_terminal and error is not None:
|
|
values["error"] = error
|
|
if not require_terminal and goal_verdict is not None:
|
|
values["goal_verdict"] = self._safe_json(goal_verdict)
|
|
allowed_sources = ["pending", "running"]
|
|
if status not in allowed_sources:
|
|
allowed_sources.append(status)
|
|
if status == "error" and "interrupted" not in allowed_sources:
|
|
allowed_sources.append("interrupted")
|
|
if require_terminal and status in ("pending", "running"):
|
|
return False
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
conditions = [RunRow.run_id == run_id]
|
|
if require_terminal:
|
|
conditions.extend([RunRow.status == status, RunRow.owner_worker_id == owner_worker_id])
|
|
else:
|
|
conditions.append(RunRow.status.in_(tuple(allowed_sources)))
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
*conditions,
|
|
)
|
|
.values(**values)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def update_run_progress(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
total_input_tokens: int | None = None,
|
|
total_output_tokens: int | None = None,
|
|
total_tokens: int | None = None,
|
|
llm_call_count: int | None = None,
|
|
lead_agent_tokens: int | None = None,
|
|
subagent_tokens: int | None = None,
|
|
middleware_tokens: int | None = None,
|
|
token_usage_by_model: dict[str, dict[str, int]] | None = None,
|
|
message_count: int | None = None,
|
|
last_ai_message: str | None = None,
|
|
first_human_message: str | None = None,
|
|
) -> None:
|
|
"""Update token usage + convenience fields while a run is still active."""
|
|
values: dict[str, Any] = {"updated_at": datetime.now(UTC)}
|
|
optional_counters = {
|
|
"total_input_tokens": total_input_tokens,
|
|
"total_output_tokens": total_output_tokens,
|
|
"total_tokens": total_tokens,
|
|
"llm_call_count": llm_call_count,
|
|
"lead_agent_tokens": lead_agent_tokens,
|
|
"subagent_tokens": subagent_tokens,
|
|
"middleware_tokens": middleware_tokens,
|
|
"message_count": message_count,
|
|
}
|
|
for key, value in optional_counters.items():
|
|
if value is not None:
|
|
values[key] = value
|
|
if token_usage_by_model is not None:
|
|
values["token_usage_by_model"] = self._safe_json(token_usage_by_model) or {}
|
|
if last_ai_message is not None:
|
|
values["last_ai_message"] = last_ai_message[:2000]
|
|
if first_human_message is not None:
|
|
values["first_human_message"] = first_human_message[:2000]
|
|
async with self._sf() as session:
|
|
await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status == "running").values(**values))
|
|
await session.commit()
|
|
|
|
async def aggregate_tokens_by_thread(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
include_active: bool = False,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, Any]:
|
|
"""Aggregate token usage for a thread.
|
|
|
|
``by_model`` is reduced in Python from each row's ``token_usage_by_model``
|
|
JSON column so subagent / middleware tokens land on the model that
|
|
actually produced them (issue #3645). Rows written before that column
|
|
existed fall back to ``RunRow.model_name`` + ``RunRow.total_tokens``,
|
|
preserving the legacy lead-only behavior instead of dropping the data.
|
|
|
|
Headline totals (``total_tokens``, ``total_input_tokens``,
|
|
``total_output_tokens``) and the ``by_caller`` bucket are summed from
|
|
their own columns and are therefore unaffected by the JSON column being
|
|
empty.
|
|
"""
|
|
statuses = ("success", "error", "running") if include_active else ("success", "error")
|
|
_completed = RunRow.status.in_(statuses)
|
|
_thread = RunRow.thread_id == thread_id
|
|
_run_operation = RunRow.operation_kind == "run"
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.aggregate_tokens_by_thread")
|
|
|
|
stmt = select(
|
|
RunRow.model_name,
|
|
RunRow.total_tokens,
|
|
RunRow.total_input_tokens,
|
|
RunRow.total_output_tokens,
|
|
RunRow.lead_agent_tokens,
|
|
RunRow.subagent_tokens,
|
|
RunRow.middleware_tokens,
|
|
RunRow.token_usage_by_model,
|
|
).where(_thread, _run_operation, _completed)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).all()
|
|
|
|
total_tokens = total_input = total_output = total_runs = 0
|
|
lead_agent = subagent = middleware = 0
|
|
by_model: dict[str, dict] = {}
|
|
for r in rows:
|
|
total_runs += 1
|
|
total_tokens += r.total_tokens
|
|
total_input += r.total_input_tokens
|
|
total_output += r.total_output_tokens
|
|
lead_agent += r.lead_agent_tokens
|
|
subagent += r.subagent_tokens
|
|
middleware += r.middleware_tokens
|
|
|
|
# ``or {}`` covers rows written before ``token_usage_by_model``
|
|
# existed (the column is NULL on a manual ALTER ADD COLUMN without
|
|
# backfill); fresh rows always carry the journal-produced dict.
|
|
usage_by_model = r.token_usage_by_model or {}
|
|
if usage_by_model:
|
|
for model, usage in usage_by_model.items():
|
|
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
|
|
entry["tokens"] += usage.get("total_tokens", 0)
|
|
entry["runs"] += 1
|
|
else:
|
|
model = r.model_name or "unknown"
|
|
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
|
|
entry["tokens"] += r.total_tokens
|
|
entry["runs"] += 1
|
|
|
|
return {
|
|
"total_tokens": total_tokens,
|
|
"total_input_tokens": total_input,
|
|
"total_output_tokens": total_output,
|
|
"total_runs": total_runs,
|
|
"by_model": by_model,
|
|
"by_caller": {
|
|
"lead_agent": lead_agent,
|
|
"subagent": subagent,
|
|
"middleware": middleware,
|
|
},
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Multi-worker run ownership methods
|
|
# ------------------------------------------------------------------
|
|
|
|
async def update_lease(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
owner_worker_id: str,
|
|
lease_expires_at: str,
|
|
allow_interrupted: bool = False,
|
|
) -> bool:
|
|
lease_dt = datetime.fromisoformat(lease_expires_at)
|
|
values: dict[str, Any] = {
|
|
"owner_worker_id": owner_worker_id,
|
|
"lease_expires_at": lease_dt,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
async with self._sf() as session:
|
|
await self._lock_owned_row(session, run_id)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.owner_worker_id == owner_worker_id,
|
|
RunRow.status.in_(("pending", "running", "interrupted") if allow_interrupted else ("pending", "running")),
|
|
self._live_lease_predicate(session),
|
|
self._live_lease_predicate(session, literal(lease_dt)),
|
|
)
|
|
.values(**values)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def renew_lease(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
owner_worker_id: str,
|
|
lease_expires_at: str,
|
|
allow_interrupted: bool = False,
|
|
) -> LeaseRenewal:
|
|
"""Renew the owner lease and read cancellation intent atomically."""
|
|
lease_dt = datetime.fromisoformat(lease_expires_at)
|
|
async with self._sf() as session:
|
|
await self._lock_owned_row(session, run_id)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.owner_worker_id == owner_worker_id,
|
|
RunRow.status.in_(("pending", "running", "interrupted") if allow_interrupted else ("pending", "running")),
|
|
self._live_lease_predicate(session),
|
|
self._live_lease_predicate(session, literal(lease_dt)),
|
|
)
|
|
.values(
|
|
lease_expires_at=lease_dt,
|
|
updated_at=datetime.now(UTC),
|
|
)
|
|
.returning(RunRow.run_id, RunRow.cancel_action)
|
|
)
|
|
row = result.first()
|
|
await session.commit()
|
|
if row is None:
|
|
return LeaseRenewal(renewed=False)
|
|
return LeaseRenewal(renewed=True, cancel_action=row.cancel_action)
|
|
|
|
async def request_cancel(self, run_id: str, *, action: str, execution_fence=None, user_id=None) -> str | None:
|
|
"""Atomically persist the first cancellation action on an active run."""
|
|
if action not in ("interrupt", "rollback"):
|
|
raise ValueError(f"Unsupported cancellation action: {action}")
|
|
now = datetime.now(UTC)
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
if execution_fence is not None:
|
|
from deerflow.persistence.extension_leases.fence import lock_execution_fence
|
|
|
|
await lock_execution_fence(session, execution_fence, user_id)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
*([RunRow.user_id == user_id] if execution_fence is not None else []),
|
|
RunRow.status.in_(("pending", "running")),
|
|
)
|
|
.values(
|
|
cancel_action=case(
|
|
(RunRow.cancel_action.is_(None), action),
|
|
else_=RunRow.cancel_action,
|
|
),
|
|
cancel_requested_at=case(
|
|
(RunRow.cancel_requested_at.is_(None), now),
|
|
else_=RunRow.cancel_requested_at,
|
|
),
|
|
updated_at=now,
|
|
change_seq=change_seq,
|
|
)
|
|
.returning(RunRow.cancel_action)
|
|
)
|
|
row = result.first()
|
|
await session.commit()
|
|
return row.cancel_action if row is not None else None
|
|
|
|
async def finalize_if_not_cancelled(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
status: str,
|
|
error: str | None = None,
|
|
stop_reason: str | None = None,
|
|
goal_verdict: dict[str, Any] | None = None,
|
|
owner_worker_id: str | None = None,
|
|
) -> StatusFinalization:
|
|
"""Atomically let completion win only before cancellation."""
|
|
values: dict[str, Any] = {
|
|
"status": status,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
if error is not None:
|
|
values["error"] = error
|
|
if stop_reason is not None:
|
|
values["stop_reason"] = stop_reason
|
|
if goal_verdict is not None:
|
|
values["goal_verdict"] = self._safe_json(goal_verdict)
|
|
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
conditions = [RunRow.run_id == run_id, RunRow.status.in_(("pending", "running")), RunRow.cancel_action.is_(None)]
|
|
if owner_worker_id is not None:
|
|
await self._lock_owned_row(session, run_id)
|
|
conditions.extend([RunRow.owner_worker_id == owner_worker_id, self._live_lease_predicate(session)])
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
*conditions,
|
|
)
|
|
.values(**values)
|
|
.returning(RunRow.run_id)
|
|
)
|
|
if result.first() is not None:
|
|
await session.commit()
|
|
return StatusFinalization(finalized=True)
|
|
|
|
cancellation_conditions = [RunRow.run_id == run_id]
|
|
if owner_worker_id is not None:
|
|
cancellation_conditions.extend([RunRow.owner_worker_id == owner_worker_id, RunRow.status.in_(("pending", "running")), self._live_lease_predicate(session)])
|
|
current = await session.execute(select(RunRow.cancel_action).where(*cancellation_conditions))
|
|
cancel_action = current.scalar_one_or_none()
|
|
await session.commit()
|
|
return StatusFinalization(
|
|
finalized=False,
|
|
cancel_action=cancel_action,
|
|
)
|
|
|
|
async def claim_for_takeover(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
grace_seconds: int,
|
|
error: str,
|
|
stop_reason: str | None = None,
|
|
) -> bool:
|
|
cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds)
|
|
values: dict[str, Any] = {
|
|
"status": "error",
|
|
"error": error,
|
|
"evidence_seal_state": "partial",
|
|
"evidence_seal_error": "Worker lease expired without a completed journal seal.",
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
if stop_reason is not None:
|
|
values["stop_reason"] = stop_reason
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
_lease_expired_or_null(RunRow.lease_expires_at, cutoff),
|
|
)
|
|
.values(**values)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def recover_unsealed_evidence(
|
|
self,
|
|
*,
|
|
before: str | None = None,
|
|
grace_seconds: int = 10,
|
|
exclude_run_ids: tuple[str, ...] = (),
|
|
) -> int:
|
|
"""Publish crash-interrupted terminal seals without changing business status.
|
|
|
|
Clock-first locking matches other run mutations. The bounded scan and
|
|
conditional update run in one transaction, so a renewed lease wins.
|
|
"""
|
|
before_dt = datetime.fromisoformat(before) if isinstance(before, str) else before or datetime.now(UTC)
|
|
cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds)
|
|
predicates = (
|
|
RunRow.operation_kind == "run",
|
|
RunRow.status.in_(("success", "error", "timeout", "interrupted")),
|
|
RunRow.evidence_seal_state == "pending",
|
|
RunRow.created_at <= before_dt,
|
|
_lease_expired_or_null(RunRow.lease_expires_at, cutoff),
|
|
RunRow.run_id.not_in(exclude_run_ids),
|
|
)
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
candidates = tuple((await session.scalars(select(RunRow.run_id).where(*predicates).order_by(RunRow.created_at, RunRow.run_id).limit(200))).all())
|
|
if not candidates:
|
|
await session.rollback()
|
|
return 0
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(RunRow.run_id.in_(candidates), *predicates)
|
|
.values(
|
|
evidence_seal_state="partial",
|
|
evidence_seal_error="Worker lease expired without a completed journal seal.",
|
|
evidence_revision=uuid4().hex,
|
|
change_seq=change_seq,
|
|
)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount
|
|
|
|
async def list_inflight_with_expired_lease(
|
|
self,
|
|
*,
|
|
before: str | None = None,
|
|
grace_seconds: int = 10,
|
|
) -> list[dict[str, Any]]:
|
|
if before is None:
|
|
before_dt = datetime.now(UTC)
|
|
elif isinstance(before, datetime):
|
|
before_dt = before
|
|
else:
|
|
before_dt = datetime.fromisoformat(before)
|
|
cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds)
|
|
stmt = (
|
|
select(RunRow)
|
|
.where(
|
|
RunRow.status.in_(("pending", "running")),
|
|
RunRow.created_at <= before_dt,
|
|
_lease_expired_or_null(RunRow.lease_expires_at, cutoff),
|
|
)
|
|
.order_by(RunRow.created_at.asc())
|
|
)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def create_thread_operation_atomic(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
thread_id: str,
|
|
owner_worker_id: str,
|
|
lease_expires_at: str | None,
|
|
operation_kind: str = "run",
|
|
multitask_strategy: str = "reject",
|
|
assistant_id: str | None = None,
|
|
user_id: str | None = None,
|
|
model_name: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
kwargs: dict[str, Any] | None = None,
|
|
created_at: str | None = None,
|
|
grace_seconds: int = 10,
|
|
idempotency_key: str | None = None,
|
|
evidence_origin: str = "unknown",
|
|
evidence_agent_id: str | None = None,
|
|
idempotency_request: dict[str, Any] | None = None,
|
|
execution_fence=None,
|
|
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
|
|
"""Atomically create a run with cross-process thread-uniqueness.
|
|
|
|
- For ``reject``: INSERT, let the partial unique index enforce
|
|
single-active-run. Returns ``(row_dict, [])`` on success, raises
|
|
``IntegrityError`` on conflict.
|
|
- For ``interrupt`` / ``rollback``: SELECT FOR UPDATE inflight
|
|
rows for the thread, cancel them (unless their lease is still valid),
|
|
then INSERT the new row — all in one transaction. Returns
|
|
``(row_dict, claimed_row_dicts)``.
|
|
|
|
Returns:
|
|
Tuple of ``(new_run_dict, claimed_run_dicts)``.
|
|
"""
|
|
from deerflow.runtime.runs.manager import ConflictError
|
|
|
|
resolved_user_id = resolve_user_id(user_id or AUTO, 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
|
|
cutoff = now - timedelta(seconds=grace_seconds)
|
|
|
|
values = {
|
|
"thread_id": thread_id,
|
|
"evidence_origin": evidence_origin,
|
|
"evidence_agent_id": evidence_agent_id,
|
|
"evidence_seal_state": "pending",
|
|
"assistant_id": assistant_id,
|
|
"user_id": resolved_user_id,
|
|
"model_name": self._normalize_model_name(model_name),
|
|
"status": "pending",
|
|
"operation_kind": operation_kind,
|
|
"multitask_strategy": multitask_strategy,
|
|
"metadata_json": self._safe_json(metadata) or {},
|
|
"kwargs_json": self._safe_json(kwargs) or {},
|
|
"owner_worker_id": owner_worker_id,
|
|
"lease_expires_at": lease_dt,
|
|
"idempotency_key": idempotency_key,
|
|
"idempotency_request_json": self._safe_json(idempotency_request),
|
|
"origin_kind": _origin_kind(operation_kind, metadata),
|
|
"created_at": created,
|
|
"updated_at": now,
|
|
}
|
|
|
|
async with self._sf() as session:
|
|
# Keep the global clock -> run-row lock order used by every other
|
|
# mutator. One position covers this atomic set of changes; run_id
|
|
# provides deterministic ordering within the position.
|
|
change_seq = await self._next_change_seq(session)
|
|
if execution_fence is not None:
|
|
from deerflow.persistence.extension_leases.fence import lock_execution_fence
|
|
|
|
await lock_execution_fence(session, execution_fence, resolved_user_id)
|
|
claimed: list[dict[str, Any]] = []
|
|
|
|
if idempotency_key is not None:
|
|
existing = (await session.execute(select(RunRow).where(RunRow.idempotency_key == idempotency_key).with_for_update())).scalar_one_or_none()
|
|
if existing is not None:
|
|
raise RunIdempotencyConflict(self._row_to_dict(existing))
|
|
|
|
if multitask_strategy in ("interrupt", "rollback"):
|
|
stmt = (
|
|
select(RunRow)
|
|
.where(
|
|
RunRow.thread_id == thread_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
)
|
|
.with_for_update()
|
|
)
|
|
result = await session.execute(stmt)
|
|
for row in result.scalars():
|
|
lease_expired = False
|
|
if row.lease_expires_at is not None:
|
|
# SQLite drops tzinfo on read despite
|
|
# ``DateTime(timezone=True)`` (see ``_row_to_dict``).
|
|
# Treat naive values as UTC — same convention as
|
|
# ``coerce_iso`` — so the Python-side comparison
|
|
# against the aware ``cutoff`` does not raise
|
|
# ``TypeError: can't compare offset-naive and
|
|
# offset-aware datetimes`` when heartbeat is enabled
|
|
# on SQLite.
|
|
row_lease = row.lease_expires_at
|
|
if row_lease.tzinfo is None:
|
|
row_lease = row_lease.replace(tzinfo=UTC)
|
|
lease_expired = row_lease < cutoff
|
|
if row_lease >= cutoff and row.owner_worker_id != owner_worker_id:
|
|
# Live run owned by another worker — we cannot
|
|
# interrupt it and the partial unique index would
|
|
# reject our INSERT anyway. Surface as
|
|
# ConflictError so the caller gets a clean signal
|
|
# instead of a retry loop on IntegrityError.
|
|
raise ConflictError(f"Thread {thread_id} already has an active run owned by another worker")
|
|
if row.operation_kind != "run" and not lease_expired:
|
|
raise ConflictError(f"Thread {thread_id} has an active checkpoint write")
|
|
row.status = "interrupted"
|
|
row.error = "Cancelled by newer run"
|
|
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))
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError as exc:
|
|
await session.rollback()
|
|
if idempotency_key is not None:
|
|
existing = (await session.execute(select(RunRow).where(RunRow.idempotency_key == idempotency_key))).scalar_one_or_none()
|
|
if existing is not None:
|
|
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
|