agent_model_contract/
usage.rs

1//! Observed token counters: absent is unreported, while Some(0) is an explicit zero.
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
6#[serde(rename_all = "camelCase")]
7pub struct TokenUsage {
8    #[serde(default, skip_serializing_if = "Option::is_none")]
9    pub input_tokens: Option<u64>,
10    #[serde(default, skip_serializing_if = "Option::is_none")]
11    pub output_tokens: Option<u64>,
12    #[serde(default, skip_serializing_if = "Option::is_none")]
13    pub input_cached_tokens: Option<u64>,
14    #[serde(default, skip_serializing_if = "Option::is_none")]
15    pub input_cache_write_tokens: Option<u64>,
16}
17
18impl TokenUsage {
19    pub fn is_reported(&self) -> bool {
20        self.input_tokens.is_some()
21            || self.output_tokens.is_some()
22            || self.input_cached_tokens.is_some()
23            || self.input_cache_write_tokens.is_some()
24    }
25
26    /// SSE usage frames contain cumulative counters, not increments.
27    pub fn merge(&mut self, next: &Self) {
28        self.input_tokens = next.input_tokens.or(self.input_tokens);
29        self.output_tokens = next.output_tokens.or(self.output_tokens);
30        self.input_cached_tokens = next.input_cached_tokens.or(self.input_cached_tokens);
31        self.input_cache_write_tokens = next
32            .input_cache_write_tokens
33            .or(self.input_cache_write_tokens);
34    }
35
36    pub fn total_tokens(&self) -> Option<u64> {
37        self.input_tokens?.checked_add(self.output_tokens?)
38    }
39}
40
41/// Number of invocations that actually reported each counter; denominator is invocation_count.
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
43#[serde(rename_all = "camelCase")]
44pub struct TokenUsageCoverage {
45    pub input_tokens: u64,
46    pub output_tokens: u64,
47    pub input_cached_tokens: u64,
48    pub input_cache_write_tokens: u64,
49}
50
51impl crate::ModelUsageSummary {
52    pub fn observe_usage(&mut self, usage: &TokenUsage) {
53        if let Some(tokens) = usage.input_tokens {
54            self.input_tokens = self.input_tokens.saturating_add(tokens);
55            self.usage_coverage.input_tokens += 1;
56        }
57        if let Some(tokens) = usage.output_tokens {
58            self.output_tokens = self.output_tokens.saturating_add(tokens);
59            self.usage_coverage.output_tokens += 1;
60        }
61        self.usage_coverage.input_cached_tokens += u64::from(usage.input_cached_tokens.is_some());
62        self.usage_coverage.input_cache_write_tokens +=
63            u64::from(usage.input_cache_write_tokens.is_some());
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use serde_json::json;
71
72    #[test]
73    fn missing_zero_and_partial_counters_survive_serialization_and_merge() {
74        let mut usage: TokenUsage =
75            serde_json::from_value(json!({"inputTokens":10,"inputCachedTokens":6})).unwrap();
76        assert_eq!(usage.total_tokens(), None);
77        usage.merge(&serde_json::from_value(json!({"outputTokens":4})).unwrap());
78        usage.merge(&TokenUsage::default());
79        assert_eq!(usage.total_tokens(), Some(14));
80        assert_eq!(
81            serde_json::to_value(&usage).unwrap(),
82            json!({"inputTokens":10,"outputTokens":4,"inputCachedTokens":6})
83        );
84        usage.merge(&serde_json::from_value(json!({"inputCachedTokens":0})).unwrap());
85        assert_eq!(usage.input_cached_tokens, Some(0));
86        assert!(!TokenUsage::default().is_reported());
87        assert_eq!(
88            serde_json::to_value(TokenUsage::default()).unwrap(),
89            json!({})
90        );
91    }
92
93    #[test]
94    fn summaries_count_only_observed_fields_and_total_overflow_is_absent() {
95        let mut summary = crate::ModelUsageSummary::default();
96        summary.observe_usage(&TokenUsage::default());
97        summary.observe_usage(&TokenUsage {
98            input_tokens: Some(0),
99            input_cached_tokens: Some(0),
100            ..Default::default()
101        });
102        assert_eq!(summary.usage_coverage.input_tokens, 1);
103        assert_eq!(summary.usage_coverage.output_tokens, 0);
104        assert_eq!(summary.usage_coverage.input_cached_tokens, 1);
105        assert_eq!(
106            TokenUsage {
107                input_tokens: Some(u64::MAX),
108                output_tokens: Some(1),
109                ..Default::default()
110            }
111            .total_tokens(),
112            None
113        );
114    }
115}
agent_model_contract/usage.rs

agent_model_contract/
usage.rs
