diff --git a/api/services/workflow_copilot/repository.py b/api/services/workflow_copilot/repository.py new file mode 100644 index 00000000000..3c8377f9cf0 --- /dev/null +++ b/api/services/workflow_copilot/repository.py @@ -0,0 +1,207 @@ +"""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 currently implements the session lifecycle + the version-CAS +(``create_session``/``get_session``/``compare_and_advance`` -- P3a Task 4). +The remaining ``Repository`` methods (checkpoints/runs/test-inputs/ +list_conversation) are P3a Task 5 and are stubbed with ``NotImplementedError`` +so the class is importable and satisfies ``isinstance(..., Repository)``. + +``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. +""" + +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, + Run, + Snapshot, + TestInput, +) +from core.workflow_copilot.models import ( + Session as DomainSession, +) +from core.workflow_copilot.state import PcState +from models.workflow_copilot import ( + CopilotConversationItem, + CopilotSession, + CopilotSessionCommit, +) +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: + raise NotImplementedError + + def get_checkpoint(self, id: str) -> tuple[Checkpoint, Snapshot]: + raise NotImplementedError + + # -- runs (P3a Task 5) -- + + def save_run(self, session_id: str, run: Run) -> None: + raise NotImplementedError + + def get_run(self, id: str) -> Run: + raise NotImplementedError + + # -- test inputs (P3a Task 5) -- + + def save_test_input(self, ti: TestInput) -> None: + raise NotImplementedError + + def get_test_input(self, id: str) -> TestInput: + raise NotImplementedError + + # -- conversation (P3a Task 5) -- + + def list_conversation(self, session_id: str) -> list[ConversationItem]: + raise NotImplementedError + + # -- 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, + ) diff --git a/api/tests/unit_tests/services/workflow_copilot/test_repository.py b/api/tests/unit_tests/services/workflow_copilot/test_repository.py new file mode 100644 index 00000000000..aeb09f5efb2 --- /dev/null +++ b/api/tests/unit_tests/services/workflow_copilot/test_repository.py @@ -0,0 +1,162 @@ +"""Tests for ``SqlCopilotRepository``'s session lifecycle + version-CAS. + +Verifies the P3a Task 4 slice (``create_session``/``get_session``/ +``compare_and_advance``) against an in-memory SQLite database: real +``Base.metadata.create_all`` DDL, a real ``UPDATE ... WHERE version=`` CAS, +and a real ``UNIQUE(session_id, seq)`` constraint backing seq authority. +The remaining ``Repository`` methods (checkpoints/runs/test-inputs/ +list_conversation) are Task 5 and are not exercised here. +""" + +from collections.abc import Iterator + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from sqlalchemy.orm import sessionmaker + +from core.workflow_copilot.errors import ConflictError, NotFoundError +from core.workflow_copilot.models import ConversationItem, EntryMode, FixContext, MutationIntent +from core.workflow_copilot.models import Session as DomainSession +from core.workflow_copilot.state import PcState +from models.base import Base +from services.workflow_copilot.repository import SqlCopilotRepository + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +APP_ID = "22222222-2222-2222-2222-222222222222" +ACCOUNT_ID = "33333333-3333-3333-3333-333333333333" + + +@pytest.fixture +def engine() -> Iterator[Engine]: + engine = create_engine("sqlite://") + Base.metadata.create_all(engine) + try: + yield engine + finally: + engine.dispose() + + +@pytest.fixture +def repo(engine: Engine) -> SqlCopilotRepository: + factory = sessionmaker(bind=engine) + return SqlCopilotRepository(factory) + + +def _make_domain_session(**overrides) -> DomainSession: + kwargs = { + "app_id": APP_ID, + "tenant_id": TENANT_ID, + "owner_account_id": ACCOUNT_ID, + "entry_mode": EntryMode.FIX, + "current_state": PcState.FIX_DIAGNOSE, + } + kwargs.update(overrides) + return DomainSession(**kwargs) + + +def test_create_session_then_get_session_round_trips(repo: SqlCopilotRepository) -> None: + domain_session = _make_domain_session() + fc = FixContext(failed_run_id="TR-1") + items = [ + ConversationItem(seq=0, kind="run-context", payload={"a": 1}, at_version=1), + ConversationItem(seq=1, kind="assistant-turn", payload={"b": 2}, at_version=1), + ] + + repo.create_session(domain_session, fc, items) + + # id was generated + version defaults to 1. + assert domain_session.id + assert domain_session.version == 1 + # next_seq advances past the seeded items. + assert fc.next_seq == 2 + + loaded_session, loaded_fc = repo.get_session(domain_session.id) + + assert loaded_session.id == domain_session.id + assert loaded_session.app_id == APP_ID + assert loaded_session.tenant_id == TENANT_ID + assert loaded_session.owner_account_id == ACCOUNT_ID + assert loaded_session.entry_mode == EntryMode.FIX + assert loaded_session.current_state == PcState.FIX_DIAGNOSE + assert loaded_session.version == 1 + assert loaded_fc == FixContext(failed_run_id="TR-1", next_seq=2) + + +def test_get_session_missing_id_raises_not_found(repo: SqlCopilotRepository) -> None: + with pytest.raises(NotFoundError): + repo.get_session("does-not-exist") + + +def test_compare_and_advance_bumps_version_and_appends_commit_and_items(repo: SqlCopilotRepository) -> None: + domain_session = _make_domain_session() + fc = FixContext(failed_run_id="TR-1") + repo.create_session(domain_session, fc, []) + + next_fc = FixContext(failed_run_id="TR-1", staged_repair=[MutationIntent(op="set_node_config")]) + new_items = [ConversationItem(seq=0, kind="diagnosis", payload={"x": 1}, at_version=2)] + + new_version = repo.compare_and_advance( + domain_session.id, base_version=1, next=PcState.FIX_PROPOSE, fc=next_fc, items=new_items + ) + + assert new_version == 2 + + loaded_session, loaded_fc = repo.get_session(domain_session.id) + assert loaded_session.version == 2 + assert loaded_session.current_state == PcState.FIX_PROPOSE + assert loaded_fc == next_fc + + +def test_compare_and_advance_stale_version_raises_conflict_and_leaves_row_unchanged( + repo: SqlCopilotRepository, +) -> None: + domain_session = _make_domain_session() + fc = FixContext(failed_run_id="TR-1") + repo.create_session(domain_session, fc, []) + + with pytest.raises(ConflictError): + repo.compare_and_advance( + domain_session.id, + base_version=999, + next=PcState.FIX_PROPOSE, + fc=FixContext(), + items=[], + ) + + loaded_session, loaded_fc = repo.get_session(domain_session.id) + assert loaded_session.version == 1 + assert loaded_session.current_state == PcState.FIX_DIAGNOSE + assert loaded_fc == fc + + +def test_compare_and_advance_missing_session_raises_not_found(repo: SqlCopilotRepository) -> None: + with pytest.raises(NotFoundError): + repo.compare_and_advance( + "does-not-exist", + base_version=1, + next=PcState.FIX_PROPOSE, + fc=FixContext(), + items=[], + ) + + +def test_compare_and_advance_duplicate_seq_raises_conflict(repo: SqlCopilotRepository) -> None: + domain_session = _make_domain_session() + fc = FixContext(failed_run_id="TR-1") + seed_items = [ConversationItem(seq=0, kind="run-context", payload={}, at_version=1)] + repo.create_session(domain_session, fc, seed_items) + + with pytest.raises(ConflictError): + repo.compare_and_advance( + domain_session.id, + base_version=1, + next=PcState.FIX_PROPOSE, + fc=FixContext(), + items=[ConversationItem(seq=0, kind="diagnosis", payload={}, at_version=2)], + ) + + # The row must not have advanced -- the whole call is one transaction. + loaded_session, _ = repo.get_session(domain_session.id) + assert loaded_session.version == 1 + assert loaded_session.current_state == PcState.FIX_DIAGNOSE