huskies/server/src/agents/mod.rs

//! Agent subsystem — types, configuration, and orchestration for coding agents.
/// Acceptance-gate checks (clippy, tests, doc coverage).
pub mod gates;
/// Agent start/stop and pipeline advancement on completion.
pub mod lifecycle;
/// Constructs the system prompt sent to coding agents.
pub mod local_prompt;
/// Merge-conflict resolution helpers.
pub mod merge;
pub(crate) mod pool;
pub(crate) mod pty;
/// Runtime backends (Claude Code, OpenAI, Gemini) that execute agent sessions.
pub mod runtime;
/// Persistent session-ID storage for agent resume support.
pub mod session_store;
/// Token-usage tracking and budget estimation.
pub mod token_usage;

/// Typed agent model enum (Sonnet/Opus/Haiku).
pub mod model;
pub use model::AgentModel;

use crate::config::AgentConfig;
use serde::{Deserialize, Serialize};

pub use lifecycle::{
    close_bug_to_archive, feature_branch_has_unmerged_changes, move_story_to_done,
    move_story_to_merge, move_story_to_qa, move_story_to_stage, reject_story_from_qa,
};
pub use pool::AgentPool;

/// Events emitted during server startup reconciliation to broadcast real-time
/// progress to connected WebSocket clients.
#[derive(Debug, Clone, Serialize)]
pub struct ReconciliationEvent {
    /// The story being reconciled, or empty string for the overall "done" event.
    pub story_id: String,
    /// Coarse status: "checking", "gates_running", "advanced", "skipped", "failed", "done"
    pub status: String,
    /// Human-readable details.
    pub message: String,
}

/// Events streamed from a running agent to SSE clients.
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentEvent {
    /// Agent status changed.
    Status {
        story_id: String,
        agent_name: String,
        status: String,
    },
    /// Raw text output from the agent process.
    Output {
        story_id: String,
        agent_name: String,
        text: String,
    },
    /// Agent produced a JSON event from `--output-format stream-json`.
    AgentJson {
        story_id: String,
        agent_name: String,
        data: serde_json::Value,
    },
    /// Agent finished.
    Done {
        story_id: String,
        agent_name: String,
        session_id: Option<String>,
    },
    /// Agent errored.
    Error {
        story_id: String,
        agent_name: String,
        message: String,
    },
    /// Thinking tokens from an extended-thinking block.
    Thinking {
        story_id: String,
        agent_name: String,
        text: String,
    },
}

#[derive(Debug, Clone, Serialize, PartialEq)]
#[serde(rename_all = "snake_case")]
/// Lifecycle state of an agent session.
pub enum AgentStatus {
    Pending,
    Running,
    Completed,
    Failed,
}

/// The execution state of a rate-limited agent session.
///
/// Replaces the legacy `throttled: bool` flag, carrying the expiry time so
/// the scheduler can decide when to allow a retry rather than skipping
/// indefinitely.
#[derive(Debug, Clone, Serialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AgentExecution {
    /// The agent hit a rate-limit and is paused until `until`.
    Throttled {
        /// UTC instant at which the rate limit expires and the agent may resume.
        until: chrono::DateTime<chrono::Utc>,
    },
}

impl AgentExecution {
    /// Return `true` if the throttle period has not yet elapsed.
    pub fn is_active(&self) -> bool {
        match self {
            Self::Throttled { until } => chrono::Utc::now() < *until,
        }
    }
}

/// Why an agent was forcibly terminated by the watchdog.
#[derive(Debug, Clone, Serialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum TerminationReason {
    /// Agent exceeded its configured `max_turns`.
    TurnLimit,
    /// Agent exceeded its configured `max_budget_usd`.
    BudgetLimit,
}

impl std::fmt::Display for AgentStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Pending => write!(f, "pending"),
            Self::Running => write!(f, "running"),
            Self::Completed => write!(f, "completed"),
            Self::Failed => write!(f, "failed"),
        }
    }
}

/// Pipeline stages for automatic story advancement.
#[derive(Debug, Clone, PartialEq)]
pub enum PipelineStage {
    /// Coding agents (coder-1, coder-2, etc.)
    Coder,
    /// QA review agent
    Qa,
    /// Mergemaster agent
    Mergemaster,
    /// Supervisors and unknown agents — no automatic advancement.
    Other,
}

/// Determine the pipeline stage from an agent name.
pub fn pipeline_stage(agent_name: &str) -> PipelineStage {
    match agent_name {
        "qa" => PipelineStage::Qa,
        "mergemaster" => PipelineStage::Mergemaster,
        name if name.starts_with("coder") => PipelineStage::Coder,
        _ => PipelineStage::Other,
    }
}

/// Map a pipeline [`Stage`] to the canonical [`PipelineStage`] for LLM agent spawning.
///
/// Returns `None` for stages where no LLM agent should be active (terminal states,
/// blocked, frozen, or unclassified merge failures requiring human intervention).
/// Returns `Some(stage)` naming the single LLM-agent type that may run on this story.
/// Used by `validate_agent_stage` and `reconcile_canonical_agents` to enforce the
/// one-agent-per-story invariant (story 1100).
pub fn canonical_pipeline_stage(s: &crate::pipeline_state::Stage) -> Option<PipelineStage> {
    use crate::pipeline_state::{MergeFailureKind, Stage};
    match s {
        Stage::Coding { .. } => Some(PipelineStage::Coder),
        Stage::Qa => Some(PipelineStage::Qa),
        Stage::Merge { .. } => Some(PipelineStage::Mergemaster),
        Stage::MergeFailure {
            kind: MergeFailureKind::ConflictDetected(_),
            ..
        } => Some(PipelineStage::Mergemaster),
        Stage::MergeFailure {
            kind: MergeFailureKind::GatesFailed(_),
            ..
        } => Some(PipelineStage::Coder),
        Stage::MergeFailureFinal { .. } => Some(PipelineStage::Mergemaster),
        Stage::Upcoming
        | Stage::Backlog
        | Stage::MergeFailure { .. }
        | Stage::Done { .. }
        | Stage::Blocked { .. }
        | Stage::Archived { .. }
        | Stage::Frozen { .. }
        | Stage::ReviewHold { .. }
        | Stage::Abandoned { .. }
        | Stage::Superseded { .. }
        | Stage::Rejected { .. } => None,
    }
}

/// Determine the pipeline stage for a configured agent.
///
/// Prefers the explicit `stage` config field (added in Bug 150) over the
/// legacy name-based heuristic so that agents with non-standard names
/// (e.g. `qa-2`, `coder-opus`) are assigned to the correct stage.
pub(crate) fn agent_config_stage(cfg: &AgentConfig) -> PipelineStage {
    match cfg.stage.as_deref() {
        Some("coder") => PipelineStage::Coder,
        Some("qa") => PipelineStage::Qa,
        Some("mergemaster") => PipelineStage::Mergemaster,
        Some(_) => PipelineStage::Other,
        None => pipeline_stage(&cfg.name),
    }
}

/// Completion report produced when acceptance gates are run.
///
/// Created automatically by the server when an agent process exits normally,
/// or via the internal `report_completion` method.
#[derive(Debug, Serialize, Clone)]
pub struct CompletionReport {
    pub summary: String,
    pub gates_passed: bool,
    pub gate_output: String,
    /// True when the coder exited with no commits but left uncommitted content in
    /// the worktree.  The pipeline advance will issue a commit-only recovery respawn
    /// rather than a normal retry, and will NOT consume a `retry_count` slot.
    #[serde(default)]
    pub needs_commit_recovery: bool,
}

/// Token usage from a Claude Code session's `result` event.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct TokenUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_creation_input_tokens: u64,
    pub cache_read_input_tokens: u64,
    pub total_cost_usd: f64,
}

impl TokenUsage {
    /// Estimate USD cost from token counts using approximate per-model pricing.
    ///
    /// Used by the watchdog to compute in-flight budget from per-message usage
    /// data in the agent log (since `total_cost_usd` is only available in the
    /// `result` event at session end).  Uses conservative (high) pricing when
    /// the model is unknown so budget limits are hit sooner rather than later.
    pub fn estimate_cost_usd(&self, model: Option<&AgentModel>) -> f64 {
        // Per-million-token pricing (input, output, cache_read, cache_create).
        let (inp, out, cr, cc) = match model {
            Some(AgentModel::Haiku) => (0.80, 4.0, 0.08, 1.00),
            Some(AgentModel::Sonnet) => (3.0, 15.0, 0.30, 3.75),
            // Opus, Other, or unknown → most expensive = conservative.
            Some(AgentModel::Opus) | Some(AgentModel::Other(_)) | None => (15.0, 75.0, 1.50, 18.75),
        };
        (self.input_tokens as f64 * inp
            + self.output_tokens as f64 * out
            + self.cache_read_input_tokens as f64 * cr
            + self.cache_creation_input_tokens as f64 * cc)
            / 1_000_000.0
    }

    /// Parse token usage from a Claude Code `result` JSON event.
    pub fn from_result_event(json: &serde_json::Value) -> Option<Self> {
        let usage = json.get("usage")?;
        Some(Self {
            input_tokens: usage
                .get("input_tokens")
                .and_then(|v| v.as_u64())
                .unwrap_or(0),
            output_tokens: usage
                .get("output_tokens")
                .and_then(|v| v.as_u64())
                .unwrap_or(0),
            cache_creation_input_tokens: usage
                .get("cache_creation_input_tokens")
                .and_then(|v| v.as_u64())
                .unwrap_or(0),
            cache_read_input_tokens: usage
                .get("cache_read_input_tokens")
                .and_then(|v| v.as_u64())
                .unwrap_or(0),
            total_cost_usd: json
                .get("total_cost_usd")
                .and_then(|v| v.as_f64())
                .unwrap_or(0.0),
        })
    }
}

#[derive(Debug, Serialize, Clone)]
/// Snapshot of a running or completed agent, exposed via the HTTP API.
pub struct AgentInfo {
    pub story_id: String,
    pub agent_name: String,
    pub status: AgentStatus,
    pub session_id: Option<String>,
    pub worktree_path: Option<String>,
    pub base_branch: Option<String>,
    pub completion: Option<CompletionReport>,
    /// UUID identifying the persistent log file for this session.
    pub log_session_id: Option<String>,
    /// Set when the agent is rate-limited; holds the UTC expiry time.
    /// `None` when the agent is not throttled.
    pub throttled: Option<chrono::DateTime<chrono::Utc>>,
    /// Set when the watchdog terminates the agent for exceeding a limit.
    pub termination_reason: Option<TerminationReason>,
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── pipeline_stage tests ──────────────────────────────────────────────────

    #[test]
    fn pipeline_stage_detects_coders() {
        assert_eq!(pipeline_stage("coder-1"), PipelineStage::Coder);
        assert_eq!(pipeline_stage("coder-2"), PipelineStage::Coder);
        assert_eq!(pipeline_stage("coder-3"), PipelineStage::Coder);
    }

    #[test]
    fn pipeline_stage_detects_qa() {
        assert_eq!(pipeline_stage("qa"), PipelineStage::Qa);
    }

    #[test]
    fn pipeline_stage_detects_mergemaster() {
        assert_eq!(pipeline_stage("mergemaster"), PipelineStage::Mergemaster);
    }

    #[test]
    fn pipeline_stage_supervisor_is_other() {
        assert_eq!(pipeline_stage("supervisor"), PipelineStage::Other);
        assert_eq!(pipeline_stage("default"), PipelineStage::Other);
        assert_eq!(pipeline_stage("unknown"), PipelineStage::Other);
    }
}