"""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