llm/usage/
session.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3
4use super::{LlmCallPurpose, ModelIdentity, TokenUsage, UsageCost, UsageSource, Usd};
5
6/// Running token totals and cost estimate across every call an agent has seen.
7#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
8pub struct SessionUsageTotals {
9    pub tokens: TokenUsage,
10    /// Sum of every priced call's estimated cost.
11    pub estimated_usd: Usd,
12    /// Cumulative estimated cost of non-cached input tokens, in USD.
13    pub estimated_input_usd: Usd,
14    /// Cumulative estimated cost of output tokens, in USD.
15    pub estimated_output_usd: Usd,
16    /// Cumulative estimated cost of cache-read tokens, in USD.
17    pub estimated_cache_read_usd: Usd,
18    /// Cumulative estimated cost of cache-creation tokens, in USD.
19    pub estimated_cache_creation_usd: Usd,
20    /// Calls with nonzero usage but no catalog pricing, which `estimated_usd`
21    /// therefore leaves out.
22    pub unpriced_calls: u64,
23}
24
25impl SessionUsageTotals {
26    pub fn add(&mut self, tokens: TokenUsage, estimated_cost: Option<UsageCost>) {
27        if tokens.is_zero() {
28            return;
29        }
30        self.tokens += tokens;
31        match estimated_cost {
32            Some(cost) => {
33                self.estimated_usd += cost.total_usd;
34                self.estimated_input_usd += cost.input_usd;
35                self.estimated_output_usd += cost.output_usd;
36                self.estimated_cache_read_usd += cost.cache_read_usd;
37                self.estimated_cache_creation_usd += cost.cache_creation_usd;
38            }
39            None => self.unpriced_calls += 1,
40        }
41    }
42
43    /// Whether `estimated_usd` accounts for every call with nonzero usage.
44    pub fn is_fully_priced(&self) -> bool {
45        self.unpriced_calls == 0
46    }
47}
48
49/// One provider usage sample, its estimated cost, and the totals after it.
50#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
51pub struct SessionUsageEvent {
52    /// Sequence assigned by the emitting tracker, including folded child samples.
53    pub sequence: u64,
54    /// Attribution of this sample, not the owner of the cumulative totals.
55    pub source: UsageSource,
56    pub purpose: LlmCallPurpose,
57    pub model: ModelIdentity,
58    pub tokens: TokenUsage,
59    pub estimated_cost: Option<UsageCost>,
60    pub totals: SessionUsageTotals,
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use serde_json::json;
67
68    #[test]
69    fn cumulative_cost_components_include_each_priced_sample_once() {
70        let mut totals = SessionUsageTotals::default();
71        let cost = UsageCost {
72            input_usd: Usd::new(0.125),
73            output_usd: Usd::new(0.25),
74            cache_read_usd: Usd::new(0.0625),
75            cache_creation_usd: Usd::new(0.0625),
76            total_usd: Usd::new(0.5),
77        };
78        totals.add(TokenUsage::new(10, 2), Some(cost));
79        totals.add(TokenUsage::new(20, 4), Some(cost));
80        totals.add(TokenUsage::new(5, 1), None);
81        totals.add(TokenUsage::default(), Some(cost));
82        totals.add(TokenUsage::default(), None);
83
84        let serialized = serde_json::to_value(&totals).unwrap();
85        assert_eq!(serialized["estimated_input_usd"], json!(0.25));
86        assert_eq!(serialized["estimated_output_usd"], json!(0.5));
87        assert_eq!(serialized["estimated_cache_read_usd"], json!(0.125));
88        assert_eq!(serialized["estimated_cache_creation_usd"], json!(0.125));
89        assert_eq!(totals.estimated_usd, Usd::new(1.0));
90        assert_eq!(totals.tokens, TokenUsage::new(35, 7));
91        assert_eq!(totals.unpriced_calls, 1);
92        assert!(!totals.is_fully_priced());
93        assert_eq!(serde_json::from_value::<SessionUsageTotals>(serialized).unwrap(), totals);
94    }
95}
llm/usage/session.rs

llm/usage/
session.rs
