refactor: clean up log usage

crates/goose-cli/src/commands/bench.rs

@@ -48,17 +48,7 @@ impl BenchAgent for BenchSession {
     }
 
     async fn get_token_usage(&self) -> Option<i32> {
-        // Get token usage from the provider
-        if let Ok(usage) = self.session.get_usage().await {
-            // Sum up total tokens across all usage entries
-            let total_tokens = usage
-                .iter()
-                .map(|u| u.usage.total_tokens.unwrap_or(0))
-                .sum();
-            Some(total_tokens)
-        } else {
-            None
-        }
+        self.session.get_total_token_usage().ok().flatten()
     }
 }
 

crates/goose-cli/src/lib.rs

@@ -1,7 +1,6 @@
 use etcetera::AppStrategyArgs;
 use once_cell::sync::Lazy;
 pub mod commands;
-pub mod log_usage;
 pub mod logging;
 pub mod session;
 

crates/goose-cli/src/session/mod.rs

@@ -16,7 +16,6 @@ use goose::agents::extension::{Envs, ExtensionConfig};
 use goose::agents::{Agent, SessionConfig};
 use goose::config::Config;
 use goose::message::{Message, MessageContent};
-use goose::providers::base::ProviderUsage;
 use goose::session;
 use mcp_core::handler::ToolError;
 use mcp_core::prompt::PromptMessage;
@@ -29,8 +28,6 @@ use std::sync::Arc;
 use std::time::Instant;
 use tokio;
 
-use crate::log_usage::log_usage;
-
 pub struct Session {
     agent: Box<dyn Agent>,
     messages: Vec<Message>,
@@ -413,19 +410,10 @@ impl Session {
             }
         }
 
-        // Log usage and cleanup
-        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()
-            );
-        }
+        println!(
+            "\nClosing session. Recorded to {}",
+            self.session_file.display()
+        );
         Ok(())
     }
 
@@ -645,8 +633,18 @@ impl Session {
         self.messages.clone()
     }
 
-    /// Get the token usage from the agent
-    pub async fn get_usage(&self) -> Result<Vec<ProviderUsage>> {
-        Ok(self.agent.usage().await)
+    /// Get the session metadata
+    pub fn get_metadata(&self) -> Result<session::SessionMetadata> {
+        if !self.session_file.exists() {
+            return Err(anyhow::anyhow!("Session file does not exist"));
+        }
+
+        session::read_metadata(&self.session_file)
+    }
+
+    // Get the session's total token usage
+    pub fn get_total_token_usage(&self) -> Result<Option<i32>> {
+        let metadata = self.get_metadata()?;
+        Ok(metadata.total_tokens)
     }
 }

crates/goose/src/agents/agent.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 
 use super::extension::{ExtensionConfig, ExtensionResult};
 use crate::message::Message;
-use crate::providers::base::{Provider, ProviderUsage};
+use crate::providers::base::Provider;
 use crate::session;
 use mcp_core::prompt::Prompt;
 use mcp_core::protocol::GetPromptResult;
@@ -47,9 +47,6 @@ pub trait Agent: Send + Sync {
     /// Pass through a JSON-RPC request to a specific extension
     async fn passthrough(&self, extension: &str, request: Value) -> ExtensionResult<Value>;
 
-    /// Get the total usage of the agent
-    async fn usage(&self) -> Vec<ProviderUsage>;
-
     /// Add custom text to be included in the system prompt
     async fn extend_system_prompt(&mut self, extension: String);
 

crates/goose/src/agents/capabilities.rs

@@ -13,7 +13,7 @@ use tracing::{debug, instrument};
 use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult};
 use crate::config::Config;
 use crate::prompt_template;
-use crate::providers::base::{Provider, ProviderUsage};
+use crate::providers::base::Provider;
 use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
 use mcp_client::transport::{SseTransport, StdioTransport, Transport};
 use mcp_core::{prompt::Prompt, Content, Tool, ToolCall, ToolError, ToolResult};
@@ -32,7 +32,6 @@ pub struct Capabilities {
     instructions: HashMap<String, String>,
     resource_capable_extensions: HashSet<String>,
     provider: Arc<Box<dyn Provider>>,
-    provider_usage: Mutex<Vec<ProviderUsage>>,
     system_prompt_override: Option<String>,
     system_prompt_extensions: Vec<String>,
 }
@@ -92,7 +91,6 @@ impl Capabilities {
             instructions: HashMap::new(),
             resource_capable_extensions: HashSet::new(),
             provider: Arc::new(provider),
-            provider_usage: Mutex::new(Vec::new()),
             system_prompt_override: None,
             system_prompt_extensions: Vec::new(),
         }
@@ -207,12 +205,6 @@ impl Capabilities {
         Arc::clone(&self.provider)
     }
 
-    /// Record provider usage
-    // TODO consider moving this off to the provider or as a form of logging
-    pub async fn record_usage(&self, usage: ProviderUsage) {
-        self.provider_usage.lock().await.push(usage);
-    }
-
     /// Get aggregated usage statistics
     pub async fn remove_extension(&mut self, name: &str) -> ExtensionResult<()> {
         let sanitized_name = normalize(name.to_string());
@@ -227,29 +219,6 @@ impl Capabilities {
         Ok(self.clients.keys().cloned().collect())
     }
 
-    pub async fn get_usage(&self) -> Vec<ProviderUsage> {
-        let provider_usage = self.provider_usage.lock().await.clone();
-        let mut usage_map: HashMap<String, ProviderUsage> = HashMap::new();
-
-        provider_usage.iter().for_each(|usage| {
-            usage_map
-                .entry(usage.model.clone())
-                .and_modify(|e| {
-                    e.usage.input_tokens = Some(
-                        e.usage.input_tokens.unwrap_or(0) + usage.usage.input_tokens.unwrap_or(0),
-                    );
-                    e.usage.output_tokens = Some(
-                        e.usage.output_tokens.unwrap_or(0) + usage.usage.output_tokens.unwrap_or(0),
-                    );
-                    e.usage.total_tokens = Some(
-                        e.usage.total_tokens.unwrap_or(0) + usage.usage.total_tokens.unwrap_or(0),
-                    );
-                })
-                .or_insert_with(|| usage.clone());
-        });
-        usage_map.into_values().collect()
-    }
-
     /// Get all tools from all clients with proper prefixing
     pub async fn get_prefixed_tools(&mut self) -> ExtensionResult<Vec<Tool>> {
         let mut tools = Vec::new();

crates/goose/src/agents/reference.rs

@@ -13,7 +13,6 @@ use crate::agents::capabilities::Capabilities;
 use crate::agents::extension::{ExtensionConfig, ExtensionResult};
 use crate::message::{Message, ToolRequest};
 use crate::providers::base::Provider;
-use crate::providers::base::ProviderUsage;
 use crate::token_counter::TokenCounter;
 use crate::{register_agent, session};
 use anyhow::{anyhow, Result};
@@ -146,7 +145,6 @@ impl Agent for ReferenceAgent {
                     &messages,
                     &tools,
                 ).await?;
-                capabilities.record_usage(usage.clone()).await;
 
                 // record usage for the session in the session file
                 if let Some(session) = session.clone() {
@@ -204,11 +202,6 @@ impl Agent for ReferenceAgent {
         }))
     }
 
-    async fn usage(&self) -> Vec<ProviderUsage> {
-        let capabilities = self.capabilities.lock().await;
-        capabilities.get_usage().await
-    }
-
     async fn extend_system_prompt(&mut self, extension: String) {
         let mut capabilities = self.capabilities.lock().await;
         capabilities.add_system_prompt_extension(extension);

crates/goose/src/agents/summarize.rs

@@ -18,7 +18,6 @@ use crate::config::Config;
 use crate::memory_condense::condense_messages;
 use crate::message::{Message, ToolRequest};
 use crate::providers::base::Provider;
-use crate::providers::base::ProviderUsage;
 use crate::providers::errors::ProviderError;
 use crate::register_agent;
 use crate::session;
@@ -243,8 +242,6 @@ impl Agent for SummarizeAgent {
                     &tools,
                 ).await {
                     Ok((response, usage)) => {
-                        capabilities.record_usage(usage.clone()).await;
-
                         // record usage for the session in the session file
                         if let Some(session) = session.clone() {
                             // TODO: track session_id in langfuse tracing
@@ -419,11 +416,6 @@ impl Agent for SummarizeAgent {
         }))
     }
 
-    async fn usage(&self) -> Vec<ProviderUsage> {
-        let capabilities = self.capabilities.lock().await;
-        capabilities.get_usage().await
-    }
-
     async fn extend_system_prompt(&mut self, extension: String) {
         let mut capabilities = self.capabilities.lock().await;
         capabilities.add_system_prompt_extension(extension);

crates/goose/src/agents/truncate.rs

@@ -17,7 +17,6 @@ use crate::agents::ToolPermissionStore;
 use crate::config::Config;
 use crate::message::{Message, ToolRequest};
 use crate::providers::base::Provider;
-use crate::providers::base::ProviderUsage;
 use crate::providers::errors::ProviderError;
 use crate::providers::toolshim::{
     augment_message_with_tool_calls, modify_system_prompt_for_tool_json, OllamaInterpreter,
@@ -258,8 +257,6 @@ impl Agent for TruncateAgent {
                             response = augment_message_with_tool_calls(&interpreter, response, &toolshim_tools).await?;
                         }
 
-                        capabilities.record_usage(usage.clone()).await;
-
                         // record usage for the session in the session file
                         if let Some(session) = session.clone() {
                             // TODO: track session_id in langfuse tracing
@@ -473,11 +470,6 @@ impl Agent for TruncateAgent {
         }))
     }
 
-    async fn usage(&self) -> Vec<ProviderUsage> {
-        let capabilities = self.capabilities.lock().await;
-        capabilities.get_usage().await
-    }
-
     async fn extend_system_prompt(&mut self, extension: String) {
         let mut capabilities = self.capabilities.lock().await;
         capabilities.add_system_prompt_extension(extension);