forge_primitives/
token_usage.rs

1//! Token/cost accounting shapes for the `usage` subject.
2//!
3//! Distinct from [`crate::usage`] (Anthropic plan/subscription
4//! utilization): these describe per-model and per-project token counts
5//! plus a notional API-equivalent cost, rolled up from the session
6//! JSONL pool. The scanner lives in `forge-agent::env::token_usage`,
7//! the redb cache in `forge-workspace::store::token_usage`.
8
9use serde::{Deserialize, Serialize};
10
11/// Per-token USD pricing for one model, sourced from the LiteLLM table.
12/// Cache-write splits by TTL tier to match the JSONL's
13/// `cache_creation.ephemeral_{1h,5m}` breakdown.
14#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
15pub struct ModelPricing {
16    pub input: f64,
17    pub output: f64,
18    pub cache_write_1h: f64,
19    pub cache_write_5m: f64,
20    pub cache_read: f64,
21}
22
23/// One row of the usage table: `label` is a model id, a folded project
24/// name, or `"TOTAL"` for the aggregate row. Token counts are split;
25/// `cost_usd` is notional (at API pricing, not a bill).
26#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
27pub struct UsageRow {
28    pub label: String,
29    pub input: u64,
30    pub cache_write_1h: u64,
31    pub cache_write_5m: u64,
32    pub cache_read: u64,
33    pub output: u64,
34    pub cost_usd: f64,
35}
36
37impl UsageRow {
38    /// Total tokens across the five token fields.
39    pub fn tokens(&self) -> u64 {
40        self.input
41            .saturating_add(self.cache_write_1h)
42            .saturating_add(self.cache_write_5m)
43            .saturating_add(self.cache_read)
44            .saturating_add(self.output)
45    }
46}
47
48/// One time window's usage: the same data grouped two ways plus the
49/// aggregate `total` row. Each `Vec` is sorted by cost descending.
50#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
51pub struct WindowUsage {
52    pub by_model: Vec<UsageRow>,
53    pub by_project: Vec<UsageRow>,
54    pub total: UsageRow,
55}
56
57/// The four rolling windows the `/usage` overlay drew. `lifetime` is
58/// a static sum; the other three are recomputed relative to now.
59#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
60pub struct UsageReport {
61    pub today: WindowUsage,
62    pub week: WindowUsage,
63    pub month: WindowUsage,
64    pub lifetime: WindowUsage,
65    /// Whether a pricing table was loaded when this report was built.
66    /// `false` (cold cache, failed or empty fetch) means every cost is
67    /// a placeholder 0, so the UI blanks costs rather than reading a
68    /// misleading `$0.00`.
69    pub pricing_available: bool,
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    fn row(label: &str) -> UsageRow {
77        UsageRow {
78            label: label.to_owned(),
79            input: 1,
80            cache_write_1h: 2,
81            cache_write_5m: 4,
82            cache_read: 8,
83            output: 16,
84            cost_usd: 1.5,
85        }
86    }
87
88    #[test]
89    fn tokens_sums_the_five_fields() {
90        assert_eq!(row("opus").tokens(), 1 + 2 + 4 + 8 + 16);
91    }
92
93    #[test]
94    fn usage_report_round_trips_through_serde() {
95        let window = WindowUsage {
96            by_model: vec![row("opus"), row("sonnet")],
97            by_project: vec![row("forge")],
98            total: row("TOTAL"),
99        };
100        let report = UsageReport {
101            today: window.clone(),
102            week: window.clone(),
103            month: window.clone(),
104            lifetime: window,
105            pricing_available: true,
106        };
107        let json = serde_json::to_string(&report).expect("serialize");
108        let back: UsageReport = serde_json::from_str(&json).expect("deserialize");
109        assert_eq!(report, back);
110    }
111}
forge_primitives/token_usage.rs

forge_primitives/
token_usage.rs
