lean_ctx/core/a2a/
cost_attribution.rs

1use chrono::Utc;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::path::PathBuf;
5use std::sync::Mutex;
6
7#[derive(Debug, Clone, Serialize, Deserialize, Default)]
8pub struct CostStore {
9    pub agents: HashMap<String, AgentCost>,
10    pub tools: HashMap<String, ToolCost>,
11    pub sessions: Vec<SessionCostSnapshot>,
12    pub updated_at: Option<String>,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize, Default)]
16pub struct AgentCost {
17    pub agent_id: String,
18    pub agent_type: String,
19    #[serde(default)]
20    pub model_key: Option<String>,
21    #[serde(default)]
22    pub pricing_match: Option<String>,
23    pub total_input_tokens: u64,
24    pub total_output_tokens: u64,
25    pub total_cached_tokens: u64,
26    pub total_calls: u64,
27    pub cost_usd: f64,
28    pub tools_used: HashMap<String, u64>,
29    pub first_seen: Option<String>,
30    pub last_seen: Option<String>,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize, Default)]
34pub struct ToolCost {
35    pub tool_name: String,
36    pub total_input_tokens: u64,
37    pub total_output_tokens: u64,
38    #[serde(default)]
39    pub total_cached_tokens: u64,
40    pub total_calls: u64,
41    pub avg_input_tokens: f64,
42    pub avg_output_tokens: f64,
43    #[serde(default)]
44    pub avg_cached_tokens: f64,
45    pub cost_usd: f64,
46    /// RFC3339 timestamp of the most recent call to this tool. Set on the
47    /// existing post-dispatch write path, so the tool-health "rot" report
48    /// (#848) gets an honest last-used signal at no new hot-path cost.
49    #[serde(default)]
50    pub last_used: Option<String>,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct SessionCostSnapshot {
55    pub timestamp: String,
56    pub agent_id: String,
57    #[serde(default)]
58    pub model_key: Option<String>,
59    pub total_input: u64,
60    pub total_output: u64,
61    #[serde(default)]
62    pub total_cached: u64,
63    pub total_saved: u64,
64    pub cost_usd: f64,
65    pub duration_secs: u64,
66}
67
68pub fn estimate_cost(model_key: Option<&str>, input: u64, output: u64, cached: u64) -> f64 {
69    let pricing = crate::core::gain::model_pricing::ModelPricing::load();
70    let quote = pricing.quote(model_key);
71    quote.cost.estimate_usd(input, output, 0, cached)
72}
73
74static COST_BUFFER: Mutex<Option<CostStore>> = Mutex::new(None);
75
76impl CostStore {
77    pub fn load() -> Self {
78        let mut guard = COST_BUFFER
79            .lock()
80            .unwrap_or_else(std::sync::PoisonError::into_inner);
81        if let Some(ref store) = *guard {
82            return store.clone();
83        }
84
85        let store = load_from_disk();
86        *guard = Some(store.clone());
87        store
88    }
89
90    pub fn record_tool_call(
91        &mut self,
92        agent_id: &str,
93        agent_type: &str,
94        tool_name: &str,
95        input_tokens: u64,
96        output_tokens: u64,
97        cached_tokens: u64,
98    ) {
99        let now = Utc::now().to_rfc3339();
100        let pricing = crate::core::gain::model_pricing::ModelPricing::load();
101        // Honors a declared model for MCP-only IDEs (`[cost.models]`/default).
102        let quote = pricing.quote_for_client(agent_type);
103        let cost = quote
104            .cost
105            .estimate_usd(input_tokens, output_tokens, 0, cached_tokens);
106
107        let agent = self
108            .agents
109            .entry(agent_id.to_string())
110            .or_insert_with(|| AgentCost {
111                agent_id: agent_id.to_string(),
112                agent_type: agent_type.to_string(),
113                first_seen: Some(now.clone()),
114                ..Default::default()
115            });
116        agent.total_input_tokens += input_tokens;
117        agent.total_output_tokens += output_tokens;
118        agent.total_cached_tokens += cached_tokens;
119        agent.total_calls += 1;
120        agent.cost_usd += cost;
121        agent.last_seen = Some(now.clone());
122        agent.model_key = Some(quote.model_key.clone());
123        agent.pricing_match = Some(format!("{:?}", quote.match_kind));
124        *agent.tools_used.entry(tool_name.to_string()).or_insert(0) += 1;
125
126        let tool = self
127            .tools
128            .entry(tool_name.to_string())
129            .or_insert_with(|| ToolCost {
130                tool_name: tool_name.to_string(),
131                ..Default::default()
132            });
133        tool.total_input_tokens += input_tokens;
134        tool.total_output_tokens += output_tokens;
135        tool.total_cached_tokens += cached_tokens;
136        tool.total_calls += 1;
137        tool.cost_usd += cost;
138        tool.last_used = Some(now.clone());
139        if tool.total_calls > 0 {
140            tool.avg_input_tokens = tool.total_input_tokens as f64 / tool.total_calls as f64;
141            tool.avg_output_tokens = tool.total_output_tokens as f64 / tool.total_calls as f64;
142            tool.avg_cached_tokens = tool.total_cached_tokens as f64 / tool.total_calls as f64;
143        }
144
145        self.updated_at = Some(now);
146    }
147
148    pub fn save(&self) -> std::io::Result<()> {
149        save_to_disk(self)?;
150        let mut guard = COST_BUFFER
151            .lock()
152            .unwrap_or_else(std::sync::PoisonError::into_inner);
153        *guard = Some(self.clone());
154        Ok(())
155    }
156
157    pub fn top_agents(&self, limit: usize) -> Vec<&AgentCost> {
158        let mut agents: Vec<_> = self.agents.values().collect();
159        agents.sort_by(|a, b| {
160            b.cost_usd
161                .partial_cmp(&a.cost_usd)
162                .unwrap_or(std::cmp::Ordering::Equal)
163        });
164        agents.truncate(limit);
165        agents
166    }
167
168    pub fn top_tools(&self, limit: usize) -> Vec<&ToolCost> {
169        let mut tools: Vec<_> = self.tools.values().collect();
170        tools.sort_by(|a, b| {
171            b.cost_usd
172                .partial_cmp(&a.cost_usd)
173                .unwrap_or(std::cmp::Ordering::Equal)
174        });
175        tools.truncate(limit);
176        tools
177    }
178
179    pub fn total_cost(&self) -> f64 {
180        self.agents.values().map(|a| a.cost_usd).sum()
181    }
182
183    pub fn total_tokens(&self) -> (u64, u64, u64) {
184        let input: u64 = self.agents.values().map(|a| a.total_input_tokens).sum();
185        let output: u64 = self.agents.values().map(|a| a.total_output_tokens).sum();
186        let cached: u64 = self.agents.values().map(|a| a.total_cached_tokens).sum();
187        (input, output, cached)
188    }
189
190    pub fn add_session_snapshot(
191        &mut self,
192        agent_id: &str,
193        input: u64,
194        output: u64,
195        saved: u64,
196        duration_secs: u64,
197    ) {
198        let model_key = self
199            .agents
200            .get(agent_id)
201            .and_then(|a| a.model_key.as_deref())
202            .map(std::string::ToString::to_string);
203        let cost = estimate_cost(model_key.as_deref(), input, output, 0);
204        self.sessions.push(SessionCostSnapshot {
205            timestamp: Utc::now().to_rfc3339(),
206            agent_id: agent_id.to_string(),
207            model_key,
208            total_input: input,
209            total_output: output,
210            total_cached: 0,
211            total_saved: saved,
212            cost_usd: cost,
213            duration_secs,
214        });
215
216        if self.sessions.len() > 500 {
217            self.sessions.drain(0..self.sessions.len() - 500);
218        }
219    }
220}
221
222fn cost_store_path() -> Option<PathBuf> {
223    crate::core::paths::state_dir()
224        .ok()
225        .map(|d| d.join("cost_attribution.json"))
226}
227
228fn load_from_disk() -> CostStore {
229    let Some(path) = cost_store_path() else {
230        return CostStore::default();
231    };
232    match std::fs::read_to_string(&path) {
233        Ok(content) => serde_json::from_str(&content).unwrap_or_default(),
234        Err(_) => CostStore::default(),
235    }
236}
237
238fn save_to_disk(store: &CostStore) -> std::io::Result<()> {
239    let Some(path) = cost_store_path() else {
240        return Err(std::io::Error::new(
241            std::io::ErrorKind::NotFound,
242            "no home dir",
243        ));
244    };
245
246    if let Some(parent) = path.parent() {
247        std::fs::create_dir_all(parent)?;
248    }
249
250    let json = serde_json::to_string(store).map_err(std::io::Error::other)?;
251    let tmp = path.with_extension("tmp");
252    std::fs::write(&tmp, &json)?;
253    std::fs::rename(&tmp, &path)?;
254    Ok(())
255}
256
257pub fn format_cost_report(store: &CostStore, limit: usize) -> String {
258    let mut lines = Vec::new();
259    let (total_in, total_out, total_cached) = store.total_tokens();
260    let total_cost = store.total_cost();
261
262    lines.push(format!(
263        "Cost Attribution Report ({} agents, {} tools)",
264        store.agents.len(),
265        store.tools.len()
266    ));
267    lines.push(format!(
268        "Total: {total_in} input + {total_out} output + {total_cached} cached tokens = ${total_cost:.4}"
269    ));
270    if let Ok(m) = std::env::var("LEAN_CTX_MODEL").or_else(|_| std::env::var("LCTX_MODEL"))
271        && !m.trim().is_empty()
272    {
273        let pricing = crate::core::gain::model_pricing::ModelPricing::load();
274        let q = pricing.quote(Some(&m));
275        lines.push(format!(
276            "Pricing: model={} ({:?}) in=${:.2}/M out=${:.2}/M cacheR=${:.3}/M",
277            q.model_key,
278            q.match_kind,
279            q.cost.input_per_m,
280            q.cost.output_per_m,
281            q.cost.cache_read_per_m
282        ));
283        // Cached vs uncached input split (GL #573): cached tokens are
284        // billed at the cache-read rate, never at the input rate.
285        let uncached_cost = total_in as f64 / 1_000_000.0 * q.cost.input_per_m;
286        let cached_cost = total_cached as f64 / 1_000_000.0 * q.cost.cache_read_per_m;
287        lines.push(format!(
288                "Input split: uncached {total_in} tok = ${uncached_cost:.4} | cached {total_cached} tok = ${cached_cost:.4} (cache-read rate)"
289            ));
290    }
291    lines.push(String::new());
292
293    let top_agents = store.top_agents(limit);
294    if !top_agents.is_empty() {
295        lines.push("Top Agents by Cost:".to_string());
296        for (i, agent) in top_agents.iter().enumerate() {
297            lines.push(format!(
298                "  {}. {} ({}) — {} calls, {} in + {} out + {} cached tok, ${:.4}{}",
299                i + 1,
300                agent.agent_id,
301                agent.agent_type,
302                agent.total_calls,
303                agent.total_input_tokens,
304                agent.total_output_tokens,
305                agent.total_cached_tokens,
306                agent.cost_usd,
307                agent
308                    .model_key
309                    .as_deref()
310                    .map(|m| format!(" [{m}]"))
311                    .unwrap_or_default()
312            ));
313        }
314        lines.push(String::new());
315    }
316
317    let top_tools = store.top_tools(limit);
318    if !top_tools.is_empty() {
319        lines.push("Top Tools by Cost:".to_string());
320        for (i, tool) in top_tools.iter().enumerate() {
321            lines.push(format!(
322                "  {}. {} — {} calls, avg {:.0} in + {:.0} out + {:.0} cached tok, ${:.4}",
323                i + 1,
324                tool.tool_name,
325                tool.total_calls,
326                tool.avg_input_tokens,
327                tool.avg_output_tokens,
328                tool.avg_cached_tokens,
329                tool.cost_usd
330            ));
331        }
332    }
333
334    lines.join("\n")
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340
341    #[test]
342    fn cost_estimation() {
343        let cost = estimate_cost(Some("fallback-blended"), 1000, 100, 500);
344        assert!(cost > 0.0);
345    }
346
347    #[test]
348    fn record_and_query() {
349        let mut store = CostStore::default();
350        store.record_tool_call("agent-1", "mcp", "ctx_read", 5000, 200, 0);
351        store.record_tool_call("agent-1", "mcp", "ctx_read", 3000, 150, 500);
352        store.record_tool_call("agent-2", "cursor", "ctx_shell", 1000, 100, 0);
353
354        assert_eq!(store.agents.len(), 2);
355        assert_eq!(store.tools.len(), 2);
356
357        let agent1 = &store.agents["agent-1"];
358        assert_eq!(agent1.total_calls, 2);
359        assert_eq!(agent1.total_input_tokens, 8000);
360        assert_eq!(agent1.total_cached_tokens, 500);
361        assert_eq!(*agent1.tools_used.get("ctx_read").unwrap(), 2);
362
363        let top = store.top_agents(5);
364        assert_eq!(top[0].agent_id, "agent-1");
365    }
366
367    #[test]
368    fn format_report() {
369        let mut store = CostStore::default();
370        store.record_tool_call("agent-a", "mcp", "ctx_read", 10000, 500, 1000);
371        store.record_tool_call("agent-b", "cursor", "ctx_shell", 2000, 100, 0);
372
373        let report = format_cost_report(&store, 5);
374        assert!(report.contains("Cost Attribution Report"));
375        assert!(report.contains("agent-a"));
376        assert!(report.contains("ctx_read"));
377    }
378
379    #[test]
380    fn session_snapshots() {
381        let mut store = CostStore::default();
382        store.add_session_snapshot("agent-a", 50000, 5000, 30000, 120);
383        assert_eq!(store.sessions.len(), 1);
384        assert!(store.sessions[0].cost_usd > 0.0);
385    }
386}
lean_ctx/core/a2a/cost_attribution.rs

lean_ctx/core/a2a/
cost_attribution.rs
