huskies/server/src/service/agents/token.rs

//! Pure token usage aggregation — no I/O, no side effects.
//!
//! Functions here take slices of `TokenUsageRecord` (already loaded by `io.rs`)
//! and compute summaries. Tests cover every branch without touching the filesystem.
use crate::agents::token_usage::TokenUsageRecord;
use std::collections::HashMap;

/// Per-agent cost breakdown entry.
#[derive(Debug, Clone, PartialEq)]
pub struct AgentTokenCost {
    pub agent_name: String,
    pub model: Option<crate::agents::AgentModel>,
    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,
}

/// Aggregated token cost for a story.
#[derive(Debug, Clone, PartialEq)]
pub struct TokenCostSummary {
    pub total_cost_usd: f64,
    pub agents: Vec<AgentTokenCost>,
}

/// Aggregate token usage records for a single story.
///
/// Records for other stories are ignored. The returned `agents` list is sorted
/// alphabetically by `agent_name` for deterministic output. Returns a zero-cost
/// summary when no records match the given `story_id`.
pub fn aggregate_for_story(records: &[TokenUsageRecord], story_id: &str) -> TokenCostSummary {
    let mut agent_map: HashMap<String, AgentTokenCost> = HashMap::new();
    let mut total_cost_usd = 0.0_f64;

    for record in records.iter().filter(|r| r.story_id == story_id) {
        total_cost_usd += record.usage.total_cost_usd;
        let entry = agent_map
            .entry(record.agent_name.clone())
            .or_insert_with(|| AgentTokenCost {
                agent_name: record.agent_name.clone(),
                model: record.model.clone(),
                input_tokens: 0,
                output_tokens: 0,
                cache_creation_input_tokens: 0,
                cache_read_input_tokens: 0,
                total_cost_usd: 0.0,
            });
        entry.input_tokens += record.usage.input_tokens;
        entry.output_tokens += record.usage.output_tokens;
        entry.cache_creation_input_tokens += record.usage.cache_creation_input_tokens;
        entry.cache_read_input_tokens += record.usage.cache_read_input_tokens;
        entry.total_cost_usd += record.usage.total_cost_usd;
    }

    let mut agents: Vec<AgentTokenCost> = agent_map.into_values().collect();
    agents.sort_by(|a, b| a.agent_name.cmp(&b.agent_name));

    TokenCostSummary {
        total_cost_usd,
        agents,
    }
}

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

    fn make_record(story_id: &str, agent: &str, cost: f64) -> TokenUsageRecord {
        TokenUsageRecord {
            story_id: story_id.to_string(),
            agent_name: agent.to_string(),
            timestamp: "2024-01-01T00:00:00Z".to_string(),
            model: None,
            usage: TokenUsage {
                input_tokens: 100,
                output_tokens: 50,
                cache_creation_input_tokens: 10,
                cache_read_input_tokens: 20,
                total_cost_usd: cost,
            },
        }
    }

    #[test]
    fn aggregate_returns_zero_when_no_records() {
        let summary = aggregate_for_story(&[], "42_story_foo");
        assert_eq!(summary.total_cost_usd, 0.0);
        assert!(summary.agents.is_empty());
    }

    #[test]
    fn aggregate_filters_to_story_id() {
        let records = vec![
            make_record("42_story_foo", "coder-1", 1.0),
            make_record("99_story_other", "coder-1", 5.0),
        ];
        let summary = aggregate_for_story(&records, "42_story_foo");
        assert!((summary.total_cost_usd - 1.0).abs() < f64::EPSILON);
        assert_eq!(summary.agents.len(), 1);
    }

    #[test]
    fn aggregate_sums_tokens_per_agent() {
        let records = vec![
            make_record("42_story_foo", "coder-1", 1.0),
            make_record("42_story_foo", "coder-1", 2.0),
        ];
        let summary = aggregate_for_story(&records, "42_story_foo");
        assert!((summary.total_cost_usd - 3.0).abs() < f64::EPSILON);
        assert_eq!(summary.agents.len(), 1);
        assert_eq!(summary.agents[0].input_tokens, 200);
        assert_eq!(summary.agents[0].output_tokens, 100);
        assert!((summary.agents[0].total_cost_usd - 3.0).abs() < f64::EPSILON);
    }

    #[test]
    fn aggregate_splits_by_agent() {
        let records = vec![
            make_record("42_story_foo", "coder-1", 1.0),
            make_record("42_story_foo", "qa", 0.5),
        ];
        let summary = aggregate_for_story(&records, "42_story_foo");
        assert!((summary.total_cost_usd - 1.5).abs() < f64::EPSILON);
        assert_eq!(summary.agents.len(), 2);
        // sorted alphabetically
        assert_eq!(summary.agents[0].agent_name, "coder-1");
        assert_eq!(summary.agents[1].agent_name, "qa");
    }

    #[test]
    fn aggregate_sorts_agents_alphabetically() {
        let records = vec![
            make_record("42_story_foo", "z-agent", 1.0),
            make_record("42_story_foo", "a-agent", 1.0),
            make_record("42_story_foo", "m-agent", 1.0),
        ];
        let summary = aggregate_for_story(&records, "42_story_foo");
        assert_eq!(summary.agents[0].agent_name, "a-agent");
        assert_eq!(summary.agents[1].agent_name, "m-agent");
        assert_eq!(summary.agents[2].agent_name, "z-agent");
    }

    #[test]
    fn aggregate_returns_zero_when_no_matching_story() {
        let records = vec![make_record("99_other", "coder-1", 5.0)];
        let summary = aggregate_for_story(&records, "42_story_foo");
        assert_eq!(summary.total_cost_usd, 0.0);
        assert!(summary.agents.is_empty());
    }

    #[test]
    fn aggregate_preserves_model_from_first_record() {
        let mut r = make_record("42_story_foo", "coder-1", 1.0);
        r.model = Some(crate::agents::AgentModel::Sonnet);
        let summary = aggregate_for_story(&[r], "42_story_foo");
        assert_eq!(
            summary.agents[0].model,
            Some(crate::agents::AgentModel::Sonnet)
        );
    }
}