1pub 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 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 CacheExclusive,
29 Gemini,
31 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 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}