AgentEvalTool/backend/agenteval/storage/repository/run.py

"""Repository for evaluation runs."""

from dataclasses import dataclass
from enum import Enum
from typing import Optional

from sqlalchemy import insert as sql_insert
from sqlalchemy import literal
from sqlalchemy.exc import IntegrityError
from sqlmodel import select

from agenteval.models import CampaignStatus, EvalResult, EvalRun, RunStatus, RunTrigger
from agenteval.storage.db import (
    CampaignDB,
    EvalResultDB,
    EvalRunDB,
    TurnDB,
    new_uuid,
    utc_now,
)
from agenteval.storage.repository.base import BaseRepository
from agenteval.storage.repository.result import result_from_db


class CampaignRunClaimStatus(str, Enum):
    """Outcome of a durable Campaign child-Run claim."""

    CLAIMED = "claimed"
    EXISTING = "existing"
    NOT_FOUND = "not_found"
    CONFLICT = "conflict"


@dataclass(frozen=True)
class CampaignRunClaimResult:
    """Typed result for the idempotent child-Run claim seam."""

    status: CampaignRunClaimStatus
    run: Optional[EvalRun] = None

    @property
    def accepted(self) -> bool:
        return self.status in (CampaignRunClaimStatus.CLAIMED, CampaignRunClaimStatus.EXISTING)

    @property
    def created(self) -> bool:
        return self.status is CampaignRunClaimStatus.CLAIMED


class RunRepository(BaseRepository[EvalRun, EvalRunDB]):
    """Repository for evaluation runs."""

    _table = EvalRunDB
    _order_by = "started_at"

    def claim_campaign_run(
        self,
        *,
        campaign_id: str,
        scenario_id: str,
        scenario_version: int,
        plan_index: int,
        occurrence_index: int,
    ) -> "CampaignRunClaimResult":
        """Claim one durable child Run occurrence while its Campaign is running.

        The conditional ``INSERT ... SELECT`` makes the Campaign status check
        part of the write.  The composite unique index turns retries or a
        competing scheduler into an ``EXISTING`` result instead of a second
        pending Run.  No external messages are sent by this operation.
        """
        if plan_index < 0 or occurrence_index < 0:
            raise ValueError("plan_index and occurrence_index must be non-negative")

        run_id = new_uuid()
        now = utc_now()
        statement = sql_insert(EvalRunDB).from_select(
            [
                "id",
                "target_id",
                "scenario_id",
                "scenario_version",
                "campaign_id",
                "campaign_plan_index",
                "campaign_occurrence_index",
                "status",
                "triggered_by",
                "started_at",
            ],
            select(
                literal(run_id),
                CampaignDB.target_id,
                literal(scenario_id),
                literal(scenario_version),
                CampaignDB.id,
                literal(plan_index),
                literal(occurrence_index),
                literal(RunStatus.PENDING.value),
                literal(RunTrigger.CAMPAIGN.value),
                literal(now),
            ).where(
                CampaignDB.id == campaign_id,
                CampaignDB.status == CampaignStatus.RUNNING.value,
                CampaignDB.target_id.is_not(None),
            ),
        )

        try:
            result = self.session.exec(statement)
            if result.rowcount == 1:
                self.session.commit()
                db = self.session.get(EvalRunDB, run_id)
                return CampaignRunClaimResult(CampaignRunClaimStatus.CLAIMED, self._from_db(db))
            self.session.rollback()
        except IntegrityError:
            self.session.rollback()
        except Exception:
            self.session.rollback()
            raise

        existing = self._get_campaign_identity(campaign_id, plan_index, occurrence_index)
        campaign = self.session.get(CampaignDB, campaign_id)
        if campaign is None:
            return CampaignRunClaimResult(CampaignRunClaimStatus.NOT_FOUND)
        if campaign.status != CampaignStatus.RUNNING.value:
            return CampaignRunClaimResult(CampaignRunClaimStatus.CONFLICT)
        if existing is not None:
            return CampaignRunClaimResult(CampaignRunClaimStatus.EXISTING, self._from_db(existing))
        return CampaignRunClaimResult(CampaignRunClaimStatus.CONFLICT)

    def _get_campaign_identity(self, campaign_id: str, plan_index: int, occurrence_index: int) -> Optional[EvalRunDB]:
        statement = select(EvalRunDB).where(
            EvalRunDB.campaign_id == campaign_id,
            EvalRunDB.campaign_plan_index == plan_index,
            EvalRunDB.campaign_occurrence_index == occurrence_index,
        )
        return self.session.exec(statement).first()

    def _to_db(self, run: EvalRun) -> EvalRunDB:
        db = EvalRunDB(
            id=run.id,
            target_id=run.target_id,
            scenario_id=run.scenario_id,
            scenario_version=run.scenario_version,
            campaign_id=run.campaign_id,
            campaign_plan_index=run.campaign_plan_index,
            campaign_occurrence_index=run.campaign_occurrence_index,
            status=run.status.value,
            triggered_by=run.triggered_by.value,
            started_at=run.started_at,
            completed_at=run.completed_at,
        )
        if run.summary is not None:
            db.set_summary(run.summary.model_dump(mode="json"))
        return db

    def _from_db(self, db: EvalRunDB) -> EvalRun:
        return EvalRun(
            id=db.id,
            target_id=db.target_id,
            scenario_id=db.scenario_id,
            scenario_version=db.scenario_version or 1,
            campaign_id=db.campaign_id,
            campaign_plan_index=db.campaign_plan_index,
            campaign_occurrence_index=db.campaign_occurrence_index,
            status=db.status,
            triggered_by=db.triggered_by or "manual",
            started_at=db.started_at,
            completed_at=db.completed_at,
            summary=db.get_summary(),
        )

    def list_by_campaign(self, campaign_id: str) -> list[EvalRun]:
        statement = select(EvalRunDB).where(EvalRunDB.campaign_id == campaign_id).order_by(EvalRunDB.started_at)
        return [self._from_db(r) for r in self.session.exec(statement).all()]

    def list_by_campaigns(self, campaign_ids: list[str]) -> dict[str, list[EvalRun]]:
        """Load child Runs for multiple Campaigns in one query."""
        grouped = {campaign_id: [] for campaign_id in campaign_ids}
        if not campaign_ids:
            return grouped
        statement = (
            select(EvalRunDB)
            .where(EvalRunDB.campaign_id.in_(campaign_ids))  # type: ignore[union-attr]
            .order_by(EvalRunDB.started_at)
        )
        for row in self.session.exec(statement).all():
            if row.campaign_id in grouped:
                grouped[row.campaign_id].append(self._from_db(row))
        return grouped

    def mark_orphans_failed(self) -> int:
        """Fail process-orphaned Runs while preserving recoverable child claims.

        Every RUNNING Run may already have sent messages and must not replay.
        A PENDING Run is preserved only when it has a complete Campaign plan
        identity and its parent Campaign is still running; the Campaign loop
        can safely resume those claims before any external send occurred.
        """
        statement = select(EvalRunDB).where(EvalRunDB.status.in_(["running", "pending"]))  # type: ignore[attr-defined]
        candidates = self.session.exec(statement).all()

        # Batch load all campaigns for pending runs to avoid N+1 queries
        campaign_ids = {db.campaign_id for db in candidates if db.campaign_id and db.status == RunStatus.PENDING.value}
        campaigns = {}
        if campaign_ids:
            campaign_stmt = select(CampaignDB).where(CampaignDB.id.in_(campaign_ids))  # type: ignore[attr-defined]
            campaigns = {c.id: c for c in self.session.exec(campaign_stmt).all()}

        failed: list[EvalRunDB] = []
        for db in candidates:
            if db.status == RunStatus.PENDING.value and self._is_recoverable_campaign_pending(db, campaigns):
                continue
            self._set_interrupted(db)
            failed.append(db)
        if failed:
            self.session.commit()
        return len(failed)

    def _is_recoverable_campaign_pending(self, db: EvalRunDB, campaigns: dict[str, CampaignDB]) -> bool:
        if db.campaign_id is None or db.campaign_plan_index is None or db.campaign_occurrence_index is None:
            return False
        campaign = campaigns.get(db.campaign_id)
        return campaign is not None and campaign.status == CampaignStatus.RUNNING.value

    def _set_interrupted(self, db: EvalRunDB) -> None:
        db.status = RunStatus.FAILED.value
        db.completed_at = db.completed_at or utc_now()
        summary = db.get_summary() or {}
        summary["error"] = {"code": "interrupted", "message": "服务重启导致评测中断"}
        db.set_summary(summary)
        self.session.add(db)

    def list_pending_campaign_children(self, campaign_id: str) -> list[EvalRun]:
        """Return durable pending claims in deterministic plan order."""
        statement = (
            select(EvalRunDB)
            .where(
                EvalRunDB.campaign_id == campaign_id,
                EvalRunDB.campaign_plan_index.is_not(None),
                EvalRunDB.campaign_occurrence_index.is_not(None),
                EvalRunDB.status == RunStatus.PENDING.value,
            )
            .order_by(EvalRunDB.campaign_plan_index, EvalRunDB.campaign_occurrence_index)
        )
        return [self._from_db(db) for db in self.session.exec(statement).all()]

    def mark_campaign_running_interrupted(self, campaign_id: str) -> list[str]:
        """Fail running child Runs from a previous process without replaying."""
        statement = select(EvalRunDB).where(
            EvalRunDB.campaign_id == campaign_id,
            EvalRunDB.campaign_plan_index.is_not(None),
            EvalRunDB.campaign_occurrence_index.is_not(None),
            EvalRunDB.status == RunStatus.RUNNING.value,
        )
        rows = list(self.session.exec(statement).all())
        for db in rows:
            self._set_interrupted(db)
        if rows:
            self.session.commit()
        return [db.id or "" for db in rows]

    def mark_pending_campaign_child_interrupted(self, run_id: str) -> Optional[EvalRun]:
        """Fail one still-pending child claim that cannot be safely recovered."""
        db = self.session.get(EvalRunDB, run_id)
        if db is None or db.status != RunStatus.PENDING.value:
            return self._from_db(db) if db is not None else None
        self._set_interrupted(db)
        self.session.commit()
        self.session.refresh(db)
        return self._from_db(db)

    def update(self, run: EvalRun) -> Optional[EvalRun]:
        existing = self.session.get(EvalRunDB, run.id)
        if not existing:
            return None
        existing.target_id = run.target_id
        existing.scenario_id = run.scenario_id
        existing.scenario_version = run.scenario_version
        existing.campaign_id = run.campaign_id
        existing.campaign_plan_index = run.campaign_plan_index
        existing.campaign_occurrence_index = run.campaign_occurrence_index
        existing.status = run.status.value
        existing.triggered_by = run.triggered_by.value
        existing.completed_at = run.completed_at
        if run.summary is not None:
            existing.set_summary(run.summary.model_dump(mode="json"))
        self.session.add(existing)
        self.session.commit()
        self.session.refresh(existing)
        return self._from_db(existing)

    def get_turns(self, run_id: str) -> list[TurnDB]:
        statement = select(TurnDB).where(TurnDB.run_id == run_id).order_by(TurnDB.sent_at)
        return list(self.session.exec(statement).all())

    def get_results(self, run_id: str) -> list[EvalResult]:
        statement = select(EvalResultDB).where(EvalResultDB.run_id == run_id)
        return [result_from_db(r) for r in self.session.exec(statement).all()]