storkit/server/src/agents/mod.rs

pub mod gates;
pub mod lifecycle;
pub mod merge;
mod pool;
mod pty;
pub mod token_usage;

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

pub use lifecycle::{
    close_bug_to_archive, feature_branch_has_unmerged_changes, move_story_to_archived,
    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")]
pub enum AgentStatus {
    Pending,
    Running,
    Completed,
    Failed,
}

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,
    }
}

/// 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,
}

/// 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 {
    /// 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)]
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>,
}

#[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);
    }
}