1use serde_json::Value;
2
3use crate::types::{
4 CacheUsage, CacheUsageStatus, CycleRecord, TaskTokenUsage, TokenUsage, UsageSource,
5};
6
7pub fn normalize_token_usage(raw_usage: &Value) -> TokenUsage {
8 normalize_token_usage_with_hints(raw_usage, None, None)
9}
10
11pub fn normalize_token_usage_with_hints(
12 raw_usage: &Value,
13 usage_source: Option<UsageSource>,
14 cache_status: Option<CacheUsageStatus>,
15) -> TokenUsage {
16 let Some(raw) = raw_usage.as_object() else {
17 let status = cache_status.unwrap_or_default();
18 return TokenUsage {
19 usage_source: usage_source.unwrap_or_default(),
20 cache_usage: CacheUsage {
21 status,
22 source: (status == CacheUsageStatus::Unsupported)
23 .then(|| "adapter_capability".to_string()),
24 ..CacheUsage::default()
25 },
26 ..TokenUsage::default()
27 };
28 };
29
30 let prompt_tokens = read_int(raw.get("prompt_tokens"));
31 let completion_tokens = read_int(raw.get("completion_tokens"));
32 let input_tokens = read_int(raw.get("input_tokens")).or(prompt_tokens);
33 let output_tokens = read_int(raw.get("output_tokens")).or(completion_tokens);
34 let total_tokens = read_int(raw.get("total_tokens")).unwrap_or_else(|| {
35 prompt_tokens.or(input_tokens).unwrap_or_default()
36 + completion_tokens.or(output_tokens).unwrap_or_default()
37 });
38 let cached_tokens = read_nested_cache_int(
39 raw_usage,
40 &[
41 &["cache_read_tokens"],
42 &["prompt_tokens_details", "cached_tokens"],
43 &["input_tokens_details", "cached_tokens"],
44 &["cache_read_input_tokens"],
45 ],
46 );
47 let reasoning_tokens = read_nested_int(
48 raw_usage,
49 &[
50 &["completion_tokens_details", "reasoning_tokens"],
51 &["output_tokens_details", "reasoning_tokens"],
52 &["reasoning_tokens"],
53 ],
54 )
55 .unwrap_or_default();
56 let cache_creation_tokens = read_nested_cache_int(
57 raw_usage,
58 &[
59 &["cache_creation_tokens"],
60 &["cache_write_tokens"],
61 &["input_tokens_details", "cache_creation_tokens"],
62 &["prompt_tokens_details", "cache_creation_tokens"],
63 &["cache_creation_input_tokens"],
64 &["cache_write_input_tokens"],
65 ],
66 );
67 let mut uncached_input_tokens = read_nested_cache_int(raw_usage, &[&["uncached_input_tokens"]]);
68 let observed_cache_metric = cached_tokens.is_some()
69 || cache_creation_tokens.is_some()
70 || uncached_input_tokens.is_some();
71 let normalized_cache_status = if observed_cache_metric {
72 CacheUsageStatus::ProviderReported
73 } else {
74 cache_status.unwrap_or_default()
75 };
76
77 if normalized_cache_status == CacheUsageStatus::ProviderReported
78 && uncached_input_tokens.is_none()
79 {
80 if has_any_key(
81 raw_usage,
82 &[
83 "cache_read_input_tokens",
84 "cache_creation_input_tokens",
85 "cache_write_input_tokens",
86 ],
87 ) {
88 uncached_input_tokens = input_tokens;
89 } else if (has_nested_path(raw_usage, &["prompt_tokens_details", "cached_tokens"])
90 || has_nested_path(raw_usage, &["input_tokens_details", "cached_tokens"]))
91 && input_tokens.is_some()
92 && cached_tokens.is_some()
93 {
94 uncached_input_tokens = Some(
95 input_tokens
96 .unwrap_or_default()
97 .saturating_sub(cached_tokens.unwrap_or_default()),
98 );
99 }
100 }
101
102 let cache_usage = CacheUsage {
103 status: normalized_cache_status,
104 read_tokens: (normalized_cache_status == CacheUsageStatus::ProviderReported)
105 .then_some(cached_tokens)
106 .flatten(),
107 write_tokens: (normalized_cache_status == CacheUsageStatus::ProviderReported)
108 .then_some(cache_creation_tokens)
109 .flatten(),
110 uncached_input_tokens: (normalized_cache_status == CacheUsageStatus::ProviderReported)
111 .then_some(uncached_input_tokens)
112 .flatten(),
113 source: match normalized_cache_status {
114 CacheUsageStatus::ProviderReported => Some("provider_usage".to_string()),
115 CacheUsageStatus::Unsupported => Some("adapter_capability".to_string()),
116 CacheUsageStatus::AccountingMissing => None,
117 },
118 };
119
120 TokenUsage {
121 prompt_tokens: prompt_tokens.or(input_tokens).unwrap_or_default(),
122 completion_tokens: completion_tokens.or(output_tokens).unwrap_or_default(),
123 total_tokens,
124 cached_tokens: cached_tokens.unwrap_or_default(),
125 reasoning_tokens,
126 input_tokens: input_tokens.unwrap_or_default(),
127 output_tokens: output_tokens.unwrap_or_default(),
128 cache_creation_tokens: cache_creation_tokens.unwrap_or_default(),
129 usage_source: usage_source.unwrap_or_else(|| infer_usage_source(raw_usage)),
130 cache_usage,
131 raw: raw_usage.clone(),
132 }
133}
134
135pub fn summarize_task_token_usage(cycles: &[CycleRecord]) -> TaskTokenUsage {
136 let mut summary = TaskTokenUsage::default();
137 for cycle in cycles {
138 summary.add_cycle(cycle.index, cycle.token_usage.clone());
139 }
140 summary
141}
142
143fn read_nested_int(source: &Value, path_options: &[&[&str]]) -> Option<u64> {
144 path_options
145 .iter()
146 .find_map(|path| nested_value(source, path).and_then(|value| read_int(Some(value))))
147}
148
149fn read_nested_cache_int(source: &Value, path_options: &[&[&str]]) -> Option<u64> {
150 path_options
151 .iter()
152 .find_map(|path| nested_value(source, path).and_then(read_cache_int))
153}
154
155fn nested_value<'a>(source: &'a Value, path: &[&str]) -> Option<&'a Value> {
156 let mut current = source;
157 for key in path {
158 current = current.as_object()?.get(*key)?;
159 }
160 Some(current)
161}
162
163fn read_cache_int(value: &Value) -> Option<u64> {
164 match value {
165 Value::Bool(_) => None,
166 Value::Number(number) => number.as_u64().or_else(|| {
167 number.as_f64().and_then(|value| {
168 (value.is_finite() && value >= 0.0 && value.fract() == 0.0).then_some(value as u64)
169 })
170 }),
171 Value::String(value) => value.trim().parse::<u64>().ok(),
172 _ => None,
173 }
174}
175
176fn read_int(value: Option<&Value>) -> Option<u64> {
177 match value? {
178 Value::Bool(_) => None,
179 Value::Number(number) => number
180 .as_u64()
181 .or_else(|| {
182 number
183 .as_i64()
184 .and_then(|value| (value >= 0).then_some(value as u64))
185 })
186 .or_else(|| number.as_f64().and_then(float_to_u64)),
187 Value::String(value) => value.trim().parse::<u64>().ok(),
188 _ => None,
189 }
190}
191
192fn float_to_u64(value: f64) -> Option<u64> {
193 value
194 .is_finite()
195 .then_some(value.trunc())
196 .and_then(|value| (value >= 0.0).then_some(value as u64))
197}
198
199fn infer_usage_source(raw_usage: &Value) -> UsageSource {
200 let Some(raw) = raw_usage.as_object() else {
201 return UsageSource::AccountingMissing;
202 };
203 if [
204 "prompt_tokens",
205 "completion_tokens",
206 "total_tokens",
207 "input_tokens",
208 "output_tokens",
209 ]
210 .iter()
211 .any(|key| raw.contains_key(*key) && read_int(raw.get(*key)).is_some())
212 {
213 UsageSource::ProviderReported
214 } else {
215 UsageSource::default()
216 }
217}
218
219fn has_any_key(source: &Value, keys: &[&str]) -> bool {
220 source
221 .as_object()
222 .is_some_and(|object| keys.iter().any(|key| object.contains_key(*key)))
223}
224
225fn has_nested_path(source: &Value, path: &[&str]) -> bool {
226 nested_value(source, path).is_some()
227}