lash_core/runtime/
usage.rs

1//! Token usage accounting: ledger entries, usage totals, reports, and diff helpers.
2//!
3//! Extracted from `runtime/mod.rs` as part of the runtime split. All items
4//! keep their original public paths via `pub use` in `mod.rs` — no API
5//! changes.
6
7use std::collections::{BTreeMap, HashMap};
8
9use crate::session_model::TokenUsage;
10use lash_sansio::PromptUsage;
11
12/// A single row in the token cost ledger. One per unique
13/// `(source, model)` pair — accumulated, not per-call.
14#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
15pub struct TokenLedgerEntry {
16    /// Caller-supplied label: `"turn"`, `"subagent"`, `"compaction"`,
17    /// `"observer"`, `"reflector"`, or any plugin-defined
18    /// string. Core doesn't interpret the value; the UI uses it for
19    /// grouping and display.
20    pub source: String,
21    /// Model identifier used for the LLM call (e.g.
22    /// `"anthropic/claude-haiku-4-5"`).
23    pub model: String,
24    /// Accumulated token counts for this `(source, model)` pair.
25    pub usage: TokenUsage,
26}
27
28#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
29pub struct UsageTotals {
30    pub input_tokens: i64,
31    pub output_tokens: i64,
32    pub cache_read_input_tokens: i64,
33    #[serde(default)]
34    pub cache_write_input_tokens: i64,
35    #[serde(default)]
36    pub reasoning_output_tokens: i64,
37    pub total_tokens: i64,
38    pub context_total_tokens: i64,
39}
40
41impl UsageTotals {
42    fn from_usage(usage: &TokenUsage) -> Self {
43        let total_tokens = usage.total();
44        Self {
45            input_tokens: usage.input_tokens,
46            output_tokens: usage.output_tokens,
47            cache_read_input_tokens: usage.cache_read_input_tokens,
48            cache_write_input_tokens: usage.cache_write_input_tokens,
49            reasoning_output_tokens: usage.reasoning_output_tokens,
50            total_tokens,
51            context_total_tokens: total_tokens,
52        }
53    }
54}
55
56#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
57pub struct UsageReportRow {
58    pub source: String,
59    pub model: String,
60    pub usage: UsageTotals,
61}
62
63#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
64pub struct SessionUsageReport {
65    pub entry_count: usize,
66    pub usage: UsageTotals,
67    pub by_source: BTreeMap<String, UsageTotals>,
68    pub by_model: BTreeMap<String, UsageTotals>,
69    pub by_source_model: Vec<UsageReportRow>,
70}
71
72impl SessionUsageReport {
73    pub fn from_entries(entries: &[TokenLedgerEntry]) -> Self {
74        let mut total = TokenUsage::default();
75        let mut by_source_usage = BTreeMap::<String, TokenUsage>::new();
76        let mut by_model_usage = BTreeMap::<String, TokenUsage>::new();
77        let mut by_source_model = Vec::with_capacity(entries.len());
78
79        for entry in entries {
80            total.add(&entry.usage);
81            by_source_usage
82                .entry(entry.source.clone())
83                .or_default()
84                .add(&entry.usage);
85            by_model_usage
86                .entry(entry.model.clone())
87                .or_default()
88                .add(&entry.usage);
89            by_source_model.push(UsageReportRow {
90                source: entry.source.clone(),
91                model: entry.model.clone(),
92                usage: UsageTotals::from_usage(&entry.usage),
93            });
94        }
95
96        Self {
97            entry_count: entries.len(),
98            usage: UsageTotals::from_usage(&total),
99            by_source: by_source_usage
100                .into_iter()
101                .map(|(key, usage)| (key, UsageTotals::from_usage(&usage)))
102                .collect(),
103            by_model: by_model_usage
104                .into_iter()
105                .map(|(key, usage)| (key, UsageTotals::from_usage(&usage)))
106                .collect(),
107            by_source_model,
108        }
109    }
110}
111
112pub fn diff_token_ledger(
113    before: &[TokenLedgerEntry],
114    after: &[TokenLedgerEntry],
115) -> Result<Vec<TokenLedgerEntry>, String> {
116    let before_index = before
117        .iter()
118        .map(|entry| ((entry.source.as_str(), entry.model.as_str()), &entry.usage))
119        .collect::<HashMap<_, _>>();
120    let after_index = after
121        .iter()
122        .map(|entry| ((entry.source.as_str(), entry.model.as_str()), &entry.usage))
123        .collect::<HashMap<_, _>>();
124
125    let mut keys = before_index
126        .keys()
127        .copied()
128        .chain(after_index.keys().copied())
129        .collect::<Vec<_>>();
130    keys.sort_unstable();
131    keys.dedup();
132
133    let mut out = Vec::new();
134    for (source, model) in keys {
135        let before_usage = before_index
136            .get(&(source, model))
137            .copied()
138            .cloned()
139            .unwrap_or_default();
140        let after_usage = after_index
141            .get(&(source, model))
142            .copied()
143            .cloned()
144            .unwrap_or_default();
145        let delta = TokenUsage {
146            input_tokens: after_usage.input_tokens - before_usage.input_tokens,
147            output_tokens: after_usage.output_tokens - before_usage.output_tokens,
148            cache_read_input_tokens: after_usage.cache_read_input_tokens
149                - before_usage.cache_read_input_tokens,
150            cache_write_input_tokens: after_usage.cache_write_input_tokens
151                - before_usage.cache_write_input_tokens,
152            reasoning_output_tokens: after_usage.reasoning_output_tokens
153                - before_usage.reasoning_output_tokens,
154        };
155        if delta.input_tokens < 0
156            || delta.output_tokens < 0
157            || delta.cache_read_input_tokens < 0
158            || delta.cache_write_input_tokens < 0
159            || delta.reasoning_output_tokens < 0
160        {
161            return Err(format!(
162                "token ledger decreased for source/model ({source}, {model})"
163            ));
164        }
165        if delta.total() == 0 {
166            continue;
167        }
168        out.push(TokenLedgerEntry {
169            source: source.to_string(),
170            model: model.to_string(),
171            usage: delta,
172        });
173    }
174    Ok(out)
175}
176
177pub fn diff_usage_reports(
178    before: &SessionUsageReport,
179    after: &SessionUsageReport,
180) -> Result<Vec<TokenLedgerEntry>, String> {
181    let before_entries = before
182        .by_source_model
183        .iter()
184        .map(|row| TokenLedgerEntry {
185            source: row.source.clone(),
186            model: row.model.clone(),
187            usage: TokenUsage {
188                input_tokens: row.usage.input_tokens,
189                output_tokens: row.usage.output_tokens,
190                cache_read_input_tokens: row.usage.cache_read_input_tokens,
191                cache_write_input_tokens: row.usage.cache_write_input_tokens,
192                reasoning_output_tokens: row.usage.reasoning_output_tokens,
193            },
194        })
195        .collect::<Vec<_>>();
196    let after_entries = after
197        .by_source_model
198        .iter()
199        .map(|row| TokenLedgerEntry {
200            source: row.source.clone(),
201            model: row.model.clone(),
202            usage: TokenUsage {
203                input_tokens: row.usage.input_tokens,
204                output_tokens: row.usage.output_tokens,
205                cache_read_input_tokens: row.usage.cache_read_input_tokens,
206                cache_write_input_tokens: row.usage.cache_write_input_tokens,
207                reasoning_output_tokens: row.usage.reasoning_output_tokens,
208            },
209        })
210        .collect::<Vec<_>>();
211    diff_token_ledger(&before_entries, &after_entries)
212}
213
214pub(super) fn merge_ledger_entry(ledger: &mut Vec<TokenLedgerEntry>, entry: TokenLedgerEntry) {
215    if entry.usage.total() == 0 {
216        return;
217    }
218    if let Some(existing) = ledger
219        .iter_mut()
220        .find(|e| e.source == entry.source && e.model == entry.model)
221    {
222        existing.usage.input_tokens += entry.usage.input_tokens;
223        existing.usage.output_tokens += entry.usage.output_tokens;
224        existing.usage.cache_read_input_tokens += entry.usage.cache_read_input_tokens;
225        existing.usage.cache_write_input_tokens += entry.usage.cache_write_input_tokens;
226        existing.usage.reasoning_output_tokens += entry.usage.reasoning_output_tokens;
227    } else {
228        ledger.push(entry);
229    }
230}
231
232pub(super) fn merge_usage_delta_entries(entries: Vec<TokenLedgerEntry>) -> Vec<TokenLedgerEntry> {
233    let mut merged = Vec::new();
234    for entry in entries {
235        merge_ledger_entry(&mut merged, entry);
236    }
237    merged
238}
239
240pub(super) fn normalize_prompt_usage(usage: &TokenUsage) -> Option<PromptUsage> {
241    let input_tokens = usage.input_tokens.max(0) as usize;
242    let output_tokens = usage.output_tokens.max(0) as usize;
243    let cache_read_input_tokens = usage.cache_read_input_tokens.max(0) as usize;
244    let cache_write_input_tokens = usage.cache_write_input_tokens.max(0) as usize;
245    if input_tokens == 0
246        && cache_read_input_tokens == 0
247        && cache_write_input_tokens == 0
248        && output_tokens == 0
249    {
250        return None;
251    }
252
253    let prompt_context_tokens = input_tokens
254        .saturating_add(cache_read_input_tokens)
255        .saturating_add(cache_write_input_tokens);
256    let context_budget_tokens = input_tokens
257        .saturating_add(output_tokens)
258        .saturating_add(cache_read_input_tokens)
259        .saturating_add(cache_write_input_tokens);
260
261    Some(PromptUsage {
262        prompt_context_tokens,
263        input_tokens,
264        cache_read_input_tokens,
265        cache_write_input_tokens,
266        context_budget_tokens,
267    })
268}
lash_core/runtime/usage.rs

lash_core/runtime/
usage.rs
