Tokens in session file should accumulate for accurate reporting of token usage (#2136)

crates/goose/src/agents/agent.rs

@@ -401,6 +401,18 @@ impl Agent {
                             metadata.total_tokens = usage.usage.total_tokens;
                             metadata.input_tokens = usage.usage.input_tokens;
                             metadata.output_tokens = usage.usage.output_tokens;
+
+                            let accumulate = |a: Option<i32>, b: Option<i32>| -> Option<i32> {
+                                match (a, b) {
+                                    (Some(x), Some(y)) => Some(x + y),
+                                    _ => a.or(b)
+                                }
+                            };
+
+                            metadata.accumulated_total_tokens = accumulate(metadata.accumulated_total_tokens, usage.usage.total_tokens);
+                            metadata.accumulated_input_tokens = accumulate(metadata.accumulated_input_tokens, usage.usage.input_tokens);
+                            metadata.accumulated_output_tokens = accumulate(metadata.accumulated_output_tokens, usage.usage.output_tokens);
+
                             // The message count is the number of messages in the session + 1 for the response
                             // The message count does not include the tool response till next iteration
                             metadata.message_count = messages.len() + 1;

crates/goose/src/session/storage.rs

@@ -31,6 +31,12 @@ pub struct SessionMetadata {
     pub input_tokens: Option<i32>,
     /// The number of output tokens used in the session. Retrieved from the provider's last usage.
     pub output_tokens: Option<i32>,
+    /// The total number of tokens used in the session. Accumulated across all messages.
+    pub accumulated_total_tokens: Option<i32>,
+    /// The number of input tokens used in the session. Accumulated across all messages.
+    pub accumulated_input_tokens: Option<i32>,
+    /// The number of output tokens used in the session. Accumulated across all messages.
+    pub accumulated_output_tokens: Option<i32>,
 }
 
 // Custom deserializer to handle old sessions without working_dir
@@ -46,6 +52,9 @@ impl<'de> Deserialize<'de> for SessionMetadata {
             total_tokens: Option<i32>,
             input_tokens: Option<i32>,
             output_tokens: Option<i32>,
+            accumulated_total_tokens: Option<i32>,
+            accumulated_input_tokens: Option<i32>,
+            accumulated_output_tokens: Option<i32>,
             working_dir: Option<PathBuf>,
         }
 
@@ -57,6 +66,9 @@ impl<'de> Deserialize<'de> for SessionMetadata {
             total_tokens: helper.total_tokens,
             input_tokens: helper.input_tokens,
             output_tokens: helper.output_tokens,
+            accumulated_total_tokens: helper.accumulated_total_tokens,
+            accumulated_input_tokens: helper.accumulated_input_tokens,
+            accumulated_output_tokens: helper.accumulated_output_tokens,
             working_dir: helper.working_dir.unwrap_or_else(get_home_dir),
         })
     }
@@ -71,6 +83,9 @@ impl SessionMetadata {
             total_tokens: None,
             input_tokens: None,
             output_tokens: None,
+            accumulated_total_tokens: None,
+            accumulated_input_tokens: None,
+            accumulated_output_tokens: None,
         }
     }
 }