vv_agent/runtime/
token_usage.rs

1use serde_json::Value;
2
3use crate::types::{
4    CacheUsage, CacheUsageStatus, ModelCallRecord, 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_count(raw.get("prompt_tokens"));
31    let completion_tokens = read_count(raw.get("completion_tokens"));
32    let native_input_tokens = read_count(raw.get("input_tokens")).or(prompt_tokens);
33    let output_tokens = read_count(raw.get("output_tokens")).or(completion_tokens);
34    let cache_read_input_tokens = read_nested_count(
35        raw_usage,
36        &[
37            &["cache_read_input_tokens"],
38            &["cache_read_tokens"],
39            &["prompt_tokens_details", "cached_tokens"],
40            &["input_tokens_details", "cached_tokens"],
41        ],
42    );
43    let reasoning_tokens = read_nested_count(
44        raw_usage,
45        &[
46            &["completion_tokens_details", "reasoning_tokens"],
47            &["output_tokens_details", "reasoning_tokens"],
48            &["reasoning_tokens"],
49        ],
50    );
51    let cache_write_input_tokens = read_nested_count(
52        raw_usage,
53        &[
54            &["cache_write_input_tokens"],
55            &["cache_creation_input_tokens"],
56            &["cache_write_tokens"],
57            &["input_tokens_details", "cache_creation_tokens"],
58            &["prompt_tokens_details", "cache_creation_tokens"],
59        ],
60    );
61    let mut uncached_input_tokens = read_nested_count(raw_usage, &[&["uncached_input_tokens"]]);
62
63    let anthropic_native = prompt_tokens.is_none()
64        && !raw.contains_key("total_tokens")
65        && uncached_input_tokens.is_none()
66        && has_any_key(
67            raw_usage,
68            &["cache_read_input_tokens", "cache_creation_input_tokens"],
69        );
70    let input_tokens = if anthropic_native {
71        native_input_tokens.and_then(|native| {
72            native
73                .checked_add(cache_read_input_tokens.unwrap_or_default())?
74                .checked_add(cache_write_input_tokens.unwrap_or_default())
75        })
76    } else {
77        native_input_tokens
78    };
79    let total_tokens = read_count(raw.get("total_tokens")).or_else(|| {
80        input_tokens.and_then(|input| output_tokens.and_then(|output| input.checked_add(output)))
81    });
82
83    let observed_cache_metric = cache_read_input_tokens.is_some()
84        || cache_write_input_tokens.is_some()
85        || uncached_input_tokens.is_some();
86    let normalized_cache_status = if observed_cache_metric {
87        CacheUsageStatus::ProviderReported
88    } else {
89        cache_status.unwrap_or_default()
90    };
91
92    if normalized_cache_status == CacheUsageStatus::ProviderReported
93        && uncached_input_tokens.is_none()
94    {
95        if anthropic_native {
96            uncached_input_tokens = native_input_tokens.and_then(|native| {
97                native.checked_add(cache_write_input_tokens.unwrap_or_default())
98            });
99        } else if let (Some(input), Some(read)) = (input_tokens, cache_read_input_tokens) {
100            uncached_input_tokens = Some(input.saturating_sub(read));
101        }
102    }
103
104    let cache_usage = CacheUsage {
105        status: normalized_cache_status,
106        read_input_tokens: (normalized_cache_status == CacheUsageStatus::ProviderReported)
107            .then_some(cache_read_input_tokens)
108            .flatten(),
109        write_input_tokens: (normalized_cache_status == CacheUsageStatus::ProviderReported)
110            .then_some(cache_write_input_tokens)
111            .flatten(),
112        uncached_input_tokens: (normalized_cache_status == CacheUsageStatus::ProviderReported)
113            .then_some(uncached_input_tokens)
114            .flatten(),
115        source: match normalized_cache_status {
116            CacheUsageStatus::ProviderReported => Some("provider_usage".to_string()),
117            CacheUsageStatus::Unsupported => Some("adapter_capability".to_string()),
118            CacheUsageStatus::AccountingMissing => None,
119        },
120    };
121
122    TokenUsage {
123        input_tokens,
124        output_tokens,
125        total_tokens,
126        reasoning_tokens,
127        usage_source: usage_source.unwrap_or_else(|| infer_usage_source(raw_usage)),
128        cache_usage,
129        provider_usage: raw.clone(),
130    }
131}
132
133pub fn summarize_task_token_usage(model_calls: &[ModelCallRecord]) -> TaskTokenUsage {
134    let mut summary = TaskTokenUsage::default();
135    for model_call in model_calls {
136        summary
137            .add_model_call(model_call.clone())
138            .expect("runtime model call ledger contains duplicate call_id");
139    }
140    summary
141}
142
143fn read_nested_count(source: &Value, path_options: &[&[&str]]) -> Option<u64> {
144    path_options
145        .iter()
146        .find_map(|path| nested_value(source, path).and_then(|value| read_count(Some(value))))
147}
148
149fn nested_value<'a>(source: &'a Value, path: &[&str]) -> Option<&'a Value> {
150    let mut current = source;
151    for key in path {
152        current = current.as_object()?.get(*key)?;
153    }
154    Some(current)
155}
156
157fn read_count(value: Option<&Value>) -> Option<u64> {
158    match value? {
159        Value::Bool(_) => None,
160        Value::Number(number) => number.as_u64().or_else(|| {
161            number.as_f64().and_then(|value| {
162                (value.is_finite() && value >= 0.0 && value.fract() == 0.0).then_some(value as u64)
163            })
164        }),
165        Value::String(value) => value.trim().parse::<u64>().ok(),
166        _ => None,
167    }
168}
169
170fn infer_usage_source(raw_usage: &Value) -> UsageSource {
171    let Some(raw) = raw_usage.as_object() else {
172        return UsageSource::AccountingMissing;
173    };
174    if [
175        "prompt_tokens",
176        "completion_tokens",
177        "total_tokens",
178        "input_tokens",
179        "output_tokens",
180    ]
181    .iter()
182    .any(|key| raw.contains_key(*key) && read_count(raw.get(*key)).is_some())
183    {
184        UsageSource::ProviderReported
185    } else {
186        UsageSource::AccountingMissing
187    }
188}
189
190fn has_any_key(source: &Value, keys: &[&str]) -> bool {
191    source
192        .as_object()
193        .is_some_and(|object| keys.iter().any(|key| object.contains_key(*key)))
194}
vv_agent/runtime/token_usage.rs

vv_agent/runtime/
token_usage.rs
