Skip to main content

agent_sessions/
accounting.rs

1use crate::{AccountingPolicy, LineErrorKind, TokenCounts};
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
6#[non_exhaustive]
7pub enum UsageAdjustment {
8    NegativeClamped { field: String, original: i64 },
9    CacheTtlCapped { original: u64, total: u64 },
10    UnreportedCounters,
11    MissingCumulativeBaseline,
12    CumulativeRegression,
13    UncertainCumulativeBaseline,
14}
15
16pub(crate) fn counts(
17    v: &Value,
18    claude: bool,
19    policy: AccountingPolicy,
20) -> Result<(TokenCounts, Vec<UsageAdjustment>), LineErrorKind> {
21    let statistics = policy == AccountingPolicy::UsageStatistics;
22    let mut corrected: Option<Value> = None;
23    let mut adjustments = Vec::new();
24    if statistics && claude {
25        for key in [
26            "input_tokens",
27            "output_tokens",
28            "cache_creation_input_tokens",
29            "cache_read_input_tokens",
30        ] {
31            if let Some(original) = v.get(key).and_then(Value::as_i64).filter(|n| *n < 0) {
32                corrected.get_or_insert_with(|| v.clone())[key] = Value::from(0);
33                adjustments.push(UsageAdjustment::NegativeClamped {
34                    field: key.into(),
35                    original,
36                });
37            }
38        }
39        if let Some(original) = v
40            .pointer("/cache_creation/ephemeral_1h_input_tokens")
41            .and_then(Value::as_i64)
42            .filter(|n| *n < 0)
43        {
44            corrected.get_or_insert_with(|| v.clone())["cache_creation"]["ephemeral_1h_input_tokens"] =
45                Value::from(0);
46            adjustments.push(UsageAdjustment::NegativeClamped {
47                field: "ephemeral_1h_input_tokens".into(),
48                original,
49            });
50        }
51        let current = corrected.as_ref().unwrap_or(v);
52        let total = current
53            .get("cache_creation_input_tokens")
54            .and_then(Value::as_u64)
55            .unwrap_or(0);
56        if let Some(original) = current
57            .pointer("/cache_creation/ephemeral_1h_input_tokens")
58            .and_then(Value::as_u64)
59            .filter(|n| *n > total)
60        {
61            corrected.get_or_insert_with(|| v.clone())["cache_creation"]["ephemeral_1h_input_tokens"] =
62                Value::from(total);
63            adjustments.push(UsageAdjustment::CacheTtlCapped { original, total });
64        }
65    }
66    let counts = crate::tokens::parse_counts(corrected.as_ref().unwrap_or(v), claude, statistics)?;
67    if counts.is_missing() {
68        adjustments.push(UsageAdjustment::UnreportedCounters);
69    }
70    Ok((counts, adjustments))
71}