Skip to main content

agent_sessions/
tokens.rs

1/// `ReadSummary::ignored_types` label for selected statistical token-count
2/// records with nonzero last usage but no cumulative total. No usage is emitted
3/// because an independent ledger cannot be inferred from that record.
4pub const CODEX_MISSING_TOTAL_USAGE: &str = "event_msg:token_count:missing_total";
5
6use crate::LineErrorKind;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
11pub struct TokenCounts {
12    pub input: Option<u64>,
13    pub output: Option<u64>,
14    pub cache_read: Option<u64>,
15    pub cache_write: Option<u64>,
16    /// Subset of cache_write, never an additional category to sum.
17    pub cache_write_1h: Option<u64>,
18    pub reasoning: Option<u64>,
19    pub reported_total: Option<u64>,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
23#[non_exhaustive]
24pub enum TokenSemantics {
25    ClaudeExclusive,
26    CodexInclusive,
27    /// Input excludes cache categories; output includes reasoning.
28    CacheExclusive,
29    /// Gemini prompt includes cached input; candidate output excludes thoughts.
30    Gemini,
31    /// Native overlap has not been established; exclusive() returns None.
32    Unknown,
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
36#[non_exhaustive]
37pub enum UsageBasis {
38    Message,
39    LastSample,
40    CumulativeDelta,
41    InitialCumulative,
42    Response,
43}
44
45impl TokenCounts {
46    pub fn is_missing(&self) -> bool {
47        self.values().iter().all(Option::is_none)
48    }
49    pub(crate) fn values(&self) -> [Option<u64>; 7] {
50        [
51            self.input,
52            self.output,
53            self.cache_read,
54            self.cache_write,
55            self.cache_write_1h,
56            self.reasoning,
57            self.reported_total,
58        ]
59    }
60    pub(crate) fn statistics_eq(self, other: Self) -> bool {
61        self.values()
62            .into_iter()
63            .zip(other.values())
64            .all(|(a, b)| a.unwrap_or(0) == b.unwrap_or(0))
65    }
66    pub(crate) fn statistics_delta(self, prev: Self) -> Self {
67        fn sub(a: Option<u64>, b: Option<u64>) -> Option<u64> {
68            a.map(|a| a.saturating_sub(b.unwrap_or(0)))
69        }
70        Self {
71            input: sub(self.input, prev.input),
72            output: sub(self.output, prev.output),
73            cache_read: sub(self.cache_read, prev.cache_read),
74            cache_write: sub(self.cache_write, prev.cache_write),
75            cache_write_1h: sub(self.cache_write_1h, prev.cache_write_1h),
76            reasoning: sub(self.reasoning, prev.reasoning),
77            reported_total: sub(self.reported_total, prev.reported_total),
78        }
79    }
80    pub(crate) fn regressed_from(&self, prev: &Self) -> bool {
81        self.values()
82            .into_iter()
83            .zip(prev.values())
84            .any(|(a, b)| matches!((a,b), (Some(a),Some(b)) if a < b))
85    }
86    pub(crate) fn delta(self, prev: Self) -> Self {
87        fn sub(a: Option<u64>, b: Option<u64>) -> Option<u64> {
88            a?.checked_sub(b?)
89        }
90        Self {
91            input: sub(self.input, prev.input),
92            output: sub(self.output, prev.output),
93            cache_read: sub(self.cache_read, prev.cache_read),
94            cache_write: sub(self.cache_write, prev.cache_write),
95            cache_write_1h: sub(self.cache_write_1h, prev.cache_write_1h),
96            reasoning: sub(self.reasoning, prev.reasoning),
97            reported_total: sub(self.reported_total, prev.reported_total),
98        }
99    }
100    /// Convert inclusive Codex counts into disjoint categories without guessing
101    /// absent counters. An unprovable subtraction stays None.
102    pub fn exclusive(self, semantics: TokenSemantics) -> Option<Self> {
103        if semantics == TokenSemantics::Unknown {
104            return None;
105        }
106        if semantics == TokenSemantics::ClaudeExclusive {
107            return Some(self);
108        }
109        if semantics == TokenSemantics::Gemini {
110            return Some(Self {
111                input: match (self.input, self.cache_read) {
112                    (Some(input), Some(cache)) => Some(input.checked_sub(cache)?),
113                    _ => None,
114                },
115                ..self
116            });
117        }
118        if semantics == TokenSemantics::CacheExclusive {
119            return Some(Self {
120                output: match (self.output, self.reasoning) {
121                    (Some(output), Some(reasoning)) => Some(output.checked_sub(reasoning)?),
122                    _ => None,
123                },
124                ..self
125            });
126        }
127        let input = match (self.input, self.cache_read, self.cache_write) {
128            (Some(i), Some(r), Some(w)) => Some(i.checked_sub(r)?.checked_sub(w)?),
129            _ => None,
130        };
131        let output = match (self.output, self.reasoning) {
132            (Some(o), Some(r)) => Some(o.checked_sub(r)?),
133            _ => None,
134        };
135        Some(Self {
136            input,
137            output,
138            ..self
139        })
140    }
141}
142
143fn count(v: &Value, key: &str) -> Result<Option<u64>, LineErrorKind> {
144    match v.get(key) {
145        None | Some(Value::Null) => Ok(None),
146        Some(n) => n
147            .as_u64()
148            .map(Some)
149            .ok_or(LineErrorKind::InvalidField(key.into())),
150    }
151}
152
153pub(crate) fn parse_counts(
154    v: &Value,
155    claude: bool,
156    allow_missing: bool,
157) -> Result<TokenCounts, LineErrorKind> {
158    if !v.is_object() {
159        return Err(LineErrorKind::InvalidField("usage".into()));
160    }
161    let cache_read = if claude {
162        count(v, "cache_read_input_tokens")?
163    } else {
164        count(v, "cached_input_tokens")?.or(count(v, "cache_read_input_tokens")?)
165    };
166    let result = TokenCounts {
167        input: count(v, "input_tokens")?,
168        output: count(v, "output_tokens")?,
169        cache_read,
170        cache_write: count(
171            v,
172            if claude {
173                "cache_creation_input_tokens"
174            } else {
175                "cache_write_input_tokens"
176            },
177        )?,
178        cache_write_1h: match v.get("cache_creation") {
179            None | Some(Value::Null) => None,
180            Some(c) if c.is_object() => count(c, "ephemeral_1h_input_tokens")?,
181            _ => return Err(LineErrorKind::InvalidField("cache_creation".into())),
182        },
183        reasoning: count(v, "reasoning_output_tokens")?,
184        reported_total: count(v, "total_tokens")?,
185    };
186    if result.is_missing() && !allow_missing {
187        return Err(LineErrorKind::MissingField("usage counters"));
188    }
189    if matches!((result.cache_write_1h,result.cache_write),(Some(h),Some(w)) if h>w) {
190        return Err(LineErrorKind::InvalidField("cache_creation".into()));
191    }
192    Ok(result)
193}