//! 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 { 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 = 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 `-`, 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 { 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"); } }