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 (
    CampaignRepository,
    CampaignWriteStatus,
    ScenarioRepository,
    TargetRepository,
)


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

    status_code: int
    detail: str


@dataclass(frozen=True)
class CampaignLifecycleError(Exception):
    """Campaign transition failure translated by the HTTP adapter."""

    status_code: int
    detail: str


def start_campaign(
    session: Session,
    campaign_id: str,
    *,
    launch: Optional[Callable[[str], 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:
        result = repo.start_if_planned(campaign_id, utc_now())
        campaign = result.campaign
        if not result.applied or campaign is None:
            return None
    if launch is not None and campaign.id:
        launch(campaign.id)
    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], 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)
    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)
    result = repo.cancel_if_active(campaign_id, utc_now())
    campaign = result.campaign
    if result.status is CampaignWriteStatus.NOT_FOUND:
        raise CampaignLifecycleError(404, "campaign not found")
    if result.status is CampaignWriteStatus.CONFLICT or campaign is None:
        raise CampaignLifecycleError(400, "campaign is not in a cancellable state")

    if stop is not None:
        stop(campaign_id)
    return campaign


def complete_campaign(
    session: Session,
    campaign_id: str,
) -> Campaign:
    """Atomically complete a running Campaign and settle exploration.

    A competing cancellation wins because the status predicate is evaluated in
    the database. Settlement and the terminal status share one transaction, so
    a failure leaves the Campaign running for a later tick to retry.
    """
    repo = CampaignRepository(session)
    result = repo.complete_if_running(campaign_id, utc_now())
    campaign = result.campaign
    if result.status is CampaignWriteStatus.NOT_FOUND:
        raise CampaignLifecycleError(404, "campaign not found")
    if result.status is CampaignWriteStatus.CONFLICT or campaign is None:
        raise CampaignLifecycleError(409, "campaign is not running")
    return campaign