claude_agent/tokens/
budget.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
4pub struct TokenBudget {
5    pub input_tokens: u64,
6    pub cache_read_tokens: u64,
7    pub cache_write_tokens: u64,
8    pub output_tokens: u64,
9}
10
11impl TokenBudget {
12    #[inline]
13    pub fn context_usage(&self) -> u64 {
14        self.input_tokens + self.cache_read_tokens + self.cache_write_tokens
15    }
16
17    #[inline]
18    pub fn total(&self) -> u64 {
19        self.context_usage() + self.output_tokens
20    }
21
22    pub fn add(&mut self, other: &TokenBudget) {
23        self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
24        self.cache_read_tokens = self
25            .cache_read_tokens
26            .saturating_add(other.cache_read_tokens);
27        self.cache_write_tokens = self
28            .cache_write_tokens
29            .saturating_add(other.cache_write_tokens);
30        self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
31    }
32
33    pub fn is_empty(&self) -> bool {
34        self.context_usage() == 0 && self.output_tokens == 0
35    }
36}
37
38impl From<&crate::types::Usage> for TokenBudget {
39    fn from(usage: &crate::types::Usage) -> Self {
40        Self {
41            input_tokens: usage.input_tokens as u64,
42            cache_read_tokens: usage.cache_read_input_tokens.unwrap_or(0) as u64,
43            cache_write_tokens: usage.cache_creation_input_tokens.unwrap_or(0) as u64,
44            output_tokens: usage.output_tokens as u64,
45        }
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn test_context_usage() {
55        let budget = TokenBudget {
56            input_tokens: 100,
57            cache_read_tokens: 200_000,
58            cache_write_tokens: 0,
59            output_tokens: 500,
60        };
61
62        assert_eq!(budget.context_usage(), 200_100);
63        assert_eq!(budget.total(), 200_600);
64    }
65
66    #[test]
67    fn test_add() {
68        let mut a = TokenBudget {
69            input_tokens: 100,
70            cache_read_tokens: 50,
71            cache_write_tokens: 25,
72            output_tokens: 200,
73        };
74        let b = TokenBudget {
75            input_tokens: 100,
76            cache_read_tokens: 50,
77            cache_write_tokens: 25,
78            output_tokens: 200,
79        };
80
81        a.add(&b);
82        assert_eq!(a.input_tokens, 200);
83        assert_eq!(a.cache_read_tokens, 100);
84    }
85}
claude_agent/tokens/budget.rs

claude_agent/tokens/
budget.rs
