1use serde::{Deserialize, Serialize};
9
10use crate::context::{ContextItem, ContextVisibility};
11
12pub const TOKEN_ESTIMATOR_VERSION: &str = "utf8-bytes-divisor-3-overhead-16-v1";
14
15pub const USAGE_NORMALIZATION_VERSION: &str = "provider-inclusive-input-v1";
17
18pub const MODEL_PROJECTION_MAX_BYTES: usize = 128 * 1024;
23
24#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
26#[serde(tag = "kind", rename_all = "snake_case")]
27pub enum MeasurementProvenance {
28 ExactSerialized {
30 boundary: String,
32 },
33 Estimated {
35 estimator: String,
37 version: String,
39 },
40 ProviderReported {
42 provider: String,
44 component: String,
46 },
47 Derived {
49 rule: String,
51 version: String,
53 },
54 Unknown,
56}
57
58#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
60pub struct ByteMeasurement {
61 pub value: u64,
63 pub provenance: MeasurementProvenance,
65}
66
67#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
69pub struct TokenMeasurement {
70 pub value: Option<u64>,
72 pub provenance: MeasurementProvenance,
74}
75
76impl TokenMeasurement {
77 pub const fn unknown() -> Self {
79 Self { value: None, provenance: MeasurementProvenance::Unknown }
80 }
81
82 pub fn provider(provider: &str, component: &str, value: Option<u64>) -> Self {
84 Self {
85 value,
86 provenance: if value.is_some() {
87 MeasurementProvenance::ProviderReported {
88 provider: provider.to_string(),
89 component: component.to_string(),
90 }
91 } else {
92 MeasurementProvenance::Unknown
93 },
94 }
95 }
96}
97
98#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
100pub struct ProviderUsageComponents {
101 pub input_tokens: Option<u64>,
103 pub output_tokens: Option<u64>,
105 pub cache_read_input_tokens: Option<u64>,
107 pub cache_creation_input_tokens: Option<u64>,
109 pub reasoning_tokens: Option<u64>,
111}
112
113impl ProviderUsageComponents {
114 pub const fn new(input_tokens: u64, output_tokens: u64) -> Self {
116 Self {
117 input_tokens: Some(input_tokens),
118 output_tokens: Some(output_tokens),
119 cache_read_input_tokens: None,
120 cache_creation_input_tokens: None,
121 reasoning_tokens: None,
122 }
123 }
124
125 pub fn merge_snapshot(&mut self, next: &Self) {
127 merge_max(&mut self.input_tokens, next.input_tokens);
128 merge_max(&mut self.output_tokens, next.output_tokens);
129 merge_max(&mut self.cache_read_input_tokens, next.cache_read_input_tokens);
130 merge_max(&mut self.cache_creation_input_tokens, next.cache_creation_input_tokens);
131 merge_max(&mut self.reasoning_tokens, next.reasoning_tokens);
132 }
133
134 pub fn normalize(&self, provider: &str, rule: ProviderUsageRule) -> ProviderUsage {
136 let inclusive_input_tokens = match rule {
137 ProviderUsageRule::AnthropicMessages => self.input_tokens.map(|input| {
138 input
139 .saturating_add(self.cache_read_input_tokens.unwrap_or(0))
140 .saturating_add(self.cache_creation_input_tokens.unwrap_or(0))
141 }),
142 ProviderUsageRule::OpenAiChat | ProviderUsageRule::OpenAiResponses => self.input_tokens,
143 };
144 ProviderUsage {
145 provider: provider.to_string(),
146 rule,
147 components: self.clone(),
148 inclusive_input_tokens: TokenMeasurement {
149 value: inclusive_input_tokens,
150 provenance: if inclusive_input_tokens.is_some() {
151 MeasurementProvenance::Derived {
152 rule: rule.label().to_string(),
153 version: USAGE_NORMALIZATION_VERSION.to_string(),
154 }
155 } else {
156 MeasurementProvenance::Unknown
157 },
158 },
159 }
160 }
161}
162
163fn merge_max(current: &mut Option<u64>, next: Option<u64>) {
164 if let Some(next) = next {
165 *current = Some(current.map_or(next, |current| current.max(next)));
166 }
167}
168
169#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
171#[serde(rename_all = "snake_case")]
172pub enum ProviderUsageRule {
173 AnthropicMessages,
175 OpenAiChat,
177 OpenAiResponses,
179}
180
181impl ProviderUsageRule {
182 pub const fn label(self) -> &'static str {
184 match self {
185 Self::AnthropicMessages => "anthropic_input_plus_cache_components",
186 Self::OpenAiChat => "openai_prompt_tokens_inclusive",
187 Self::OpenAiResponses => "openai_responses_input_tokens_inclusive",
188 }
189 }
190}
191
192#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
194pub struct ProviderUsage {
195 pub provider: String,
197 pub rule: ProviderUsageRule,
199 pub components: ProviderUsageComponents,
201 pub inclusive_input_tokens: TokenMeasurement,
203}
204
205#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
207pub struct ModelProjectionMessage {
208 pub role: String,
210 pub content: String,
212}
213
214#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
216pub struct ContextReductionReceipt {
217 pub item_id: String,
219 pub method: String,
221 pub version: String,
223 pub before_bytes: u64,
225 pub after_bytes: u64,
227 pub lossy: bool,
229}
230
231#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
233pub struct ContextItemSnapshot {
234 pub id: String,
236 #[serde(default, skip_serializing_if = "Option::is_none")]
238 pub artifact_handle: Option<String>,
239 pub state: ContextVisibility,
241 pub reason_code: String,
243 pub reason: String,
245}
246
247impl From<&ContextItem> for ContextItemSnapshot {
248 fn from(item: &ContextItem) -> Self {
249 Self {
250 id: item.id.clone(),
251 artifact_handle: item.artifact_handle.clone(),
252 state: item.visibility.clone(),
253 reason_code: item.reason_code.clone(),
254 reason: item.reason.clone(),
255 }
256 }
257}
258
259pub fn snapshot_context(items: &[ContextItem]) -> Vec<ContextItemSnapshot> {
261 items.iter().map(ContextItemSnapshot::from).collect()
262}
263
264#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
266pub struct ProviderRequestAccounting {
267 pub turn_id: String,
269 pub request_id: String,
271 pub attempt: u32,
273 pub provider: String,
275 pub model: String,
277 pub serialized_bytes: ByteMeasurement,
279 pub estimated_input_tokens: TokenMeasurement,
281 pub provider_usage: Option<ProviderUsage>,
283 pub context: Vec<ContextItemSnapshot>,
285 #[serde(default)]
287 pub shadow_receipts: Vec<ContextReductionReceipt>,
288 #[serde(skip)]
293 pub model_projection: Vec<ModelProjectionMessage>,
294}
295
296impl ProviderRequestAccounting {
297 pub fn from_serialized_request(
299 turn_id: impl Into<String>, request_id: impl Into<String>, attempt: u32, provider: &str, model: &str,
300 bytes: &[u8], context: Vec<ContextItemSnapshot>,
301 ) -> Self {
302 let byte_count = bytes.len() as u64;
303 Self {
304 turn_id: turn_id.into(),
305 request_id: request_id.into(),
306 attempt,
307 provider: provider.to_string(),
308 model: model.to_string(),
309 serialized_bytes: ByteMeasurement {
310 value: byte_count,
311 provenance: MeasurementProvenance::ExactSerialized { boundary: "provider_request_body".to_string() },
312 },
313 estimated_input_tokens: TokenMeasurement {
314 value: Some(estimate_serialized_tokens(byte_count)),
315 provenance: MeasurementProvenance::Estimated {
316 estimator: "utf8_bytes_divisor_3_plus_item_overhead".to_string(),
317 version: TOKEN_ESTIMATOR_VERSION.to_string(),
318 },
319 },
320 provider_usage: None,
321 context,
322 shadow_receipts: Vec::new(),
323 model_projection: Vec::new(),
324 }
325 }
326
327 pub fn with_model_projection(mut self, projection: Vec<ModelProjectionMessage>) -> Self {
330 let mut bytes = 0usize;
331 self.model_projection = projection
332 .into_iter()
333 .filter_map(|mut message| {
334 let remaining = MODEL_PROJECTION_MAX_BYTES.saturating_sub(bytes);
335 if remaining == 0 {
336 return None;
337 }
338 message.content = truncate_utf8(&message.content, remaining);
339 bytes = bytes
340 .saturating_add(message.role.len())
341 .saturating_add(message.content.len());
342 Some(message)
343 })
344 .collect();
345 self
346 }
347}
348
349fn truncate_utf8(value: &str, max_bytes: usize) -> String {
350 if value.len() <= max_bytes {
351 return value.to_string();
352 }
353 let mut end = max_bytes;
354 while end > 0 && !value.is_char_boundary(end) {
355 end -= 1;
356 }
357 value[..end].to_string()
358}
359
360pub const fn estimate_serialized_tokens(bytes: u64) -> u64 {
362 bytes.div_ceil(3).saturating_add(16)
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368
369 #[test]
370 fn provider_components_preserve_unknown_and_measured_zero() {
371 let components = ProviderUsageComponents {
372 input_tokens: Some(0),
373 output_tokens: None,
374 cache_read_input_tokens: None,
375 cache_creation_input_tokens: Some(0),
376 reasoning_tokens: None,
377 };
378 assert_eq!(components.input_tokens, Some(0));
379 assert_eq!(components.output_tokens, None);
380 assert_eq!(components.cache_creation_input_tokens, Some(0));
381 }
382
383 #[test]
384 fn repeated_stream_snapshots_are_not_added_twice() {
385 let mut total = ProviderUsageComponents::new(12, 3);
386 total.merge_snapshot(&ProviderUsageComponents::new(12, 3));
387 total.merge_snapshot(&ProviderUsageComponents::new(15, 4));
388 assert_eq!(total.input_tokens, Some(15));
389 assert_eq!(total.output_tokens, Some(4));
390 }
391
392 #[test]
393 fn provider_rules_normalize_anthropic_and_openai_inputs_differently() {
394 let components = ProviderUsageComponents {
395 input_tokens: Some(100),
396 output_tokens: Some(10),
397 cache_read_input_tokens: Some(20),
398 cache_creation_input_tokens: Some(5),
399 reasoning_tokens: Some(2),
400 };
401 assert_eq!(
402 components
403 .normalize("anthropic", ProviderUsageRule::AnthropicMessages)
404 .inclusive_input_tokens
405 .value,
406 Some(125)
407 );
408 assert_eq!(
409 components
410 .normalize("openai", ProviderUsageRule::OpenAiChat)
411 .inclusive_input_tokens
412 .value,
413 Some(100)
414 );
415 }
416
417 #[test]
418 fn request_accounting_records_exact_bytes_and_context_once() {
419 let item = ContextItem {
420 id: "item-1".to_string(),
421 kind: crate::context::ContextItemKind::Transcript,
422 label: "turn".to_string(),
423 source_path: None,
424 scope: ".".to_string(),
425 content_hash: None,
426 artifact_handle: None,
427 byte_count: 4,
428 content: None,
429 token_estimate: 18,
430 visibility: ContextVisibility::Visible,
431 reason_code: "recent_transcript".to_string(),
432 reason: "recent transcript entry".to_string(),
433 };
434 let context = snapshot_context(std::slice::from_ref(&item));
435 let accounting = ProviderRequestAccounting::from_serialized_request(
436 "turn_1",
437 "turn_1:request:0",
438 1,
439 "provider",
440 "model",
441 b"{}",
442 context,
443 );
444 assert_eq!(accounting.serialized_bytes.value, 2);
445 assert_eq!(accounting.estimated_input_tokens.value, Some(17));
446 assert_eq!(accounting.context.len(), 1);
447 assert_eq!(accounting.context[0].reason_code, "recent_transcript");
448 }
449}