1use crate::models::{PerProviderUsage, ProviderActiveDays, UsageResult};
10use crate::pricing::CostSource;
11use crate::usage::StoredCosts;
12use serde_json::Value;
13use std::borrow::Cow;
14
15#[derive(Clone, Copy)]
17enum ProviderPricing<'a> {
18 Litellm,
20 GrokGauge,
23 OpenCode(&'a crate::constants::FastHashMap<String, f64>),
25 CursorEstimate,
27 Hermes(&'a crate::constants::FastHashMap<String, f64>),
29}
30
31impl ProviderPricing<'_> {
32 fn source_for(&self, model: &str) -> CostSource {
34 let stored =
35 |m: &crate::constants::FastHashMap<String, f64>| m.get(model).copied().unwrap_or(0.0);
36 match self {
37 Self::Litellm => CostSource::Litellm,
38 Self::GrokGauge => CostSource::GrokGauge,
39 Self::OpenCode(m) => CostSource::OpenCodeStored(stored(m)),
40 Self::CursorEstimate => CostSource::OpenCodeStored(0.0),
41 Self::Hermes(m) => CostSource::HermesStored(stored(m)),
42 }
43 }
44}
45
46#[derive(Default, Clone)]
57pub struct UsageRow {
58 pub model: String, pub display_model: String, pub input_tokens: i64,
65 pub output_tokens: i64,
67 pub reasoning_tokens: i64,
69 pub cache_read: i64,
71 pub cache_creation: i64,
73 pub total: i64,
75 pub cost: f64,
77}
78
79impl UsageRow {
80 #[inline]
84 pub fn output_with_reasoning(&self) -> i64 {
85 self.output_tokens + self.reasoning_tokens
86 }
87}
88
89#[derive(Default)]
91pub struct UsageTotals {
92 pub input_tokens: i64,
94 pub output_tokens: i64,
96 pub reasoning_tokens: i64,
98 pub cache_read: i64,
100 pub cache_creation: i64,
102 pub total: i64,
104 pub cost: f64,
106}
107
108impl UsageTotals {
109 pub fn accumulate(&mut self, row: &UsageRow) {
111 self.input_tokens += row.input_tokens;
112 self.output_tokens += row.output_tokens;
113 self.reasoning_tokens += row.reasoning_tokens;
114 self.cache_read += row.cache_read;
115 self.cache_creation += row.cache_creation;
116 self.total += row.total;
117 self.cost += row.cost;
118 }
119
120 #[inline]
122 pub fn output_with_reasoning(&self) -> i64 {
123 self.output_tokens + self.reasoning_tokens
124 }
125}
126
127#[derive(Default, Clone)]
131pub struct ProviderStats {
132 pub total_tokens: i64,
134 pub total_cost: f64,
136 pub days_count: usize,
138}
139
140impl ProviderStats {
141 fn accumulate_row(&mut self, row: &UsageRow) {
143 self.total_tokens += row.total;
144 self.total_cost += row.cost;
145 }
146}
147
148pub type UsageProviderTotals = crate::models::ProviderTotals<ProviderStats>;
150
151#[derive(Default)]
153pub struct UsageSummary {
154 pub rows: Vec<UsageRow>,
156 pub totals: UsageTotals,
158 pub provider_totals: UsageProviderTotals,
160}
161
162pub fn calculate_provider_totals_from_per_provider(
171 per_provider: &PerProviderUsage,
172 provider_days: &ProviderActiveDays,
173 pricing_map: &crate::pricing::ModelPricingMap,
174 stored_costs: &StoredCosts,
175) -> UsageProviderTotals {
176 let mut totals = UsageProviderTotals::default();
177
178 totals.claude.days_count = provider_days.claude;
179 totals.codex.days_count = provider_days.codex;
180 totals.copilot.days_count = provider_days.copilot;
181 totals.gemini.days_count = provider_days.gemini;
182 totals.grok.days_count = provider_days.grok;
183 totals.opencode.days_count = provider_days.opencode;
184 totals.cursor.days_count = provider_days.cursor;
185 totals.hermes.days_count = provider_days.hermes;
186 totals.overall.days_count = provider_days.total;
187
188 accumulate_provider(
189 &mut totals.claude,
190 &per_provider.claude,
191 pricing_map,
192 ProviderPricing::Litellm,
193 );
194 accumulate_provider(
195 &mut totals.codex,
196 &per_provider.codex,
197 pricing_map,
198 ProviderPricing::Litellm,
199 );
200 accumulate_provider(
201 &mut totals.copilot,
202 &per_provider.copilot,
203 pricing_map,
204 ProviderPricing::Litellm,
205 );
206 accumulate_provider(
207 &mut totals.gemini,
208 &per_provider.gemini,
209 pricing_map,
210 ProviderPricing::Litellm,
211 );
212 accumulate_provider(
213 &mut totals.grok,
214 &per_provider.grok,
215 pricing_map,
216 ProviderPricing::GrokGauge,
217 );
218 accumulate_provider(
219 &mut totals.opencode,
220 &per_provider.opencode,
221 pricing_map,
222 ProviderPricing::OpenCode(&stored_costs.opencode),
223 );
224 accumulate_provider(
225 &mut totals.cursor,
226 &per_provider.cursor,
227 pricing_map,
228 ProviderPricing::CursorEstimate,
229 );
230 accumulate_provider(
231 &mut totals.hermes,
232 &per_provider.hermes,
233 pricing_map,
234 ProviderPricing::Hermes(&stored_costs.hermes),
235 );
236
237 totals.overall.total_tokens = totals.claude.total_tokens
247 + totals.codex.total_tokens
248 + totals.copilot.total_tokens
249 + totals.gemini.total_tokens
250 + totals.grok.total_tokens
251 + totals.opencode.total_tokens
252 + totals.cursor.total_tokens
253 + totals.hermes.total_tokens;
254 totals.overall.total_cost = totals.claude.total_cost
255 + totals.codex.total_cost
256 + totals.copilot.total_cost
257 + totals.gemini.total_cost
258 + totals.grok.total_cost
259 + totals.opencode.total_cost
260 + totals.cursor.total_cost
261 + totals.hermes.total_cost;
262
263 totals
264}
265
266fn accumulate_provider(
269 stats: &mut ProviderStats,
270 usage: &UsageResult,
271 pricing_map: &crate::pricing::ModelPricingMap,
272 pricing: ProviderPricing,
273) {
274 for (model, raw_usage) in usage {
275 let row = extract_usage_row(model, raw_usage, pricing_map, pricing.source_for(model));
276 stats.accumulate_row(&row);
277 }
278}
279
280pub fn build_usage_summary(
288 usage_data: &UsageResult,
289 per_provider: &PerProviderUsage,
290 provider_days: &ProviderActiveDays,
291 pricing_map: &crate::pricing::ModelPricingMap,
292 stored_costs: &StoredCosts,
293) -> UsageSummary {
294 if usage_data.is_empty() {
295 return UsageSummary::default();
296 }
297
298 let mut summary = UsageSummary::default();
299
300 summary.rows.reserve(usage_data.len());
302
303 for (model, usage) in usage_data.iter() {
305 let (cost, matched_model) =
306 crate::usage::resolve_merged_model_cost(model, per_provider, pricing_map, stored_costs)
307 .unwrap_or_else(|| price_usage(model, usage, pricing_map, CostSource::Litellm));
308 let row = build_usage_row(model, usage, cost, matched_model);
309 summary.rows.push(row);
310 }
311
312 summary.rows.sort_by(|a, b| {
314 a.cost
315 .partial_cmp(&b.cost)
316 .unwrap_or(std::cmp::Ordering::Equal)
317 .then_with(|| a.model.cmp(&b.model))
318 });
319
320 for row in &summary.rows {
321 summary.totals.accumulate(row);
322 }
323
324 summary.provider_totals = calculate_provider_totals_from_per_provider(
325 per_provider,
326 provider_days,
327 pricing_map,
328 stored_costs,
329 );
330 summary
331}
332
333fn extract_usage_row(
341 model: &str,
342 usage: &Value,
343 pricing_map: &crate::pricing::ModelPricingMap,
344 source: CostSource,
345) -> UsageRow {
346 use crate::pricing::resolve_model_cost;
347 use crate::utils::extract_token_counts;
348
349 let counts = extract_token_counts(usage);
351 let (cost, matched_model) = resolve_model_cost(model, &counts, pricing_map, source);
352 build_usage_row_from_counts(model, &counts, cost, matched_model)
353}
354
355fn price_usage(
357 model: &str,
358 usage: &Value,
359 pricing_map: &crate::pricing::ModelPricingMap,
360 source: CostSource,
361) -> (f64, Option<String>) {
362 use crate::pricing::resolve_model_cost;
363 use crate::utils::extract_token_counts;
364
365 let counts = extract_token_counts(usage);
366 resolve_model_cost(model, &counts, pricing_map, source)
367}
368
369fn build_usage_row(
371 model: &str,
372 usage: &Value,
373 cost: f64,
374 matched_model: Option<String>,
375) -> UsageRow {
376 let counts = crate::utils::extract_token_counts(usage);
377 build_usage_row_from_counts(model, &counts, cost, matched_model)
378}
379
380fn build_usage_row_from_counts(
383 model: &str,
384 counts: &crate::utils::TokenCounts,
385 cost: f64,
386 matched_model: Option<String>,
387) -> UsageRow {
388 let display_model = if let Some(matched) = &matched_model {
390 Cow::Owned(format!("{} ({})", model, matched))
391 } else {
392 Cow::Borrowed(model)
393 };
394
395 UsageRow {
396 model: model.to_string(),
397 display_model: display_model.into_owned(),
398 input_tokens: counts.input_tokens,
399 output_tokens: counts.output_tokens,
400 reasoning_tokens: counts.reasoning_tokens,
401 cache_read: counts.cache_read,
402 cache_creation: counts.cache_creation,
403 total: counts.total,
404 cost,
405 }
406}
407
408fn base_model_key(model: &str) -> &str {
414 model.split_once('/').map(|(_, rest)| rest).unwrap_or(model)
415}
416
417pub fn merge_rows_by_base_model(rows: &[UsageRow]) -> Vec<UsageRow> {
429 use std::collections::HashMap;
430
431 let mut groups: HashMap<&str, Vec<&UsageRow>> = HashMap::new();
432 for row in rows {
433 groups
434 .entry(base_model_key(&row.model))
435 .or_default()
436 .push(row);
437 }
438
439 let mut merged: Vec<UsageRow> = Vec::with_capacity(groups.len());
440 for (key, members) in groups {
441 let mut acc = UsageRow {
442 model: key.to_string(),
443 display_model: key.to_string(),
444 ..UsageRow::default()
445 };
446 for m in members {
447 acc.input_tokens += m.input_tokens;
448 acc.output_tokens += m.output_tokens;
449 acc.reasoning_tokens += m.reasoning_tokens;
450 acc.cache_read += m.cache_read;
451 acc.cache_creation += m.cache_creation;
452 acc.total += m.total;
453 acc.cost += m.cost;
454 }
455 merged.push(acc);
456 }
457
458 merged.sort_by(|a, b| {
460 a.cost
461 .partial_cmp(&b.cost)
462 .unwrap_or(std::cmp::Ordering::Equal)
463 .then_with(|| a.model.cmp(&b.model))
464 });
465 merged
466}
467
468#[cfg(test)]
469mod tests {
470 use super::*;
471 use crate::pricing::{ModelPricing, ModelPricingMap, clear_pricing_cache};
472 use serde_json::json;
473
474 #[test]
475 fn merged_rows_include_grok_source_cost() {
476 clear_pricing_cache();
477 let mut raw_pricing = std::collections::HashMap::new();
478 raw_pricing.insert(
479 "shared-model".to_string(),
480 ModelPricing {
481 input_cost_per_token: 0.01,
482 ..Default::default()
483 },
484 );
485 let pricing_map = ModelPricingMap::new(raw_pricing);
486 let mut usage_data = UsageResult::default();
487 usage_data.insert("shared-model".to_string(), json!({"input_tokens": 200}));
488 let mut per_provider = PerProviderUsage::default();
489 per_provider
490 .claude
491 .insert("shared-model".to_string(), json!({"input_tokens": 100}));
492 per_provider
493 .grok
494 .insert("shared-model".to_string(), json!({"input_tokens": 100}));
495
496 let summary = build_usage_summary(
497 &usage_data,
498 &per_provider,
499 &ProviderActiveDays::default(),
500 &pricing_map,
501 &StoredCosts::default(),
502 );
503
504 assert!((summary.rows[0].cost - 2.0).abs() < 1e-9);
505 }
506
507 #[test]
508 fn merged_rows_price_opencode_fallback_only_for_opencode_tokens() {
509 clear_pricing_cache();
510
511 let mut raw_pricing = std::collections::HashMap::new();
512 raw_pricing.insert(
513 "shared".to_string(),
514 ModelPricing {
515 input_cost_per_token: 0.01,
516 ..Default::default()
517 },
518 );
519 let pricing_map = ModelPricingMap::new(raw_pricing);
520
521 let mut usage_data = UsageResult::default();
522 usage_data.insert("shared-pro".to_string(), json!({"input_tokens": 200}));
523
524 let mut per_provider = PerProviderUsage::default();
525 per_provider
526 .claude
527 .insert("shared-pro".to_string(), json!({"input_tokens": 100}));
528 per_provider
529 .opencode
530 .insert("shared-pro".to_string(), json!({"input_tokens": 100}));
531
532 let mut stored_costs = StoredCosts::default();
533 stored_costs.opencode.insert("shared-pro".to_string(), 7.0);
534
535 let summary = build_usage_summary(
536 &usage_data,
537 &per_provider,
538 &ProviderActiveDays::default(),
539 &pricing_map,
540 &stored_costs,
541 );
542
543 assert_eq!(summary.rows.len(), 1);
544 assert!((summary.rows[0].cost - 8.0).abs() < 1e-9);
545 assert_eq!(summary.rows[0].display_model, "shared-pro (shared)");
546 }
547
548 #[test]
549 fn cursor_row_uses_exact_litellm_price_and_ignores_legacy_stored_cost() {
550 clear_pricing_cache();
551
552 let mut raw_pricing = std::collections::HashMap::new();
554 raw_pricing.insert(
555 "gemini-2.5-pro".to_string(),
556 ModelPricing {
557 input_cost_per_token: 0.01,
558 ..Default::default()
559 },
560 );
561 let pricing_map = ModelPricingMap::new(raw_pricing);
562
563 let mut usage_data = UsageResult::default();
564 usage_data.insert("gemini-2.5-pro".to_string(), json!({"input_tokens": 1000}));
565
566 let mut per_provider = PerProviderUsage::default();
567 per_provider
568 .cursor
569 .insert("gemini-2.5-pro".to_string(), json!({"input_tokens": 1000}));
570
571 let mut stored_costs = StoredCosts::default();
573 stored_costs
574 .cursor
575 .insert("gemini-2.5-pro".to_string(), 0.3425);
576
577 let summary = build_usage_summary(
578 &usage_data,
579 &per_provider,
580 &ProviderActiveDays::default(),
581 &pricing_map,
582 &stored_costs,
583 );
584
585 assert_eq!(summary.rows.len(), 1);
586 assert!((summary.rows[0].cost - 10.0).abs() < 1e-9);
587 }
588
589 #[test]
590 fn stored_costs_do_not_cross_contaminate_on_name_collision() {
591 clear_pricing_cache();
592 let pricing_map = ModelPricingMap::new(std::collections::HashMap::new());
595
596 let mut usage_data = UsageResult::default();
598 usage_data.insert("collide".to_string(), json!({"input_tokens": 10}));
599
600 let mut per_provider = PerProviderUsage::default();
601 per_provider
602 .opencode
603 .insert("collide".to_string(), json!({"input_tokens": 5}));
604 per_provider
605 .cursor
606 .insert("collide".to_string(), json!({"input_tokens": 5}));
607
608 let mut stored_costs = StoredCosts::default();
609 stored_costs.opencode.insert("collide".to_string(), 5.0);
610 stored_costs.cursor.insert("collide".to_string(), 3.0);
611
612 let summary = build_usage_summary(
613 &usage_data,
614 &per_provider,
615 &ProviderActiveDays::default(),
616 &pricing_map,
617 &stored_costs,
618 );
619
620 assert_eq!(summary.rows.len(), 1);
621 assert!((summary.rows[0].cost - 5.0).abs() < 1e-9);
622 assert!((summary.provider_totals.opencode.total_cost - 5.0).abs() < 1e-9);
623 assert!(summary.provider_totals.cursor.total_cost.abs() < 1e-9);
624 }
625
626 fn row(model: &str, input: i64, total: i64, cost: f64) -> UsageRow {
627 UsageRow {
628 model: model.to_string(),
629 display_model: model.to_string(),
630 input_tokens: input,
631 total,
632 cost,
633 ..UsageRow::default()
634 }
635 }
636
637 #[test]
638 fn merge_collapses_prefixed_and_bare_names_and_sums() {
639 let rows = vec![
640 row("openai/gpt-5.5", 100, 100, 0.20),
641 row("azure/gpt-5.5", 200, 200, 3.00),
642 row("gpt-5.5", 300, 300, 5.00),
643 ];
644
645 let merged = merge_rows_by_base_model(&rows);
646
647 assert_eq!(merged.len(), 1);
648 let m = &merged[0];
649 assert_eq!(m.model, "gpt-5.5");
650 assert_eq!(m.display_model, "gpt-5.5");
651 assert_eq!(m.input_tokens, 600);
652 assert_eq!(m.total, 600);
653 assert!((m.cost - 8.20).abs() < 1e-9);
654 }
655
656 #[test]
657 fn merge_keeps_different_versions_apart() {
658 let rows = vec![
661 row("openai/gpt-5.5", 10, 10, 1.0),
662 row("azure/gpt-5.5", 30, 30, 3.0),
663 row("openai/gpt-5.4", 20, 20, 2.0),
664 ];
665
666 let merged = merge_rows_by_base_model(&rows);
667
668 assert_eq!(merged.len(), 2);
669 let five_five = merged.iter().find(|r| r.model == "gpt-5.5").unwrap();
671 assert_eq!(five_five.display_model, "gpt-5.5");
672 assert_eq!(five_five.total, 40);
673 let five_four = merged.iter().find(|r| r.model == "gpt-5.4").unwrap();
675 assert_eq!(five_four.display_model, "gpt-5.4");
676 assert_eq!(five_four.total, 20);
677 }
678
679 #[test]
680 fn merge_strips_prefix_from_single_row() {
681 let mut only = row("deepseek/deepseek-v4-pro", 5, 5, 1.5);
684 only.display_model = "deepseek/deepseek-v4-pro (deepseek-v4)".to_string();
685
686 let merged = merge_rows_by_base_model(std::slice::from_ref(&only));
687
688 assert_eq!(merged.len(), 1);
689 assert_eq!(merged[0].model, "deepseek-v4-pro");
690 assert_eq!(merged[0].display_model, "deepseek-v4-pro");
691 assert_eq!(merged[0].total, 5);
692 }
693
694 #[test]
695 fn merge_only_strips_first_slash_segment() {
696 assert_eq!(base_model_key("a/b/c"), "b/c");
697 assert_eq!(base_model_key("gpt-5.5"), "gpt-5.5");
698 assert_eq!(base_model_key("openai/gpt-5.5"), "gpt-5.5");
699 }
700
701 #[test]
702 fn merge_reorders_by_ascending_cost() {
703 let rows = vec![
704 row("openai/gpt-5.5", 1, 1, 9.0),
705 row("azure/gpt-5.5", 1, 1, 9.0),
706 row("cheap-model", 1, 1, 0.01),
707 ];
708
709 let merged = merge_rows_by_base_model(&rows);
710
711 assert_eq!(merged.len(), 2);
712 assert_eq!(merged[0].model, "cheap-model");
714 assert_eq!(merged[1].model, "gpt-5.5");
715 assert!((merged[1].cost - 18.0).abs() < 1e-9);
716 }
717}