bamboo_domain/
token_usage.rs

1//! Shared token usage value object used across domain, application, and infrastructure layers.
2
3use serde::{Deserialize, Serialize};
4
5/// Largest token counter that can be represented losslessly by the signed
6/// 64-bit integer columns used by durable metrics stores.
7///
8/// Keeping this boundary beside the shared value object lets runtime budgets
9/// and persistence apply one policy instead of allowing `u64 as i64` wraparound
10/// or storage-specific divergence.
11pub const MAX_DURABLE_TOKEN_COUNT: u64 = i64::MAX as u64;
12
13/// Token consumption statistics for a single LLM call or aggregated period.
14///
15/// This is a stable, cross-layer value object. Every crate that needs to
16/// represent "how many tokens were used" should use this type (or re-export it)
17/// instead of defining a local duplicate.
18#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
19pub struct TokenUsage {
20    /// Tokens in the LLM request (prompt / input).
21    pub prompt_tokens: u64,
22    /// Tokens in the LLM response (completion / output).
23    pub completion_tokens: u64,
24    /// Total tokens consumed (normally prompt + completion).
25    pub total_tokens: u64,
26}
27
28impl TokenUsage {
29    /// Accumulate another usage snapshot into this one.
30    pub fn add_assign(&mut self, other: TokenUsage) {
31        self.prompt_tokens = self.prompt_tokens.saturating_add(other.prompt_tokens);
32        self.completion_tokens = self
33            .completion_tokens
34            .saturating_add(other.completion_tokens);
35        self.total_tokens = self.total_tokens.saturating_add(other.total_tokens);
36    }
37
38    /// Recompute `total_tokens` from the two component fields.
39    pub fn recompute_total(&mut self) {
40        self.total_tokens = self.prompt_tokens.saturating_add(self.completion_tokens);
41    }
42
43    /// Normalize all counters to the lossless signed-64 durable range.
44    ///
45    /// Components saturate independently and the total is recomputed from the
46    /// saturated components, then saturated to the same boundary. This makes
47    /// the policy deterministic even if an upstream supplied an inconsistent
48    /// total and keeps round, session, and runtime views reconcilable.
49    pub fn clamped_for_durable_metrics(mut self) -> Self {
50        self.prompt_tokens = self.prompt_tokens.min(MAX_DURABLE_TOKEN_COUNT);
51        self.completion_tokens = self.completion_tokens.min(MAX_DURABLE_TOKEN_COUNT);
52        self.total_tokens = self
53            .prompt_tokens
54            .saturating_add(self.completion_tokens)
55            .min(MAX_DURABLE_TOKEN_COUNT);
56        self
57    }
58
59    /// Accumulate usage with the same saturation policy used by durable stores.
60    pub fn add_assign_durable(&mut self, other: TokenUsage) {
61        self.prompt_tokens = self
62            .prompt_tokens
63            .saturating_add(other.prompt_tokens)
64            .min(MAX_DURABLE_TOKEN_COUNT);
65        self.completion_tokens = self
66            .completion_tokens
67            .saturating_add(other.completion_tokens)
68            .min(MAX_DURABLE_TOKEN_COUNT);
69        self.total_tokens = self
70            .prompt_tokens
71            .saturating_add(self.completion_tokens)
72            .min(MAX_DURABLE_TOKEN_COUNT);
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn default_is_zero() {
82        let usage = TokenUsage::default();
83        assert_eq!(usage.prompt_tokens, 0);
84        assert_eq!(usage.completion_tokens, 0);
85        assert_eq!(usage.total_tokens, 0);
86    }
87
88    #[test]
89    fn add_assign_accumulates() {
90        let mut usage1 = TokenUsage {
91            prompt_tokens: 100,
92            completion_tokens: 50,
93            total_tokens: 150,
94        };
95        let usage2 = TokenUsage {
96            prompt_tokens: 200,
97            completion_tokens: 100,
98            total_tokens: 300,
99        };
100        usage1.add_assign(usage2);
101        assert_eq!(usage1.prompt_tokens, 300);
102        assert_eq!(usage1.completion_tokens, 150);
103        assert_eq!(usage1.total_tokens, 450);
104    }
105
106    #[test]
107    fn recompute_total_uses_saturating_add() {
108        let mut usage = TokenUsage {
109            prompt_tokens: u64::MAX - 5,
110            completion_tokens: u64::MAX - 9,
111            total_tokens: 0,
112        };
113        usage.recompute_total();
114        assert_eq!(usage.total_tokens, u64::MAX);
115    }
116
117    #[test]
118    fn durable_metrics_policy_clamps_components_and_recomputes_total() {
119        let usage = TokenUsage {
120            prompt_tokens: u64::MAX,
121            completion_tokens: u64::MAX,
122            total_tokens: 1,
123        }
124        .clamped_for_durable_metrics();
125
126        assert_eq!(usage.prompt_tokens, MAX_DURABLE_TOKEN_COUNT);
127        assert_eq!(usage.completion_tokens, MAX_DURABLE_TOKEN_COUNT);
128        assert_eq!(usage.total_tokens, MAX_DURABLE_TOKEN_COUNT);
129    }
130
131    #[test]
132    fn durable_accumulation_saturates_without_wraparound() {
133        let mut usage = TokenUsage {
134            prompt_tokens: MAX_DURABLE_TOKEN_COUNT - 2,
135            completion_tokens: MAX_DURABLE_TOKEN_COUNT - 3,
136            total_tokens: MAX_DURABLE_TOKEN_COUNT,
137        };
138        usage.add_assign_durable(TokenUsage {
139            prompt_tokens: 10,
140            completion_tokens: 20,
141            total_tokens: 30,
142        });
143
144        assert_eq!(usage.prompt_tokens, MAX_DURABLE_TOKEN_COUNT);
145        assert_eq!(usage.completion_tokens, MAX_DURABLE_TOKEN_COUNT);
146        assert_eq!(usage.total_tokens, MAX_DURABLE_TOKEN_COUNT);
147    }
148}
bamboo_domain/token_usage.rs

bamboo_domain/
token_usage.rs
