lean_ctx/tools/
ctx_cost.rs

1use crate::core::a2a::cost_attribution::{format_cost_report, CostStore};
2use crate::core::gain::GainEngine;
3
4pub fn handle(action: &str, agent_id: Option<&str>, limit: Option<usize>) -> String {
5    let engine = GainEngine::load();
6    let lim = limit.unwrap_or(10);
7
8    match action {
9        "report" | "status" => format_cost_report(&engine.costs, lim),
10        "agent" => handle_agent_detail(&engine.costs, agent_id),
11        "tools" => handle_tool_breakdown(&engine.costs, lim),
12        "json" => serde_json::to_string_pretty(&engine.costs).unwrap_or_else(|_| "{}".to_string()),
13        "reset" => handle_reset(),
14        _ => format!("Unknown action '{action}'. Available: report, agent, tools, json, reset"),
15    }
16}
17
18fn handle_agent_detail(store: &CostStore, agent_id: Option<&str>) -> String {
19    let Some(aid) = agent_id else {
20        let agents = store.top_agents(20);
21        if agents.is_empty() {
22            return "No agent cost data recorded yet.".to_string();
23        }
24        let mut lines = vec![format!("All agents ({}):", agents.len())];
25        for a in &agents {
26            lines.push(format!(
27                "  {} ({}) — {} calls, ${:.4}",
28                a.agent_id, a.agent_type, a.total_calls, a.cost_usd
29            ));
30        }
31        return lines.join("\n");
32    };
33
34    match store.agents.get(aid) {
35        Some(agent) => {
36            let mut lines = vec![
37                format!("Agent: {} ({})", agent.agent_id, agent.agent_type),
38                format!("  Calls: {}", agent.total_calls),
39                format!("  Input tokens: {}", agent.total_input_tokens),
40                format!("  Output tokens: {}", agent.total_output_tokens),
41                format!("  Cached tokens: {}", agent.total_cached_tokens),
42                format!("  Estimated cost: ${:.4}", agent.cost_usd),
43            ];
44            if !agent.tools_used.is_empty() {
45                lines.push("  Tools used:".to_string());
46                let mut tools: Vec<_> = agent.tools_used.iter().collect();
47                tools.sort_by_key(|item| std::cmp::Reverse(*item.1));
48                for (name, count) in tools {
49                    lines.push(format!("    {name}: {count} calls"));
50                }
51            }
52            lines.join("\n")
53        }
54        None => format!("No cost data found for agent '{aid}'"),
55    }
56}
57
58fn handle_tool_breakdown(store: &CostStore, limit: usize) -> String {
59    let tools = store.top_tools(limit);
60    if tools.is_empty() {
61        return "No tool cost data recorded yet.".to_string();
62    }
63
64    let mut lines = vec![format!("Tool Cost Breakdown ({} tools):", tools.len())];
65    for (i, tool) in tools.iter().enumerate() {
66        lines.push(format!(
67            "  {}. {} — {} calls, avg {:.0} in + {:.0} out + {:.0} cached tok, ${:.4}",
68            i + 1,
69            tool.tool_name,
70            tool.total_calls,
71            tool.avg_input_tokens,
72            tool.avg_output_tokens,
73            tool.avg_cached_tokens,
74            tool.cost_usd
75        ));
76    }
77    lines.join("\n")
78}
79
80fn handle_reset() -> String {
81    let store = CostStore::default();
82    match store.save() {
83        Ok(()) => "Cost attribution data has been reset.".to_string(),
84        Err(e) => format!("Error resetting cost data: {e}"),
85    }
86}
lean_ctx/tools/ctx_cost.rs

lean_ctx/tools/
ctx_cost.rs
