diff --git a/server/src/matrix/commands.rs b/server/src/matrix/commands.rs index 85cec8c..09314a5 100644 --- a/server/src/matrix/commands.rs +++ b/server/src/matrix/commands.rs @@ -8,7 +8,7 @@ use crate::agents::{AgentPool, AgentStatus}; use crate::config::ProjectConfig; use matrix_sdk::ruma::OwnedRoomId; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::Path; use std::sync::{Arc, Mutex}; @@ -98,6 +98,11 @@ pub fn commands() -> &'static [BotCommand] { description: "Show live system and agent process dashboard (`htop`, `htop 10m`, `htop stop`)", handler: handle_htop_fallback, }, + BotCommand { + name: "cost", + description: "Show token spend: 24h total, top stories, breakdown by agent type, and all-time total", + handler: handle_cost, + }, ] } @@ -457,6 +462,97 @@ fn handle_htop_fallback(_ctx: &CommandContext) -> Option { None } +/// Show token spend: 24h total, top 5 stories, agent-type breakdown, and +/// all-time total. +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. +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() +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -1002,4 +1098,168 @@ mod tests { let result = try_cmd_addressed("Timmy", "@timmy:homeserver.local", "@timmy GIT"); assert!(result.is_some(), "GIT should match case-insensitively"); } + + // -- cost command ------------------------------------------------------- + + 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, + usage: make_usage(cost), + } + } + + fn cost_cmd_with_root(root: &std::path::Path) -> Option { + let agents = test_agents(); + let ambient_rooms = test_ambient_rooms(); + let room_id = make_room_id("!test:example.com"); + 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, + is_addressed: true, + }; + try_handle_command(&dispatch, "@timmy cost") + } + + #[test] + fn cost_command_is_registered() { + 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 = 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 = 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"); + } }