generated from nhcarrigan/template
feat: massive overhaul to manage costs (#103)
### Explanation _No response_ ### Issue Closes #102 ### Attestations - [ ] I have read and agree to the [Code of Conduct](https://docs.nhcarrigan.com/community/coc/) - [ ] I have read and agree to the [Community Guidelines](https://docs.nhcarrigan.com/community/guide/). - [ ] My contribution complies with the [Contributor Covenant](https://docs.nhcarrigan.com/dev/covenant/). ### Dependencies - [ ] I have pinned the dependencies to a specific patch version. ### Style - [ ] I have run the linter and resolved any errors. - [ ] My pull request uses an appropriate title, matching the conventional commit standards. - [ ] My scope of feat/fix/chore/etc. correctly matches the nature of changes in my pull request. ### Tests - [ ] My contribution adds new code, and I have added tests to cover it. - [ ] My contribution modifies existing code, and I have updated the tests to reflect these changes. - [ ] All new and existing tests pass locally with my changes. - [ ] Code coverage remains at or above the configured threshold. ### Documentation _No response_ ### Versioning _No response_ Reviewed-on: #103 Co-authored-by: Naomi Carrigan <commits@nhcarrigan.com> Co-committed-by: Naomi Carrigan <commits@nhcarrigan.com>
This commit was merged in pull request #103.
This commit is contained in:
@@ -1935,6 +1935,7 @@ pub fn check_achievements(
|
||||
let search_count: u64 = search_tools
|
||||
.iter()
|
||||
.filter_map(|tool| stats.tools_usage.get(*tool))
|
||||
.map(|t| t.call_count)
|
||||
.sum();
|
||||
if search_count >= 50 && progress.unlock(AchievementId::Explorer) {
|
||||
newly_unlocked.push(AchievementId::Explorer);
|
||||
@@ -1988,25 +1989,25 @@ pub fn check_achievements(
|
||||
// TODO: Track different Claude models used
|
||||
|
||||
// Tool mastery achievements
|
||||
if let Some(bash_count) = stats.tools_usage.get("Bash") {
|
||||
if *bash_count >= 50 && progress.unlock(AchievementId::BashMaster) {
|
||||
if let Some(bash_stats) = stats.tools_usage.get("Bash") {
|
||||
if bash_stats.call_count >= 50 && progress.unlock(AchievementId::BashMaster) {
|
||||
newly_unlocked.push(AchievementId::BashMaster);
|
||||
}
|
||||
}
|
||||
if let Some(read_count) = stats.tools_usage.get("Read") {
|
||||
if *read_count >= 100 && progress.unlock(AchievementId::FileExplorer) {
|
||||
if let Some(read_stats) = stats.tools_usage.get("Read") {
|
||||
if read_stats.call_count >= 100 && progress.unlock(AchievementId::FileExplorer) {
|
||||
newly_unlocked.push(AchievementId::FileExplorer);
|
||||
}
|
||||
}
|
||||
if let Some(grep_count) = stats.tools_usage.get("Grep") {
|
||||
if *grep_count >= 50 && progress.unlock(AchievementId::SearchExpert) {
|
||||
if let Some(grep_stats) = stats.tools_usage.get("Grep") {
|
||||
if grep_stats.call_count >= 50 && progress.unlock(AchievementId::SearchExpert) {
|
||||
newly_unlocked.push(AchievementId::SearchExpert);
|
||||
}
|
||||
}
|
||||
|
||||
// Git Guru - check git command usage in Bash
|
||||
if let Some(bash_count) = stats.tools_usage.get("Bash") {
|
||||
if *bash_count >= 10 && progress.unlock(AchievementId::GitGuru) {
|
||||
if let Some(bash_stats) = stats.tools_usage.get("Bash") {
|
||||
if bash_stats.call_count >= 10 && progress.unlock(AchievementId::GitGuru) {
|
||||
// TODO: More specific git command tracking
|
||||
newly_unlocked.push(AchievementId::GitGuru);
|
||||
}
|
||||
@@ -2055,28 +2056,28 @@ pub fn check_achievements(
|
||||
}
|
||||
|
||||
// More tool mastery achievements
|
||||
if let Some(edit_count) = stats.tools_usage.get("Edit") {
|
||||
if *edit_count >= 100 && progress.unlock(AchievementId::EditMaster) {
|
||||
if let Some(edit_stats) = stats.tools_usage.get("Edit") {
|
||||
if edit_stats.call_count >= 100 && progress.unlock(AchievementId::EditMaster) {
|
||||
newly_unlocked.push(AchievementId::EditMaster);
|
||||
}
|
||||
}
|
||||
if let Some(write_count) = stats.tools_usage.get("Write") {
|
||||
if *write_count >= 50 && progress.unlock(AchievementId::WriteMaster) {
|
||||
if let Some(write_stats) = stats.tools_usage.get("Write") {
|
||||
if write_stats.call_count >= 50 && progress.unlock(AchievementId::WriteMaster) {
|
||||
newly_unlocked.push(AchievementId::WriteMaster);
|
||||
}
|
||||
}
|
||||
if let Some(glob_count) = stats.tools_usage.get("Glob") {
|
||||
if *glob_count >= 100 && progress.unlock(AchievementId::GlobMaster) {
|
||||
if let Some(glob_stats) = stats.tools_usage.get("Glob") {
|
||||
if glob_stats.call_count >= 100 && progress.unlock(AchievementId::GlobMaster) {
|
||||
newly_unlocked.push(AchievementId::GlobMaster);
|
||||
}
|
||||
}
|
||||
if let Some(task_count) = stats.tools_usage.get("Task") {
|
||||
if *task_count >= 50 && progress.unlock(AchievementId::TaskMaster) {
|
||||
if let Some(task_stats) = stats.tools_usage.get("Task") {
|
||||
if task_stats.call_count >= 50 && progress.unlock(AchievementId::TaskMaster) {
|
||||
newly_unlocked.push(AchievementId::TaskMaster);
|
||||
}
|
||||
}
|
||||
if let Some(web_count) = stats.tools_usage.get("WebFetch") {
|
||||
if *web_count >= 20 && progress.unlock(AchievementId::WebFetcher) {
|
||||
if let Some(web_stats) = stats.tools_usage.get("WebFetch") {
|
||||
if web_stats.call_count >= 20 && progress.unlock(AchievementId::WebFetcher) {
|
||||
newly_unlocked.push(AchievementId::WebFetcher);
|
||||
}
|
||||
}
|
||||
@@ -2085,7 +2086,7 @@ pub fn check_achievements(
|
||||
.tools_usage
|
||||
.iter()
|
||||
.filter(|(name, _)| name.starts_with("mcp__"))
|
||||
.map(|(_, count)| count)
|
||||
.map(|(_, tool_stats)| tool_stats.call_count)
|
||||
.sum();
|
||||
if mcp_count >= 50 && progress.unlock(AchievementId::McpExplorer) {
|
||||
newly_unlocked.push(AchievementId::McpExplorer);
|
||||
@@ -2323,6 +2324,11 @@ mod tests {
|
||||
morning_sessions: 0,
|
||||
night_sessions: 0,
|
||||
last_session_date: None,
|
||||
context_tokens_used: 0,
|
||||
context_window_limit: 200_000,
|
||||
context_utilisation_percent: 0.0,
|
||||
potential_cache_hits: 0,
|
||||
potential_cache_savings_tokens: 0,
|
||||
achievements: AchievementProgress::new(),
|
||||
}
|
||||
}
|
||||
@@ -2733,12 +2739,21 @@ mod tests {
|
||||
// check_achievements tests - Tool Usage
|
||||
// =====================
|
||||
|
||||
// Helper function to create a ToolTokenStats with just call count for tests
|
||||
fn tool_stats(call_count: u64) -> crate::stats::ToolTokenStats {
|
||||
crate::stats::ToolTokenStats {
|
||||
call_count,
|
||||
estimated_input_tokens: 0,
|
||||
estimated_output_tokens: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_check_achievements_first_tool() {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Read".to_string(), 1);
|
||||
stats.tools_usage.insert("Read".to_string(), tool_stats(1));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::FirstTool));
|
||||
@@ -2749,11 +2764,11 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Read".to_string(), 1);
|
||||
stats.tools_usage.insert("Write".to_string(), 1);
|
||||
stats.tools_usage.insert("Edit".to_string(), 1);
|
||||
stats.tools_usage.insert("Bash".to_string(), 1);
|
||||
stats.tools_usage.insert("Grep".to_string(), 1);
|
||||
stats.tools_usage.insert("Read".to_string(), tool_stats(1));
|
||||
stats.tools_usage.insert("Write".to_string(), tool_stats(1));
|
||||
stats.tools_usage.insert("Edit".to_string(), tool_stats(1));
|
||||
stats.tools_usage.insert("Bash".to_string(), tool_stats(1));
|
||||
stats.tools_usage.insert("Grep".to_string(), tool_stats(1));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::Toolsmith));
|
||||
@@ -2765,7 +2780,7 @@ mod tests {
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
for i in 0..10 {
|
||||
stats.tools_usage.insert(format!("Tool{}", i), 1);
|
||||
stats.tools_usage.insert(format!("Tool{}", i), tool_stats(1));
|
||||
}
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
@@ -2777,7 +2792,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Bash".to_string(), 50);
|
||||
stats.tools_usage.insert("Bash".to_string(), tool_stats(50));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::BashMaster));
|
||||
@@ -2788,7 +2803,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Read".to_string(), 100);
|
||||
stats.tools_usage.insert("Read".to_string(), tool_stats(100));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::FileExplorer));
|
||||
@@ -2799,7 +2814,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Grep".to_string(), 50);
|
||||
stats.tools_usage.insert("Grep".to_string(), tool_stats(50));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::SearchExpert));
|
||||
@@ -2810,7 +2825,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Edit".to_string(), 100);
|
||||
stats.tools_usage.insert("Edit".to_string(), tool_stats(100));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::EditMaster));
|
||||
@@ -2821,7 +2836,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Write".to_string(), 50);
|
||||
stats.tools_usage.insert("Write".to_string(), tool_stats(50));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::WriteMaster));
|
||||
@@ -2832,7 +2847,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Glob".to_string(), 100);
|
||||
stats.tools_usage.insert("Glob".to_string(), tool_stats(100));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::GlobMaster));
|
||||
@@ -2843,7 +2858,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Task".to_string(), 50);
|
||||
stats.tools_usage.insert("Task".to_string(), tool_stats(50));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::TaskMaster));
|
||||
@@ -2854,7 +2869,7 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("WebFetch".to_string(), 20);
|
||||
stats.tools_usage.insert("WebFetch".to_string(), tool_stats(20));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::WebFetcher));
|
||||
@@ -2865,8 +2880,8 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("mcp__github__create_issue".to_string(), 25);
|
||||
stats.tools_usage.insert("mcp__notion__search".to_string(), 25);
|
||||
stats.tools_usage.insert("mcp__github__create_issue".to_string(), tool_stats(25));
|
||||
stats.tools_usage.insert("mcp__notion__search".to_string(), tool_stats(25));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::McpExplorer));
|
||||
@@ -2881,8 +2896,8 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Grep".to_string(), 30);
|
||||
stats.tools_usage.insert("Glob".to_string(), 20);
|
||||
stats.tools_usage.insert("Grep".to_string(), tool_stats(30));
|
||||
stats.tools_usage.insert("Glob".to_string(), tool_stats(20));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::Explorer));
|
||||
@@ -2893,9 +2908,9 @@ mod tests {
|
||||
let mut stats = create_test_stats();
|
||||
let mut progress = AchievementProgress::new();
|
||||
|
||||
stats.tools_usage.insert("Grep".to_string(), 200);
|
||||
stats.tools_usage.insert("Glob".to_string(), 200);
|
||||
stats.tools_usage.insert("Task".to_string(), 100);
|
||||
stats.tools_usage.insert("Grep".to_string(), tool_stats(200));
|
||||
stats.tools_usage.insert("Glob".to_string(), tool_stats(200));
|
||||
stats.tools_usage.insert("Task".to_string(), tool_stats(100));
|
||||
|
||||
let newly = check_achievements(&stats, &mut progress);
|
||||
assert!(newly.contains(&AchievementId::MasterSearcher));
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tauri::AppHandle;
|
||||
|
||||
use crate::commands::record_session;
|
||||
use crate::config::ClaudeStartOptions;
|
||||
use crate::stats::UsageStats;
|
||||
use crate::wsl_bridge::WslBridge;
|
||||
@@ -53,7 +54,12 @@ impl BridgeManager {
|
||||
.or_insert_with(|| WslBridge::new_with_conversation_id(conversation_id.to_string()));
|
||||
|
||||
// Start the Claude process
|
||||
bridge.start(app, options)?;
|
||||
bridge.start(app.clone(), options)?;
|
||||
|
||||
// Record session start for cost tracking
|
||||
tauri::async_runtime::spawn(async move {
|
||||
record_session(&app).await;
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -556,6 +556,97 @@ pub async fn rename_path(old_path: String, new_path: String) -> Result<(), Strin
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ==================== Cost Tracking Commands ====================
|
||||
|
||||
const COST_HISTORY_STORE_KEY: &str = "cost_history";
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn get_cost_summary(app: AppHandle, days: u32) -> Result<crate::cost_tracking::CostSummary, String> {
|
||||
let history = load_cost_history(&app).await;
|
||||
Ok(history.get_summary(days))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn get_cost_alerts(app: AppHandle) -> Result<Vec<crate::cost_tracking::CostAlert>, String> {
|
||||
let mut history = load_cost_history(&app).await;
|
||||
let alerts = history.check_alerts();
|
||||
|
||||
// Save updated alert state
|
||||
save_cost_history(&app, &history).await?;
|
||||
|
||||
Ok(alerts)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn set_cost_alert_thresholds(
|
||||
app: AppHandle,
|
||||
daily: Option<f64>,
|
||||
weekly: Option<f64>,
|
||||
monthly: Option<f64>,
|
||||
) -> Result<(), String> {
|
||||
let mut history = load_cost_history(&app).await;
|
||||
history.set_alert_thresholds(daily, weekly, monthly);
|
||||
save_cost_history(&app, &history).await
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn export_cost_csv(app: AppHandle, days: u32) -> Result<String, String> {
|
||||
let history = load_cost_history(&app).await;
|
||||
Ok(history.export_csv(days))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn get_today_cost(app: AppHandle) -> Result<f64, String> {
|
||||
let history = load_cost_history(&app).await;
|
||||
Ok(history.get_today_cost())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn get_week_cost(app: AppHandle) -> Result<f64, String> {
|
||||
let history = load_cost_history(&app).await;
|
||||
Ok(history.get_week_cost())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn get_month_cost(app: AppHandle) -> Result<f64, String> {
|
||||
let history = load_cost_history(&app).await;
|
||||
Ok(history.get_month_cost())
|
||||
}
|
||||
|
||||
/// Add cost to history (called internally when stats are updated)
|
||||
pub async fn record_cost(app: &AppHandle, input_tokens: u64, output_tokens: u64, cost_usd: f64) {
|
||||
let mut history = load_cost_history(app).await;
|
||||
history.add_cost(input_tokens, output_tokens, cost_usd);
|
||||
let _ = save_cost_history(app, &history).await;
|
||||
}
|
||||
|
||||
/// Record a new session
|
||||
pub async fn record_session(app: &AppHandle) {
|
||||
let mut history = load_cost_history(app).await;
|
||||
history.increment_sessions();
|
||||
let _ = save_cost_history(app, &history).await;
|
||||
}
|
||||
|
||||
async fn load_cost_history(app: &AppHandle) -> crate::cost_tracking::CostHistory {
|
||||
let store = match app.store("hikari-cost-history.json") {
|
||||
Ok(s) => s,
|
||||
Err(_) => return crate::cost_tracking::CostHistory::new(),
|
||||
};
|
||||
|
||||
match store.get(COST_HISTORY_STORE_KEY) {
|
||||
Some(value) => serde_json::from_value(value.clone()).unwrap_or_default(),
|
||||
None => crate::cost_tracking::CostHistory::new(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn save_cost_history(app: &AppHandle, history: &crate::cost_tracking::CostHistory) -> Result<(), String> {
|
||||
let store = app.store("hikari-cost-history.json").map_err(|e| e.to_string())?;
|
||||
let value = serde_json::to_value(history).map_err(|e| e.to_string())?;
|
||||
store.set(COST_HISTORY_STORE_KEY, value);
|
||||
store.save().map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -96,6 +96,22 @@ pub struct HikariConfig {
|
||||
// Custom theme colors
|
||||
#[serde(default)]
|
||||
pub custom_theme_colors: CustomThemeColors,
|
||||
|
||||
// Token budget settings
|
||||
#[serde(default)]
|
||||
pub budget_enabled: bool,
|
||||
|
||||
#[serde(default)]
|
||||
pub session_token_budget: Option<u64>,
|
||||
|
||||
#[serde(default)]
|
||||
pub session_cost_budget: Option<f64>,
|
||||
|
||||
#[serde(default = "default_budget_action")]
|
||||
pub budget_action: BudgetAction,
|
||||
|
||||
#[serde(default = "default_budget_warning_threshold")]
|
||||
pub budget_warning_threshold: f32,
|
||||
}
|
||||
|
||||
impl Default for HikariConfig {
|
||||
@@ -123,6 +139,11 @@ impl Default for HikariConfig {
|
||||
profile_avatar_path: None,
|
||||
profile_bio: None,
|
||||
custom_theme_colors: CustomThemeColors::default(),
|
||||
budget_enabled: false,
|
||||
session_token_budget: None,
|
||||
session_cost_budget: None,
|
||||
budget_action: BudgetAction::Warn,
|
||||
budget_warning_threshold: 0.8,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -147,6 +168,22 @@ fn default_font_size() -> u32 {
|
||||
14
|
||||
}
|
||||
|
||||
fn default_budget_action() -> BudgetAction {
|
||||
BudgetAction::Warn
|
||||
}
|
||||
|
||||
fn default_budget_warning_threshold() -> f32 {
|
||||
0.8
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BudgetAction {
|
||||
#[default]
|
||||
Warn,
|
||||
Block,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum Theme {
|
||||
@@ -205,6 +242,11 @@ mod tests {
|
||||
assert!(config.profile_avatar_path.is_none());
|
||||
assert!(config.profile_bio.is_none());
|
||||
assert_eq!(config.custom_theme_colors, CustomThemeColors::default());
|
||||
assert!(!config.budget_enabled);
|
||||
assert!(config.session_token_budget.is_none());
|
||||
assert!(config.session_cost_budget.is_none());
|
||||
assert_eq!(config.budget_action, BudgetAction::Warn);
|
||||
assert!((config.budget_warning_threshold - 0.8).abs() < f32::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -232,6 +274,11 @@ mod tests {
|
||||
profile_avatar_path: None,
|
||||
profile_bio: Some("A test bio".to_string()),
|
||||
custom_theme_colors: CustomThemeColors::default(),
|
||||
budget_enabled: true,
|
||||
session_token_budget: Some(100000),
|
||||
session_cost_budget: Some(1.50),
|
||||
budget_action: BudgetAction::Block,
|
||||
budget_warning_threshold: 0.75,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&config).unwrap();
|
||||
|
||||
@@ -0,0 +1,376 @@
|
||||
use chrono::{Datelike, Local, NaiveDate, Weekday};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Represents a single day's cost data
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct DailyCost {
|
||||
pub date: String, // ISO date string (YYYY-MM-DD)
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cost_usd: f64,
|
||||
pub messages_sent: u64,
|
||||
pub sessions_count: u64,
|
||||
}
|
||||
|
||||
/// Historical cost tracking data
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct CostHistory {
|
||||
/// Daily costs indexed by date string (YYYY-MM-DD)
|
||||
pub daily_costs: HashMap<String, DailyCost>,
|
||||
/// Cost alert thresholds
|
||||
pub daily_alert_threshold: Option<f64>,
|
||||
pub weekly_alert_threshold: Option<f64>,
|
||||
pub monthly_alert_threshold: Option<f64>,
|
||||
/// Whether alerts have been triggered today
|
||||
pub daily_alert_triggered: bool,
|
||||
pub weekly_alert_triggered: bool,
|
||||
pub monthly_alert_triggered: bool,
|
||||
pub last_alert_reset_date: Option<String>,
|
||||
}
|
||||
|
||||
impl CostHistory {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Get today's date as a string
|
||||
fn today_str() -> String {
|
||||
Local::now().format("%Y-%m-%d").to_string()
|
||||
}
|
||||
|
||||
/// Get the start of the current week (Monday)
|
||||
fn week_start() -> NaiveDate {
|
||||
let today = Local::now().date_naive();
|
||||
let days_since_monday = today.weekday().num_days_from_monday();
|
||||
today - chrono::Duration::days(days_since_monday as i64)
|
||||
}
|
||||
|
||||
/// Get the start of the current month
|
||||
fn month_start() -> NaiveDate {
|
||||
let today = Local::now().date_naive();
|
||||
NaiveDate::from_ymd_opt(today.year(), today.month(), 1).unwrap_or(today)
|
||||
}
|
||||
|
||||
/// Add cost for today
|
||||
pub fn add_cost(&mut self, input_tokens: u64, output_tokens: u64, cost_usd: f64) {
|
||||
let today = Self::today_str();
|
||||
|
||||
// Reset alert flags if it's a new day
|
||||
if self.last_alert_reset_date.as_ref() != Some(&today) {
|
||||
self.daily_alert_triggered = false;
|
||||
// Reset weekly on Monday
|
||||
if Local::now().weekday() == Weekday::Mon {
|
||||
self.weekly_alert_triggered = false;
|
||||
}
|
||||
// Reset monthly on the 1st
|
||||
if Local::now().day() == 1 {
|
||||
self.monthly_alert_triggered = false;
|
||||
}
|
||||
self.last_alert_reset_date = Some(today.clone());
|
||||
}
|
||||
|
||||
let daily = self.daily_costs.entry(today).or_default();
|
||||
daily.input_tokens += input_tokens;
|
||||
daily.output_tokens += output_tokens;
|
||||
daily.cost_usd += cost_usd;
|
||||
daily.messages_sent += 1;
|
||||
}
|
||||
|
||||
/// Increment session count for today
|
||||
pub fn increment_sessions(&mut self) {
|
||||
let today = Self::today_str();
|
||||
let daily = self.daily_costs.entry(today.clone()).or_insert_with(|| DailyCost {
|
||||
date: today,
|
||||
..Default::default()
|
||||
});
|
||||
daily.sessions_count += 1;
|
||||
}
|
||||
|
||||
/// Get today's cost
|
||||
pub fn get_today_cost(&self) -> f64 {
|
||||
self.daily_costs
|
||||
.get(&Self::today_str())
|
||||
.map(|d| d.cost_usd)
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
|
||||
/// Get this week's cost (Monday to Sunday)
|
||||
pub fn get_week_cost(&self) -> f64 {
|
||||
let week_start = Self::week_start();
|
||||
self.daily_costs
|
||||
.values()
|
||||
.filter(|d| {
|
||||
NaiveDate::parse_from_str(&d.date, "%Y-%m-%d")
|
||||
.map(|date| date >= week_start)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.map(|d| d.cost_usd)
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Get this month's cost
|
||||
pub fn get_month_cost(&self) -> f64 {
|
||||
let month_start = Self::month_start();
|
||||
self.daily_costs
|
||||
.values()
|
||||
.filter(|d| {
|
||||
NaiveDate::parse_from_str(&d.date, "%Y-%m-%d")
|
||||
.map(|date| date >= month_start)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.map(|d| d.cost_usd)
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Get cost summary for a date range
|
||||
pub fn get_summary(&self, days: u32) -> CostSummary {
|
||||
let today = Local::now().date_naive();
|
||||
let start_date = today - chrono::Duration::days(days as i64 - 1);
|
||||
|
||||
let mut total_input_tokens = 0u64;
|
||||
let mut total_output_tokens = 0u64;
|
||||
let mut total_cost = 0.0f64;
|
||||
let mut total_messages = 0u64;
|
||||
let mut total_sessions = 0u64;
|
||||
let mut daily_breakdown = Vec::new();
|
||||
|
||||
for i in 0..days {
|
||||
let date = start_date + chrono::Duration::days(i as i64);
|
||||
let date_str = date.format("%Y-%m-%d").to_string();
|
||||
|
||||
if let Some(daily) = self.daily_costs.get(&date_str) {
|
||||
total_input_tokens += daily.input_tokens;
|
||||
total_output_tokens += daily.output_tokens;
|
||||
total_cost += daily.cost_usd;
|
||||
total_messages += daily.messages_sent;
|
||||
total_sessions += daily.sessions_count;
|
||||
daily_breakdown.push(daily.clone());
|
||||
} else {
|
||||
daily_breakdown.push(DailyCost {
|
||||
date: date_str,
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
CostSummary {
|
||||
period_days: days,
|
||||
total_input_tokens,
|
||||
total_output_tokens,
|
||||
total_cost,
|
||||
total_messages,
|
||||
total_sessions,
|
||||
average_daily_cost: if days > 0 { total_cost / days as f64 } else { 0.0 },
|
||||
daily_breakdown,
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if any alert thresholds are exceeded and return which ones
|
||||
pub fn check_alerts(&mut self) -> Vec<CostAlert> {
|
||||
let mut alerts = Vec::new();
|
||||
|
||||
if let Some(threshold) = self.daily_alert_threshold {
|
||||
let today_cost = self.get_today_cost();
|
||||
if today_cost >= threshold && !self.daily_alert_triggered {
|
||||
self.daily_alert_triggered = true;
|
||||
alerts.push(CostAlert {
|
||||
alert_type: AlertType::Daily,
|
||||
threshold,
|
||||
current_cost: today_cost,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(threshold) = self.weekly_alert_threshold {
|
||||
let week_cost = self.get_week_cost();
|
||||
if week_cost >= threshold && !self.weekly_alert_triggered {
|
||||
self.weekly_alert_triggered = true;
|
||||
alerts.push(CostAlert {
|
||||
alert_type: AlertType::Weekly,
|
||||
threshold,
|
||||
current_cost: week_cost,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(threshold) = self.monthly_alert_threshold {
|
||||
let month_cost = self.get_month_cost();
|
||||
if month_cost >= threshold && !self.monthly_alert_triggered {
|
||||
self.monthly_alert_triggered = true;
|
||||
alerts.push(CostAlert {
|
||||
alert_type: AlertType::Monthly,
|
||||
threshold,
|
||||
current_cost: month_cost,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
alerts
|
||||
}
|
||||
|
||||
/// Set alert thresholds
|
||||
pub fn set_alert_thresholds(
|
||||
&mut self,
|
||||
daily: Option<f64>,
|
||||
weekly: Option<f64>,
|
||||
monthly: Option<f64>,
|
||||
) {
|
||||
self.daily_alert_threshold = daily;
|
||||
self.weekly_alert_threshold = weekly;
|
||||
self.monthly_alert_threshold = monthly;
|
||||
}
|
||||
|
||||
/// Clean up old data (keep last N days)
|
||||
#[allow(dead_code)]
|
||||
pub fn cleanup_old_data(&mut self, keep_days: u32) {
|
||||
let cutoff = Local::now().date_naive() - chrono::Duration::days(keep_days as i64);
|
||||
self.daily_costs.retain(|date_str, _| {
|
||||
NaiveDate::parse_from_str(date_str, "%Y-%m-%d")
|
||||
.map(|date| date >= cutoff)
|
||||
.unwrap_or(false)
|
||||
});
|
||||
}
|
||||
|
||||
/// Export to CSV format
|
||||
pub fn export_csv(&self, days: u32) -> String {
|
||||
let summary = self.get_summary(days);
|
||||
let mut csv = String::from("Date,Input Tokens,Output Tokens,Cost (USD),Messages,Sessions\n");
|
||||
|
||||
for daily in &summary.daily_breakdown {
|
||||
csv.push_str(&format!(
|
||||
"{},{},{},{:.4},{},{}\n",
|
||||
daily.date,
|
||||
daily.input_tokens,
|
||||
daily.output_tokens,
|
||||
daily.cost_usd,
|
||||
daily.messages_sent,
|
||||
daily.sessions_count
|
||||
));
|
||||
}
|
||||
|
||||
// Add totals row
|
||||
csv.push_str(&format!(
|
||||
"TOTAL,{},{},{:.4},{},{}\n",
|
||||
summary.total_input_tokens,
|
||||
summary.total_output_tokens,
|
||||
summary.total_cost,
|
||||
summary.total_messages,
|
||||
summary.total_sessions
|
||||
));
|
||||
|
||||
csv
|
||||
}
|
||||
}
|
||||
|
||||
/// Cost summary for a period
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CostSummary {
|
||||
pub period_days: u32,
|
||||
pub total_input_tokens: u64,
|
||||
pub total_output_tokens: u64,
|
||||
pub total_cost: f64,
|
||||
pub total_messages: u64,
|
||||
pub total_sessions: u64,
|
||||
pub average_daily_cost: f64,
|
||||
pub daily_breakdown: Vec<DailyCost>,
|
||||
}
|
||||
|
||||
/// Alert types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub enum AlertType {
|
||||
Daily,
|
||||
Weekly,
|
||||
Monthly,
|
||||
}
|
||||
|
||||
/// Cost alert notification
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CostAlert {
|
||||
pub alert_type: AlertType,
|
||||
pub threshold: f64,
|
||||
pub current_cost: f64,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_add_cost() {
|
||||
let mut history = CostHistory::new();
|
||||
history.add_cost(1000, 500, 0.05);
|
||||
|
||||
let today_cost = history.get_today_cost();
|
||||
assert!((today_cost - 0.05).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_accumulate_daily_cost() {
|
||||
let mut history = CostHistory::new();
|
||||
history.add_cost(1000, 500, 0.05);
|
||||
history.add_cost(2000, 1000, 0.10);
|
||||
|
||||
let today_cost = history.get_today_cost();
|
||||
assert!((today_cost - 0.15).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_summary() {
|
||||
let mut history = CostHistory::new();
|
||||
history.add_cost(1000, 500, 0.05);
|
||||
|
||||
let summary = history.get_summary(7);
|
||||
assert_eq!(summary.period_days, 7);
|
||||
assert!((summary.total_cost - 0.05).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_daily_alert() {
|
||||
let mut history = CostHistory::new();
|
||||
history.set_alert_thresholds(Some(0.10), None, None);
|
||||
|
||||
history.add_cost(1000, 500, 0.05);
|
||||
let alerts = history.check_alerts();
|
||||
assert!(alerts.is_empty());
|
||||
|
||||
history.add_cost(1000, 500, 0.06);
|
||||
let alerts = history.check_alerts();
|
||||
assert_eq!(alerts.len(), 1);
|
||||
assert_eq!(alerts[0].alert_type, AlertType::Daily);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_alert_only_triggers_once() {
|
||||
let mut history = CostHistory::new();
|
||||
history.set_alert_thresholds(Some(0.10), None, None);
|
||||
|
||||
history.add_cost(1000, 500, 0.15);
|
||||
let alerts = history.check_alerts();
|
||||
assert_eq!(alerts.len(), 1);
|
||||
|
||||
// Second check should not trigger again
|
||||
let alerts = history.check_alerts();
|
||||
assert!(alerts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_export_csv() {
|
||||
let mut history = CostHistory::new();
|
||||
history.add_cost(1000, 500, 0.05);
|
||||
|
||||
let csv = history.export_csv(1);
|
||||
assert!(csv.contains("Date,Input Tokens"));
|
||||
assert!(csv.contains("TOTAL"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_increment_sessions() {
|
||||
let mut history = CostHistory::new();
|
||||
history.increment_sessions();
|
||||
history.increment_sessions();
|
||||
|
||||
let summary = history.get_summary(1);
|
||||
assert_eq!(summary.total_sessions, 2);
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@ mod bridge_manager;
|
||||
mod clipboard;
|
||||
mod commands;
|
||||
mod config;
|
||||
mod cost_tracking;
|
||||
mod git;
|
||||
mod notifications;
|
||||
mod quick_actions;
|
||||
@@ -10,6 +11,7 @@ mod sessions;
|
||||
mod snippets;
|
||||
mod stats;
|
||||
mod temp_manager;
|
||||
mod tool_cache;
|
||||
mod tray;
|
||||
mod types;
|
||||
mod vbs_notification;
|
||||
@@ -159,6 +161,14 @@ pub fn run() {
|
||||
delete_file,
|
||||
delete_directory,
|
||||
rename_path,
|
||||
// Cost tracking commands
|
||||
get_cost_summary,
|
||||
get_cost_alerts,
|
||||
set_cost_alert_thresholds,
|
||||
export_cost_csv,
|
||||
get_today_cost,
|
||||
get_week_cost,
|
||||
get_month_cost,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
|
||||
+657
-24
@@ -5,6 +5,110 @@ use std::collections::HashMap;
|
||||
use std::time::Instant;
|
||||
use tauri_plugin_store::StoreExt;
|
||||
|
||||
/// Per-tool token usage statistics
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct ToolTokenStats {
|
||||
pub call_count: u64,
|
||||
pub estimated_input_tokens: u64,
|
||||
pub estimated_output_tokens: u64,
|
||||
}
|
||||
|
||||
impl ToolTokenStats {
|
||||
#[allow(dead_code)]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn increment_call(&mut self) {
|
||||
self.call_count += 1;
|
||||
}
|
||||
|
||||
pub fn add_tokens(&mut self, input: u64, output: u64) {
|
||||
self.estimated_input_tokens += input;
|
||||
self.estimated_output_tokens += output;
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn total_tokens(&self) -> u64 {
|
||||
self.estimated_input_tokens + self.estimated_output_tokens
|
||||
}
|
||||
}
|
||||
|
||||
/// Warning levels for context window utilisation
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ContextWarning {
|
||||
/// 50-74% utilisation - conversation is getting long
|
||||
Moderate,
|
||||
/// 75-89% utilisation - consider summarising
|
||||
High,
|
||||
/// 90%+ utilisation - approaching limit
|
||||
Critical,
|
||||
}
|
||||
|
||||
/// Budget status indicating whether user is within their limits
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BudgetStatus {
|
||||
/// Within budget, no concerns
|
||||
Ok,
|
||||
/// Approaching budget limit (warning threshold reached)
|
||||
Warning {
|
||||
budget_type: BudgetType,
|
||||
percent_used: f32,
|
||||
},
|
||||
/// Budget exceeded
|
||||
Exceeded { budget_type: BudgetType },
|
||||
}
|
||||
|
||||
/// Type of budget limit
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BudgetType {
|
||||
Token,
|
||||
Cost,
|
||||
}
|
||||
|
||||
impl ContextWarning {
|
||||
#[allow(dead_code)]
|
||||
pub fn message(&self) -> &'static str {
|
||||
match self {
|
||||
ContextWarning::Moderate => "Context window is 50%+ full. Consider starting a new conversation for better performance.",
|
||||
ContextWarning::High => "Context window is 75%+ full. Responses may degrade. Consider summarising or starting fresh.",
|
||||
ContextWarning::Critical => "Context window is nearly full (90%+)! Start a new conversation to avoid errors.",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the context window limit (in tokens) for a given model
|
||||
fn get_context_window_limit(model: &str) -> u64 {
|
||||
match model {
|
||||
// Claude 4.5 family - 200K standard context
|
||||
"claude-opus-4-5-20251101"
|
||||
| "claude-sonnet-4-5-20250929"
|
||||
| "claude-haiku-4-5-20251001" => 200_000,
|
||||
// Claude 4.x family - 200K standard context
|
||||
"claude-opus-4-1-20250805"
|
||||
| "claude-opus-4-20250514"
|
||||
| "claude-sonnet-4-20250514" => 200_000,
|
||||
// Claude 3.x family
|
||||
"claude-3-7-sonnet-20250219"
|
||||
| "claude-3-5-sonnet-20241022"
|
||||
| "claude-3-5-sonnet-20240620"
|
||||
| "claude-3-5-haiku-20241022"
|
||||
| "claude-3-opus-20240229"
|
||||
| "claude-3-sonnet-20240229"
|
||||
| "claude-3-haiku-20240307" => 200_000,
|
||||
// Default to 200K for unknown Claude models
|
||||
_ if model.starts_with("claude") => 200_000,
|
||||
// For non-Claude models (Ollama, etc.), use a conservative default
|
||||
_ => 128_000,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct UsageStats {
|
||||
pub total_input_tokens: u64,
|
||||
@@ -24,8 +128,8 @@ pub struct UsageStats {
|
||||
pub session_files_edited: u64,
|
||||
pub files_created: u64,
|
||||
pub session_files_created: u64,
|
||||
pub tools_usage: HashMap<String, u64>,
|
||||
pub session_tools_usage: HashMap<String, u64>,
|
||||
pub tools_usage: HashMap<String, ToolTokenStats>,
|
||||
pub session_tools_usage: HashMap<String, ToolTokenStats>,
|
||||
pub session_duration_seconds: u64,
|
||||
#[serde(skip)]
|
||||
pub session_start: Option<Instant>,
|
||||
@@ -38,6 +142,15 @@ pub struct UsageStats {
|
||||
pub night_sessions: u64, // Sessions started after 10 PM
|
||||
pub last_session_date: Option<String>, // ISO date string for streak tracking
|
||||
|
||||
// Context window tracking
|
||||
pub context_tokens_used: u64,
|
||||
pub context_window_limit: u64,
|
||||
pub context_utilisation_percent: f32,
|
||||
|
||||
// Cache analytics (tracks potential savings from repeated tool calls)
|
||||
pub potential_cache_hits: u64,
|
||||
pub potential_cache_savings_tokens: u64,
|
||||
|
||||
// Achievement tracking
|
||||
#[serde(skip)]
|
||||
pub achievements: AchievementProgress,
|
||||
@@ -61,6 +174,114 @@ impl UsageStats {
|
||||
self.session_cost_usd += cost;
|
||||
|
||||
self.model = Some(model.to_string());
|
||||
|
||||
// Update context window tracking
|
||||
self.update_context_tracking(model);
|
||||
}
|
||||
|
||||
pub fn update_context_tracking(&mut self, model: &str) {
|
||||
// Get context window limit for the current model
|
||||
self.context_window_limit = get_context_window_limit(model);
|
||||
|
||||
// Context tokens = input tokens (the prompt/context sent to the model)
|
||||
// We track cumulative session input as a proxy for context growth
|
||||
self.context_tokens_used = self.session_input_tokens;
|
||||
|
||||
// Calculate utilisation percentage
|
||||
if self.context_window_limit > 0 {
|
||||
self.context_utilisation_percent =
|
||||
(self.context_tokens_used as f32 / self.context_window_limit as f32) * 100.0;
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn get_context_warning(&self) -> Option<ContextWarning> {
|
||||
if self.context_utilisation_percent >= 90.0 {
|
||||
Some(ContextWarning::Critical)
|
||||
} else if self.context_utilisation_percent >= 75.0 {
|
||||
Some(ContextWarning::High)
|
||||
} else if self.context_utilisation_percent >= 50.0 {
|
||||
Some(ContextWarning::Moderate)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn estimate_remaining_tokens(&self) -> u64 {
|
||||
self.context_window_limit
|
||||
.saturating_sub(self.context_tokens_used)
|
||||
}
|
||||
|
||||
/// Check budget status given current usage and budget settings
|
||||
#[allow(dead_code)]
|
||||
pub fn check_budget(
|
||||
&self,
|
||||
budget_enabled: bool,
|
||||
token_budget: Option<u64>,
|
||||
cost_budget: Option<f64>,
|
||||
warning_threshold: f32,
|
||||
) -> BudgetStatus {
|
||||
if !budget_enabled {
|
||||
return BudgetStatus::Ok;
|
||||
}
|
||||
|
||||
let session_tokens = self.session_input_tokens + self.session_output_tokens;
|
||||
|
||||
// Check token budget
|
||||
if let Some(limit) = token_budget {
|
||||
if session_tokens >= limit {
|
||||
return BudgetStatus::Exceeded {
|
||||
budget_type: BudgetType::Token,
|
||||
};
|
||||
}
|
||||
let percent_used = session_tokens as f32 / limit as f32;
|
||||
if percent_used >= warning_threshold {
|
||||
return BudgetStatus::Warning {
|
||||
budget_type: BudgetType::Token,
|
||||
percent_used: percent_used * 100.0,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Check cost budget
|
||||
if let Some(limit) = cost_budget {
|
||||
if self.session_cost_usd >= limit {
|
||||
return BudgetStatus::Exceeded {
|
||||
budget_type: BudgetType::Cost,
|
||||
};
|
||||
}
|
||||
let percent_used = (self.session_cost_usd / limit) as f32;
|
||||
if percent_used >= warning_threshold {
|
||||
return BudgetStatus::Warning {
|
||||
budget_type: BudgetType::Cost,
|
||||
percent_used: percent_used * 100.0,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
BudgetStatus::Ok
|
||||
}
|
||||
|
||||
/// Get remaining token budget (None if no budget set)
|
||||
#[allow(dead_code)]
|
||||
pub fn get_remaining_token_budget(&self, token_budget: Option<u64>) -> Option<u64> {
|
||||
token_budget.map(|limit| {
|
||||
let used = self.session_input_tokens + self.session_output_tokens;
|
||||
limit.saturating_sub(used)
|
||||
})
|
||||
}
|
||||
|
||||
/// Get remaining cost budget (None if no budget set)
|
||||
#[allow(dead_code)]
|
||||
pub fn get_remaining_cost_budget(&self, cost_budget: Option<f64>) -> Option<f64> {
|
||||
cost_budget.map(|limit| {
|
||||
if limit > self.session_cost_usd {
|
||||
limit - self.session_cost_usd
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub fn reset_session(&mut self) {
|
||||
@@ -76,6 +297,13 @@ impl UsageStats {
|
||||
self.session_start = Some(Instant::now());
|
||||
self.achievements.start_session();
|
||||
|
||||
// Reset context window tracking
|
||||
self.context_tokens_used = 0;
|
||||
self.context_utilisation_percent = 0.0;
|
||||
|
||||
// Note: Cache analytics are NOT reset here - they're cumulative across sessions
|
||||
// to show total potential savings over time
|
||||
|
||||
// Track session start for achievements
|
||||
self.track_session_start();
|
||||
}
|
||||
@@ -139,11 +367,32 @@ impl UsageStats {
|
||||
}
|
||||
|
||||
pub fn increment_tool_usage(&mut self, tool_name: &str) {
|
||||
*self.tools_usage.entry(tool_name.to_string()).or_insert(0) += 1;
|
||||
*self
|
||||
.session_tools_usage
|
||||
self.tools_usage
|
||||
.entry(tool_name.to_string())
|
||||
.or_insert(0) += 1;
|
||||
.or_default()
|
||||
.increment_call();
|
||||
self.session_tools_usage
|
||||
.entry(tool_name.to_string())
|
||||
.or_default()
|
||||
.increment_call();
|
||||
}
|
||||
|
||||
pub fn add_tool_tokens(&mut self, tool_name: &str, input_tokens: u64, output_tokens: u64) {
|
||||
self.tools_usage
|
||||
.entry(tool_name.to_string())
|
||||
.or_default()
|
||||
.add_tokens(input_tokens, output_tokens);
|
||||
self.session_tools_usage
|
||||
.entry(tool_name.to_string())
|
||||
.or_default()
|
||||
.add_tokens(input_tokens, output_tokens);
|
||||
}
|
||||
|
||||
/// Record a potential cache hit (when the same tool call is made twice)
|
||||
#[allow(dead_code)]
|
||||
pub fn add_potential_cache_hit(&mut self, tokens_saved: u64) {
|
||||
self.potential_cache_hits += 1;
|
||||
self.potential_cache_savings_tokens += tokens_saved;
|
||||
}
|
||||
|
||||
pub fn get_session_duration(&mut self) -> u64 {
|
||||
@@ -184,6 +433,11 @@ impl UsageStats {
|
||||
morning_sessions: self.morning_sessions,
|
||||
night_sessions: self.night_sessions,
|
||||
last_session_date: self.last_session_date.clone(),
|
||||
context_tokens_used: self.context_tokens_used,
|
||||
context_window_limit: self.context_window_limit,
|
||||
context_utilisation_percent: self.context_utilisation_percent,
|
||||
potential_cache_hits: self.potential_cache_hits,
|
||||
potential_cache_savings_tokens: self.potential_cache_savings_tokens,
|
||||
achievements: AchievementProgress::new(), // Dummy for copy
|
||||
};
|
||||
check_achievements(&stats_copy, &mut self.achievements)
|
||||
@@ -206,20 +460,22 @@ fn is_consecutive_day(prev_date: &str, current_date: &str) -> bool {
|
||||
}
|
||||
}
|
||||
|
||||
// Pricing as of January 2025
|
||||
// https://www.anthropic.com/pricing
|
||||
fn calculate_cost(input_tokens: u64, output_tokens: u64, model: &str) -> f64 {
|
||||
// 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 {
|
||||
let (input_price_per_million, output_price_per_million) = match model {
|
||||
// Opus 4.5
|
||||
"claude-opus-4-5-20251101" => (15.0, 75.0),
|
||||
// Current generation (Claude 4.5)
|
||||
"claude-opus-4-5-20251101" => (5.0, 25.0),
|
||||
"claude-sonnet-4-5-20250929" => (3.0, 15.0),
|
||||
"claude-haiku-4-5-20251001" => (1.0, 5.0),
|
||||
|
||||
// Opus 4
|
||||
// Previous generation (Claude 4.x)
|
||||
"claude-opus-4-1-20250805" => (15.0, 75.0),
|
||||
"claude-opus-4-20250514" => (15.0, 75.0),
|
||||
|
||||
// Sonnet 4
|
||||
"claude-sonnet-4-20250514" => (3.0, 15.0),
|
||||
|
||||
// Previous generation models
|
||||
// Legacy (Claude 3.x)
|
||||
"claude-3-7-sonnet-20250219" => (3.0, 15.0),
|
||||
"claude-3-5-sonnet-20241022" => (3.0, 15.0),
|
||||
"claude-3-5-sonnet-20240620" => (3.0, 15.0),
|
||||
"claude-3-5-haiku-20241022" => (1.0, 5.0),
|
||||
@@ -252,7 +508,7 @@ pub struct PersistedStats {
|
||||
pub code_blocks_generated: u64,
|
||||
pub files_edited: u64,
|
||||
pub files_created: u64,
|
||||
pub tools_usage: HashMap<String, u64>,
|
||||
pub tools_usage: HashMap<String, ToolTokenStats>,
|
||||
pub sessions_started: u64,
|
||||
pub consecutive_days: u64,
|
||||
pub total_days_used: u64,
|
||||
@@ -372,8 +628,10 @@ mod tests {
|
||||
#[test]
|
||||
fn test_cost_calculation_opus_45() {
|
||||
let cost = calculate_cost(1000, 2000, "claude-opus-4-5-20251101");
|
||||
// Same pricing as Opus 4
|
||||
assert!((cost - 0.165).abs() < 0.0001);
|
||||
// Opus 4.5 pricing: $5/MTok input, $25/MTok output
|
||||
// 1000 input tokens = $0.005, 2000 output tokens = $0.05
|
||||
// Total = $0.055
|
||||
assert!((cost - 0.055).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -512,10 +770,33 @@ mod tests {
|
||||
stats.increment_tool_usage("Read");
|
||||
stats.increment_tool_usage("Write");
|
||||
|
||||
assert_eq!(stats.tools_usage.get("Read"), Some(&2));
|
||||
assert_eq!(stats.tools_usage.get("Write"), Some(&1));
|
||||
assert_eq!(stats.session_tools_usage.get("Read"), Some(&2));
|
||||
assert_eq!(stats.session_tools_usage.get("Write"), Some(&1));
|
||||
assert_eq!(stats.tools_usage.get("Read").map(|t| t.call_count), Some(2));
|
||||
assert_eq!(stats.tools_usage.get("Write").map(|t| t.call_count), Some(1));
|
||||
assert_eq!(stats.session_tools_usage.get("Read").map(|t| t.call_count), Some(2));
|
||||
assert_eq!(stats.session_tools_usage.get("Write").map(|t| t.call_count), Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_tool_tokens() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.increment_tool_usage("Read");
|
||||
stats.add_tool_tokens("Read", 100, 50);
|
||||
stats.add_tool_tokens("Read", 200, 100);
|
||||
|
||||
let read_stats = stats.tools_usage.get("Read").unwrap();
|
||||
assert_eq!(read_stats.call_count, 1);
|
||||
assert_eq!(read_stats.estimated_input_tokens, 300);
|
||||
assert_eq!(read_stats.estimated_output_tokens, 150);
|
||||
assert_eq!(read_stats.total_tokens(), 450);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_token_stats_default() {
|
||||
let tool_stats = ToolTokenStats::new();
|
||||
assert_eq!(tool_stats.call_count, 0);
|
||||
assert_eq!(tool_stats.estimated_input_tokens, 0);
|
||||
assert_eq!(tool_stats.estimated_output_tokens, 0);
|
||||
assert_eq!(tool_stats.total_tokens(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -590,7 +871,11 @@ mod tests {
|
||||
files_created: 5,
|
||||
tools_usage: {
|
||||
let mut map = HashMap::new();
|
||||
map.insert("Read".to_string(), 50);
|
||||
map.insert("Read".to_string(), ToolTokenStats {
|
||||
call_count: 50,
|
||||
estimated_input_tokens: 5000,
|
||||
estimated_output_tokens: 2500,
|
||||
});
|
||||
map
|
||||
},
|
||||
sessions_started: 10,
|
||||
@@ -608,7 +893,8 @@ mod tests {
|
||||
assert_eq!(stats.total_output_tokens, 20000);
|
||||
assert_eq!(stats.total_cost_usd, 5.50);
|
||||
assert_eq!(stats.messages_exchanged, 100);
|
||||
assert_eq!(stats.tools_usage.get("Read"), Some(&50));
|
||||
assert_eq!(stats.tools_usage.get("Read").map(|t| t.call_count), Some(50));
|
||||
assert_eq!(stats.tools_usage.get("Read").map(|t| t.estimated_input_tokens), Some(5000));
|
||||
assert_eq!(stats.consecutive_days, 7);
|
||||
assert_eq!(stats.morning_sessions, 3);
|
||||
assert_eq!(stats.last_session_date, Some("2024-06-15".to_string()));
|
||||
@@ -672,4 +958,351 @@ mod tests {
|
||||
assert!(json.contains("stats"));
|
||||
assert!(json.contains("total_input_tokens"));
|
||||
}
|
||||
|
||||
// =====================
|
||||
// Context Window Tracking tests
|
||||
// =====================
|
||||
|
||||
#[test]
|
||||
fn test_context_window_limit_claude_4() {
|
||||
assert_eq!(get_context_window_limit("claude-opus-4-5-20251101"), 200_000);
|
||||
assert_eq!(get_context_window_limit("claude-opus-4-20250514"), 200_000);
|
||||
assert_eq!(get_context_window_limit("claude-sonnet-4-20250514"), 200_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_window_limit_claude_35() {
|
||||
assert_eq!(
|
||||
get_context_window_limit("claude-3-5-sonnet-20241022"),
|
||||
200_000
|
||||
);
|
||||
assert_eq!(
|
||||
get_context_window_limit("claude-3-5-sonnet-20240620"),
|
||||
200_000
|
||||
);
|
||||
assert_eq!(
|
||||
get_context_window_limit("claude-3-5-haiku-20241022"),
|
||||
200_000
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_window_limit_unknown_claude() {
|
||||
assert_eq!(
|
||||
get_context_window_limit("claude-some-future-model"),
|
||||
200_000
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_window_limit_non_claude() {
|
||||
assert_eq!(get_context_window_limit("gpt-4"), 128_000);
|
||||
assert_eq!(get_context_window_limit("llama-3"), 128_000);
|
||||
assert_eq!(get_context_window_limit("unknown-model"), 128_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_tracking_update() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(50_000, 10_000, "claude-sonnet-4-20250514");
|
||||
|
||||
assert_eq!(stats.context_tokens_used, 50_000);
|
||||
assert_eq!(stats.context_window_limit, 200_000);
|
||||
assert!((stats.context_utilisation_percent - 25.0).abs() < 0.1);
|
||||
}
|
||||
|
||||
#[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");
|
||||
|
||||
assert_eq!(stats.context_tokens_used, 100_000);
|
||||
assert!((stats.context_utilisation_percent - 50.0).abs() < 0.1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_warning_none() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_utilisation_percent = 40.0;
|
||||
assert!(stats.get_context_warning().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_warning_moderate() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_utilisation_percent = 55.0;
|
||||
assert_eq!(stats.get_context_warning(), Some(ContextWarning::Moderate));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_warning_high() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_utilisation_percent = 80.0;
|
||||
assert_eq!(stats.get_context_warning(), Some(ContextWarning::High));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_warning_critical() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_utilisation_percent = 95.0;
|
||||
assert_eq!(stats.get_context_warning(), Some(ContextWarning::Critical));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_remaining_tokens() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_tokens_used = 50_000;
|
||||
stats.context_window_limit = 200_000;
|
||||
|
||||
assert_eq!(stats.estimate_remaining_tokens(), 150_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_remaining_tokens_at_limit() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_tokens_used = 200_000;
|
||||
stats.context_window_limit = 200_000;
|
||||
|
||||
assert_eq!(stats.estimate_remaining_tokens(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_remaining_tokens_over_limit() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.context_tokens_used = 250_000;
|
||||
stats.context_window_limit = 200_000;
|
||||
|
||||
assert_eq!(stats.estimate_remaining_tokens(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_reset_on_session_reset() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.add_usage(100_000, 20_000, "claude-sonnet-4-20250514");
|
||||
|
||||
assert!(stats.context_tokens_used > 0);
|
||||
assert!(stats.context_utilisation_percent > 0.0);
|
||||
|
||||
stats.reset_session();
|
||||
|
||||
assert_eq!(stats.context_tokens_used, 0);
|
||||
assert_eq!(stats.context_utilisation_percent, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_warning_message() {
|
||||
assert_eq!(
|
||||
ContextWarning::Moderate.message(),
|
||||
"Context window is 50%+ full. Consider starting a new conversation for better performance."
|
||||
);
|
||||
assert_eq!(
|
||||
ContextWarning::High.message(),
|
||||
"Context window is 75%+ full. Responses may degrade. Consider summarising or starting fresh."
|
||||
);
|
||||
assert_eq!(
|
||||
ContextWarning::Critical.message(),
|
||||
"Context window is nearly full (90%+)! Start a new conversation to avoid errors."
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_context_warning_serialization() {
|
||||
let warning = ContextWarning::Critical;
|
||||
let json = serde_json::to_string(&warning).expect("Failed to serialize");
|
||||
assert_eq!(json, "\"critical\"");
|
||||
|
||||
let warning = ContextWarning::Moderate;
|
||||
let json = serde_json::to_string(&warning).expect("Failed to serialize");
|
||||
assert_eq!(json, "\"moderate\"");
|
||||
}
|
||||
|
||||
// =====================
|
||||
// Budget Tracking tests
|
||||
// =====================
|
||||
|
||||
#[test]
|
||||
fn test_budget_disabled_returns_ok() {
|
||||
let stats = UsageStats::new();
|
||||
let status = stats.check_budget(false, Some(1000), Some(1.0), 0.8);
|
||||
assert_eq!(status, BudgetStatus::Ok);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_no_limits_returns_ok() {
|
||||
let stats = UsageStats::new();
|
||||
let status = stats.check_budget(true, None, None, 0.8);
|
||||
assert_eq!(status, BudgetStatus::Ok);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_within_limit() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_input_tokens = 500;
|
||||
stats.session_output_tokens = 300;
|
||||
|
||||
let status = stats.check_budget(true, Some(10000), None, 0.8);
|
||||
assert_eq!(status, BudgetStatus::Ok);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_warning() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_input_tokens = 4500;
|
||||
stats.session_output_tokens = 4000;
|
||||
|
||||
let status = stats.check_budget(true, Some(10000), None, 0.8);
|
||||
match status {
|
||||
BudgetStatus::Warning {
|
||||
budget_type,
|
||||
percent_used,
|
||||
} => {
|
||||
assert_eq!(budget_type, BudgetType::Token);
|
||||
assert!(percent_used >= 80.0);
|
||||
}
|
||||
_ => panic!("Expected Warning status"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_exceeded() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_input_tokens = 6000;
|
||||
stats.session_output_tokens = 5000;
|
||||
|
||||
let status = stats.check_budget(true, Some(10000), None, 0.8);
|
||||
assert_eq!(
|
||||
status,
|
||||
BudgetStatus::Exceeded {
|
||||
budget_type: BudgetType::Token
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cost_budget_within_limit() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_cost_usd = 0.50;
|
||||
|
||||
let status = stats.check_budget(true, None, Some(5.0), 0.8);
|
||||
assert_eq!(status, BudgetStatus::Ok);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cost_budget_warning() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_cost_usd = 4.25;
|
||||
|
||||
let status = stats.check_budget(true, None, Some(5.0), 0.8);
|
||||
match status {
|
||||
BudgetStatus::Warning {
|
||||
budget_type,
|
||||
percent_used,
|
||||
} => {
|
||||
assert_eq!(budget_type, BudgetType::Cost);
|
||||
assert!(percent_used >= 80.0);
|
||||
}
|
||||
_ => panic!("Expected Warning status"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cost_budget_exceeded() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_cost_usd = 5.50;
|
||||
|
||||
let status = stats.check_budget(true, None, Some(5.0), 0.8);
|
||||
assert_eq!(
|
||||
status,
|
||||
BudgetStatus::Exceeded {
|
||||
budget_type: BudgetType::Cost
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_takes_priority() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_input_tokens = 12000;
|
||||
stats.session_output_tokens = 0;
|
||||
stats.session_cost_usd = 0.01;
|
||||
|
||||
// Token budget exceeded, cost budget OK
|
||||
let status = stats.check_budget(true, Some(10000), Some(5.0), 0.8);
|
||||
assert_eq!(
|
||||
status,
|
||||
BudgetStatus::Exceeded {
|
||||
budget_type: BudgetType::Token
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remaining_token_budget() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_input_tokens = 3000;
|
||||
stats.session_output_tokens = 2000;
|
||||
|
||||
assert_eq!(stats.get_remaining_token_budget(Some(10000)), Some(5000));
|
||||
assert_eq!(stats.get_remaining_token_budget(None), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remaining_token_budget_exceeded() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_input_tokens = 8000;
|
||||
stats.session_output_tokens = 5000;
|
||||
|
||||
assert_eq!(stats.get_remaining_token_budget(Some(10000)), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remaining_cost_budget() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_cost_usd = 2.50;
|
||||
|
||||
let remaining = stats.get_remaining_cost_budget(Some(5.0));
|
||||
assert!(remaining.is_some());
|
||||
assert!((remaining.unwrap() - 2.50).abs() < 0.001);
|
||||
assert_eq!(stats.get_remaining_cost_budget(None), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remaining_cost_budget_exceeded() {
|
||||
let mut stats = UsageStats::new();
|
||||
stats.session_cost_usd = 6.0;
|
||||
|
||||
let remaining = stats.get_remaining_cost_budget(Some(5.0));
|
||||
assert!(remaining.is_some());
|
||||
assert!((remaining.unwrap() - 0.0).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_status_serialization() {
|
||||
let status = BudgetStatus::Warning {
|
||||
budget_type: BudgetType::Token,
|
||||
percent_used: 85.5,
|
||||
};
|
||||
let json = serde_json::to_string(&status).expect("Failed to serialize");
|
||||
assert!(json.contains("warning"));
|
||||
assert!(json.contains("token"));
|
||||
|
||||
let status = BudgetStatus::Exceeded {
|
||||
budget_type: BudgetType::Cost,
|
||||
};
|
||||
let json = serde_json::to_string(&status).expect("Failed to serialize");
|
||||
assert!(json.contains("exceeded"));
|
||||
assert!(json.contains("cost"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_type_serialization() {
|
||||
let token = BudgetType::Token;
|
||||
let json = serde_json::to_string(&token).expect("Failed to serialize");
|
||||
assert_eq!(json, "\"token\"");
|
||||
|
||||
let cost = BudgetType::Cost;
|
||||
let json = serde_json::to_string(&cost).expect("Failed to serialize");
|
||||
assert_eq!(json, "\"cost\"");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,266 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::collections::HashMap;
|
||||
use std::hash::{Hash, Hasher};
|
||||
|
||||
/// Tools that could benefit from caching
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CacheableTool {
|
||||
Read,
|
||||
Glob,
|
||||
Grep,
|
||||
}
|
||||
|
||||
impl CacheableTool {
|
||||
#[allow(dead_code)]
|
||||
pub fn from_name(name: &str) -> Option<Self> {
|
||||
match name {
|
||||
"Read" => Some(Self::Read),
|
||||
"Glob" => Some(Self::Glob),
|
||||
"Grep" => Some(Self::Grep),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics about potential cache savings
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
|
||||
pub struct CacheAnalytics {
|
||||
/// Number of tool calls that could have been cache hits
|
||||
pub potential_cache_hits: u64,
|
||||
/// Estimated tokens that could have been saved
|
||||
pub potential_savings_tokens: u64,
|
||||
/// Tracks unique tool invocations: hash -> (tool_name, call_count)
|
||||
#[serde(skip)]
|
||||
recent_invocations: HashMap<u64, (String, u64)>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
impl CacheAnalytics {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Compute a hash key from tool name and input
|
||||
fn compute_key(tool_name: &str, input: &serde_json::Value) -> u64 {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
tool_name.hash(&mut hasher);
|
||||
input.to_string().hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Track a tool invocation for analytics
|
||||
/// Returns true if this was a repeated invocation (potential cache hit)
|
||||
pub fn track_invocation(
|
||||
&mut self,
|
||||
tool_name: &str,
|
||||
input: &serde_json::Value,
|
||||
estimated_tokens: u64,
|
||||
) -> bool {
|
||||
// Only track cacheable tools
|
||||
if CacheableTool::from_name(tool_name).is_none() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let key = Self::compute_key(tool_name, input);
|
||||
|
||||
if let Some((_, count)) = self.recent_invocations.get_mut(&key) {
|
||||
*count += 1;
|
||||
// This is a repeat - could have been a cache hit
|
||||
self.potential_cache_hits += 1;
|
||||
self.potential_savings_tokens += estimated_tokens;
|
||||
true
|
||||
} else {
|
||||
self.recent_invocations
|
||||
.insert(key, (tool_name.to_string(), 1));
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the number of unique tool invocations being tracked
|
||||
pub fn unique_invocations(&self) -> usize {
|
||||
self.recent_invocations.len()
|
||||
}
|
||||
|
||||
/// Get invocations that were called more than once
|
||||
pub fn repeated_invocations(&self) -> Vec<(&str, u64)> {
|
||||
self.recent_invocations
|
||||
.values()
|
||||
.filter(|(_, count)| *count > 1)
|
||||
.map(|(name, count)| (name.as_str(), *count))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Clear session analytics (keep totals)
|
||||
pub fn clear_session(&mut self) {
|
||||
self.recent_invocations.clear();
|
||||
}
|
||||
|
||||
/// Fully reset all analytics
|
||||
pub fn reset(&mut self) {
|
||||
self.potential_cache_hits = 0;
|
||||
self.potential_savings_tokens = 0;
|
||||
self.recent_invocations.clear();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_cacheable_tool_from_name() {
|
||||
assert_eq!(CacheableTool::from_name("Read"), Some(CacheableTool::Read));
|
||||
assert_eq!(CacheableTool::from_name("Glob"), Some(CacheableTool::Glob));
|
||||
assert_eq!(CacheableTool::from_name("Grep"), Some(CacheableTool::Grep));
|
||||
assert_eq!(CacheableTool::from_name("Bash"), None);
|
||||
assert_eq!(CacheableTool::from_name("Edit"), None);
|
||||
assert_eq!(CacheableTool::from_name("Write"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_first_invocation_not_cache_hit() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input = json!({"file_path": "/home/test/file.txt"});
|
||||
|
||||
let is_repeat = analytics.track_invocation("Read", &input, 100);
|
||||
|
||||
assert!(!is_repeat);
|
||||
assert_eq!(analytics.potential_cache_hits, 0);
|
||||
assert_eq!(analytics.potential_savings_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_second_invocation_is_cache_hit() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input = json!({"file_path": "/home/test/file.txt"});
|
||||
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
let is_repeat = analytics.track_invocation("Read", &input, 100);
|
||||
|
||||
assert!(is_repeat);
|
||||
assert_eq!(analytics.potential_cache_hits, 1);
|
||||
assert_eq!(analytics.potential_savings_tokens, 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_inputs_not_cache_hit() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input1 = json!({"file_path": "/home/test/file1.txt"});
|
||||
let input2 = json!({"file_path": "/home/test/file2.txt"});
|
||||
|
||||
analytics.track_invocation("Read", &input1, 100);
|
||||
let is_repeat = analytics.track_invocation("Read", &input2, 100);
|
||||
|
||||
assert!(!is_repeat);
|
||||
assert_eq!(analytics.potential_cache_hits, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_non_cacheable_tool_ignored() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input = json!({"command": "ls -la"});
|
||||
|
||||
let is_repeat = analytics.track_invocation("Bash", &input, 100);
|
||||
analytics.track_invocation("Bash", &input, 100);
|
||||
|
||||
assert!(!is_repeat);
|
||||
assert_eq!(analytics.potential_cache_hits, 0);
|
||||
assert_eq!(analytics.unique_invocations(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_repeated_invocations() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input = json!({"file_path": "/home/test/file.txt"});
|
||||
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
|
||||
assert_eq!(analytics.potential_cache_hits, 2);
|
||||
assert_eq!(analytics.potential_savings_tokens, 200);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unique_invocations_count() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
|
||||
analytics.track_invocation("Read", &json!({"file_path": "/file1.txt"}), 100);
|
||||
analytics.track_invocation("Read", &json!({"file_path": "/file2.txt"}), 100);
|
||||
analytics.track_invocation("Glob", &json!({"pattern": "*.rs"}), 50);
|
||||
|
||||
assert_eq!(analytics.unique_invocations(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_repeated_invocations_list() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
|
||||
// file1 read twice
|
||||
analytics.track_invocation("Read", &json!({"file_path": "/file1.txt"}), 100);
|
||||
analytics.track_invocation("Read", &json!({"file_path": "/file1.txt"}), 100);
|
||||
|
||||
// file2 read once
|
||||
analytics.track_invocation("Read", &json!({"file_path": "/file2.txt"}), 100);
|
||||
|
||||
// glob run 3 times
|
||||
analytics.track_invocation("Glob", &json!({"pattern": "*.rs"}), 50);
|
||||
analytics.track_invocation("Glob", &json!({"pattern": "*.rs"}), 50);
|
||||
analytics.track_invocation("Glob", &json!({"pattern": "*.rs"}), 50);
|
||||
|
||||
let repeated = analytics.repeated_invocations();
|
||||
assert_eq!(repeated.len(), 2); // file1 and glob pattern
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear_session() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input = json!({"file_path": "/file.txt"});
|
||||
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
|
||||
assert_eq!(analytics.potential_cache_hits, 1);
|
||||
assert_eq!(analytics.unique_invocations(), 1);
|
||||
|
||||
analytics.clear_session();
|
||||
|
||||
assert_eq!(analytics.potential_cache_hits, 1); // Preserved
|
||||
assert_eq!(analytics.unique_invocations(), 0); // Cleared
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_reset() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
let input = json!({"file_path": "/file.txt"});
|
||||
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
analytics.track_invocation("Read", &input, 100);
|
||||
|
||||
analytics.reset();
|
||||
|
||||
assert_eq!(analytics.potential_cache_hits, 0);
|
||||
assert_eq!(analytics.potential_savings_tokens, 0);
|
||||
assert_eq!(analytics.unique_invocations(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_serialization() {
|
||||
let mut analytics = CacheAnalytics::new();
|
||||
analytics.potential_cache_hits = 10;
|
||||
analytics.potential_savings_tokens = 500;
|
||||
|
||||
let json = serde_json::to_string(&analytics).expect("Failed to serialize");
|
||||
let deserialized: CacheAnalytics =
|
||||
serde_json::from_str(&json).expect("Failed to deserialize");
|
||||
|
||||
assert_eq!(deserialized.potential_cache_hits, 10);
|
||||
assert_eq!(deserialized.potential_savings_tokens, 500);
|
||||
// recent_invocations is skipped in serialization
|
||||
assert_eq!(deserialized.unique_invocations(), 0);
|
||||
}
|
||||
}
|
||||
@@ -176,6 +176,14 @@ pub struct StateChangeEvent {
|
||||
pub conversation_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Cost information for a message
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MessageCost {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cost_usd: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OutputEvent {
|
||||
pub line_type: String,
|
||||
@@ -183,6 +191,8 @@ pub struct OutputEvent {
|
||||
pub tool_name: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub conversation_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cost: Option<MessageCost>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -354,10 +364,31 @@ mod tests {
|
||||
content: "Test output".to_string(),
|
||||
tool_name: None,
|
||||
conversation_id: None,
|
||||
cost: None,
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&event).unwrap();
|
||||
assert!(serialized.contains("\"line_type\":\"assistant\""));
|
||||
assert!(serialized.contains("\"content\":\"Test output\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_output_event_with_cost() {
|
||||
let event = OutputEvent {
|
||||
line_type: "assistant".to_string(),
|
||||
content: "Test output".to_string(),
|
||||
tool_name: None,
|
||||
conversation_id: Some("conv-123".to_string()),
|
||||
cost: Some(MessageCost {
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
cost_usd: 0.005,
|
||||
}),
|
||||
};
|
||||
|
||||
let serialized = serde_json::to_string(&event).unwrap();
|
||||
assert!(serialized.contains("\"cost\":"));
|
||||
assert!(serialized.contains("\"input_tokens\":100"));
|
||||
assert!(serialized.contains("\"output_tokens\":50"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,12 +9,13 @@ use tempfile::NamedTempFile;
|
||||
use std::os::windows::process::CommandExt;
|
||||
|
||||
use crate::achievements::{get_achievement_info, AchievementUnlockedEvent};
|
||||
use crate::commands::record_cost;
|
||||
use crate::config::ClaudeStartOptions;
|
||||
use crate::stats::{StatsUpdateEvent, UsageStats};
|
||||
use crate::stats::{calculate_cost, StatsUpdateEvent, UsageStats};
|
||||
use crate::types::{
|
||||
CharacterState, ClaudeMessage, ConnectionEvent, ConnectionStatus, ContentBlock, OutputEvent,
|
||||
PermissionPromptEvent, QuestionOption, SessionEvent, StateChangeEvent, UserQuestionEvent,
|
||||
WorkingDirectoryEvent,
|
||||
CharacterState, ClaudeMessage, ConnectionEvent, ConnectionStatus, ContentBlock, MessageCost,
|
||||
OutputEvent, PermissionPromptEvent, QuestionOption, SessionEvent, StateChangeEvent,
|
||||
UserQuestionEvent, WorkingDirectoryEvent,
|
||||
};
|
||||
use parking_lot::RwLock;
|
||||
|
||||
@@ -534,6 +535,7 @@ fn handle_stderr(
|
||||
content: line,
|
||||
tool_name: None,
|
||||
conversation_id: conversation_id.clone(),
|
||||
cost: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -586,17 +588,57 @@ fn process_json_line(
|
||||
let mut state = CharacterState::Typing;
|
||||
let mut tool_name = None;
|
||||
|
||||
// Collect all tool names in this message for token attribution
|
||||
let tools_in_message: Vec<String> = message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::ToolUse { name, .. } => Some(name.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Track message cost for display
|
||||
let mut message_cost: Option<MessageCost> = None;
|
||||
|
||||
// Only update stats if we have usage information
|
||||
if let Some(usage) = &message.usage {
|
||||
if let Some(model) = &message.model {
|
||||
// Calculate cost for historical tracking
|
||||
let cost_usd = calculate_cost(usage.input_tokens, usage.output_tokens, model);
|
||||
|
||||
// Store cost for later use in output events
|
||||
message_cost = Some(MessageCost {
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cost_usd,
|
||||
});
|
||||
|
||||
// Batch all stats updates in a single write lock
|
||||
{
|
||||
let mut stats_guard = stats.write();
|
||||
stats_guard.increment_messages();
|
||||
stats_guard.add_usage(usage.input_tokens, usage.output_tokens, model);
|
||||
stats_guard.get_session_duration();
|
||||
|
||||
// Attribute tokens to tools if any tools were used in this message
|
||||
if !tools_in_message.is_empty() {
|
||||
let per_tool_input = usage.input_tokens / tools_in_message.len() as u64;
|
||||
let per_tool_output = usage.output_tokens / tools_in_message.len() as u64;
|
||||
for tool in &tools_in_message {
|
||||
stats_guard.add_tool_tokens(tool, per_tool_input, per_tool_output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Record to historical cost tracking
|
||||
let app_clone = app.clone();
|
||||
let input = usage.input_tokens;
|
||||
let output = usage.output_tokens;
|
||||
tauri::async_runtime::spawn(async move {
|
||||
record_cost(&app_clone, input, output, cost_usd).await;
|
||||
});
|
||||
|
||||
// Don't emit here - we'll emit on Result message instead
|
||||
// This reduces the frequency of updates
|
||||
} else {
|
||||
@@ -635,6 +677,7 @@ fn process_json_line(
|
||||
content: desc,
|
||||
tool_name: Some(name.clone()),
|
||||
conversation_id: conversation_id.clone(),
|
||||
cost: None, // Tool use doesn't have separate cost
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -652,6 +695,7 @@ fn process_json_line(
|
||||
content: text.clone(),
|
||||
tool_name: None,
|
||||
conversation_id: conversation_id.clone(),
|
||||
cost: message_cost.clone(), // Include cost with assistant text
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -664,6 +708,7 @@ fn process_json_line(
|
||||
content: format!("[Thinking] {}", thinking),
|
||||
tool_name: None,
|
||||
conversation_id: conversation_id.clone(),
|
||||
cost: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -723,9 +768,20 @@ fn process_json_line(
|
||||
stats_guard.model.clone().unwrap_or_else(|| "claude-opus-4-20250514".to_string())
|
||||
};
|
||||
|
||||
// Calculate cost for historical tracking
|
||||
let cost_usd = calculate_cost(usage_info.input_tokens, usage_info.output_tokens, &model);
|
||||
|
||||
let mut stats_guard = stats.write();
|
||||
stats_guard.add_usage(usage_info.input_tokens, usage_info.output_tokens, &model);
|
||||
println!("Result message tokens - input: {}, output: {}", usage_info.input_tokens, usage_info.output_tokens);
|
||||
|
||||
// Record to historical cost tracking
|
||||
let app_clone = app.clone();
|
||||
let input = usage_info.input_tokens;
|
||||
let output = usage_info.output_tokens;
|
||||
tauri::async_runtime::spawn(async move {
|
||||
record_cost(&app_clone, input, output, cost_usd).await;
|
||||
});
|
||||
}
|
||||
|
||||
// Always emit updated stats on result message (less frequent)
|
||||
@@ -797,6 +853,7 @@ fn process_json_line(
|
||||
content: text.clone(),
|
||||
tool_name: None,
|
||||
conversation_id: conversation_id.clone(),
|
||||
cost: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user