ag_ui/
token_usage.rs

1//! Per-model token accounting carried by `RUN_FINISHED` and `RUN_ERROR`.
2
3use serde::{Deserialize, Serialize};
4
5/// A numeric-only token usage summary for one `(provider, model)` pair.
6///
7/// Deliberately carries nothing identifying or content-bearing — no prompts,
8/// completions, message text, or thread/run/user ids. Only provider and model
9/// labels plus counts, so usage can be logged and aggregated in places where
10/// conversation content must not go.
11///
12/// Every count is optional and every count is a non-negative integer: `None`
13/// means *the provider did not report it*, which is distinct from zero.
14#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
15#[serde(rename_all = "camelCase")]
16#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
17#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
18pub struct TokenUsage {
19    /// The inference provider, for example `"anthropic"`.
20    #[serde(default, skip_serializing_if = "Option::is_none")]
21    pub provider: Option<String>,
22    /// The model id, for example `"claude-opus-5"`.
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub model: Option<String>,
25    /// Tokens consumed by the prompt.
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub input_tokens: Option<u64>,
28    /// Tokens produced by the model.
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub output_tokens: Option<u64>,
31    /// Total tokens, as reported by the provider (not necessarily the sum of
32    /// the other fields).
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub total_tokens: Option<u64>,
35    /// Output tokens spent on reasoning.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub reasoning_tokens: Option<u64>,
38    /// Input tokens served from the provider's prompt cache.
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub cached_input_tokens: Option<u64>,
41}
42
43impl TokenUsage {
44    /// An empty usage entry.
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Whether any count is present. A labels-only entry reports nothing and
50    /// should be omitted rather than emitted.
51    pub const fn has_counts(&self) -> bool {
52        self.input_tokens.is_some()
53            || self.output_tokens.is_some()
54            || self.total_tokens.is_some()
55            || self.reasoning_tokens.is_some()
56            || self.cached_input_tokens.is_some()
57    }
58}
59
60/// Sums per-call [`TokenUsage`] entries into one entry per `(provider, model)`
61/// pair, in order of first appearance.
62///
63/// A count stays `None` when no member of a group reported it, keeping "not
64/// reported" distinct from zero.
65///
66/// ```
67/// # use ag_ui::{TokenUsage, aggregate_token_usage};
68/// let calls = vec![
69///     TokenUsage { model: Some("m".into()), input_tokens: Some(10), ..Default::default() },
70///     TokenUsage { model: Some("m".into()), input_tokens: Some(5), ..Default::default() },
71/// ];
72/// let total = aggregate_token_usage(&calls);
73/// assert_eq!(total.len(), 1);
74/// assert_eq!(total[0].input_tokens, Some(15));
75/// assert_eq!(total[0].output_tokens, None);
76/// ```
77pub fn aggregate_token_usage(entries: &[TokenUsage]) -> Vec<TokenUsage> {
78    let mut grouped: Vec<TokenUsage> = Vec::new();
79
80    for entry in entries {
81        let index = match grouped
82            .iter()
83            .position(|g| g.provider == entry.provider && g.model == entry.model)
84        {
85            Some(index) => index,
86            None => {
87                grouped.push(TokenUsage {
88                    provider: entry.provider.clone(),
89                    model: entry.model.clone(),
90                    ..Default::default()
91                });
92                grouped.len() - 1
93            }
94        };
95
96        let target = &mut grouped[index];
97        add_into(&mut target.input_tokens, entry.input_tokens);
98        add_into(&mut target.output_tokens, entry.output_tokens);
99        add_into(&mut target.total_tokens, entry.total_tokens);
100        add_into(&mut target.reasoning_tokens, entry.reasoning_tokens);
101        add_into(&mut target.cached_input_tokens, entry.cached_input_tokens);
102    }
103
104    grouped
105}
106
107/// Adds `value` into `target`, leaving `target` untouched when nothing was
108/// reported. Saturating, so a hostile producer cannot panic the aggregation.
109fn add_into(target: &mut Option<u64>, value: Option<u64>) {
110    if let Some(value) = value {
111        *target = Some(target.unwrap_or(0).saturating_add(value));
112    }
113}
ag_ui/token_usage.rs

ag_ui/
token_usage.rs
