feat: implement async token counter with network resilience and performance optimizations (#3111)

Co-authored-by: jack <>
Co-authored-by: Salman Mohammed <smohammed@squareup.com>

crates/goose/src/context_mgmt/common.rs

@@ -2,7 +2,11 @@ use std::sync::Arc;
 
 use mcp_core::Tool;
 
-use crate::{message::Message, providers::base::Provider, token_counter::TokenCounter};
+use crate::{
+    message::Message,
+    providers::base::Provider,
+    token_counter::{AsyncTokenCounter, TokenCounter},
+};
 
 const ESTIMATE_FACTOR: f32 = 0.7;
 const SYSTEM_PROMPT_TOKEN_OVERHEAD: usize = 3_000;
@@ -28,6 +32,19 @@ pub fn get_messages_token_counts(token_counter: &TokenCounter, messages: &[Messa
         .collect()
 }
 
+/// Async version of get_messages_token_counts for better performance
+pub fn get_messages_token_counts_async(
+    token_counter: &AsyncTokenCounter,
+    messages: &[Message],
+) -> Vec<usize> {
+    // Calculate current token count of each message, use count_chat_tokens to ensure we
+    // capture the full content of the message, include ToolRequests and ToolResponses
+    messages
+        .iter()
+        .map(|msg| token_counter.count_chat_tokens("", std::slice::from_ref(msg), &[]))
+        .collect()
+}
+
 // These are not being used now but could be useful in the future
 
 #[allow(dead_code)]
@@ -55,3 +72,23 @@ pub fn get_token_counts(
         messages: messages_token_count,
     }
 }
+
+/// Async version of get_token_counts for better performance
+#[allow(dead_code)]
+pub fn get_token_counts_async(
+    token_counter: &AsyncTokenCounter,
+    messages: &mut [Message],
+    system_prompt: &str,
+    tools: &mut Vec<Tool>,
+) -> ChatTokenCounts {
+    // Take into account the system prompt (includes goosehints), and our tools input
+    let system_prompt_token_count = token_counter.count_tokens(system_prompt);
+    let tools_token_count = token_counter.count_tokens_for_tools(tools.as_slice());
+    let messages_token_count = get_messages_token_counts_async(token_counter, messages);
+
+    ChatTokenCounts {
+        system: system_prompt_token_count,
+        tools: tools_token_count,
+        messages: messages_token_count,
+    }
+}

crates/goose/src/context_mgmt/summarize.rs

@@ -1,7 +1,7 @@
-use super::common::get_messages_token_counts;
+use super::common::{get_messages_token_counts, get_messages_token_counts_async};
 use crate::message::{Message, MessageContent};
 use crate::providers::base::Provider;
-use crate::token_counter::TokenCounter;
+use crate::token_counter::{AsyncTokenCounter, TokenCounter};
 use anyhow::Result;
 use mcp_core::Role;
 use std::sync::Arc;
@@ -159,6 +159,59 @@ pub async fn summarize_messages(
     ))
 }
 
+/// Async version using AsyncTokenCounter for better performance
+pub async fn summarize_messages_async(
+    provider: Arc<dyn Provider>,
+    messages: &[Message],
+    token_counter: &AsyncTokenCounter,
+    context_limit: usize,
+) -> Result<(Vec<Message>, Vec<usize>), anyhow::Error> {
+    let chunk_size = context_limit / 3; // 33% of the context window.
+    let summary_prompt_tokens = token_counter.count_tokens(SUMMARY_PROMPT);
+    let mut accumulated_summary = Vec::new();
+
+    // Preprocess messages to handle tool response edge case.
+    let (preprocessed_messages, removed_messages) = preprocess_messages(messages);
+
+    // Get token counts for each message.
+    let token_counts = get_messages_token_counts_async(token_counter, &preprocessed_messages);
+
+    // Tokenize and break messages into chunks.
+    let mut current_chunk: Vec<Message> = Vec::new();
+    let mut current_chunk_tokens = 0;
+
+    for (message, message_tokens) in preprocessed_messages.iter().zip(token_counts.iter()) {
+        if current_chunk_tokens + message_tokens > chunk_size - summary_prompt_tokens {
+            // Summarize the current chunk with the accumulated summary.
+            accumulated_summary =
+                summarize_combined_messages(&provider, &accumulated_summary, &current_chunk)
+                    .await?;
+
+            // Reset for the next chunk.
+            current_chunk.clear();
+            current_chunk_tokens = 0;
+        }
+
+        // Add message to the current chunk.
+        current_chunk.push(message.clone());
+        current_chunk_tokens += message_tokens;
+    }
+
+    // Summarize the final chunk if it exists.
+    if !current_chunk.is_empty() {
+        accumulated_summary =
+            summarize_combined_messages(&provider, &accumulated_summary, &current_chunk).await?;
+    }
+
+    // Add back removed messages.
+    let final_summary = reintegrate_removed_messages(&accumulated_summary, &removed_messages);
+
+    Ok((
+        final_summary.clone(),
+        get_messages_token_counts_async(token_counter, &final_summary),
+    ))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;