release/1.4.1

Cargo.toml

@@ -4,7 +4,7 @@ resolver = "2"
 
 [workspace.package]
 edition = "2021"
-version = "1.3.0"
+version = "1.4.1"
 authors = ["Block <[email protected]>"]
 license = "Apache-2.0"
 repository = "https://github.com/block/goose"

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

@@ -620,7 +620,7 @@ async fn process_message_streaming(
 
                                     // For now, auto-summarize in web mode
                                     // TODO: Implement proper UI for context handling
-                                    let (summarized_messages, _) =
+                                    let (summarized_messages, _, _) =
                                         agent.summarize_context(messages.messages()).await?;
                                     {
                                         let mut session_msgs = session_messages.lock().await;

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

@@ -162,7 +162,7 @@ impl Session {
         message_suffix: &str,
     ) -> Result<()> {
         // Summarize messages to fit within context length
-        let (summarized_messages, _) = agent.summarize_context(messages.messages()).await?;
+        let (summarized_messages, _, _) = agent.summarize_context(messages.messages()).await?;
         let msg = format!("Context maxed out\n{}\n{}", "-".repeat(50), message_suffix);
         output::render_text(&msg, Some(Color::Yellow), true);
         *messages = summarized_messages;
@@ -719,15 +719,15 @@ impl Session {
                         let provider = self.agent.provider().await?;
 
                         // Call the summarize_context method which uses the summarize_messages function
-                        let (summarized_messages, _) = self
+                        let (summarized_messages, _token_counts, summarization_usage) = self
                             .agent
                             .summarize_context(self.messages.messages())
                             .await?;
 
                         // Update the session messages with the summarized ones
                         self.messages = summarized_messages;
 
-                        // Persist the summarized messages
+                        // Persist the summarized messages and update session metadata with new token counts
                         if let Some(session_file) = &self.session_file {
                             let working_dir = std::env::current_dir().ok();
                             session::persist_messages_with_schedule_id(
@@ -738,6 +738,46 @@ impl Session {
                                 working_dir,
                             )
                             .await?;
+
+                            // Update session metadata with the new token counts from summarization
+                            if let Some(usage) = summarization_usage {
+                                let session_file_path = session::storage::get_path(
+                                    session::storage::Identifier::Path(session_file.to_path_buf()),
+                                )?;
+                                let mut metadata =
+                                    session::storage::read_metadata(&session_file_path)?;
+
+                                // Update token counts with the summarization usage
+                                // Use output tokens as total since that's what's actually in the context going forward
+                                let summary_tokens = usage.usage.output_tokens.unwrap_or(0);
+                                metadata.total_tokens = Some(summary_tokens);
+                                metadata.input_tokens = None; // Clear input tokens since we now have a summary
+                                metadata.output_tokens = Some(summary_tokens);
+                                metadata.message_count = self.messages.len();
+
+                                // Update accumulated tokens (add the summarization cost)
+                                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,
+                                );
+
+                                session::storage::update_metadata(&session_file_path, &metadata)
+                                    .await?;
+                            }
                         }
 
                         output::hide_thinking();

crates/goose-server/src/routes/context.rs

@@ -67,7 +67,7 @@ async fn manage_context(
             .await
             .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
     } else if request.manage_action == "summarize" {
-        (processed_messages, token_counts) = agent
+        (processed_messages, token_counts, _) = agent
             .summarize_context(&request.messages)
             .await
             .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

crates/goose/Cargo.toml

@@ -122,11 +122,3 @@ path = "examples/agent.rs"
 [[example]]
 name = "databricks_oauth"
 path = "examples/databricks_oauth.rs"
-
-[[example]]
-name = "async_token_counter_demo"
-path = "examples/async_token_counter_demo.rs"
-
-[[bench]]
-name = "tokenization_benchmark"
-harness = false