diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 2a3584a6..380eb71c 100644 --- a/crates/goose/src/agents/agent.rs +++ b/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, b: Option| -> Option { + 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; diff --git a/crates/goose/src/session/storage.rs b/crates/goose/src/session/storage.rs index 4387c657..a6ade795 100644 --- a/crates/goose/src/session/storage.rs +++ b/crates/goose/src/session/storage.rs @@ -31,6 +31,12 @@ pub struct SessionMetadata { pub input_tokens: Option, /// The number of output tokens used in the session. Retrieved from the provider's last usage. pub output_tokens: Option, + /// The total number of tokens used in the session. Accumulated across all messages. + pub accumulated_total_tokens: Option, + /// The number of input tokens used in the session. Accumulated across all messages. + pub accumulated_input_tokens: Option, + /// The number of output tokens used in the session. Accumulated across all messages. + pub accumulated_output_tokens: Option, } // Custom deserializer to handle old sessions without working_dir @@ -46,6 +52,9 @@ impl<'de> Deserialize<'de> for SessionMetadata { total_tokens: Option, input_tokens: Option, output_tokens: Option, + accumulated_total_tokens: Option, + accumulated_input_tokens: Option, + accumulated_output_tokens: Option, working_dir: Option, } @@ -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, } } }