story-kit: merge 302_story_bot_cost_command_shows_total_and_per_story_token_spend

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<String> {
     None
 }
 
+/// Show token spend: 24h total, top 5 stories, agent-type breakdown, and
+/// all-time total.
+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.
+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<String> {
+        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");
+    }
 }