dify/api/services/workflow_copilot/repository.py

"""SQL-backed implementation of the workflow copilot ``Repository`` port.

Implements ``core.workflow_copilot.ports.Repository`` against the tables in
``models.workflow_copilot`` (P3a Task 1) via a caller-provided
``sessionmaker``. Every public method opens exactly one transaction:
``with self._session_factory() as session, session.begin(): ...``.

This module implements the session lifecycle + the version-CAS
(``create_session``/``get_session``/``compare_and_advance`` -- P3a Task 4) and
the remaining ``Repository`` surface -- checkpoints/runs/test-inputs/
list_conversation (P3a Task 5).

``compare_and_advance`` is the single concurrency primitive: it is a real
``UPDATE ... WHERE id = ? AND version = ?`` -- if no row matches, the caller
either named a session that doesn't exist (``NotFoundError``) or raced
another writer and lost (``ConflictError``, stale ``base_version``). The
conversation items' ``UNIQUE(session_id, seq)`` DB constraint is the seq
authority: a collision surfaces as ``IntegrityError`` from the driver, which
this method maps to ``ConflictError`` so callers never see a raw SQLAlchemy
exception.
"""

import dataclasses

from sqlalchemy import select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import sessionmaker

from core.workflow_copilot.errors import ConflictError, NotFoundError
from core.workflow_copilot.models import (
    Checkpoint,
    ConversationItem,
    EntryMode,
    FixContext,
    NodeOutput,
    Run,
    Snapshot,
    TestInput,
)
from core.workflow_copilot.models import (
    Session as DomainSession,
)
from core.workflow_copilot.state import PcState
from models.workflow_copilot import (
    CopilotCheckpoint,
    CopilotConversationItem,
    CopilotRun,
    CopilotSession,
    CopilotSessionCommit,
    CopilotSnapshot,
    CopilotTestInput,
)
from services.workflow_copilot.serde import fix_context_from_dict, fix_context_to_dict

__all__ = ["SqlCopilotRepository"]


class SqlCopilotRepository:
    """SQL-backed ``Repository`` for the workflow copilot engine."""

    def __init__(self, session_factory: sessionmaker) -> None:
        self._session_factory = session_factory

    # -- sessions --

    def create_session(self, session: DomainSession, initial_fc: FixContext, items: list[ConversationItem]) -> None:
        with self._session_factory() as db_session, db_session.begin():
            row = CopilotSession(
                app_id=session.app_id,
                tenant_id=session.tenant_id,
                owner_account_id=session.owner_account_id,
                entry_mode=str(session.entry_mode),
                current_state=str(session.current_state),
            )
            if session.id:
                row.id = session.id
            db_session.add(row)
            db_session.flush()

            session.id = row.id
            session.version = 1

            db_session.add_all(self._to_conversation_row(row.id, item) for item in items)

            # Mutate next_seq before serializing the commit's context so the
            # persisted state and the caller's in-process `initial_fc` never
            # diverge -- a `get_session` immediately after `create_session`
            # must see the same next_seq the caller now holds.
            initial_fc.next_seq = len(items)

            db_session.add(
                CopilotSessionCommit(
                    session_id=row.id,
                    version=1,
                    state=str(session.current_state),
                    context=fix_context_to_dict(initial_fc),
                    actor="",
                )
            )

    def get_session(self, id: str) -> tuple[DomainSession, FixContext]:
        with self._session_factory() as db_session, db_session.begin():
            row = db_session.get(CopilotSession, id)
            if row is None:
                raise NotFoundError(f"session {id} not found")

            stmt = (
                select(CopilotSessionCommit)
                .where(CopilotSessionCommit.session_id == id)
                .order_by(CopilotSessionCommit.version.desc())
                .limit(1)
            )
            commit = db_session.execute(stmt).scalar_one_or_none()
            if commit is None:
                raise NotFoundError(f"no commits found for session {id}")

            fc = fix_context_from_dict(commit.context)
            return self._to_domain_session(row), fc

    def compare_and_advance(
        self,
        session_id: str,
        base_version: int,
        next: PcState,
        fc: FixContext,
        items: list[ConversationItem],
    ) -> int:
        new_version = base_version + 1

        with self._session_factory() as db_session, db_session.begin():
            result = db_session.execute(
                update(CopilotSession)
                .where(CopilotSession.id == session_id, CopilotSession.version == base_version)
                .values(version=CopilotSession.version + 1, current_state=str(next))
                .execution_options(synchronize_session=False)
            )

            if result.rowcount == 0:
                if db_session.get(CopilotSession, session_id) is None:
                    raise NotFoundError(f"session {session_id} not found")
                raise ConflictError(f"stale base_version {base_version} for session {session_id}")

            db_session.add(
                CopilotSessionCommit(
                    session_id=session_id,
                    version=new_version,
                    state=str(next),
                    context=fix_context_to_dict(fc),
                    actor="",
                )
            )

            db_session.add_all(self._to_conversation_row(session_id, item) for item in items)

            try:
                db_session.flush()
            except IntegrityError as exc:
                raise ConflictError(f"duplicate conversation item seq for session {session_id}") from exc

        return new_version

    # -- checkpoints (P3a Task 5) --

    def create_checkpoint(self, cp: Checkpoint, snap: Snapshot) -> None:
        with self._session_factory() as db_session, db_session.begin():
            snap_row = CopilotSnapshot(
                session_id=snap.session_id,
                hash=snap.hash,
                graph=snap.graph,
            )
            if snap.id:
                snap_row.id = snap.id
            db_session.add(snap_row)
            db_session.flush()
            snap.id = snap_row.id

            cp_row = CopilotCheckpoint(
                session_id=cp.session_id,
                state=str(cp.state),
                snapshot_id=snap_row.id,
            )
            if cp.id:
                cp_row.id = cp.id
            db_session.add(cp_row)
            db_session.flush()
            cp.id = cp_row.id
            cp.snapshot_id = snap_row.id

    def get_checkpoint(self, id: str) -> tuple[Checkpoint, Snapshot]:
        with self._session_factory() as db_session, db_session.begin():
            cp_row = db_session.get(CopilotCheckpoint, id)
            if cp_row is None:
                raise NotFoundError(f"checkpoint {id} not found")

            snap_row = db_session.get(CopilotSnapshot, cp_row.snapshot_id)
            if snap_row is None:
                raise NotFoundError(f"snapshot {cp_row.snapshot_id} not found")

            return self._to_domain_checkpoint(cp_row), self._to_domain_snapshot(snap_row)

    # -- runs (P3a Task 5) --

    def save_run(self, session_id: str, run: Run) -> None:
        with self._session_factory() as db_session, db_session.begin():
            row = CopilotRun(
                session_id=session_id,
                kind=run.kind,
                dify_run_id=run.dify_run_id,
                status=run.status,
                per_node=[dataclasses.asdict(n) for n in run.per_node],
                culprit_node_id=run.culprit_node_id,
                inputs_ref=run.inputs_ref,
                tokens=run.tokens,
                elapsed_ms=run.elapsed_ms,
                immutable=run.immutable,
            )
            # Id-preservation contract (see `ports.Repository.save_run`):
            # a caller-supplied non-empty `run.id` must survive unchanged.
            if run.id:
                row.id = run.id
            db_session.add(row)
            db_session.flush()
            run.id = row.id

    def get_run(self, id: str) -> Run:
        with self._session_factory() as db_session, db_session.begin():
            row = db_session.get(CopilotRun, id)
            if row is None:
                raise NotFoundError(f"run {id} not found")

            return self._to_domain_run(row)

    # -- test inputs (P3a Task 5) --

    def save_test_input(self, ti: TestInput) -> None:
        with self._session_factory() as db_session, db_session.begin():
            row = CopilotTestInput(
                session_id=ti.session_id,
                source=ti.source,
                inputs=ti.inputs,
                start_schema_hash=ti.start_schema_hash,
            )
            if ti.id:
                row.id = ti.id
            db_session.add(row)
            db_session.flush()
            ti.id = row.id

    def get_test_input(self, id: str) -> TestInput:
        with self._session_factory() as db_session, db_session.begin():
            row = db_session.get(CopilotTestInput, id)
            if row is None:
                raise NotFoundError(f"test input {id} not found")

            return self._to_domain_test_input(row)

    # -- conversation (P3a Task 5) --

    def list_conversation(self, session_id: str) -> list[ConversationItem]:
        with self._session_factory() as db_session, db_session.begin():
            stmt = (
                select(CopilotConversationItem)
                .where(CopilotConversationItem.session_id == session_id)
                .order_by(CopilotConversationItem.seq)
            )
            rows = db_session.execute(stmt).scalars().all()
            return [self._to_domain_conversation_item(row) for row in rows]

    # -- mappers --

    @staticmethod
    def _to_domain_session(row: CopilotSession) -> DomainSession:
        return DomainSession(
            id=row.id,
            app_id=row.app_id,
            tenant_id=row.tenant_id,
            owner_account_id=row.owner_account_id,
            entry_mode=EntryMode(row.entry_mode),
            current_state=PcState(row.current_state),
            version=row.version,
        )

    @staticmethod
    def _to_conversation_row(session_id: str, item: ConversationItem) -> CopilotConversationItem:
        return CopilotConversationItem(
            session_id=session_id,
            seq=item.seq,
            kind=item.kind,
            payload=item.payload,
            at_version=item.at_version,
        )

    @staticmethod
    def _to_domain_checkpoint(row: CopilotCheckpoint) -> Checkpoint:
        return Checkpoint(
            id=row.id,
            session_id=row.session_id,
            state=PcState(row.state),
            snapshot_id=row.snapshot_id,
        )

    @staticmethod
    def _to_domain_snapshot(row: CopilotSnapshot) -> Snapshot:
        return Snapshot(
            id=row.id,
            session_id=row.session_id,
            hash=row.hash,
            graph=row.graph,
        )

    @staticmethod
    def _to_domain_run(row: CopilotRun) -> Run:
        return Run(
            id=row.id,
            kind=row.kind,
            dify_run_id=row.dify_run_id or "",
            status=row.status,
            per_node=[NodeOutput(**d) for d in (row.per_node or [])],
            culprit_node_id=row.culprit_node_id or "",
            inputs_ref=row.inputs_ref or "",
            tokens=row.tokens or 0,
            elapsed_ms=row.elapsed_ms or 0,
            immutable=row.immutable,
        )

    @staticmethod
    def _to_domain_test_input(row: CopilotTestInput) -> TestInput:
        return TestInput(
            id=row.id,
            session_id=row.session_id,
            source=row.source,
            inputs=row.inputs,
            start_schema_hash=row.start_schema_hash,
        )

    @staticmethod
    def _to_domain_conversation_item(row: CopilotConversationItem) -> ConversationItem:
        return ConversationItem(
            seq=row.seq,
            kind=row.kind,
            payload=row.payload,
            at_version=row.at_version,
        )