diff --git a/crates/goose-cli/src/log_usage.rs b/crates/goose-cli/src/log_usage.rs index 7be55cd9c..56435f02d 100644 --- a/crates/goose-cli/src/log_usage.rs +++ b/crates/goose-cli/src/log_usage.rs @@ -1,4 +1,4 @@ -use etcetera::{choose_app_strategy, AppStrategy}; +use etcetera::AppStrategy; use goose::providers::base::ProviderUsage; #[derive(Debug, serde::Serialize, serde::Deserialize)] @@ -7,51 +7,51 @@ struct SessionLog { usage: Vec, } -pub fn log_usage(session_file: String, usage: Vec) { +pub fn log_usage( + home_dir: etcetera::app_strategy::Xdg, + session_file: String, + usage: Vec, +) { let log = SessionLog { session_file, usage, }; - // Ensure log directory exists - if let Ok(home_dir) = choose_app_strategy(crate::APP_STRATEGY.clone()) { - // choose_app_strategy().state_dir() - // - macOS/Linux: ~/.local/state/goose/logs/ - // - Windows: ~\AppData\Roaming\Block\goose\data\logs - // - Windows has no convention for state_dir, use data_dir instead - let log_dir = home_dir - .in_state_dir("logs") - .unwrap_or_else(|| home_dir.in_data_dir("logs")); + // choose_app_strategy().state_dir() + // - macOS/Linux: ~/.local/state/goose/logs/ + // - Windows: ~\AppData\Roaming\Block\goose\data\logs + // - Windows has no convention for state_dir, use data_dir instead + let log_dir = home_dir + .in_state_dir("logs") + .unwrap_or_else(|| home_dir.in_data_dir("logs")); + + if let Err(e) = std::fs::create_dir_all(&log_dir) { + eprintln!("Failed to create log directory: {}", e); + return; + } - if let Err(e) = std::fs::create_dir_all(&log_dir) { - eprintln!("Failed to create log directory: {}", e); + let log_file = log_dir.join("goose.log"); + + let serialized = match serde_json::to_string(&log) { + Ok(s) => s, + Err(e) => { + eprintln!("Failed to serialize usage log: {}", e); return; } + }; - let log_file = log_dir.join("goose.log"); - let serialized = match serde_json::to_string(&log) { - Ok(s) => s, - Err(e) => { - eprintln!("Failed to serialize usage log: {}", e); - return; - } - }; - - // Append to log file - if let Err(e) = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(log_file) - .and_then(|mut file| { - std::io::Write::write_all(&mut file, serialized.as_bytes())?; - std::io::Write::write_all(&mut file, b"\n")?; - Ok(()) - }) - { - eprintln!("Failed to write to usage log file: {}", e); - } - } else { - eprintln!("Failed to write to usage log file: Failed to determine home directory"); + // Append to log file + if let Err(e) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(log_file) + .and_then(|mut file| { + std::io::Write::write_all(&mut file, serialized.as_bytes())?; + std::io::Write::write_all(&mut file, b"\n")?; + Ok(()) + }) + { + eprintln!("Failed to write to usage log file: {}", e); } } @@ -62,46 +62,45 @@ mod tests { use crate::log_usage::{log_usage, SessionLog}; - pub fn run_with_tmp_dir T, T>(func: F) -> T { + #[test] + fn test_session_logging() { use tempfile::tempdir; + // Create a temporary directory let temp_dir = tempdir().unwrap(); - let temp_dir_path = temp_dir.path().to_path_buf(); - - temp_env::with_vars([("HOME", Some(temp_dir_path.as_os_str()))], func) - } + let temp_home = temp_dir.path().to_path_buf(); - #[test] - fn test_session_logging() { - run_with_tmp_dir(|| { - let home_dir = choose_app_strategy(crate::APP_STRATEGY.clone()).unwrap(); - let log_file = home_dir - .in_state_dir("logs") - .unwrap_or_else(|| home_dir.in_data_dir("logs")) - .join("goose.log"); - - log_usage( - "path.txt".to_string(), - vec![ProviderUsage::new( - "model".to_string(), - Usage::new(Some(10), Some(20), Some(30)), - )], - ); - - // Check if log file exists and contains the expected content - assert!(log_file.exists(), "Log file should exist"); - - let log_content = std::fs::read_to_string(&log_file).unwrap(); - let log: SessionLog = serde_json::from_str(&log_content).unwrap(); - - assert!(log.session_file.contains("path.txt")); - assert_eq!(log.usage[0].usage.input_tokens, Some(10)); - assert_eq!(log.usage[0].usage.output_tokens, Some(20)); - assert_eq!(log.usage[0].usage.total_tokens, Some(30)); - assert_eq!(log.usage[0].model, "model"); - - // Remove the log file after test - std::fs::remove_file(&log_file).ok(); - }) + // Temporarily set `HOME` to the temp directory + std::env::set_var("HOME", temp_home.as_os_str()); + let home_dir = choose_app_strategy(crate::APP_STRATEGY.clone()).unwrap(); + let log_file = home_dir + .in_state_dir("logs") + .unwrap_or_else(|| home_dir.in_data_dir("logs")) + .join("goose.log"); + + log_usage( + home_dir, + "path.txt".to_string(), + vec![ProviderUsage::new( + "model".to_string(), + Usage::new(Some(10), Some(20), Some(30)), + )], + ); + + // Check if log file exists and contains the expected content + assert!(log_file.exists(), "Log file should exist"); + + let log_content = std::fs::read_to_string(&log_file).unwrap(); + let log: SessionLog = serde_json::from_str(&log_content).unwrap(); + + assert!(log.session_file.contains("path.txt")); + assert_eq!(log.usage[0].usage.input_tokens, Some(10)); + assert_eq!(log.usage[0].usage.output_tokens, Some(20)); + assert_eq!(log.usage[0].usage.total_tokens, Some(30)); + assert_eq!(log.usage[0].model, "model"); + + // Remove the log file after test + std::fs::remove_file(&log_file).ok(); + std::env::remove_var("HOME"); } } diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 2a65ef058..d359a297a 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -8,6 +8,7 @@ mod thinking; pub use builder::build_session; use anyhow::Result; +use etcetera::choose_app_strategy; use goose::agents::extension::{Envs, ExtensionConfig}; use goose::agents::Agent; use goose::message::{Message, MessageContent}; @@ -168,12 +169,18 @@ impl Session { } // Log usage and cleanup - let usage = self.agent.usage().await; - log_usage(self.session_file.to_string_lossy().to_string(), usage); - println!( - "\nClosing session. Recorded to {}", - self.session_file.display() - ); + if let Ok(home_dir) = choose_app_strategy(crate::APP_STRATEGY.clone()) { + let usage = self.agent.usage().await; + log_usage( + home_dir, + self.session_file.to_string_lossy().to_string(), + usage, + ); + println!( + "\nClosing session. Recorded to {}", + self.session_file.display() + ); + } Ok(()) }