usage data types

crates/goose-types/src/lib.rs

@@ -11,9 +11,11 @@ use unicode_normalization::UnicodeNormalization;
 pub mod canonical;
 pub mod conversation;
 pub mod model;
+pub mod usage;
 
 pub use conversation::Conversation;
 pub use model::{Config, ConfigError, ConfigParamError, EnvConfig, ModelConfig, ThinkingEffort};
+pub use usage::{ProviderUsage, Usage};
 
 pub type ToolResult<T> = Result<T, ErrorData>;
 

crates/goose-types/src/usage.rs

@@ -0,0 +1,143 @@
+use serde::{Deserialize, Serialize};
+use std::ops::{Add, AddAssign};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ProviderUsage {
+    pub model: String,
+    pub usage: Usage,
+}
+
+impl ProviderUsage {
+    pub fn new(model: String, usage: Usage) -> Self {
+        Self { model, usage }
+    }
+
+    pub fn combine_with(&self, other: &ProviderUsage) -> ProviderUsage {
+        ProviderUsage {
+            model: self.model.clone(),
+            usage: self.usage + other.usage,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy)]
+pub struct Usage {
+    pub input_tokens: Option<i32>,
+    pub output_tokens: Option<i32>,
+    pub total_tokens: Option<i32>,
+    pub cache_read_input_tokens: Option<i32>,
+    pub cache_write_input_tokens: Option<i32>,
+}
+
+fn sum_optionals<T>(a: Option<T>, b: Option<T>) -> Option<T>
+where
+    T: Add<Output = T> + Default,
+{
+    match (a, b) {
+        (Some(x), Some(y)) => Some(x + y),
+        (Some(x), None) => Some(x + T::default()),
+        (None, Some(y)) => Some(T::default() + y),
+        (None, None) => None,
+    }
+}
+
+impl Add for Usage {
+    type Output = Self;
+
+    fn add(self, other: Self) -> Self {
+        Self::new(
+            sum_optionals(self.input_tokens, other.input_tokens),
+            sum_optionals(self.output_tokens, other.output_tokens),
+            sum_optionals(self.total_tokens, other.total_tokens),
+        )
+        .with_cache_tokens(
+            sum_optionals(self.cache_read_input_tokens, other.cache_read_input_tokens),
+            sum_optionals(
+                self.cache_write_input_tokens,
+                other.cache_write_input_tokens,
+            ),
+        )
+    }
+}
+
+impl AddAssign for Usage {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
+}
+
+impl Usage {
+    pub fn new(
+        input_tokens: Option<i32>,
+        output_tokens: Option<i32>,
+        total_tokens: Option<i32>,
+    ) -> Self {
+        let calculated_total = if total_tokens.is_none() {
+            match (input_tokens, output_tokens) {
+                (Some(input), Some(output)) => Some(input + output),
+                (Some(input), None) => Some(input),
+                (None, Some(output)) => Some(output),
+                (None, None) => None,
+            }
+        } else {
+            total_tokens
+        };
+
+        Self {
+            input_tokens,
+            output_tokens,
+            total_tokens: calculated_total,
+            cache_read_input_tokens: None,
+            cache_write_input_tokens: None,
+        }
+    }
+
+    pub fn with_cache_tokens(
+        mut self,
+        cache_read_input_tokens: Option<i32>,
+        cache_write_input_tokens: Option<i32>,
+    ) -> Self {
+        self.cache_read_input_tokens = cache_read_input_tokens;
+        self.cache_write_input_tokens = cache_write_input_tokens;
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn usage_new_calculates_total_when_missing() {
+        let usage = Usage::new(Some(10), Some(20), None);
+        assert_eq!(usage.total_tokens, Some(30));
+    }
+
+    #[test]
+    fn usage_add_combines_token_counts() {
+        let usage_a =
+            Usage::new(Some(100), Some(20), Some(120)).with_cache_tokens(Some(10), Some(5));
+        let usage_b = Usage::new(Some(50), Some(8), Some(58)).with_cache_tokens(Some(4), Some(1));
+
+        let combined = usage_a + usage_b;
+
+        assert_eq!(combined.input_tokens, Some(150));
+        assert_eq!(combined.output_tokens, Some(28));
+        assert_eq!(combined.total_tokens, Some(178));
+        assert_eq!(combined.cache_read_input_tokens, Some(14));
+        assert_eq!(combined.cache_write_input_tokens, Some(6));
+    }
+
+    #[test]
+    fn provider_usage_combines_with_same_model() {
+        let usage_a = ProviderUsage::new("model-a".to_string(), Usage::new(Some(1), Some(2), None));
+        let usage_b = ProviderUsage::new("model-b".to_string(), Usage::new(Some(3), Some(4), None));
+
+        let combined = usage_a.combine_with(&usage_b);
+
+        assert_eq!(combined.model, "model-a");
+        assert_eq!(combined.usage.input_tokens, Some(4));
+        assert_eq!(combined.usage.output_tokens, Some(6));
+        assert_eq!(combined.usage.total_tokens, Some(10));
+    }
+}

crates/goose/src/context_mgmt/mod.rs

@@ -6,6 +6,7 @@ use crate::prompt_template::render_template;
 use crate::providers::base::{stream_from_single_message, MessageStream};
 use crate::providers::base::{Provider, ProviderUsage};
 use crate::providers::errors::ProviderError;
+use crate::providers::usage_estimator::ensure_usage_tokens;
 use crate::{config::Config, token_counter::create_token_counter};
 use anyhow::Result;
 use indoc::indoc;
@@ -319,10 +320,15 @@ async fn do_compact(
             Ok((mut response, mut provider_usage)) => {
                 response.role = Role::User;
 
-                provider_usage
-                    .ensure_tokens(&system_prompt, &summarization_request, &response, &[])
-                    .await
-                    .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?;
+                ensure_usage_tokens(
+                    &mut provider_usage,
+                    &system_prompt,
+                    &summarization_request,
+                    &response,
+                    &[],
+                )
+                .await
+                .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?;
 
                 return Ok((response, provider_usage));
             }

crates/goose/src/providers/base.rs

@@ -20,12 +20,12 @@ use crate::conversation::Conversation;
 use crate::model::ModelConfig;
 use crate::permission::PermissionConfirmation;
 use crate::utils::safe_truncate;
+pub use goose_types::usage::{ProviderUsage, Usage};
 use rmcp::model::Tool;
 use utoipa::ToSchema;
 
 use once_cell::sync::Lazy;
 use regex::Regex;
-use std::ops::{Add, AddAssign};
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::LazyLock;
@@ -678,129 +678,6 @@ impl ConfigKey {
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ProviderUsage {
-    pub model: String,
-    pub usage: Usage,
-}
-
-impl ProviderUsage {
-    pub fn new(model: String, usage: Usage) -> Self {
-        Self { model, usage }
-    }
-
-    /// Ensures this ProviderUsage has token counts, estimating them if necessary
-    pub async fn ensure_tokens(
-        &mut self,
-        system_prompt: &str,
-        request_messages: &[Message],
-        response: &Message,
-        tools: &[Tool],
-    ) -> Result<(), ProviderError> {
-        crate::providers::usage_estimator::ensure_usage_tokens(
-            self,
-            system_prompt,
-            request_messages,
-            response,
-            tools,
-        )
-        .await
-        .map_err(|e| ProviderError::ExecutionError(format!("Failed to ensure usage tokens: {}", e)))
-    }
-
-    /// Combine this ProviderUsage with another, adding their token counts
-    /// Uses the model from this ProviderUsage
-    pub fn combine_with(&self, other: &ProviderUsage) -> ProviderUsage {
-        ProviderUsage {
-            model: self.model.clone(),
-            usage: self.usage + other.usage,
-        }
-    }
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, Default, Copy)]
-pub struct Usage {
-    pub input_tokens: Option<i32>,
-    pub output_tokens: Option<i32>,
-    pub total_tokens: Option<i32>,
-    pub cache_read_input_tokens: Option<i32>,
-    pub cache_write_input_tokens: Option<i32>,
-}
-
-fn sum_optionals<T>(a: Option<T>, b: Option<T>) -> Option<T>
-where
-    T: Add<Output = T> + Default,
-{
-    match (a, b) {
-        (Some(x), Some(y)) => Some(x + y),
-        (Some(x), None) => Some(x + T::default()),
-        (None, Some(y)) => Some(T::default() + y),
-        (None, None) => None,
-    }
-}
-
-impl Add for Usage {
-    type Output = Self;
-
-    fn add(self, other: Self) -> Self {
-        Self::new(
-            sum_optionals(self.input_tokens, other.input_tokens),
-            sum_optionals(self.output_tokens, other.output_tokens),
-            sum_optionals(self.total_tokens, other.total_tokens),
-        )
-        .with_cache_tokens(
-            sum_optionals(self.cache_read_input_tokens, other.cache_read_input_tokens),
-            sum_optionals(
-                self.cache_write_input_tokens,
-                other.cache_write_input_tokens,
-            ),
-        )
-    }
-}
-
-impl AddAssign for Usage {
-    fn add_assign(&mut self, rhs: Self) {
-        *self = *self + rhs;
-    }
-}
-
-impl Usage {
-    pub fn new(
-        input_tokens: Option<i32>,
-        output_tokens: Option<i32>,
-        total_tokens: Option<i32>,
-    ) -> Self {
-        let calculated_total = if total_tokens.is_none() {
-            match (input_tokens, output_tokens) {
-                (Some(input), Some(output)) => Some(input + output),
-                (Some(input), None) => Some(input),
-                (None, Some(output)) => Some(output),
-                (None, None) => None,
-            }
-        } else {
-            total_tokens
-        };
-
-        Self {
-            input_tokens,
-            output_tokens,
-            total_tokens: calculated_total,
-            cache_read_input_tokens: None,
-            cache_write_input_tokens: None,
-        }
-    }
-
-    pub fn with_cache_tokens(
-        mut self,
-        cache_read_input_tokens: Option<i32>,
-        cache_write_input_tokens: Option<i32>,
-    ) -> Self {
-        self.cache_read_input_tokens = cache_read_input_tokens;
-        self.cache_write_input_tokens = cache_write_input_tokens;
-        self
-    }
-}
-
 pub(crate) fn current_working_dir() -> PathBuf {
     std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
 }