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}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
30#[non_exhaustive]
31pub enum UsageBasis {
32 Message,
33 LastSample,
34 CumulativeDelta,
35 InitialCumulative,
36 Response,
37}
38
39impl TokenCounts {
40 pub fn is_missing(&self) -> bool {
41 self.values().iter().all(Option::is_none)
42 }
43 pub(crate) fn values(&self) -> [Option<u64>; 7] {
44 [
45 self.input,
46 self.output,
47 self.cache_read,
48 self.cache_write,
49 self.cache_write_1h,
50 self.reasoning,
51 self.reported_total,
52 ]
53 }
54 pub(crate) fn statistics_eq(self, other: Self) -> bool {
55 self.values()
56 .into_iter()
57 .zip(other.values())
58 .all(|(a, b)| a.unwrap_or(0) == b.unwrap_or(0))
59 }
60 pub(crate) fn statistics_delta(self, prev: Self) -> Self {
61 fn sub(a: Option<u64>, b: Option<u64>) -> Option<u64> {
62 a.map(|a| a.saturating_sub(b.unwrap_or(0)))
63 }
64 Self {
65 input: sub(self.input, prev.input),
66 output: sub(self.output, prev.output),
67 cache_read: sub(self.cache_read, prev.cache_read),
68 cache_write: sub(self.cache_write, prev.cache_write),
69 cache_write_1h: sub(self.cache_write_1h, prev.cache_write_1h),
70 reasoning: sub(self.reasoning, prev.reasoning),
71 reported_total: sub(self.reported_total, prev.reported_total),
72 }
73 }
74 pub(crate) fn regressed_from(&self, prev: &Self) -> bool {
75 self.values()
76 .into_iter()
77 .zip(prev.values())
78 .any(|(a, b)| matches!((a,b), (Some(a),Some(b)) if a < b))
79 }
80 pub(crate) fn delta(self, prev: Self) -> Self {
81 fn sub(a: Option<u64>, b: Option<u64>) -> Option<u64> {
82 a?.checked_sub(b?)
83 }
84 Self {
85 input: sub(self.input, prev.input),
86 output: sub(self.output, prev.output),
87 cache_read: sub(self.cache_read, prev.cache_read),
88 cache_write: sub(self.cache_write, prev.cache_write),
89 cache_write_1h: sub(self.cache_write_1h, prev.cache_write_1h),
90 reasoning: sub(self.reasoning, prev.reasoning),
91 reported_total: sub(self.reported_total, prev.reported_total),
92 }
93 }
94 pub fn exclusive(self, semantics: TokenSemantics) -> Option<Self> {
97 if semantics == TokenSemantics::ClaudeExclusive {
98 return Some(self);
99 }
100 let input = match (self.input, self.cache_read, self.cache_write) {
101 (Some(i), Some(r), Some(w)) => Some(i.checked_sub(r)?.checked_sub(w)?),
102 _ => None,
103 };
104 let output = match (self.output, self.reasoning) {
105 (Some(o), Some(r)) => Some(o.checked_sub(r)?),
106 _ => None,
107 };
108 Some(Self {
109 input,
110 output,
111 ..self
112 })
113 }
114}
115
116fn count(v: &Value, key: &str) -> Result<Option<u64>, LineErrorKind> {
117 match v.get(key) {
118 None | Some(Value::Null) => Ok(None),
119 Some(n) => n
120 .as_u64()
121 .map(Some)
122 .ok_or(LineErrorKind::InvalidField(key.into())),
123 }
124}
125
126pub(crate) fn parse_counts(
127 v: &Value,
128 claude: bool,
129 allow_missing: bool,
130) -> Result<TokenCounts, LineErrorKind> {
131 if !v.is_object() {
132 return Err(LineErrorKind::InvalidField("usage".into()));
133 }
134 let cache_read = if claude {
135 count(v, "cache_read_input_tokens")?
136 } else {
137 count(v, "cached_input_tokens")?.or(count(v, "cache_read_input_tokens")?)
138 };
139 let result = TokenCounts {
140 input: count(v, "input_tokens")?,
141 output: count(v, "output_tokens")?,
142 cache_read,
143 cache_write: count(
144 v,
145 if claude {
146 "cache_creation_input_tokens"
147 } else {
148 "cache_write_input_tokens"
149 },
150 )?,
151 cache_write_1h: match v.get("cache_creation") {
152 None | Some(Value::Null) => None,
153 Some(c) if c.is_object() => count(c, "ephemeral_1h_input_tokens")?,
154 _ => return Err(LineErrorKind::InvalidField("cache_creation".into())),
155 },
156 reasoning: count(v, "reasoning_output_tokens")?,
157 reported_total: count(v, "total_tokens")?,
158 };
159 if result.is_missing() && !allow_missing {
160 return Err(LineErrorKind::MissingField("usage counters"));
161 }
162 if matches!((result.cache_write_1h,result.cache_write),(Some(h),Some(w)) if h>w) {
163 return Err(LineErrorKind::InvalidField("cache_creation".into()));
164 }
165 Ok(result)
166}