AgentEvalTool/backend/agenteval/storage/repository.py

"""Repository layer for database access."""

from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Generic, Optional, TypeVar

from sqlalchemy import insert as sql_insert
from sqlalchemy import literal
from sqlalchemy import update as sql_update
from sqlalchemy.exc import IntegrityError
from sqlmodel import Session, select

from agenteval.exploration.models import ExplorationMessage, ExplorationSession, ExplorationSessionStatus
from agenteval.models import (
    Campaign,
    CampaignStatus,
    CampaignSummary,
    Case,
    EvalResult,
    EvalRun,
    EvalTarget,
    RunStatus,
    RunTrigger,
    Scenario,
)
from agenteval.services.model_configs import ModelConfigService
from agenteval.storage.db import (
    CampaignAnalysisDB,
    CampaignDB,
    CampaignPeriodComparisonDB,
    EvalResultDB,
    EvalRunDB,
    EvalTargetDB,
    ExplorationMessageDB,
    ExplorationSessionDB,
    ScenarioDB,
    TurnDB,
    get_session,
    new_uuid,
    utc_now,
)
from agenteval.storage.model_config_repository import ScenarioModelBindingRepository

M = TypeVar("M")  # domain model
DB = TypeVar("DB")  # persisted table row


class BaseRepository(Generic[M, DB]):
    """Shared CRUD skeleton for id-keyed entity repositories.

    Subclasses declare the table (``_table``) and the ``list_all`` ordering
    column name (``_order_by``, newest-first), and implement the ``_to_db`` /
    ``_from_db`` converter pair. The converters are instance methods so a
    subclass whose ``_from_db`` needs cross-table reads (e.g. Scenario's model
    bindings) can reach ``self.session``. Entities with bespoke create/update
    (binding validation, versioning) override just those methods.
    """

    _table: type
    _order_by: str

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def _to_db(self, obj: M) -> DB:
        raise NotImplementedError

    def _from_db(self, db: DB) -> M:
        raise NotImplementedError

    def list_all(self) -> list[M]:
        column = getattr(self._table, self._order_by)
        statement = select(self._table).order_by(column.desc())
        return [self._from_db(r) for r in self.session.exec(statement).all()]

    def get(self, entity_id: str) -> Optional[M]:
        db = self.session.get(self._table, entity_id)
        return self._from_db(db) if db else None

    def create(self, obj: M) -> M:
        db = self._to_db(obj)
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return self._from_db(db)

    def delete(self, entity_id: str) -> bool:
        db = self.session.get(self._table, entity_id)
        if not db:
            return False
        self.session.delete(db)
        self.session.commit()
        return True


def _result_to_db(result: EvalResult) -> EvalResultDB:
    return EvalResultDB(
        id=result.id,
        run_id=result.run_id,
        case_id=result.case_id,
        turn_id=result.turn_id,
        rule_type=result.rule_type,
        passed=result.passed,
        score=result.score,
        reason=result.reason,
    )


def _result_from_db(db: EvalResultDB) -> EvalResult:
    return EvalResult(
        id=db.id,
        run_id=db.run_id,
        case_id=db.case_id,
        turn_id=db.turn_id,
        rule_type=db.rule_type,
        passed=db.passed,
        score=db.score,
        reason=db.reason,
    )


class TargetRepository(BaseRepository[EvalTarget, EvalTargetDB]):
    """Repository for evaluation targets."""

    _table = EvalTargetDB
    _order_by = "created_at"

    def _to_db(self, target: EvalTarget) -> EvalTargetDB:
        db = EvalTargetDB(
            id=target.id,
            name=target.name,
            description=target.description,
            platform=target.platform.value,
            channel_type=target.channel_type.value,
            status=target.status.value,
            created_at=target.created_at,
            updated_at=target.updated_at or utc_now(),
        )
        db.set_config(target.channel_config)
        return db

    def _from_db(self, db: EvalTargetDB) -> EvalTarget:
        return EvalTarget(
            id=db.id,
            name=db.name,
            description=db.description,
            platform=db.platform,
            channel_type=db.channel_type,
            channel_config=db.get_config(),
            status=db.status,
            created_at=db.created_at,
            updated_at=db.updated_at,
        )

    def update(self, target: EvalTarget) -> Optional[EvalTarget]:
        existing = self.session.get(EvalTargetDB, target.id)
        if not existing:
            return None
        existing.name = target.name
        existing.description = target.description
        existing.platform = target.platform.value
        existing.channel_type = target.channel_type.value
        existing.status = target.status.value
        existing.set_config(target.channel_config)
        existing.updated_at = utc_now()
        self.session.add(existing)
        self.session.commit()
        self.session.refresh(existing)
        return self._from_db(existing)


class ScenarioRepository(BaseRepository[Scenario, ScenarioDB]):
    """Repository for evaluation scenarios."""

    _table = ScenarioDB
    _order_by = "created_at"

    def _to_db(self, scenario: Scenario) -> ScenarioDB:
        db = ScenarioDB(
            id=scenario.id,
            name=scenario.name,
            description=scenario.description,
            created_at=scenario.created_at,
            updated_at=scenario.updated_at or utc_now(),
        )
        db.set_tags(scenario.tags)
        # mode="json" 与 update() 的考纲比较保持同一序列化形态，避免假升版
        db.set_cases([case.model_dump(mode="json") for case in scenario.cases])
        db.set_llm_config(scenario.llm_config)
        return db

    def _from_db(self, db: ScenarioDB) -> Scenario:
        bindings = ScenarioModelBindingRepository(self.session).get_for_scenario(db.id or "")
        return Scenario(
            id=db.id,
            name=db.name,
            description=db.description,
            tags=db.get_tags(),
            cases=[Case(**case) for case in db.get_cases()],
            model_bindings=bindings,
            llm_config=db.get_llm_config(),
            version=db.version or 1,
            created_at=db.created_at,
            updated_at=db.updated_at,
        )

    def create(self, scenario: Scenario) -> Scenario:
        db = self._to_db(scenario)
        bindings = {purpose.value: config_id for purpose, config_id in scenario.model_bindings.items()}
        try:
            ModelConfigService(self.session).validate_bindings(bindings)
            self.session.add(db)
            self.session.flush()
            ScenarioModelBindingRepository(self.session).replace_for_scenario(db.id or "", bindings)
            self.session.commit()
            self.session.refresh(db)
        except Exception:
            self.session.rollback()
            raise
        return self._from_db(db)

    def update(self, scenario: Scenario) -> Optional[Scenario]:
        existing = self.session.get(ScenarioDB, scenario.id)
        if not existing:
            return None
        bindings = {purpose.value: config_id for purpose, config_id in scenario.model_bindings.items()}
        try:
            ModelConfigService(self.session).validate_bindings(bindings)
            # 考纲字段（cases / model_bindings / llm_config）变更才升版（ADR-0001）；
            # 版本由系统维护，忽略 scenario.version 的外部传入值。
            new_cases = [case.model_dump(mode="json") for case in scenario.cases]
            old_bindings = ScenarioModelBindingRepository(self.session).get_for_scenario(existing.id or "")
            syllabus_changed = (
                existing.get_cases() != new_cases
                or existing.get_llm_config() != scenario.llm_config
                or old_bindings != bindings
            )
            if syllabus_changed:
                existing.version = (existing.version or 1) + 1
            existing.name = scenario.name
            existing.description = scenario.description
            existing.set_tags(scenario.tags)
            existing.set_cases(new_cases)
            existing.set_llm_config(scenario.llm_config)
            existing.updated_at = utc_now()
            self.session.add(existing)
            ScenarioModelBindingRepository(self.session).replace_for_scenario(existing.id or "", bindings)
            self.session.commit()
            self.session.refresh(existing)
        except Exception:
            self.session.rollback()
            raise
        return self._from_db(existing)

    def delete(self, scenario_id: str) -> bool:
        db = self.session.get(ScenarioDB, scenario_id)
        if not db:
            return False
        try:
            ScenarioModelBindingRepository(self.session).delete_for_scenario(scenario_id)
            self.session.delete(db)
            self.session.commit()
        except Exception:
            self.session.rollback()
            raise
        return True

    def name_map(self) -> dict[str, str]:
        """scenario_id → 名称映射：报告 / 时间线 / 列表等读路径共用的场景名取法。"""
        return {sid: name for sid, name in self.session.exec(select(ScenarioDB.id, ScenarioDB.name)).all()}


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 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()
        failed: list[EvalRunDB] = []
        for db in candidates:
            if db.status == RunStatus.PENDING.value and self._is_recoverable_campaign_pending(db):
                continue
            self._set_interrupted(db)
            failed.append(db)
        if failed:
            self.session.commit()
        return len(failed)

    def _is_recoverable_campaign_pending(self, db: EvalRunDB) -> bool:
        if (
            db.campaign_id is None
            or db.campaign_plan_index is None
            or db.campaign_occurrence_index is None
        ):
            return False
        campaign = self.session.get(CampaignDB, 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()]


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)


class AsyncJobRepository(Generic[DB]):
    """Base class for async job repositories (analysis, comparison).

    Provides common pattern: get_by_campaign, mark_orphans_failed.
    Subclasses implement upsert with their specific fields.
    """

    _table: type[DB]

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def get_by_campaign(self, campaign_id: str) -> Optional[DB]:
        statement = select(self._table).where(self._table.campaign_id == campaign_id)  # type: ignore[attr-defined]
        return self.session.exec(statement).first()

    def mark_orphans_failed(self, error_message: str) -> int:
        """服务启动时清理：把滞留的 generating 行标记为 failed。

        异步任务是进程内 asyncio 任务，服务重启后不会恢复；不清理则这些
        行永远停留在 generating（僵尸状态）。
        """
        rows = self.session.exec(
            select(self._table).where(self._table.status == "generating")  # type: ignore[attr-defined]
        ).all()
        for row in rows:
            row.status = "failed"  # type: ignore[attr-defined]
            row.error = error_message  # type: ignore[attr-defined]
            row.updated_at = utc_now()  # type: ignore[attr-defined]
            self.session.add(row)
        if rows:
            self.session.commit()
        return len(rows)


class CampaignAnalysisRepository(AsyncJobRepository[CampaignAnalysisDB]):
    """Repository for campaign analysis rows (one per campaign, upserted)."""

    _table = CampaignAnalysisDB

    def enqueue(self, campaign_id: str, *, triggered_by: str = "manual") -> CampaignAnalysisDB:
        """Persist an analysis job before launching its process-local task."""
        row = self.get_by_campaign(campaign_id)
        if row is None:
            row = CampaignAnalysisDB(campaign_id=campaign_id)
        row.status = "queued"
        row.result = None
        row.error = None
        row.triggered_by = triggered_by
        row.updated_at = utc_now()
        self.session.add(row)
        self.session.commit()
        self.session.refresh(row)
        return row

    def list_queued(self) -> list[CampaignAnalysisDB]:
        return list(self.session.exec(select(CampaignAnalysisDB).where(CampaignAnalysisDB.status == "queued")).all())

    def upsert(
        self,
        campaign_id: str,
        *,
        status: str,
        result: Optional[dict] = None,
        model_config_id: Optional[str] = None,
        error: Optional[str] = None,
        triggered_by: str = "manual",
    ) -> CampaignAnalysisDB:
        row = self.get_by_campaign(campaign_id)
        if row is None:
            row = CampaignAnalysisDB(campaign_id=campaign_id)
        row.status = status
        if result is not None:
            row.set_result(result)
        else:
            row.result = None
        row.model_config_id = model_config_id
        row.error = error
        row.triggered_by = triggered_by
        row.updated_at = utc_now()
        self.session.add(row)
        self.session.commit()
        self.session.refresh(row)
        return row

    def mark_orphans_failed(self) -> int:
        """服务启动时清理：把滞留的 generating 分析行标记为 failed。"""
        return super().mark_orphans_failed("服务重启导致分析生成中断")


class CampaignPeriodComparisonRepository(AsyncJobRepository[CampaignPeriodComparisonDB]):
    """Repository for period-comparison rows (one per campaign, upserted)."""

    _table = CampaignPeriodComparisonDB

    def upsert(
        self,
        campaign_id: str,
        *,
        status: str,
        baseline_campaign_id: Optional[str] = None,
        result: Optional[dict] = None,
        model_config_id: Optional[str] = None,
        error: Optional[str] = None,
        triggered_by: str = "manual",
    ) -> CampaignPeriodComparisonDB:
        row = self.get_by_campaign(campaign_id)
        if row is None:
            row = CampaignPeriodComparisonDB(
                campaign_id=campaign_id,
                baseline_campaign_id=baseline_campaign_id or "",
            )
        row.status = status
        if baseline_campaign_id is not None:
            row.baseline_campaign_id = baseline_campaign_id
        if result is not None:
            row.set_result(result)
        else:
            row.result = None
        row.model_config_id = model_config_id
        row.error = error
        row.triggered_by = triggered_by
        row.updated_at = utc_now()
        self.session.add(row)
        self.session.commit()
        self.session.refresh(row)
        return row

    def mark_orphans_failed(self) -> int:
        """服务启动时清理：把滞留的 generating 周期对比行标记为 failed。"""
        return super().mark_orphans_failed("服务重启导致周期对比生成中断")


class ResultRepository:
    """Repository for evaluation results."""

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def save_turn(self, turn) -> TurnDB:
        db = TurnDB(
            id=turn.id,
            run_id=turn.run_id,
            case_id=turn.case_id,
            round_index=turn.round_index,
            question_msg_id=turn.question_msg_id,
            sent_at=turn.sent_at,
            received_at=turn.received_at,
            latency_ms=turn.latency_ms,
        )
        db.set_sent_message(turn.sent_message)
        db.set_reply(turn.reply)
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return db

    def update_turn_exchange(
        self,
        turn_id: str,
        *,
        question_msg_id: Optional[str],
        reply: Optional[dict],
        received_at: Optional[datetime],
        latency_ms: Optional[int],
    ) -> Optional[TurnDB]:
        """Attach exchange facts to an already-persisted Turn.

        This is deliberately narrower than updating a whole Turn: the sent
        message, case identity, ordering and send timestamp remain untouched.
        A caller can therefore commit the sent fact before polling and safely
        complete the same ledger row after a reply, timeout or poll failure.
        """
        db = self.session.get(TurnDB, turn_id)
        if db is None:
            return None
        db.question_msg_id = question_msg_id
        db.set_reply(reply)
        db.received_at = received_at
        db.latency_ms = latency_ms
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return db

    def save_result(self, result: EvalResult) -> EvalResult:
        db = _result_to_db(result)
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        return _result_from_db(db)


class ExplorationSessionRepository(BaseRepository[ExplorationSession, ExplorationSessionDB]):
    """Repository for virtual-user exploration sessions (探索会话)."""

    _table = ExplorationSessionDB
    _order_by = "created_at"

    def _copy_mutable(self, db: ExplorationSessionDB, session_obj: ExplorationSession) -> None:
        db.goal = session_obj.goal
        db.status = session_obj.status.value
        db.triggered_by = session_obj.triggered_by.value
        db.turn_count = session_obj.turn_count
        db.error = session_obj.error
        db.closed_at = session_obj.closed_at
        db.set_persona(session_obj.persona)
        if session_obj.seed_ref is not None:
            db.set_seed_ref(session_obj.seed_ref)
        if session_obj.experience is not None:
            db.set_experience(session_obj.experience)
        if session_obj.judge_review is not None:
            db.set_judge_review(session_obj.judge_review)

    def _to_db(self, session_obj: ExplorationSession) -> ExplorationSessionDB:
        db = ExplorationSessionDB(
            id=session_obj.id,
            campaign_id=session_obj.campaign_id,
            target_id=session_obj.target_id,
            created_at=session_obj.created_at,
        )
        self._copy_mutable(db, session_obj)
        return db

    def _from_db(self, db: ExplorationSessionDB) -> ExplorationSession:
        return ExplorationSession(
            id=db.id,
            campaign_id=db.campaign_id,
            target_id=db.target_id,
            persona=db.get_persona(),
            goal=db.goal,
            seed_ref=db.get_seed_ref(),
            status=db.status,
            triggered_by=db.triggered_by,
            experience=db.get_experience(),
            judge_review=db.get_judge_review(),
            turn_count=db.turn_count,
            error=db.error,
            created_at=db.created_at,
            closed_at=db.closed_at,
        )

    def update(self, session_obj: ExplorationSession) -> Optional[ExplorationSession]:
        existing = self.session.get(ExplorationSessionDB, session_obj.id)
        if not existing:
            return None
        self._copy_mutable(existing, session_obj)
        self.session.add(existing)
        self.session.commit()
        self.session.refresh(existing)
        return self._from_db(existing)

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

    def expire_running_sessions(self, campaign_id: str) -> int:
        """Expire every still-running exploration session of a finalized campaign.

        Returns the number of sessions expired. Completed/failed sessions keep
        their evidence untouched.
        """
        expired = 0
        for session_obj in self.list_by_campaign(campaign_id):
            if session_obj.status != ExplorationSessionStatus.RUNNING:
                continue
            session_obj.status = ExplorationSessionStatus.EXPIRED
            session_obj.closed_at = utc_now()
            self.update(session_obj)
            expired += 1
        return expired


class ExplorationMessageRepository:
    """Append-only repository for exploration session chat rows."""

    def __init__(self, session: Optional[Session] = None):
        self.session = session or get_session()

    def save_message(self, message: ExplorationMessage) -> ExplorationMessage:
        db = ExplorationMessageDB(
            id=message.id,
            session_id=message.session_id,
            round_index=message.round_index,
            role=message.role,
            content=message.content,
            latency_ms=message.latency_ms,
            created_at=message.created_at,
        )
        self.session.add(db)
        self.session.commit()
        self.session.refresh(db)
        message.id = db.id
        return message

    def list_by_session(self, session_id: str) -> list[ExplorationMessage]:
        statement = (
            select(ExplorationMessageDB)
            .where(ExplorationMessageDB.session_id == session_id)
            .order_by(ExplorationMessageDB.created_at)
        )
        return [
            ExplorationMessage(
                id=r.id,
                session_id=r.session_id,
                round_index=r.round_index,
                role=r.role,
                content=r.content,
                latency_ms=r.latency_ms,
                created_at=r.created_at,
            )
            for r in self.session.exec(statement).all()
        ]