claude_code_acp/types/
session.rs

1//! Session-related types for token usage and statistics
2
3use serde::{Deserialize, Serialize};
4
5/// Token usage statistics
6///
7/// Tracks the number of tokens used in a session or query.
8/// Can be parsed from SDK's `ResultMessage.usage` field.
9#[derive(Debug, Clone, Default, Serialize, Deserialize)]
10pub struct TokenUsage {
11    /// Number of input tokens
12    pub input_tokens: u64,
13
14    /// Number of output tokens
15    pub output_tokens: u64,
16
17    /// Number of tokens read from cache
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub cache_read_input_tokens: Option<u64>,
20
21    /// Number of tokens written to cache
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub cache_creation_input_tokens: Option<u64>,
24}
25
26impl TokenUsage {
27    /// Create a new empty token usage
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    /// Parse from SDK's usage JSON value
33    ///
34    /// # Arguments
35    ///
36    /// * `usage` - The `usage` field from `ResultMessage` or `AssistantMessage`
37    pub fn from_sdk_usage(usage: &serde_json::Value) -> Self {
38        Self {
39            input_tokens: usage["input_tokens"].as_u64().unwrap_or(0),
40            output_tokens: usage["output_tokens"].as_u64().unwrap_or(0),
41            cache_read_input_tokens: usage["cache_read_input_tokens"].as_u64(),
42            cache_creation_input_tokens: usage["cache_creation_input_tokens"].as_u64(),
43        }
44    }
45
46    /// Add another usage to this one
47    pub fn add(&mut self, other: &TokenUsage) {
48        self.input_tokens += other.input_tokens;
49        self.output_tokens += other.output_tokens;
50
51        if let Some(v) = other.cache_read_input_tokens {
52            *self.cache_read_input_tokens.get_or_insert(0) += v;
53        }
54        if let Some(v) = other.cache_creation_input_tokens {
55            *self.cache_creation_input_tokens.get_or_insert(0) += v;
56        }
57    }
58
59    /// Get total token count (input + output)
60    pub fn total(&self) -> u64 {
61        self.input_tokens + self.output_tokens
62    }
63
64    /// Check if any tokens were used
65    pub fn is_empty(&self) -> bool {
66        self.input_tokens == 0 && self.output_tokens == 0
67    }
68}
69
70/// Session statistics
71#[derive(Debug, Clone, Default)]
72pub struct SessionStats {
73    /// Number of active sessions
74    pub active_sessions: usize,
75
76    /// Total token usage across all sessions
77    pub total_usage: TokenUsage,
78
79    /// Total cost in USD
80    pub total_cost_usd: f64,
81}
82
83impl SessionStats {
84    /// Create new empty stats
85    pub fn new() -> Self {
86        Self::default()
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use serde_json::json;
94
95    #[test]
96    fn test_token_usage_default() {
97        let usage = TokenUsage::default();
98        assert_eq!(usage.input_tokens, 0);
99        assert_eq!(usage.output_tokens, 0);
100        assert!(usage.cache_read_input_tokens.is_none());
101        assert!(usage.is_empty());
102        assert_eq!(usage.total(), 0);
103    }
104
105    #[test]
106    fn test_token_usage_from_sdk() {
107        let sdk_usage = json!({
108            "input_tokens": 1000,
109            "output_tokens": 500,
110            "cache_read_input_tokens": 200,
111            "cache_creation_input_tokens": 100
112        });
113
114        let usage = TokenUsage::from_sdk_usage(&sdk_usage);
115        assert_eq!(usage.input_tokens, 1000);
116        assert_eq!(usage.output_tokens, 500);
117        assert_eq!(usage.cache_read_input_tokens, Some(200));
118        assert_eq!(usage.cache_creation_input_tokens, Some(100));
119        assert_eq!(usage.total(), 1500);
120        assert!(!usage.is_empty());
121    }
122
123    #[test]
124    fn test_token_usage_add() {
125        let mut usage1 = TokenUsage {
126            input_tokens: 100,
127            output_tokens: 50,
128            cache_read_input_tokens: Some(10),
129            cache_creation_input_tokens: None,
130        };
131
132        let usage2 = TokenUsage {
133            input_tokens: 200,
134            output_tokens: 100,
135            cache_read_input_tokens: Some(20),
136            cache_creation_input_tokens: Some(5),
137        };
138
139        usage1.add(&usage2);
140
141        assert_eq!(usage1.input_tokens, 300);
142        assert_eq!(usage1.output_tokens, 150);
143        assert_eq!(usage1.cache_read_input_tokens, Some(30));
144        assert_eq!(usage1.cache_creation_input_tokens, Some(5));
145    }
146
147    #[test]
148    fn test_token_usage_serialization() {
149        let usage = TokenUsage {
150            input_tokens: 100,
151            output_tokens: 50,
152            cache_read_input_tokens: None,
153            cache_creation_input_tokens: None,
154        };
155
156        let json = serde_json::to_string(&usage).unwrap();
157        let parsed: TokenUsage = serde_json::from_str(&json).unwrap();
158
159        assert_eq!(parsed.input_tokens, 100);
160        assert_eq!(parsed.output_tokens, 50);
161    }
162}
claude_code_acp/types/session.rs

claude_code_acp/types/
session.rs
