agent_teams/models/
token.rs

1//! Token usage and cost types for agent session analysis.
2//!
3//! These types are used across multiple features (checkpoint, TUI, session
4//! discovery) and are therefore NOT feature-gated.
5
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8
9/// Token usage statistics from an agent session.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11#[serde(rename_all = "camelCase")]
12pub struct TokenUsage {
13    /// Number of input tokens consumed.
14    pub input_tokens: u64,
15    /// Number of output tokens generated.
16    pub output_tokens: u64,
17    /// Tokens read from cache (if applicable).
18    #[serde(default, skip_serializing_if = "Option::is_none")]
19    pub cache_read_tokens: Option<u64>,
20    /// Tokens written to cache (if applicable).
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub cache_write_tokens: Option<u64>,
23}
24
25impl TokenUsage {
26    /// Total tokens (input + output).
27    pub fn total(&self) -> u64 {
28        self.input_tokens + self.output_tokens
29    }
30
31    /// Merge another TokenUsage into this one (additive).
32    pub fn merge(&mut self, other: &TokenUsage) {
33        self.input_tokens += other.input_tokens;
34        self.output_tokens += other.output_tokens;
35        match (self.cache_read_tokens, other.cache_read_tokens) {
36            (Some(a), Some(b)) => self.cache_read_tokens = Some(a + b),
37            (None, Some(b)) => self.cache_read_tokens = Some(b),
38            _ => {}
39        }
40        match (self.cache_write_tokens, other.cache_write_tokens) {
41            (Some(a), Some(b)) => self.cache_write_tokens = Some(a + b),
42            (None, Some(b)) => self.cache_write_tokens = Some(b),
43            _ => {}
44        }
45    }
46}
47
48/// A single tool call record from an agent session (extended tier data).
49#[derive(Debug, Clone, Serialize, Deserialize)]
50#[serde(rename_all = "camelCase")]
51pub struct ToolCallRecord {
52    /// Tool name (e.g. "Read", "Write", "Bash", "Edit").
53    pub tool_name: String,
54    /// Truncated input summary (max 200 chars).
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub input_summary: Option<String>,
57    /// When the tool was called.
58    #[serde(default, skip_serializing_if = "Option::is_none")]
59    pub timestamp: Option<DateTime<Utc>>,
60}
61
62/// Aggregated cost summary across one or more sessions.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(rename_all = "camelCase")]
65pub struct CostSummary {
66    /// Total token usage across all sessions.
67    pub total_usage: TokenUsage,
68    /// Number of sessions aggregated.
69    pub session_count: usize,
70    /// Per-agent breakdown: agent_name → TokenUsage.
71    #[serde(default, skip_serializing_if = "Vec::is_empty")]
72    pub per_agent: Vec<AgentTokenUsage>,
73    /// Estimated cost in USD (approximate, based on public pricing).
74    pub estimated_cost_usd: f64,
75}
76
77/// Token usage attributed to a specific agent.
78#[derive(Debug, Clone, Serialize, Deserialize)]
79#[serde(rename_all = "camelCase")]
80pub struct AgentTokenUsage {
81    /// Agent name.
82    pub agent_name: String,
83    /// Token usage for this agent.
84    pub usage: TokenUsage,
85}
86
87/// Maximum length for prompt_summary field.
88pub const MAX_PROMPT_SUMMARY_LEN: usize = 2000;
89
90/// Maximum length for tool call input_summary field.
91pub const MAX_TOOL_INPUT_SUMMARY_LEN: usize = 200;
92
93/// Truncate a string to a maximum length, appending "..." if truncated.
94pub fn truncate_string(s: &str, max_len: usize) -> String {
95    if s.len() <= max_len {
96        s.to_string()
97    } else {
98        // Find the nearest valid char boundary at or before the target byte index
99        let target = max_len.saturating_sub(3);
100        let boundary = s[..=target.min(s.len().saturating_sub(1))]
101            .char_indices()
102            .map(|(i, _)| i)
103            .last()
104            .unwrap_or(0);
105        let truncated = &s[..boundary];
106        format!("{truncated}...")
107    }
108}
109
110// Approximate pricing (per 1M tokens, USD) — Claude Sonnet 4.5 as reference
111const INPUT_PRICE_PER_M: f64 = 3.0;
112const OUTPUT_PRICE_PER_M: f64 = 15.0;
113const CACHE_READ_PRICE_PER_M: f64 = 0.30;
114const CACHE_WRITE_PRICE_PER_M: f64 = 3.75;
115
116/// Estimate USD cost from token usage using approximate public pricing.
117pub fn estimate_cost(usage: &TokenUsage) -> f64 {
118    let input_cost = usage.input_tokens as f64 / 1_000_000.0 * INPUT_PRICE_PER_M;
119    let output_cost = usage.output_tokens as f64 / 1_000_000.0 * OUTPUT_PRICE_PER_M;
120    let cache_read_cost = usage.cache_read_tokens.unwrap_or(0) as f64 / 1_000_000.0 * CACHE_READ_PRICE_PER_M;
121    let cache_write_cost = usage.cache_write_tokens.unwrap_or(0) as f64 / 1_000_000.0 * CACHE_WRITE_PRICE_PER_M;
122    input_cost + output_cost + cache_read_cost + cache_write_cost
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn truncate_string_works() {
131        assert_eq!(truncate_string("hello", 10), "hello");
132        assert_eq!(truncate_string("hello world", 8), "hello...");
133        assert_eq!(truncate_string("", 5), "");
134        assert_eq!(truncate_string("ab", 2), "ab");
135    }
136
137    #[test]
138    fn token_usage_total() {
139        let usage = TokenUsage {
140            input_tokens: 1000,
141            output_tokens: 500,
142            cache_read_tokens: None,
143            cache_write_tokens: None,
144        };
145        assert_eq!(usage.total(), 1500);
146    }
147
148    #[test]
149    fn token_usage_merge() {
150        let mut a = TokenUsage {
151            input_tokens: 1000,
152            output_tokens: 500,
153            cache_read_tokens: Some(200),
154            cache_write_tokens: None,
155        };
156        let b = TokenUsage {
157            input_tokens: 2000,
158            output_tokens: 300,
159            cache_read_tokens: Some(100),
160            cache_write_tokens: Some(50),
161        };
162        a.merge(&b);
163        assert_eq!(a.input_tokens, 3000);
164        assert_eq!(a.output_tokens, 800);
165        assert_eq!(a.cache_read_tokens, Some(300));
166        assert_eq!(a.cache_write_tokens, Some(50));
167    }
168
169    #[test]
170    fn estimate_cost_basic() {
171        let usage = TokenUsage {
172            input_tokens: 1_000_000,
173            output_tokens: 1_000_000,
174            cache_read_tokens: None,
175            cache_write_tokens: None,
176        };
177        let cost = estimate_cost(&usage);
178        // $3 input + $15 output = $18
179        assert!((cost - 18.0).abs() < 0.01);
180    }
181}
agent_teams/models/token.rs

agent_teams/models/
token.rs
