feat: add statistics

2026-01-19 17:01:26 -08:00
parent a8f98406e1
commit 1ce43dcff8
12 changed files with 687 additions and 17 deletions
@@ -2,6 +2,7 @@ use tauri::{AppHandle, State};
 use tauri_plugin_store::StoreExt;

 use crate::config::{ClaudeStartOptions, HikariConfig};
+use crate::stats::UsageStats;
 use crate::wsl_bridge::SharedBridge;

 const CONFIG_STORE_KEY: &str = "config";
@@ -72,3 +73,9 @@ pub async fn save_config(app: AppHandle, config: HikariConfig) -> Result<(), Str

    Ok(())
 }
+
+#[tauri::command]
+pub async fn get_usage_stats(bridge: State<'_, SharedBridge>) -> Result<UsageStats, String> {
+    let bridge = bridge.lock();
+    Ok(bridge.get_stats())
+}
@@ -1,6 +1,7 @@
 mod commands;
 mod config;
 mod notifications;
+mod stats;
 mod types;
 mod wsl_bridge;
 mod wsl_notifications;
@@ -35,6 +36,7 @@ pub fn run() {
            select_wsl_directory,
            get_config,
            save_config,
+            get_usage_stats,
            send_windows_notification,
            send_simple_notification,
            send_windows_toast,
@@ -0,0 +1,182 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::time::Instant;
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct UsageStats {
+    pub total_input_tokens: u64,
+    pub total_output_tokens: u64,
+    pub total_cost_usd: f64,
+    pub session_input_tokens: u64,
+    pub session_output_tokens: u64,
+    pub session_cost_usd: f64,
+    pub model: Option<String>,
+
+    // New fields
+    pub messages_exchanged: u64,
+    pub session_messages_exchanged: u64,
+    pub code_blocks_generated: u64,
+    pub session_code_blocks_generated: u64,
+    pub files_edited: u64,
+    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 session_duration_seconds: u64,
+    #[serde(skip)]
+    pub session_start: Option<Instant>,
+}
+
+impl UsageStats {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn add_usage(&mut self, input_tokens: u64, output_tokens: u64, model: &str) {
+        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);
+        self.total_cost_usd += cost;
+        self.session_cost_usd += cost;
+
+        self.model = Some(model.to_string());
+    }
+
+    pub fn reset_session(&mut self) {
+        self.session_input_tokens = 0;
+        self.session_output_tokens = 0;
+        self.session_cost_usd = 0.0;
+        self.session_messages_exchanged = 0;
+        self.session_code_blocks_generated = 0;
+        self.session_files_edited = 0;
+        self.session_files_created = 0;
+        self.session_tools_usage.clear();
+        self.session_duration_seconds = 0;
+        self.session_start = Some(Instant::now());
+    }
+
+    pub fn increment_messages(&mut self) {
+        self.messages_exchanged += 1;
+        self.session_messages_exchanged += 1;
+    }
+
+    pub fn increment_code_blocks(&mut self) {
+        self.code_blocks_generated += 1;
+        self.session_code_blocks_generated += 1;
+    }
+
+    pub fn increment_files_edited(&mut self) {
+        self.files_edited += 1;
+        self.session_files_edited += 1;
+    }
+
+    pub fn increment_files_created(&mut self) {
+        self.files_created += 1;
+        self.session_files_created += 1;
+    }
+
+    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.entry(tool_name.to_string()).or_insert(0) += 1;
+    }
+
+    pub fn get_session_duration(&mut self) -> u64 {
+        // Only update if more than 1 second has passed to reduce calculations
+        if let Some(start) = self.session_start {
+            let elapsed = start.elapsed().as_secs();
+            if elapsed > self.session_duration_seconds {
+                self.session_duration_seconds = elapsed;
+            }
+        }
+        self.session_duration_seconds
+    }
+}
+
+// Pricing as of January 2025
+// https://www.anthropic.com/pricing
+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),
+
+        // Opus 4
+        "claude-opus-4-20250514" => (15.0, 75.0),
+
+        // Sonnet 4
+        "claude-sonnet-4-20250514" => (3.0, 15.0),
+
+        // Previous generation models
+        "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),
+        "claude-3-opus-20240229" => (15.0, 75.0),
+        "claude-3-sonnet-20240229" => (3.0, 15.0),
+        "claude-3-haiku-20240307" => (0.25, 1.25),
+
+        // Default to Sonnet pricing if model unknown
+        _ => (3.0, 15.0),
+    };
+
+    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
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StatsUpdateEvent {
+    pub stats: UsageStats,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_cost_calculation_sonnet() {
+        let cost = calculate_cost(1000, 2000, "claude-sonnet-4-20250514");
+        // 1000 input * $3/M = $0.003
+        // 2000 output * $15/M = $0.030
+        // Total = $0.033
+        assert!((cost - 0.033).abs() < 0.0001);
+    }
+
+    #[test]
+    fn test_cost_calculation_opus() {
+        let cost = calculate_cost(1000, 2000, "claude-opus-4-20250514");
+        // 1000 input * $15/M = $0.015
+        // 2000 output * $75/M = $0.150
+        // Total = $0.165
+        assert!((cost - 0.165).abs() < 0.0001);
+    }
+
+    #[test]
+    fn test_usage_stats_accumulation() {
+        let mut stats = UsageStats::new();
+        stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");
+
+        assert_eq!(stats.total_input_tokens, 1000);
+        assert_eq!(stats.total_output_tokens, 2000);
+        assert_eq!(stats.session_input_tokens, 1000);
+        assert_eq!(stats.session_output_tokens, 2000);
+        assert!((stats.total_cost_usd - 0.033).abs() < 0.0001);
+    }
+
+    #[test]
+    fn test_session_reset() {
+        let mut stats = UsageStats::new();
+        stats.add_usage(1000, 2000, "claude-sonnet-4-20250514");
+        stats.reset_session();
+
+        assert_eq!(stats.total_input_tokens, 1000);
+        assert_eq!(stats.total_output_tokens, 2000);
+        assert_eq!(stats.session_input_tokens, 0);
+        assert_eq!(stats.session_output_tokens, 0);
+        assert_eq!(stats.session_cost_usd, 0.0);
+        assert!(stats.total_cost_usd > 0.0);
+    }
+}
@@ -1,5 +1,11 @@
 use serde::{Deserialize, Serialize};

+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct UsageInfo {
+    pub input_tokens: u64,
+    pub output_tokens: u64,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
 #[serde(rename_all = "snake_case")]
 pub enum CharacterState {
@@ -87,6 +93,8 @@ pub enum ClaudeMessage {
        num_turns: Option<u32>,
        #[serde(default)]
        permission_denials: Option<Vec<PermissionDenial>>,
+        #[serde(default)]
+        usage: Option<UsageInfo>,
    },
 }

@@ -97,6 +105,8 @@ pub struct AssistantMessageContent {
    pub model: Option<String>,
    #[serde(default)]
    pub stop_reason: Option<String>,
+    #[serde(default)]
+    pub usage: Option<UsageInfo>,
 }

 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -10,6 +10,8 @@ use tempfile::NamedTempFile;
 use std::os::windows::process::CommandExt;

 use crate::config::ClaudeStartOptions;
+use crate::stats::{UsageStats, StatsUpdateEvent};
+use parking_lot::RwLock;
 use crate::types::{CharacterState, ClaudeMessage, ConnectionStatus, ContentBlock, StateChangeEvent, OutputEvent, PermissionPromptEvent};

 const SEARCH_TOOLS: [&str; 5] = ["Read", "Glob", "Grep", "WebSearch", "WebFetch"];
@@ -72,6 +74,7 @@ pub struct WslBridge {
    working_directory: String,
    session_id: Option<String>,
    mcp_config_file: Option<NamedTempFile>,
+    stats: Arc<RwLock<UsageStats>>,
 }

 impl WslBridge {
@@ -82,6 +85,7 @@ impl WslBridge {
            working_directory: String::new(),
            session_id: None,
            mcp_config_file: None,
+            stats: Arc::new(RwLock::new(UsageStats::new())),
        }
    }

@@ -249,10 +253,14 @@ impl WslBridge {
        self.stdin = stdin;
        self.process = Some(child);

+        // Reset session stats when starting new session
+        self.stats.write().reset_session();
+
        if let Some(stdout) = stdout {
            let app_clone = app.clone();
+            let stats_clone = self.stats.clone();
            thread::spawn(move || {
-                handle_stdout(stdout, app_clone);
+                handle_stdout(stdout, app_clone, stats_clone);
            });
        }

@@ -311,6 +319,18 @@ impl WslBridge {
    pub fn get_working_directory(&self) -> &str {
        &self.working_directory
    }
+
+    pub fn get_stats(&self) -> UsageStats {
+        self.stats.read().clone()
+    }
+
+    pub fn update_stats(&mut self, input_tokens: u64, output_tokens: u64, model: &str) {
+        self.stats.write().add_usage(input_tokens, output_tokens, model);
+    }
+
+    pub fn reset_session_stats(&mut self) {
+        self.stats.write().reset_session();
+    }
 }

 impl Default for WslBridge {
@@ -319,13 +339,13 @@ impl Default for WslBridge {
    }
 }

-fn handle_stdout(stdout: std::process::ChildStdout, app: AppHandle) {
+fn handle_stdout(stdout: std::process::ChildStdout, app: AppHandle, stats: Arc<RwLock<UsageStats>>) {
    let reader = BufReader::new(stdout);

    for line in reader.lines() {
        match line {
            Ok(line) if !line.is_empty() => {
-                if let Err(e) = process_json_line(&line, &app) {
+                if let Err(e) = process_json_line(&line, &app, &stats) {
                    eprintln!("Error processing line: {}", e);
                }
            }
@@ -358,7 +378,7 @@ fn handle_stderr(stderr: std::process::ChildStderr, app: AppHandle) {
    }
 }

-fn process_json_line(line: &str, app: &AppHandle) -> Result<(), String> {
+fn process_json_line(line: &str, app: &AppHandle, stats: &Arc<RwLock<UsageStats>>) -> Result<(), String> {
    let message: ClaudeMessage = serde_json::from_str(line)
        .map_err(|e| format!("Failed to parse JSON: {} - Line: {}", e, line))?;

@@ -379,12 +399,47 @@ fn process_json_line(line: &str, app: &AppHandle) -> Result<(), String> {
            let mut state = CharacterState::Typing;
            let mut tool_name = None;

+            // Only update stats if we have usage information
+            if let Some(usage) = &message.usage {
+                if let Some(model) = &message.model {
+                    // 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();
+                    }
+
+                    // Don't emit here - we'll emit on Result message instead
+                    // This reduces the frequency of updates
+                } else {
+                    // Just increment message count if no usage info
+                    stats.write().increment_messages();
+                }
+            } else {
+                // Just increment message count if no usage info
+                stats.write().increment_messages();
+            }
+
            for block in &message.content {
                match block {
                    ContentBlock::ToolUse { name, input, .. } => {
                        tool_name = Some(name.clone());
                        state = get_tool_state(name);

+                        // Batch tool tracking updates
+                        {
+                            let mut stats_guard = stats.write();
+                            stats_guard.increment_tool_usage(name);
+
+                            // Track file operations
+                            match name.as_str() {
+                                "Edit" => stats_guard.increment_files_edited(),
+                                "Write" => stats_guard.increment_files_created(),
+                                _ => {}
+                            }
+                        }
+
                        let desc = format_tool_description(name, input);
                        let _ = app.emit("claude:output", OutputEvent {
                            line_type: "tool".to_string(),
@@ -393,6 +448,12 @@ fn process_json_line(line: &str, app: &AppHandle) -> Result<(), String> {
                        });
                    }
                    ContentBlock::Text { text } => {
+                        // Count code blocks in the text
+                        let code_blocks = text.matches("```").count() / 2;
+                        for _ in 0..code_blocks {
+                            stats.write().increment_code_blocks();
+                        }
+
                        let _ = app.emit("claude:output", OutputEvent {
                            line_type: "assistant".to_string(),
                            content: text.clone(),
@@ -440,13 +501,25 @@ fn process_json_line(line: &str, app: &AppHandle) -> Result<(), String> {
            }
        }

-        ClaudeMessage::Result { subtype, result, permission_denials, .. } => {
+        ClaudeMessage::Result { subtype, result, permission_denials, usage: _, .. } => {
            let state = if subtype == "success" {
                CharacterState::Success
            } else {
                CharacterState::Error
            };

+            // Always emit updated stats on result message (less frequent)
+            // This includes the latest session duration
+            {
+                stats.write().get_session_duration();
+            }
+
+            let current_stats = stats.read().clone();
+            let stats_event = StatsUpdateEvent {
+                stats: current_stats,
+            };
+            let _ = app.emit("claude:stats", stats_event);
+
            // Only emit error results - success content is already sent via Assistant message
            if subtype != "success" {
                if let Some(text) = result {
@@ -481,6 +554,8 @@ fn process_json_line(line: &str, app: &AppHandle) -> Result<(), String> {
        }

        ClaudeMessage::User { .. } => {
+            // Increment message count for user messages
+            stats.write().increment_messages();
            emit_state_change(app, CharacterState::Thinking, None);
        }
    }