Add input and output token tracking to SessionMetadata (#1979)

crates/goose/src/agents/reference.rs

@@ -169,6 +169,8 @@ impl Agent for ReferenceAgent {
                     let mut metadata = session::read_metadata(&session_file)?;
                     metadata.working_dir = session.working_dir;
                     metadata.total_tokens = usage.usage.total_tokens;
+                    metadata.input_tokens = usage.usage.input_tokens;
+                    metadata.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/agents/summarize.rs

@@ -266,6 +266,8 @@ impl Agent for SummarizeAgent {
                             let mut metadata = session::read_metadata(&session_file)?;
                             metadata.working_dir = session.working_dir;
                             metadata.total_tokens = usage.usage.total_tokens;
+                            metadata.input_tokens = usage.usage.input_tokens;
+                            metadata.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/agents/truncate.rs

@@ -280,6 +280,8 @@ impl Agent for TruncateAgent {
                             let mut metadata = session::read_metadata(&session_file)?;
                             metadata.working_dir = session.working_dir;
                             metadata.total_tokens = usage.usage.total_tokens;
+                            metadata.input_tokens = usage.usage.input_tokens;
+                            metadata.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

@@ -27,6 +27,10 @@ pub struct SessionMetadata {
     pub message_count: usize,
     /// The total number of tokens used in the session. Retrieved from the provider's last usage.
     pub total_tokens: Option<i32>,
+    /// The number of input tokens used in the session. Retrieved from the provider's last usage.
+    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>,
 }
 
 // Custom deserializer to handle old sessions without working_dir
@@ -40,6 +44,8 @@ impl<'de> Deserialize<'de> for SessionMetadata {
             description: String,
             message_count: usize,
             total_tokens: Option<i32>,
+            input_tokens: Option<i32>,
+            output_tokens: Option<i32>,
             working_dir: Option<PathBuf>,
         }
 
@@ -49,6 +55,8 @@ impl<'de> Deserialize<'de> for SessionMetadata {
             description: helper.description,
             message_count: helper.message_count,
             total_tokens: helper.total_tokens,
+            input_tokens: helper.input_tokens,
+            output_tokens: helper.output_tokens,
             working_dir: helper.working_dir.unwrap_or_else(get_home_dir),
         })
     }
@@ -61,6 +69,8 @@ impl SessionMetadata {
             description: String::new(),
             message_count: 0,
             total_tokens: None,
+            input_tokens: None,
+            output_tokens: None,
         }
     }
 }