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}
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    /// Convert inclusive Codex counts into disjoint categories without guessing
95    /// absent counters. An unprovable subtraction stays None.
96    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}