huskies/server/src/http/mcp/diagnostics/usage.rs

//! MCP token-usage reporting tool (`tool_get_token_usage`).

use serde_json::{Value, json};

use crate::http::context::AppContext;

pub(crate) fn tool_get_token_usage(args: &Value, ctx: &AppContext) -> Result<String, String> {
    let root = ctx.state.get_project_root()?;
    let filter_story = args.get("story_id").and_then(|v| v.as_str());

    let all_records = crate::agents::token_usage::read_all(&root)?;
    let records: Vec<_> = all_records
        .into_iter()
        .filter(|r| filter_story.is_none_or(|s| r.story_id == s))
        .collect();

    let total_cost: f64 = records.iter().map(|r| r.usage.total_cost_usd).sum();
    let total_input: u64 = records.iter().map(|r| r.usage.input_tokens).sum();
    let total_output: u64 = records.iter().map(|r| r.usage.output_tokens).sum();
    let total_cache_create: u64 = records
        .iter()
        .map(|r| r.usage.cache_creation_input_tokens)
        .sum();
    let total_cache_read: u64 = records
        .iter()
        .map(|r| r.usage.cache_read_input_tokens)
        .sum();

    serde_json::to_string_pretty(&json!({
        "records": records.iter().map(|r| json!({
            "story_id": r.story_id,
            "agent_name": r.agent_name,
            "timestamp": r.timestamp,
            "input_tokens": r.usage.input_tokens,
            "output_tokens": r.usage.output_tokens,
            "cache_creation_input_tokens": r.usage.cache_creation_input_tokens,
            "cache_read_input_tokens": r.usage.cache_read_input_tokens,
            "total_cost_usd": r.usage.total_cost_usd,
        })).collect::<Vec<_>>(),
        "totals": {
            "records": records.len(),
            "input_tokens": total_input,
            "output_tokens": total_output,
            "cache_creation_input_tokens": total_cache_create,
            "cache_read_input_tokens": total_cache_read,
            "total_cost_usd": total_cost,
        }
    }))
    .map_err(|e| format!("Serialization error: {e}"))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::test_helpers::test_ctx;

    #[test]
    fn tool_get_token_usage_empty_returns_zero_totals() {
        let tmp = tempfile::tempdir().unwrap();
        let ctx = test_ctx(tmp.path());
        let result = tool_get_token_usage(&json!({}), &ctx).unwrap();
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["records"].as_array().unwrap().len(), 0);
        assert_eq!(parsed["totals"]["records"], 0);
        assert_eq!(parsed["totals"]["total_cost_usd"], 0.0);
    }

    #[test]
    fn tool_get_token_usage_returns_written_records() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        let ctx = test_ctx(root);

        let usage = crate::agents::TokenUsage {
            input_tokens: 100,
            output_tokens: 200,
            cache_creation_input_tokens: 5000,
            cache_read_input_tokens: 10000,
            total_cost_usd: 1.57,
        };
        let record =
            crate::agents::token_usage::build_record("42_story_foo", "coder-1", None, usage);
        crate::agents::token_usage::append_record(root, &record).unwrap();

        let result = tool_get_token_usage(&json!({}), &ctx).unwrap();
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["records"].as_array().unwrap().len(), 1);
        assert_eq!(parsed["records"][0]["story_id"], "42_story_foo");
        assert_eq!(parsed["records"][0]["agent_name"], "coder-1");
        assert_eq!(parsed["records"][0]["input_tokens"], 100);
        assert_eq!(parsed["totals"]["records"], 1);
        assert!((parsed["totals"]["total_cost_usd"].as_f64().unwrap() - 1.57).abs() < f64::EPSILON);
    }

    #[test]
    fn tool_get_token_usage_filters_by_story_id() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        let ctx = test_ctx(root);

        let usage = crate::agents::TokenUsage {
            input_tokens: 50,
            output_tokens: 60,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
            total_cost_usd: 0.5,
        };
        let r1 =
            crate::agents::token_usage::build_record("10_story_a", "coder-1", None, usage.clone());
        let r2 = crate::agents::token_usage::build_record("20_story_b", "coder-2", None, usage);
        crate::agents::token_usage::append_record(root, &r1).unwrap();
        crate::agents::token_usage::append_record(root, &r2).unwrap();

        let result = tool_get_token_usage(&json!({"story_id": "10_story_a"}), &ctx).unwrap();
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["records"].as_array().unwrap().len(), 1);
        assert_eq!(parsed["records"][0]["story_id"], "10_story_a");
        assert_eq!(parsed["totals"]["records"], 1);
    }
}