huskies/server/src/chat/commands/cost.rs

//! Handler for the `cost` command.

use std::collections::HashMap;

use super::status::story_short_label;
use super::CommandContext;

/// Show token spend: 24h total, top 5 stories, agent-type breakdown, and
/// all-time total.
pub(super) fn handle_cost(ctx: &CommandContext) -> Option<String> {
    let records = match crate::agents::token_usage::read_all(ctx.project_root) {
        Ok(r) => r,
        Err(e) => return Some(format!("Failed to read token usage: {e}")),
    };

    if records.is_empty() {
        return Some("**Token Spend**\n\nNo usage records found.".to_string());
    }

    let now = chrono::Utc::now();
    let cutoff = now - chrono::Duration::hours(24);

    // Partition into 24h window and all-time
    let mut recent = Vec::new();
    let mut all_time_cost = 0.0;
    for r in &records {
        all_time_cost += r.usage.total_cost_usd;
        if let Ok(ts) = chrono::DateTime::parse_from_rfc3339(&r.timestamp)
            && ts >= cutoff
        {
            recent.push(r);
        }
    }

    // 24h total
    let recent_cost: f64 = recent.iter().map(|r| r.usage.total_cost_usd).sum();

    let mut out = String::from("**Token Spend**\n\n");
    out.push_str(&format!("**Last 24h:** ${:.2}\n", recent_cost));
    out.push_str(&format!("**All-time:** ${:.2}\n\n", all_time_cost));

    // Top 5 most expensive stories (last 24h)
    let mut story_costs: HashMap<&str, f64> = HashMap::new();
    for r in &recent {
        *story_costs.entry(r.story_id.as_str()).or_default() += r.usage.total_cost_usd;
    }
    let mut story_list: Vec<(&str, f64)> = story_costs.into_iter().collect();
    story_list.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    story_list.truncate(5);

    out.push_str("**Top Stories (24h)**\n");
    if story_list.is_empty() {
        out.push_str("  *(none)*\n");
    } else {
        for (story_id, cost) in &story_list {
            let label = story_short_label(story_id, None);
            out.push_str(&format!("  • {label} — ${cost:.2}\n"));
        }
    }
    out.push('\n');

    // Breakdown by agent type (last 24h)
    // Agent names follow pattern "coder-1", "qa-1", "mergemaster" — extract
    // the type as everything before the last '-' digit, or the full name.
    let mut type_costs: HashMap<String, f64> = HashMap::new();
    for r in &recent {
        let agent_type = extract_agent_type(&r.agent_name);
        *type_costs.entry(agent_type).or_default() += r.usage.total_cost_usd;
    }
    let mut type_list: Vec<(String, f64)> = type_costs.into_iter().collect();
    type_list.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    out.push_str("**By Agent Type (24h)**\n");
    if type_list.is_empty() {
        out.push_str("  *(none)*\n");
    } else {
        for (agent_type, cost) in &type_list {
            out.push_str(&format!("  • {agent_type} — ${cost:.2}\n"));
        }
    }

    Some(out)
}

/// Extract the agent type from an agent name.
///
/// Agent names like "coder-1", "qa-2", "mergemaster" map to types "coder",
/// "qa", "mergemaster". If the name ends with `-<digits>`, strip the suffix.
pub(super) fn extract_agent_type(agent_name: &str) -> String {
    if let Some(pos) = agent_name.rfind('-') {
        let suffix = &agent_name[pos + 1..];
        if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_digit()) {
            return agent_name[..pos].to_string();
        }
    }
    agent_name.to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agents::AgentPool;
    use std::sync::Arc;

    fn write_token_records(root: &std::path::Path, records: &[crate::agents::token_usage::TokenUsageRecord]) {
        for r in records {
            crate::agents::token_usage::append_record(root, r).unwrap();
        }
    }

    fn make_usage(cost: f64) -> crate::agents::TokenUsage {
        crate::agents::TokenUsage {
            input_tokens: 100,
            output_tokens: 200,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            total_cost_usd: cost,
        }
    }

    fn make_record(story_id: &str, agent_name: &str, cost: f64, hours_ago: i64) -> crate::agents::token_usage::TokenUsageRecord {
        let ts = (chrono::Utc::now() - chrono::Duration::hours(hours_ago)).to_rfc3339();
        crate::agents::token_usage::TokenUsageRecord {
            story_id: story_id.to_string(),
            agent_name: agent_name.to_string(),
            timestamp: ts,
            model: None,
            usage: make_usage(cost),
        }
    }

    fn cost_cmd_with_root(root: &std::path::Path) -> Option<String> {
        use super::super::{CommandDispatch, try_handle_command};
        use std::collections::HashSet;
        use std::sync::Mutex;

        let agents = Arc::new(AgentPool::new_test(3000));
        let ambient_rooms = Arc::new(Mutex::new(HashSet::new()));
        let room_id = "!test:example.com".to_string();
        let dispatch = CommandDispatch {
            bot_name: "Timmy",
            bot_user_id: "@timmy:homeserver.local",
            project_root: root,
            agents: &agents,
            ambient_rooms: &ambient_rooms,
            room_id: &room_id,
        };
        try_handle_command(&dispatch, "@timmy cost")
    }

    #[test]
    fn cost_command_is_registered() {
        use super::super::commands;
        let found = commands().iter().any(|c| c.name == "cost");
        assert!(found, "cost command must be in the registry");
    }

    #[test]
    fn cost_command_appears_in_help() {
        let result = super::super::tests::try_cmd_addressed("Timmy", "@timmy:homeserver.local", "@timmy help");
        let output = result.unwrap();
        assert!(output.contains("cost"), "help should list cost command: {output}");
    }

    #[test]
    fn cost_command_no_records() {
        let tmp = tempfile::TempDir::new().unwrap();
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        assert!(output.contains("No usage records found"), "should show empty message: {output}");
    }

    #[test]
    fn cost_command_shows_24h_total() {
        let tmp = tempfile::TempDir::new().unwrap();
        write_token_records(tmp.path(), &[
            make_record("42_story_foo", "coder-1", 1.50, 2),
            make_record("42_story_foo", "coder-1", 0.50, 5),
        ]);
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        assert!(output.contains("**Last 24h:** $2.00"), "should show 24h total: {output}");
    }

    #[test]
    fn cost_command_excludes_old_from_24h() {
        let tmp = tempfile::TempDir::new().unwrap();
        write_token_records(tmp.path(), &[
            make_record("42_story_foo", "coder-1", 1.00, 2),   // within 24h
            make_record("43_story_bar", "coder-1", 5.00, 48),  // older
        ]);
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        assert!(output.contains("**Last 24h:** $1.00"), "should only count recent: {output}");
        assert!(output.contains("**All-time:** $6.00"), "all-time should include everything: {output}");
    }

    #[test]
    fn cost_command_shows_top_stories() {
        let tmp = tempfile::TempDir::new().unwrap();
        write_token_records(tmp.path(), &[
            make_record("42_story_foo", "coder-1", 3.00, 1),
            make_record("43_story_bar", "coder-1", 1.00, 1),
            make_record("42_story_foo", "qa-1", 2.00, 1),
        ]);
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        assert!(output.contains("Top Stories"), "should have top stories section: {output}");
        // Story 42 ($5.00) should appear before story 43 ($1.00)
        let pos_42 = output.find("42").unwrap();
        let pos_43 = output.find("43").unwrap();
        assert!(pos_42 < pos_43, "story 42 should appear before 43 (sorted by cost): {output}");
    }

    #[test]
    fn cost_command_limits_to_5_stories() {
        let tmp = tempfile::TempDir::new().unwrap();
        let mut records = Vec::new();
        for i in 1..=7 {
            records.push(make_record(&format!("{i}_story_s{i}"), "coder-1", i as f64, 1));
        }
        write_token_records(tmp.path(), &records);
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        // The top 5 most expensive are stories 7,6,5,4,3. Stories 1 and 2 should be excluded.
        let top_section = output.split("**By Agent Type").next().unwrap();
        assert!(!top_section.contains("• 1 —"), "story 1 should not be in top 5: {output}");
        assert!(!top_section.contains("• 2 —"), "story 2 should not be in top 5: {output}");
    }

    #[test]
    fn cost_command_shows_agent_type_breakdown() {
        let tmp = tempfile::TempDir::new().unwrap();
        write_token_records(tmp.path(), &[
            make_record("42_story_foo", "coder-1", 2.00, 1),
            make_record("42_story_foo", "qa-1", 1.50, 1),
            make_record("42_story_foo", "mergemaster", 0.50, 1),
        ]);
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        assert!(output.contains("By Agent Type"), "should have agent type section: {output}");
        assert!(output.contains("coder"), "should show coder type: {output}");
        assert!(output.contains("qa"), "should show qa type: {output}");
        assert!(output.contains("mergemaster"), "should show mergemaster type: {output}");
    }

    #[test]
    fn cost_command_shows_all_time_total() {
        let tmp = tempfile::TempDir::new().unwrap();
        write_token_records(tmp.path(), &[
            make_record("42_story_foo", "coder-1", 1.00, 2),
            make_record("43_story_bar", "coder-1", 9.00, 100),
        ]);
        let output = cost_cmd_with_root(tmp.path()).unwrap();
        assert!(output.contains("**All-time:** $10.00"), "should show all-time total: {output}");
    }

    #[test]
    fn cost_command_case_insensitive() {
        let result = super::super::tests::try_cmd_addressed("Timmy", "@timmy:homeserver.local", "@timmy COST");
        assert!(result.is_some(), "COST should match case-insensitively");
    }

    // -- extract_agent_type -------------------------------------------------

    #[test]
    fn extract_agent_type_strips_numeric_suffix() {
        assert_eq!(extract_agent_type("coder-1"), "coder");
        assert_eq!(extract_agent_type("qa-2"), "qa");
    }

    #[test]
    fn extract_agent_type_keeps_non_numeric_suffix() {
        assert_eq!(extract_agent_type("mergemaster"), "mergemaster");
        assert_eq!(extract_agent_type("coder-alpha"), "coder-alpha");
    }
}