AgentEvalTool/backend/agenteval/evaluation/campaign_lifecycle.py

"""Durable Campaign creation lifecycle.

This module owns the create seam: reference validation and Campaign
persistence happen before the process-local scheduler is launched. The
database therefore remains authoritative if task startup fails or the process
stops between commit and launch.
"""

from dataclasses import dataclass
from typing import Callable, Optional

from sqlmodel import Session

from agenteval.models import Campaign, CampaignPlanEntry, CampaignStatus, ExplorationBudgetConfig, ExplorationSeeds
from agenteval.storage.db import utc_now
from agenteval.storage.model_config_repository import ModelConfigRepository
from agenteval.storage.repository import (
    CampaignAnalysisRepository,
    CampaignPeriodComparisonRepository,
    CampaignRepository,
    ExplorationSessionRepository,
    RunRepository,
    ScenarioRepository,
    TargetRepository,
)


@dataclass(frozen=True)
class CampaignCreateError(Exception):
    """Validation failure translated by the HTTP adapter."""

    status_code: int
    detail: str


@dataclass(frozen=True)
class CampaignCancelError(Exception):
    """Cancellation failure translated by the HTTP adapter."""

    status_code: int
    detail: str


@dataclass(frozen=True)
class CampaignRecoverySummary:
    """Durable work reconciled during process startup."""

    interrupted_runs: int = 0
    interrupted_analysis: int = 0
    resumed_campaigns: int = 0
    resumed_analysis: int = 0


def start_campaign(
    session: Session,
    campaign_id: str,
    *,
    launch: Optional[Callable[[str, Session], object]] = None,
) -> Optional[Campaign]:
    """Start a planned Campaign through a conditional lifecycle write."""
    repo = CampaignRepository(session)
    campaign = repo.get(campaign_id)
    if campaign is None or campaign.status in (
        CampaignStatus.COMPLETED,
        CampaignStatus.CANCELLED,
        CampaignStatus.FAILED,
    ):
        return None
    if campaign.status is CampaignStatus.PLANNED:
        status, campaign = repo.start_if_planned(campaign_id, utc_now())
        if status != "applied" or campaign is None:
            return None
    if launch is not None and campaign.id:
        launch(campaign.id, session)
    return campaign


def create_campaign(
    session: Session,
    *,
    name: str,
    target_id: str,
    window_seconds: int,
    time_scale: float,
    plan: list[CampaignPlanEntry],
    analysis_model_config_id: Optional[str] = None,
    exploration_seeds: Optional[ExplorationSeeds] = None,
    exploration_budget: Optional[ExplorationBudgetConfig] = None,
    launch: Optional[Callable[[str, Session], object]] = None,
) -> Campaign:
    """Validate, commit and then launch one running Campaign.

    ``CampaignRepository.create`` is the transaction seam. The optional
    ``launch`` callback is invoked only after that commit, which makes this
    function straightforward to test with a fake scheduler and ensures a
    launch failure cannot erase the durable Campaign row.
    """
    if TargetRepository(session).get(target_id) is None:
        raise CampaignCreateError(404, "target not found")

    scenario_repo = ScenarioRepository(session)
    for entry in plan:
        if scenario_repo.get(entry.scenario_id) is None:
            raise CampaignCreateError(404, f"scenario not found: {entry.scenario_id}")

    if analysis_model_config_id is not None and ModelConfigRepository(session).get(analysis_model_config_id) is None:
        raise CampaignCreateError(400, "analysis model config not found")

    if exploration_seeds is not None and not exploration_seeds.personas and not exploration_seeds.goals:
        exploration_seeds = None

    campaign = Campaign(
        name=name,
        target_id=target_id,
        window_seconds=window_seconds,
        time_scale=time_scale,
        plan=plan,
        status=CampaignStatus.RUNNING,
        started_at=utc_now(),
        analysis_model_config_id=analysis_model_config_id,
        exploration_seeds=exploration_seeds,
        exploration_budget=exploration_budget,
    )
    campaign = CampaignRepository(session).create(campaign)

    # Deliberately after the repository commit. Startup recovery can relaunch
    # this Campaign if the process dies before the callback runs.
    if launch is not None and campaign.id:
        launch(campaign.id, session)
    return campaign


def cancel_campaign(
    session: Session,
    campaign_id: str,
    *,
    stop: Optional[Callable[[str], object]] = None,
) -> Campaign:
    """Cancel, settle exploration, then signal the process-local scheduler."""
    repo = CampaignRepository(session)
    status, campaign = repo.cancel_if_active(campaign_id, utc_now())
    if status == "not_found":
        raise CampaignCancelError(404, "campaign not found")
    if status == "conflict" or campaign is None:
        raise CampaignCancelError(400, "campaign is not in a cancellable state")

    ExplorationSessionRepository(session).expire_running_sessions(campaign_id)
    if stop is not None:
        stop(campaign_id)
    return campaign


def recover_campaign_runtime(session: Session, *, tick_seconds: float = 1.0) -> CampaignRecoverySummary:
    """Reconcile all durable Campaign work and relaunch safe tasks.

    The database is inspected and repaired before process-local tasks are
    launched. Child Run reconciliation remains inside each resumed Campaign
    loop, so a pending claim is resumed there while an orphaned running claim
    is marked interrupted without replaying messages.
    """
    interrupted_runs = RunRepository(session).mark_orphans_failed()
    interrupted_analysis = CampaignAnalysisRepository(session).mark_orphans_failed()
    interrupted_analysis += CampaignPeriodComparisonRepository(session).mark_orphans_failed()

    from agenteval.evaluation.analysis import resume_queued_campaign_analysis
    from agenteval.evaluation.campaign_runner import resume_running_campaigns

    resumed_campaigns = resume_running_campaigns(session, tick_seconds=tick_seconds)
    resumed_analysis = resume_queued_campaign_analysis(session)
    return CampaignRecoverySummary(
        interrupted_runs=interrupted_runs,
        interrupted_analysis=interrupted_analysis,
        resumed_campaigns=resumed_campaigns,
        resumed_analysis=resumed_analysis,
    )


def complete_campaign(
    session: Session,
    campaign_id: str,
    *,
    settle: Optional[Callable[[str, Session], object]] = None,
) -> Campaign:
    """Complete only a still-running Campaign, then settle exploration.

    A competing cancellation wins because the status predicate is evaluated in
    the database. Settlement is deliberately performed after the CAS commit;
    a failed settlement can be retried without reverting the terminal status.
    """
    repo = CampaignRepository(session)
    status, campaign = repo.complete_if_running(campaign_id, utc_now())
    if status == "not_found":
        raise CampaignCancelError(404, "campaign not found")
    if status == "conflict" or campaign is None:
        raise CampaignCancelError(409, "campaign is not running")
    if settle is not None:
        settle(campaign_id, session)
    else:
        ExplorationSessionRepository(session).expire_running_sessions(campaign_id)
    return campaign