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

lash_core/runtime/
usage.rs
