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

"""Repository for evaluation campaigns (评估活动)."""

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

from sqlalchemy import update as sql_update

from agenteval.exploration.models import ExplorationSessionStatus
from agenteval.models import Campaign, CampaignStatus, CampaignSummary
from agenteval.storage.db import CampaignDB, ExplorationSessionDB
from agenteval.storage.repository.base import BaseRepository


class CampaignWriteStatus(str, Enum):
    """Outcome of a conditional Campaign lifecycle write."""

    APPLIED = "applied"
    NOT_FOUND = "not_found"
    CONFLICT = "conflict"


@dataclass(frozen=True)
class CampaignWriteResult:
    """Typed result returned by Campaign compare-and-set operations."""

    status: CampaignWriteStatus
    campaign: Optional[Campaign] = None

    @property
    def applied(self) -> bool:
        return self.status is CampaignWriteStatus.APPLIED


class CampaignRepository(BaseRepository[Campaign, CampaignDB]):
    """Repository for evaluation campaigns (评估活动)."""

    _table = CampaignDB
    _order_by = "created_at"

    def _to_db(self, campaign: Campaign) -> CampaignDB:
        db = CampaignDB(
            id=campaign.id,
            name=campaign.name,
            target_id=campaign.target_id,
            window_seconds=campaign.window_seconds,
            time_scale=campaign.time_scale,
            status=campaign.status.value,
            started_at=campaign.started_at,
            completed_at=campaign.completed_at,
            created_at=campaign.created_at,
            analysis_model_config_id=campaign.analysis_model_config_id,
            last_patrolled_at=campaign.last_patrolled_at,
        )
        db.set_plan([entry.model_dump(mode="json") for entry in campaign.plan])
        if campaign.summary:
            db.set_summary(campaign.summary.model_dump(mode="json"))
        if campaign.exploration_seeds is not None:
            db.set_exploration_seeds(campaign.exploration_seeds.model_dump(mode="json"))
        if campaign.exploration_budget is not None:
            db.set_exploration_budget(campaign.exploration_budget.model_dump(mode="json"))
        return db

    def _from_db(self, db: CampaignDB) -> Campaign:
        return Campaign(
            id=db.id,
            name=db.name,
            target_id=db.target_id,
            window_seconds=db.window_seconds,
            time_scale=db.time_scale,
            plan=db.get_plan(),
            status=db.status,
            started_at=db.started_at,
            completed_at=db.completed_at,
            created_at=db.created_at,
            summary=db.get_summary(),
            analysis_model_config_id=db.analysis_model_config_id,
            exploration_seeds=db.get_exploration_seeds(),
            exploration_budget=db.get_exploration_budget(),
            last_patrolled_at=db.last_patrolled_at,
        )

    def update(self, campaign: Campaign) -> Optional[Campaign]:
        existing = self.session.get(CampaignDB, campaign.id)
        if not existing:
            return None
        existing.name = campaign.name
        existing.target_id = campaign.target_id
        existing.window_seconds = campaign.window_seconds
        existing.time_scale = campaign.time_scale
        existing.set_plan([entry.model_dump(mode="json") for entry in campaign.plan])
        existing.status = campaign.status.value
        existing.started_at = campaign.started_at
        existing.completed_at = campaign.completed_at
        existing.analysis_model_config_id = campaign.analysis_model_config_id
        existing.last_patrolled_at = campaign.last_patrolled_at
        if campaign.summary is not None:
            existing.set_summary(campaign.summary.model_dump(mode="json"))
        if campaign.exploration_seeds is not None:
            existing.set_exploration_seeds(campaign.exploration_seeds.model_dump(mode="json"))
        if campaign.exploration_budget is not None:
            existing.set_exploration_budget(campaign.exploration_budget.model_dump(mode="json"))
        self.session.add(existing)
        self.session.commit()
        self.session.refresh(existing)
        return self._from_db(existing)

    def touch_patrol_watermark(self, campaign_id: str, at: datetime) -> None:
        """窄口径原子更新：只写巡检水位，不覆写并发的 status / summary 变更。"""
        db = self.session.get(CampaignDB, campaign_id)
        if not db:
            return
        db.last_patrolled_at = at
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)

    def mark_cancelled(self, campaign_id: str, at: datetime) -> Optional[Campaign]:
        """窄口径终态迁移：只写 status + completed_at，不抹掉水位与调度进度。"""
        db = self.session.get(CampaignDB, campaign_id)
        if not db:
            return None
        db.status = CampaignStatus.CANCELLED.value
        db.completed_at = at
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return self._from_db(db)

    def _compare_and_set_status(
        self,
        campaign_id: str,
        *,
        expected_statuses: tuple[CampaignStatus, ...],
        target_status: CampaignStatus,
        at: datetime,
        settle_exploration: bool = False,
    ) -> CampaignWriteResult:
        """Apply one conditional Campaign transition and optional settlement.

        Final Campaign status and expiration of running exploration sessions
        share one transaction. A settlement failure therefore rolls the
        Campaign back to its prior runnable state.
        """

        values: dict[str, object] = {"status": target_status.value}
        if target_status is CampaignStatus.RUNNING:
            values["started_at"] = at
        if target_status in (CampaignStatus.CANCELLED, CampaignStatus.COMPLETED, CampaignStatus.FAILED):
            values["completed_at"] = at

        statement = (
            sql_update(CampaignDB)
            .where(
                CampaignDB.id == campaign_id,
                CampaignDB.status.in_([status.value for status in expected_statuses]),
            )
            .values(**values)
        )
        try:
            result = self.session.exec(statement)
            if result.rowcount != 1:
                self.session.rollback()
                db = self.session.get(CampaignDB, campaign_id)
                status = CampaignWriteStatus.NOT_FOUND if db is None else CampaignWriteStatus.CONFLICT
                return CampaignWriteResult(status=status, campaign=self._from_db(db) if db else None)

            if settle_exploration:
                self.session.exec(
                    sql_update(ExplorationSessionDB)
                    .where(
                        ExplorationSessionDB.campaign_id == campaign_id,
                        ExplorationSessionDB.status == ExplorationSessionStatus.RUNNING.value,
                    )
                    .values(status=ExplorationSessionStatus.EXPIRED.value, closed_at=at)
                )
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise

        self.session.expire_all()
        db = self.session.get(CampaignDB, campaign_id)
        return CampaignWriteResult(
            status=CampaignWriteStatus.APPLIED,
            campaign=self._from_db(db) if db else None,
        )

    def cancel_if_active(self, campaign_id: str, at: datetime) -> CampaignWriteResult:
        """Atomically cancel and expire running exploration sessions."""

        return self._compare_and_set_status(
            campaign_id,
            expected_statuses=(CampaignStatus.PLANNED, CampaignStatus.RUNNING),
            target_status=CampaignStatus.CANCELLED,
            at=at,
            settle_exploration=True,
        )

    def complete_if_running(self, campaign_id: str, at: datetime) -> CampaignWriteResult:
        """Atomically complete and expire running exploration sessions."""

        return self._compare_and_set_status(
            campaign_id,
            expected_statuses=(CampaignStatus.RUNNING,),
            target_status=CampaignStatus.COMPLETED,
            at=at,
            settle_exploration=True,
        )

    def start_if_planned(self, campaign_id: str, at: datetime) -> CampaignWriteResult:
        """Atomically stamp a planned Campaign as running."""

        return self._compare_and_set_status(
            campaign_id,
            expected_statuses=(CampaignStatus.PLANNED,),
            target_status=CampaignStatus.RUNNING,
            at=at,
        )

    def save_scheduler_state(self, campaign_id: str, summary: CampaignSummary) -> None:
        """窄口径调度进度持久化：只写 summary 列，不覆写并发的状态 / 水位变更。"""
        db = self.session.get(CampaignDB, campaign_id)
        if not db:
            return
        db.set_summary(summary.model_dump(mode="json"))
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)