Token counting reliability + summarization integration.

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

crates/goose/src/agents/agent.rs

@@ -843,7 +843,13 @@ impl Agent {
         &self,
         messages: &[Message],
         session: &Option<SessionConfig>,
-    ) -> Result<Option<(Conversation, String)>> {
+    ) -> Result<
+        Option<(
+            Conversation,
+            String,
+            Option<crate::providers::base::ProviderUsage>,
+        )>,
+    > {
         // Try to get session metadata for more accurate token counts
         let session_metadata = if let Some(session_config) = session {
             match session::storage::get_path(session_config.id.clone()) {
@@ -865,21 +871,23 @@ impl Agent {
         if compact_result.compacted {
             let compacted_messages = compact_result.messages;
 
-            // Create compaction notification message
-            let compaction_msg = if let (Some(before), Some(after)) =
-                (compact_result.tokens_before, compact_result.tokens_after)
-            {
-                format!(
-                    "Auto-compacted context: {} → {} tokens ({:.0}% reduction)\n\n",
-                    before,
-                    after,
-                    (1.0 - (after as f64 / before as f64)) * 100.0
-                )
-            } else {
-                "Auto-compacted context to reduce token usage\n\n".to_string()
-            };
+            // Get threshold from config to include in message
+            let config = crate::config::Config::global();
+            let threshold = config
+                .get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
+                .unwrap_or(0.8); // Default to 80%
+            let threshold_percentage = (threshold * 100.0) as u32;
+
+            let compaction_msg = format!(
+                "Exceeded auto-compact threshold of {}%. Context has been summarized and reduced.\n\n",
+                threshold_percentage
+            );
 
-            return Ok(Some((compacted_messages, compaction_msg)));
+            return Ok(Some((
+                compacted_messages,
+                compaction_msg,
+                compact_result.summarization_usage,
+            )));
         }
 
         Ok(None)
@@ -893,16 +901,16 @@ impl Agent {
         cancel_token: Option<CancellationToken>,
     ) -> Result<BoxStream<'_, Result<AgentEvent>>> {
         // Handle auto-compaction before processing
-        let (messages, compaction_msg) = match self
+        let (messages, compaction_msg, _summarization_usage) = match self
             .handle_auto_compaction(unfixed_conversation.messages(), &session)
             .await?
         {
-            Some((compacted_messages, msg)) => (compacted_messages, Some(msg)),
+            Some((compacted_messages, msg, usage)) => (compacted_messages, Some(msg), usage),
             None => {
                 let context = self
                     .prepare_reply_context(unfixed_conversation, &session)
                     .await?;
-                (context.messages, None)
+                (context.messages, None, None)
             }
         };