generated from nhcarrigan/template
fix: ensure permission/stats persist until explicit disconnect (#110)
Also includes cached tokens in cost calculations to provide more accurate billing estimates. Reviewed-on: #110 Co-authored-by: Naomi Carrigan <commits@nhcarrigan.com> Co-committed-by: Naomi Carrigan <commits@nhcarrigan.com>
This commit was merged in pull request #110.
This commit is contained in:
+71
-22
@@ -154,6 +154,16 @@ pub struct UsageStats {
|
||||
// Achievement tracking
|
||||
#[serde(skip)]
|
||||
pub achievements: AchievementProgress,
|
||||
|
||||
// Track current in-flight request for cost estimation on interrupt
|
||||
#[serde(skip)]
|
||||
pub current_request_input: Option<String>,
|
||||
#[serde(skip)]
|
||||
pub current_request_output_chars: u64,
|
||||
#[serde(skip)]
|
||||
pub current_request_thinking_chars: u64,
|
||||
#[serde(skip)]
|
||||
pub current_request_tools: Vec<String>,
|
||||
}
|
||||
|
||||
impl UsageStats {
|
||||
@@ -163,13 +173,26 @@ impl UsageStats {
|
||||
stats
|
||||
}
|
||||
|
||||
pub fn add_usage(&mut self, input_tokens: u64, output_tokens: u64, model: &str) {
|
||||
pub fn add_usage(
|
||||
&mut self,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
model: &str,
|
||||
cache_creation_tokens: Option<u64>,
|
||||
cache_read_tokens: Option<u64>,
|
||||
) {
|
||||
self.total_input_tokens += input_tokens;
|
||||
self.total_output_tokens += output_tokens;
|
||||
self.session_input_tokens += input_tokens;
|
||||
self.session_output_tokens += output_tokens;
|
||||
|
||||
let cost = calculate_cost(input_tokens, output_tokens, model);
|
||||
let cost = calculate_cost(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
model,
|
||||
cache_creation_tokens,
|
||||
cache_read_tokens,
|
||||
);
|
||||
self.total_cost_usd += cost;
|
||||
self.session_cost_usd += cost;
|
||||
|
||||
@@ -439,6 +462,10 @@ impl UsageStats {
|
||||
potential_cache_hits: self.potential_cache_hits,
|
||||
potential_cache_savings_tokens: self.potential_cache_savings_tokens,
|
||||
achievements: AchievementProgress::new(), // Dummy for copy
|
||||
current_request_input: None, // Don't copy tracking fields
|
||||
current_request_output_chars: 0,
|
||||
current_request_thinking_chars: 0,
|
||||
current_request_tools: Vec::new(),
|
||||
};
|
||||
check_achievements(&stats_copy, &mut self.achievements)
|
||||
}
|
||||
@@ -462,7 +489,14 @@ fn is_consecutive_day(prev_date: &str, current_date: &str) -> bool {
|
||||
|
||||
// Pricing as of February 2026
|
||||
// https://platform.claude.com/docs/en/about-claude/models/overview
|
||||
pub fn calculate_cost(input_tokens: u64, output_tokens: u64, model: &str) -> f64 {
|
||||
// Cache pricing: https://platform.claude.com/docs/en/build-with-claude/prompt-caching
|
||||
pub fn calculate_cost(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
model: &str,
|
||||
cache_creation_tokens: Option<u64>,
|
||||
cache_read_tokens: Option<u64>,
|
||||
) -> f64 {
|
||||
let (input_price_per_million, output_price_per_million) = match model {
|
||||
// Current generation (Claude 4.5)
|
||||
"claude-opus-4-5-20251101" => (5.0, 25.0),
|
||||
@@ -487,10 +521,25 @@ pub fn calculate_cost(input_tokens: u64, output_tokens: u64, model: &str) -> f64
|
||||
_ => (3.0, 15.0),
|
||||
};
|
||||
|
||||
// Regular input/output tokens
|
||||
let input_cost = (input_tokens as f64 / 1_000_000.0) * input_price_per_million;
|
||||
let output_cost = (output_tokens as f64 / 1_000_000.0) * output_price_per_million;
|
||||
|
||||
input_cost + output_cost
|
||||
// Cache write tokens (cache creation) cost 1.25x the base input price
|
||||
let cache_write_cost = if let Some(cache_creation) = cache_creation_tokens {
|
||||
(cache_creation as f64 / 1_000_000.0) * input_price_per_million * 1.25
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// Cache read tokens cost 0.1x (10%) the base input price
|
||||
let cache_read_cost = if let Some(cache_read) = cache_read_tokens {
|
||||
(cache_read as f64 / 1_000_000.0) * input_price_per_million * 0.1
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
input_cost + output_cost + cache_write_cost + cache_read_cost
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -609,7 +658,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_cost_calculation_sonnet() {
|
||||
let cost = calculate_cost(1000, 2000, "claude-sonnet-4-20250514");
|
||||
let cost = calculate_cost(1000, 2000, "claude-sonnet-4-20250514", None, None);
|
||||
// 1000 input * $3/M = $0.003
|
||||
// 2000 output * $15/M = $0.030
|
||||
// Total = $0.033
|
||||
@@ -618,7 +667,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_cost_calculation_opus() {
|
||||
let cost = calculate_cost(1000, 2000, "claude-opus-4-20250514");
|
||||
let cost = calculate_cost(1000, 2000, "claude-opus-4-20250514", None, None);
|
||||
// 1000 input * $15/M = $0.015
|
||||
// 2000 output * $75/M = $0.150
|
||||
// Total = $0.165
|
||||
@@ -627,7 +676,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_cost_calculation_opus_45() {
|
||||
let cost = calculate_cost(1000, 2000, "claude-opus-4-5-20251101");
|
||||
let cost = calculate_cost(1000, 2000, "claude-opus-4-5-20251101", None, None);
|
||||
// Opus 4.5 pricing: $5/MTok input, $25/MTok output
|
||||
// 1000 input tokens = $0.005, 2000 output tokens = $0.05
|
||||
// Total = $0.055
|
||||
@@ -636,7 +685,7 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_cost_calculation_haiku() {
|
||||
let cost = calculate_cost(1000, 2000, "claude-3-5-haiku-20241022");
|
||||
let cost = calculate_cost(1000, 2000, "claude-3-5-haiku-20241022", None, None);
|
||||
// 1000 input * $1/M = $0.001
|
||||
// 2000 output * $5/M = $0.010
|
||||
// Total = $0.011
|
||||
@@ -645,14 +694,14 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn test_cost_calculation_unknown_defaults_to_sonnet() {
|
||||
let cost = calculate_cost(1000, 2000, "some-unknown-model");
|
||||
let cost = calculate_cost(1000, 2000, "some-unknown-model", None, None);
|
||||
// Should default to Sonnet pricing
|
||||
assert!((cost - 0.033).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cost_calculation_legacy_sonnet() {
|
||||
let cost = calculate_cost(1000, 2000, "claude-3-5-sonnet-20241022");
|
||||
let cost = calculate_cost(1000, 2000, "claude-3-5-sonnet-20241022", None, None);
|
||||
// Same as Sonnet 4 pricing
|
||||
assert!((cost - 0.033).abs() < 0.0001);
|
||||
}
|
||||
@@ -660,7 +709,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_usage_stats_accumulation() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(1000, 2000, "claude-sonnet-4-20250514", None, None);
|
||||
|
||||
assert_eq!(stats.total_input_tokens, 1000);
|
||||
assert_eq!(stats.total_output_tokens, 2000);
|
||||
@@ -672,8 +721,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_usage_stats_multiple_accumulations() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(1000, 1000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(500, 500, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(1000, 1000, "claude-sonnet-4-20250514", None, None);
|
||||
stats.add_usage(500, 500, "claude-sonnet-4-20250514", None, None);
|
||||
|
||||
assert_eq!(stats.total_input_tokens, 1500);
|
||||
assert_eq!(stats.total_output_tokens, 1500);
|
||||
@@ -684,17 +733,17 @@ mod tests {
|
||||
#[test]
|
||||
fn test_usage_stats_model_updated() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(1000, 1000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(1000, 1000, "claude-sonnet-4-20250514", None, None);
|
||||
assert_eq!(stats.model, Some("claude-sonnet-4-20250514".to_string()));
|
||||
|
||||
stats.add_usage(500, 500, "claude-opus-4-20250514");
|
||||
stats.add_usage(500, 500, "claude-opus-4-20250514", None, None);
|
||||
assert_eq!(stats.model, Some("claude-opus-4-20250514".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_reset() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(1000, 2000, "claude-sonnet-4-20250514", None, None);
|
||||
stats.reset_session();
|
||||
|
||||
assert_eq!(stats.total_input_tokens, 1000);
|
||||
@@ -921,7 +970,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_usage_stats_serialization() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(1000, 2000, "claude-sonnet-4-20250514", None, None);
|
||||
stats.increment_messages();
|
||||
|
||||
// UsageStats should be serializable (for events)
|
||||
@@ -950,7 +999,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_stats_update_event_serialization() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(100, 200, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(100, 200, "claude-sonnet-4-20250514", None, None);
|
||||
|
||||
let event = StatsUpdateEvent { stats };
|
||||
let json = serde_json::to_string(&event).expect("Failed to serialize");
|
||||
@@ -1004,7 +1053,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_context_tracking_update() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514", None, None);
|
||||
|
||||
assert_eq!(stats.context_tokens_used, 50_000);
|
||||
assert_eq!(stats.context_window_limit, 200_000);
|
||||
@@ -1014,8 +1063,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_context_tracking_accumulates() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514", None, None);
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514", None, None);
|
||||
|
||||
assert_eq!(stats.context_tokens_used, 100_000);
|
||||
assert!((stats.context_utilisation_percent - 50.0).abs() < 0.1);
|
||||
@@ -1079,7 +1128,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_context_reset_on_session_reset() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(100_000, 20_000, "claude-sonnet-4-20250514");
|
||||
stats.add_usage(100_000, 20_000, "claude-sonnet-4-20250514", None, None);
|
||||
|
||||
assert!(stats.context_tokens_used > 0);
|
||||
assert!(stats.context_utilisation_percent > 0.0);
|
||||
|
||||
Reference in New Issue
Block a user