app/copilot/repository.py

from __future__ import annotations

from datetime import datetime, timezone

from sqlalchemy import select
from sqlalchemy.exc import IntegrityError

from app.copilot.errors import CopilotRunConflictError, CopilotRunNotFoundError
from app.copilot.models import CopilotRunRecord
from app.security.database import SecurityDatabase


class CopilotRepository:
    def __init__(self, database: SecurityDatabase) -> None:
        self.database = database

    async def create(self, record: CopilotRunRecord) -> tuple[CopilotRunRecord, bool]:
        try:
            async with self.database.tenant_session(
                workspace_id=record.workspace_id, user_id=record.user_id
            ) as session:
                session.add(record)
                await session.commit()
                await session.refresh(record)
                return record, True
        except IntegrityError:
            existing = await self.get_by_idempotency(
                record.workspace_id, record.user_id, record.idempotency_key
            )
            if existing is None:
                raise
            if existing.request_fingerprint != record.request_fingerprint:
                raise CopilotRunConflictError(
                    "Idempotency-Key is already associated with another Copilot request."
                )
            return existing, False

    async def get_by_idempotency(
        self, workspace_id: str, user_id: str, key: str
    ) -> CopilotRunRecord | None:
        async with self.database.tenant_session(
            workspace_id=workspace_id, user_id=user_id
        ) as session:
            return await session.scalar(
                select(CopilotRunRecord).where(
                    CopilotRunRecord.workspace_id == workspace_id,
                    CopilotRunRecord.idempotency_key == key,
                )
            )

    async def get(
        self, workspace_id: str, user_id: str, run_id: str, *, lock: bool = False
    ) -> CopilotRunRecord:
        async with self.database.tenant_session(
            workspace_id=workspace_id, user_id=user_id
        ) as session:
            statement = select(CopilotRunRecord).where(
                CopilotRunRecord.id == run_id,
                CopilotRunRecord.workspace_id == workspace_id,
            )
            if lock:
                statement = statement.with_for_update()
            record = await session.scalar(statement)
            if record is None:
                raise CopilotRunNotFoundError("Copilot run was not found in this workspace.")
            return record

    async def list(
        self, workspace_id: str, user_id: str, *, offset: int, limit: int
    ) -> list[CopilotRunRecord]:
        async with self.database.tenant_session(
            workspace_id=workspace_id, user_id=user_id
        ) as session:
            return list(
                (
                    await session.scalars(
                        select(CopilotRunRecord)
                        .where(CopilotRunRecord.workspace_id == workspace_id)
                        .order_by(CopilotRunRecord.created_at.desc())
                        .offset(offset)
                        .limit(limit)
                    )
                ).all()
            )

    async def claim_execution(
        self,
        workspace_id: str,
        user_id: str,
        run_id: str,
        *,
        confirmed: bool,
    ) -> CopilotRunRecord:
        async with self.database.tenant_session(
            workspace_id=workspace_id, user_id=user_id
        ) as session:
            record = await session.scalar(
                select(CopilotRunRecord)
                .where(
                    CopilotRunRecord.id == run_id,
                    CopilotRunRecord.workspace_id == workspace_id,
                )
                .with_for_update()
            )
            if record is None:
                raise CopilotRunNotFoundError("Copilot run was not found in this workspace.")
            if record.status != "plan_ready":
                raise CopilotRunConflictError("Only a plan-ready Copilot run can be executed.")
            now = datetime.now(timezone.utc)
            record.status = "executing"
            record.updated_at = now
            if confirmed and record.confirmed_at is None:
                record.confirmed_at = now
            await session.commit()
            await session.refresh(record)
            return record

    async def cancel_before_execution(
        self,
        workspace_id: str,
        user_id: str,
        run_id: str,
    ) -> tuple[CopilotRunRecord, bool]:
        async with self.database.tenant_session(
            workspace_id=workspace_id, user_id=user_id
        ) as session:
            record = await session.scalar(
                select(CopilotRunRecord)
                .where(
                    CopilotRunRecord.id == run_id,
                    CopilotRunRecord.workspace_id == workspace_id,
                )
                .with_for_update()
            )
            if record is None:
                raise CopilotRunNotFoundError("Copilot run was not found in this workspace.")
            if record.status in {"completed", "partial", "failed", "cancelled"}:
                return record, False
            if record.status == "executing":
                raise CopilotRunConflictError(
                    "This action batch is already executing; cancel its durable child job directly."
                )
            now = datetime.now(timezone.utc)
            record.status = "cancelled"
            record.current_action_id = None
            record.summary = "Copilot run cancelled before execution."
            record.updated_at = now
            record.completed_at = now
            await session.commit()
            await session.refresh(record)
            return record, True

    async def update(
        self,
        workspace_id: str,
        user_id: str,
        run_id: str,
        *,
        status: str,
        current_action_id: str | None = None,
        results: list[dict[str, object]] | None = None,
        summary: str | None = None,
        error_code: str | None = None,
        error_message: str | None = None,
        confirmed: bool = False,
        terminal: bool = False,
    ) -> CopilotRunRecord:
        async with self.database.tenant_session(
            workspace_id=workspace_id, user_id=user_id
        ) as session:
            record = await session.scalar(
                select(CopilotRunRecord)
                .where(
                    CopilotRunRecord.id == run_id,
                    CopilotRunRecord.workspace_id == workspace_id,
                )
                .with_for_update()
            )
            if record is None:
                raise CopilotRunNotFoundError("Copilot run was not found in this workspace.")
            now = datetime.now(timezone.utc)
            record.status = status
            record.current_action_id = current_action_id
            if results is not None:
                record.results_json = results
            record.summary = summary
            record.error_code = error_code
            record.error_message = error_message
            if confirmed and record.confirmed_at is None:
                record.confirmed_at = now
            record.updated_at = now
            if terminal:
                record.completed_at = now
            await session.commit()
            await session.refresh(record)
            return record