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 aid = match agent_id {
20        Some(id) => id,
21        None => {
22            let agents = store.top_agents(20);
23            if agents.is_empty() {
24                return "No agent cost data recorded yet.".to_string();
25            }
26            let mut lines = vec![format!("All agents ({}):", agents.len())];
27            for a in &agents {
28                lines.push(format!(
29                    "  {} ({}) — {} calls, ${:.4}",
30                    a.agent_id, a.agent_type, a.total_calls, a.cost_usd
31                ));
32            }
33            return lines.join("\n");
34        }
35    };
36
37    match store.agents.get(aid) {
38        Some(agent) => {
39            let mut lines = vec![
40                format!("Agent: {} ({})", agent.agent_id, agent.agent_type),
41                format!("  Calls: {}", agent.total_calls),
42                format!("  Input tokens: {}", agent.total_input_tokens),
43                format!("  Output tokens: {}", agent.total_output_tokens),
44                format!("  Estimated cost: ${:.4}", agent.cost_usd),
45            ];
46            if !agent.tools_used.is_empty() {
47                lines.push("  Tools used:".to_string());
48                let mut tools: Vec<_> = agent.tools_used.iter().collect();
49                tools.sort_by(|a, b| b.1.cmp(a.1));
50                for (name, count) in tools {
51                    lines.push(format!("    {name}: {count} calls"));
52                }
53            }
54            lines.join("\n")
55        }
56        None => format!("No cost data found for agent '{aid}'"),
57    }
58}
59
60fn handle_tool_breakdown(store: &CostStore, limit: usize) -> String {
61    let tools = store.top_tools(limit);
62    if tools.is_empty() {
63        return "No tool cost data recorded yet.".to_string();
64    }
65
66    let mut lines = vec![format!("Tool Cost Breakdown ({} tools):", tools.len())];
67    for (i, tool) in tools.iter().enumerate() {
68        lines.push(format!(
69            "  {}. {} — {} calls, avg {:.0} in + {:.0} out tok, ${:.4}",
70            i + 1,
71            tool.tool_name,
72            tool.total_calls,
73            tool.avg_input_tokens,
74            tool.avg_output_tokens,
75            tool.cost_usd
76        ));
77    }
78    lines.join("\n")
79}
80
81fn handle_reset() -> String {
82    let store = CostStore::default();
83    match store.save() {
84        Ok(()) => "Cost attribution data has been reset.".to_string(),
85        Err(e) => format!("Error resetting cost data: {e}"),
86    }
87}
lean_ctx/tools/ctx_cost.rs

lean_ctx/tools/
ctx_cost.rs
