tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 8 Packages Projects Releases Wiki Activity

Compaction overhaul (#5186)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: David Katz <dkatz@squareup.com>
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
This commit is contained in:
Douwe Osinga
2025-10-15 18:39:12 -04:00
committed by GitHub
parent ef9e68f2c6
commit 7c1f5dcb2f
56 changed files with 3069 additions and 4465 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+73 -74
View File
@@ -33,7 +33,7 @@ use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::agents::types::SessionConfig;
use crate::agents::types::{FrontendTool, ToolResultReceiver};
use crate::config::{get_enabled_extensions, get_extension_by_name, Config};
use crate::context_mgmt::auto_compact;
use crate::context_mgmt::{check_and_compact_messages, DEFAULT_COMPACTION_THRESHOLD};
use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation};
use crate::mcp_utils::ToolResult;
use crate::permission::permission_inspector::PermissionInspector;
@@ -112,7 +112,7 @@ pub enum AgentEvent {
Message(Message),
McpNotification((String, ServerNotification)),
ModelChange { model: String, mode: String },
HistoryReplaced(Vec<Message>),
HistoryReplaced(Conversation),
}
impl Default for Agent {
@@ -902,60 +902,6 @@ impl Agent {
}
}
/// Handle auto-compaction logic and return compacted messages if needed
async fn handle_auto_compaction(
&self,
messages: &[Message],
session: &Option<SessionConfig>,
) -> 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 {
SessionManager::get_session(&session_config.id, false)
.await
.ok()
} else {
None
};
let compact_result = auto_compact::check_and_compact_messages(
self,
messages,
None,
session_metadata.as_ref(),
)
.await?;
if compact_result.compacted {
let compacted_messages = compact_result.messages;
// 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,
compact_result.summarization_usage,
)));
}
Ok(None)
}
#[instrument(skip(self, unfixed_conversation, session), fields(user_message))]
pub async fn reply(
&self,
@@ -963,25 +909,66 @@ impl Agent {
session: Option<SessionConfig>,
cancel_token: Option<CancellationToken>,
) -> Result<BoxStream<'_, Result<AgentEvent>>> {
let compaction_result = self
.handle_auto_compaction(unfixed_conversation.messages(), &session)
.await?;
// Try to get session metadata for more accurate token counts
let session_metadata = if let Some(session_config) = &session {
SessionManager::get_session(&session_config.id, false)
.await
.ok()
} else {
None
};
let (did_compact, compacted_conversation, compaction_error) =
match check_and_compact_messages(
self,
unfixed_conversation.messages(),
false,
false,
None,
session_metadata.as_ref(),
)
.await
{
Ok((did_compact, conversation, _removed_indices, _summarization_usage)) => {
(did_compact, conversation, None)
}
Err(e) => (false, unfixed_conversation.clone(), Some(e)),
};
if did_compact {
// 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(DEFAULT_COMPACTION_THRESHOLD);
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
);
if let Some((conversation, compaction_message, _summarization_usage)) = compaction_result {
Ok(Box::pin(async_stream::try_stream! {
// TODO(Douwe): send this before we actually compact:
yield AgentEvent::Message(
Message::assistant().with_summarization_requested(compaction_message)
Message::assistant().with_conversation_compacted(compaction_msg)
);
yield AgentEvent::HistoryReplaced(conversation.messages().clone());
yield AgentEvent::HistoryReplaced(compacted_conversation.clone());
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &conversation).await?
SessionManager::replace_conversation(&session_to_store.id, &compacted_conversation).await?
}
let mut reply_stream = self.reply_internal(conversation, session, cancel_token).await?;
let mut reply_stream = self.reply_internal(compacted_conversation, session, cancel_token).await?;
while let Some(event) = reply_stream.next().await {
yield event?;
}
}))
} else if let Some(error) = compaction_error {
Ok(Box::pin(async_stream::try_stream! {
yield AgentEvent::Message(Message::assistant().with_text(
format!("Ran into this error trying to auto-compact: {error}.\n\nPlease try again or create a new session")
));
}))
} else {
self.reply_internal(unfixed_conversation, session, cancel_token)
.await
@@ -1113,6 +1100,7 @@ impl Agent {
let mut no_tools_called = true;
let mut messages_to_add = Conversation::default();
let mut tools_updated = false;
let mut did_recovery_compact_this_iteration = false;
while let Some(next) = stream.next().await {
if is_token_cancelled(&cancel_token) {
@@ -1306,28 +1294,37 @@ impl Agent {
messages_to_add.push(final_message_tool_resp);
}
}
Err(ProviderError::ContextLengthExceeded(error_msg)) => {
Err(ProviderError::ContextLengthExceeded(_error_msg)) => {
info!("Context length exceeded, attempting compaction");
match auto_compact::perform_compaction(self, conversation.messages()).await {
Ok(compact_result) => {
conversation = compact_result.messages;
// Get session metadata if available
let session_metadata_for_compact = if let Some(ref session_config) = session {
SessionManager::get_session(&session_config.id, false).await.ok()
} else {
None
};
match check_and_compact_messages(self, conversation.messages(), true, true, None, session_metadata_for_compact.as_ref()).await {
Ok((_did_compact, compacted_conversation, _removed_indices, _usage)) => {
conversation = compacted_conversation;
did_recovery_compact_this_iteration = true;
yield AgentEvent::Message(
Message::assistant().with_summarization_requested(
Message::assistant().with_conversation_compacted(
"Context limit reached. Conversation has been automatically compacted to continue."
)
);
yield AgentEvent::HistoryReplaced(conversation.messages().to_vec());
yield AgentEvent::HistoryReplaced(conversation.clone());
if let Some(session_to_store) = &session {
SessionManager::replace_conversation(&session_to_store.id, &conversation).await?
}
continue;
}
Err(_) => {
yield AgentEvent::Message(Message::assistant().with_context_length_exceeded(
format!("Context length exceeded and cannot summarize: {}. Unable to continue.", error_msg)
));
Err(e) => {
error!("Error: {}", e);
yield AgentEvent::Message(Message::assistant().with_text(
format!("Ran into this error trying to compact: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")
));
break;
}
}
@@ -1358,6 +1355,8 @@ impl Agent {
yield AgentEvent::Message(message);
exit_chat = true;
}
} else if did_recovery_compact_this_iteration {
// Avoid setting exit_chat; continue from last user message in the conversation
} else {
match self.handle_retry_logic(&mut conversation, &session, &initial_messages).await {
Ok(should_retry) => {
crates/goose/src/agents/context.rs
-130
View File
@@ -1,130 +0,0 @@
use anyhow::Ok;
use crate::conversation::message::{Message, MessageMetadata};
use crate::conversation::Conversation;
use crate::token_counter::create_async_token_counter;
use crate::context_mgmt::summarize::summarize_messages;
use crate::context_mgmt::truncate::{truncate_messages, OldestFirstTruncation};
use crate::context_mgmt::{estimate_target_context_limit, get_messages_token_counts_async};
use super::super::agents::Agent;
impl Agent {
/// Public API to truncate oldest messages so that the conversation's token count is within the allowed context limit.
pub async fn truncate_context(
&self,
messages: &[Message], // last message is a user msg that led to assistant message with_context_length_exceeded
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
let provider = self.provider().await?;
let token_counter = create_async_token_counter()
.await
.map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?;
let target_context_limit = estimate_target_context_limit(provider);
let token_counts = get_messages_token_counts_async(&token_counter, messages);
let (mut new_messages, mut new_token_counts) = truncate_messages(
messages,
&token_counts,
target_context_limit,
&OldestFirstTruncation,
)?;
// Only add an assistant message if we have room for it and it won't cause another overflow
let assistant_message = Message::assistant().with_text("I had run into a context length exceeded error so I truncated some of the oldest messages in our conversation.");
let assistant_tokens =
token_counter.count_chat_tokens("", std::slice::from_ref(&assistant_message), &[]);
let current_total: usize = new_token_counts.iter().sum();
if current_total + assistant_tokens <= target_context_limit {
new_messages.push(assistant_message);
new_token_counts.push(assistant_tokens);
} else {
// If we can't fit the assistant message, at least log what happened
tracing::warn!("Cannot add truncation notice message due to context limits. Current: {}, Assistant: {}, Limit: {}",
current_total, assistant_tokens, target_context_limit);
}
Ok((new_messages, new_token_counts))
}
/// Public API to summarize the conversation so that its token count is within the allowed context limit.
/// Returns the summarized messages, token counts, and the ProviderUsage from summarization
pub async fn summarize_context(
&self,
messages: &[Message], // last message is a user msg that led to assistant message with_context_length_exceeded
) -> Result<
(
Conversation,
Vec<usize>,
Option<crate::providers::base::ProviderUsage>,
),
anyhow::Error,
> {
let provider = self.provider().await?;
let summary_result = summarize_messages(provider.clone(), messages).await?;
let (summary_message, summarization_usage) = match summary_result {
Some((summary_message, provider_usage)) => (summary_message, Some(provider_usage)),
None => {
// No summary was generated (empty input)
tracing::warn!("Summarization failed. Returning empty messages.");
return Ok((Conversation::empty(), vec![], None));
}
};
// Create the final message list with updated visibility metadata:
// 1. Original messages become user_visible but not agent_visible
// 2. Summary message becomes agent_visible but not user_visible
// 3. Assistant messages to continue the conversation remain both user_visible and agent_visible
let mut final_messages = Vec::new();
let mut final_token_counts = Vec::new();
// Add all original messages with updated visibility (preserve user_visible, set agent_visible=false)
for msg in messages.iter().cloned() {
let updated_metadata = msg.metadata.with_agent_invisible();
let updated_msg = msg.with_metadata(updated_metadata);
final_messages.push(updated_msg);
// Token count doesn't matter for agent_visible=false messages, but we'll use 0
final_token_counts.push(0);
}
// Add the compaction marker (user_visible=true, agent_visible=false)
let compaction_marker = Message::assistant()
.with_summarization_requested("Conversation compacted and summarized")
.with_metadata(MessageMetadata::user_only());
let compaction_marker_tokens: usize = 0; // Not counted since agent_visible=false
final_messages.push(compaction_marker);
final_token_counts.push(compaction_marker_tokens);
// Add the summary message (agent_visible=true, user_visible=false)
let summary_msg = summary_message.with_metadata(MessageMetadata::agent_only());
// For token counting purposes, we use the output tokens (the actual summary content)
// since that's what will be in the context going forward
let summary_tokens = summarization_usage
.as_ref()
.and_then(|usage| usage.usage.output_tokens)
.unwrap_or(0) as usize;
final_messages.push(summary_msg);
final_token_counts.push(summary_tokens);
// Add an assistant message to continue the conversation (agent_visible=true, user_visible=false)
let assistant_message = Message::assistant()
.with_text(
"The previous message contains a summary that was prepared because a context limit was reached.
Do not mention that you read a summary or that conversation summarization occurred
Just continue the conversation naturally based on the summarized context"
)
.with_metadata(MessageMetadata::agent_only());
let assistant_message_tokens: usize = 0; // Not counted since it's for agent context only
final_messages.push(assistant_message);
final_token_counts.push(assistant_message_tokens);
Ok((
Conversation::new_unvalidated(final_messages),
final_token_counts,
summarization_usage,
))
}
}
crates/goose/src/agents/mod.rs
-1
View File
@@ -1,5 +1,4 @@
mod agent;
mod context;
pub mod extension;
pub mod extension_malware_check;
pub mod extension_manager;
crates/goose/src/agents/reply_parts.rs
+23 -16
View File
@@ -119,39 +119,46 @@ impl Agent {
let toolshim_tools = toolshim_tools.to_owned();
let provider = provider.clone();
let mut stream = if provider.supports_streaming() {
// Capture errors during stream creation and return them as part of the stream
// so they can be handled by the existing error handling logic in the agent
let stream_result = if provider.supports_streaming() {
debug!("WAITING_LLM_STREAM_START");
let msg_stream = provider
let result = provider
.stream(
system_prompt.as_str(),
messages_for_provider.messages(),
&tools,
)
.await?;
.await;
debug!("WAITING_LLM_STREAM_END");
msg_stream
result
} else {
debug!("WAITING_LLM_START");
let (message, mut usage) = provider
let complete_result = provider
.complete(
system_prompt.as_str(),
messages_for_provider.messages(),
&tools,
)
.await?;
.await;
debug!("WAITING_LLM_END");
// Ensure we have token counts for non-streaming case
usage
.ensure_tokens(
system_prompt.as_str(),
messages_for_provider.messages(),
&message,
&tools,
)
.await?;
match complete_result {
Ok((message, usage)) => Ok(stream_from_single_message(message, usage)),
Err(e) => Err(e),
}
};
stream_from_single_message(message, usage)
// If there was an error creating the stream, return a stream that yields that error
let mut stream = match stream_result {
Ok(s) => s,
Err(e) => {
// Return a stream that immediately yields the error
// This allows the error to be caught by existing error handling in agent.rs
return Ok(Box::pin(try_stream! {
yield Err(e)?;
}));
}
};
Ok(Box::pin(try_stream! {
crates/goose/src/agents/subagent_handler.rs
+8 -7
View File
@@ -127,7 +127,7 @@ fn get_agent_messages(
}
}
let mut session_messages =
let mut conversation =
Conversation::new_unvalidated(
vec![Message::user().with_text(text_instruction.clone())],
);
@@ -141,15 +141,16 @@ fn get_agent_messages(
};
let mut stream = agent
.reply(session_messages.clone(), Some(session_config), None)
.reply(conversation.clone(), Some(session_config), None)
.await
.map_err(|e| anyhow!("Failed to get reply from agent: {}", e))?;
while let Some(message_result) = stream.next().await {
match message_result {
Ok(AgentEvent::Message(msg)) => session_messages.push(msg),
Ok(AgentEvent::McpNotification(_))
| Ok(AgentEvent::ModelChange { .. })
| Ok(AgentEvent::HistoryReplaced(_)) => {}
Ok(AgentEvent::Message(msg)) => conversation.push(msg),
Ok(AgentEvent::McpNotification(_)) | Ok(AgentEvent::ModelChange { .. }) => {}
Ok(AgentEvent::HistoryReplaced(updated_conversation)) => {
conversation = updated_conversation;
}
Err(e) => {
tracing::error!("Error receiving message from subagent: {}", e);
break;
@@ -157,6 +158,6 @@ fn get_agent_messages(
}
}
Ok(session_messages)
Ok(conversation)
})
}
crates/goose/src/context_mgmt/auto_compact.rs
-766
View File
@@ -1,766 +0,0 @@
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::{
agents::Agent, config::Config, context_mgmt::get_messages_token_counts_async,
token_counter::create_async_token_counter,
};
use anyhow::Result;
use tracing::{debug, info};
/// Result of auto-compaction check
#[derive(Debug)]
pub struct AutoCompactResult {
/// Whether compaction was performed
pub compacted: bool,
/// The messages after potential compaction
pub messages: Conversation,
/// Provider usage from summarization (if compaction occurred)
/// This contains the actual token counts after compaction
pub summarization_usage: Option<crate::providers::base::ProviderUsage>,
}
/// Result of checking if compaction is needed
#[derive(Debug)]
pub struct CompactionCheckResult {
/// Whether compaction is needed
pub needs_compaction: bool,
/// Current token count
pub current_tokens: usize,
/// Context limit being used
pub context_limit: usize,
/// Current usage ratio (0.0 to 1.0)
pub usage_ratio: f64,
/// Remaining tokens before compaction threshold
pub remaining_tokens: usize,
/// Percentage until compaction threshold (0.0 to 100.0)
pub percentage_until_compaction: f64,
}
/// Check if messages need compaction without performing the compaction
///
/// This function analyzes the current token usage and returns detailed information
/// about whether compaction is needed and how close we are to the threshold.
/// It prioritizes actual token counts from session metadata when available,
/// falling back to estimated counts if needed.
///
/// # Arguments
/// * `agent` - The agent to use for context management
/// * `messages` - The current message history
/// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config)
/// * `session_metadata` - Optional session metadata containing actual token counts
///
/// # Returns
/// * `CompactionCheckResult` containing detailed information about compaction needs
pub async fn check_compaction_needed(
agent: &Agent,
messages: &[Message],
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::Session>,
) -> Result<CompactionCheckResult> {
// Get threshold from config or use override
let config = Config::global();
let threshold = threshold_override.unwrap_or_else(|| {
config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(0.8) // Default to 80%
});
let provider = agent.provider().await?;
let context_limit = provider.get_model_config().context_limit();
let (current_tokens, token_source) = match session_metadata.and_then(|m| m.total_tokens) {
Some(tokens) => (tokens as usize, "session metadata"),
None => {
let token_counter = create_async_token_counter()
.await
.map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?;
let token_counts = get_messages_token_counts_async(&token_counter, messages);
(token_counts.iter().sum(), "estimated")
}
};
// Calculate usage ratio
let usage_ratio = current_tokens as f64 / context_limit as f64;
// Calculate threshold token count and remaining tokens
let threshold_tokens = (context_limit as f64 * threshold) as usize;
let remaining_tokens = threshold_tokens.saturating_sub(current_tokens);
// Calculate percentage until compaction (how much more we can use before hitting threshold)
let percentage_until_compaction = if usage_ratio < threshold {
(threshold - usage_ratio) * 100.0
} else {
0.0
};
// Check if compaction is needed (disabled if threshold is invalid)
let needs_compaction = if threshold <= 0.0 || threshold >= 1.0 {
false
} else {
usage_ratio > threshold
};
debug!(
"Compaction check: {} / {} tokens ({:.1}%), threshold: {:.1}%, needs compaction: {}, source: {}",
current_tokens,
context_limit,
usage_ratio * 100.0,
threshold * 100.0,
needs_compaction,
token_source
);
Ok(CompactionCheckResult {
needs_compaction,
current_tokens,
context_limit,
usage_ratio,
remaining_tokens,
percentage_until_compaction,
})
}
/// Perform compaction on messages without checking thresholds
///
/// This function directly performs compaction on the provided messages.
/// If the most recent message is a user message, it will be preserved by removing it
/// before compaction and adding it back afterwards.
///
/// # Arguments
/// * `agent` - The agent to use for context management
/// * `messages` - The current message history
///
/// # Returns
/// * `AutoCompactResult` containing the compacted messages and metadata
pub async fn perform_compaction(agent: &Agent, messages: &[Message]) -> Result<AutoCompactResult> {
info!("Performing message compaction");
// Check if the most recent message is a user message
let (messages_to_compact, preserved_user_message) = if let Some(last_message) = messages.last()
{
if matches!(last_message.role, rmcp::model::Role::User) {
// Remove the last user message before compaction
(&messages[..messages.len() - 1], Some(last_message.clone()))
} else {
(messages, None)
}
} else {
(messages, None)
};
// Perform the compaction on messages excluding the preserved user message
let (mut compacted_messages, _, summarization_usage) =
agent.summarize_context(messages_to_compact).await?;
// Add back the preserved user message if it exists
if let Some(user_message) = preserved_user_message {
compacted_messages.push(user_message);
}
Ok(AutoCompactResult {
compacted: true,
messages: compacted_messages,
summarization_usage,
})
}
/// Check if messages need compaction and compact them if necessary
///
/// This is a convenience wrapper function that combines checking and compaction.
/// Uses perform_compaction internally to handle the actual compaction process.
///
/// # Arguments
/// * `agent` - The agent to use for context management
/// * `messages` - The current message history
/// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config)
/// * `session_metadata` - Optional session metadata containing actual token counts
///
/// # Returns
/// * `AutoCompactResult` containing the potentially compacted messages and metadata
pub async fn check_and_compact_messages(
agent: &Agent,
messages: &[Message],
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::Session>,
) -> Result<AutoCompactResult> {
// First check if compaction is needed
let check_result =
check_compaction_needed(agent, messages, threshold_override, session_metadata).await?;
// If no compaction is needed, return early
if !check_result.needs_compaction {
debug!(
"No compaction needed (usage: {:.1}% <= {:.1}% threshold)",
check_result.usage_ratio * 100.0,
check_result.percentage_until_compaction
);
return Ok(AutoCompactResult {
compacted: false,
messages: Conversation::new_unvalidated(messages.to_vec()),
summarization_usage: None,
});
}
info!(
"Auto-compacting messages (usage: {:.1}%)",
check_result.usage_ratio * 100.0
);
perform_compaction(agent, messages).await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent};
use crate::session::extension_data;
use crate::{
agents::Agent,
model::ModelConfig,
providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage},
providers::errors::ProviderError,
};
use chrono::Utc;
use rmcp::model::{AnnotateAble, RawTextContent, Role, Tool};
use std::sync::Arc;
#[derive(Clone)]
struct MockProvider {
model_config: ModelConfig,
}
#[async_trait::async_trait]
impl Provider for MockProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete_with_model(
&self,
_model_config: &ModelConfig,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
// Return a short summary message
Ok((
Message::new(
Role::Assistant,
Utc::now().timestamp(),
vec![MessageContent::Text(
RawTextContent {
text: "Summary of conversation".to_string(),
meta: None,
}
.no_annotation(),
)],
),
ProviderUsage::new("mock".to_string(), Usage::default()),
))
}
}
fn create_test_message(text: &str) -> Message {
Message::new(
Role::User,
Utc::now().timestamp(),
vec![MessageContent::text(text.to_string())],
)
}
fn create_test_session_metadata(
message_count: usize,
working_dir: &str,
) -> crate::session::Session {
use crate::conversation::Conversation;
use std::path::PathBuf;
let mut conversation = Conversation::default();
for i in 0..message_count {
conversation.push(create_test_message(format!("message {}", i).as_str()));
}
crate::session::Session {
id: "test_session".to_string(),
working_dir: PathBuf::from(working_dir),
description: "Test session".to_string(),
created_at: Default::default(),
updated_at: Default::default(),
schedule_id: Some("test_job".to_string()),
recipe: None,
total_tokens: Some(100),
input_tokens: Some(50),
output_tokens: Some(50),
accumulated_total_tokens: Some(100),
accumulated_input_tokens: Some(50),
accumulated_output_tokens: Some(50),
extension_data: extension_data::ExtensionData::new(),
conversation: Some(conversation),
message_count,
user_recipe_values: None,
}
}
#[tokio::test]
async fn test_check_compaction_needed() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(Some(100_000)),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create small messages that won't trigger compaction
let messages = vec![create_test_message("Hello"), create_test_message("World")];
let result = check_compaction_needed(&agent, &messages, Some(0.3), None)
.await
.unwrap();
assert!(!result.needs_compaction);
assert!(result.current_tokens > 0);
assert!(result.context_limit > 0);
assert!(result.usage_ratio < 0.3);
assert!(result.remaining_tokens > 0);
assert!(result.percentage_until_compaction > 0.0);
}
#[tokio::test]
async fn test_check_compaction_needed_disabled() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(Some(100_000)),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let messages = vec![create_test_message("Hello")];
// Test with threshold 0 (disabled)
let result = check_compaction_needed(&agent, &messages, Some(0.0), None)
.await
.unwrap();
assert!(!result.needs_compaction);
// Test with threshold 1.0 (disabled)
let result = check_compaction_needed(&agent, &messages, Some(1.0), None)
.await
.unwrap();
assert!(!result.needs_compaction);
}
#[tokio::test]
async fn test_auto_compact_disabled() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(Some(10_000)),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let messages = vec![create_test_message("Hello"), create_test_message("World")];
// Test with threshold 0 (disabled)
let result = check_and_compact_messages(&agent, &messages, Some(0.0), None)
.await
.unwrap();
assert!(!result.compacted);
assert_eq!(result.messages.len(), messages.len());
assert!(result.summarization_usage.is_none());
// Test with threshold 1.0 (disabled)
let result = check_and_compact_messages(&agent, &messages, Some(1.0), None)
.await
.unwrap();
assert!(!result.compacted);
}
#[tokio::test]
async fn test_auto_compact_below_threshold() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(Some(100_000)), // Increased to ensure overhead doesn't dominate
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create small messages that won't trigger compaction
let messages = vec![create_test_message("Hello"), create_test_message("World")];
let result = check_and_compact_messages(&agent, &messages, Some(0.3), None)
.await
.unwrap();
assert!(!result.compacted);
assert_eq!(result.messages.len(), messages.len());
}
#[tokio::test]
async fn test_auto_compact_above_threshold() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(30_000.into()), // Smaller context limit to make threshold easier to hit
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create messages that will exceed 30% of the context limit
// With 30k context limit, 30% is 9k tokens
let mut messages = Vec::new();
// Create much longer messages with more content to reach the threshold
for i in 0..300 {
messages.push(create_test_message(&format!(
"This is message number {} with significantly more content to increase token count substantially. \
We need to ensure that our total token usage exceeds 30% of the available context \
limit after accounting for system prompt and tools overhead. This message contains \
multiple sentences to increase the token count substantially. Adding even more text here \
to make sure we have enough tokens. Lorem ipsum dolor sit amet, consectetur adipiscing elit, \
sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, \
quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute \
irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. \
Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit \
anim id est laborum. Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium \
doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi \
architecto beatae vitae dicta sunt explicabo.",
i
)));
}
let result = check_and_compact_messages(&agent, &messages, Some(0.3), None)
.await
.unwrap();
if !result.compacted {
eprintln!("Test failed - compaction not triggered");
}
assert!(result.compacted);
assert!(result.summarization_usage.is_some());
// Verify that summarization usage contains token counts
if let Some(usage) = &result.summarization_usage {
assert!(usage.usage.total_tokens.is_some());
let after = usage.usage.total_tokens.unwrap_or(0) as usize;
assert!(
after > 0,
"Token count after compaction should be greater than 0"
);
}
// After visibility implementation, we keep all messages plus summary
// Original messages become user_visible only, summary becomes agent_visible only
assert!(result.messages.len() > messages.len());
}
#[tokio::test]
async fn test_auto_compact_respects_config() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(Some(30_000)), // Smaller context limit to make threshold easier to hit
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create enough messages to trigger compaction with low threshold
let mut messages = Vec::new();
// With 30k context limit, after overhead we have ~27k usable tokens
// 10% of 27k = 2.7k tokens, so we need messages that exceed that
for i in 0..200 {
messages.push(create_test_message(&format!(
"Message {} with enough content to ensure we exceed 10% of the context limit. \
Adding more content to increase token count substantially. This message contains \
multiple sentences to increase the token count. We need to ensure that our total \
token usage exceeds 10% of the available context limit after accounting for \
system prompt and tools overhead.",
i
)));
}
// Set config value
let config = Config::global();
config
.set_param("GOOSE_AUTO_COMPACT_THRESHOLD", serde_json::Value::from(0.1))
.unwrap();
// Should use config value when no override provided
let result = check_and_compact_messages(&agent, &messages, None, None)
.await
.unwrap();
// Debug info if not compacted
if !result.compacted {
eprintln!("Test failed - compaction not triggered");
}
// With such a low threshold (10%), it should compact
assert!(result.compacted);
// Clean up config
config
.set_param("GOOSE_AUTO_COMPACT_THRESHOLD", serde_json::Value::from(0.3))
.unwrap();
}
#[tokio::test]
async fn test_auto_compact_uses_session_metadata() {
use crate::session::Session;
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(10_000.into()),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create some test messages
let messages = vec![
create_test_message("First message"),
create_test_message("Second message"),
];
// Create session with specific token counts
#[allow(clippy::field_reassign_with_default)]
let mut session = Session::default();
{
session.total_tokens = Some(8000); // High token count to trigger compaction
session.accumulated_total_tokens = Some(15000); // Even higher accumulated count
session.input_tokens = Some(5000);
session.output_tokens = Some(3000);
}
// Test with session - should use total_tokens for compaction (not accumulated)
let result_with_metadata = check_compaction_needed(
&agent,
&messages,
Some(0.3), // 30% threshold
Some(&session),
)
.await
.unwrap();
// With 8000 tokens and context limit around 10000, should trigger compaction
assert!(result_with_metadata.needs_compaction);
assert_eq!(result_with_metadata.current_tokens, 8000);
// Test without session metadata - should use estimated tokens
let result_without_metadata = check_compaction_needed(
&agent,
&messages,
Some(0.3), // 30% threshold
None,
)
.await
.unwrap();
// Without metadata, should use much lower estimated token count
assert!(!result_without_metadata.needs_compaction);
assert!(result_without_metadata.current_tokens < 8000);
// Test with session that has only accumulated tokens (no total_tokens)
let mut session_metadata_no_total = Session::default();
#[allow(clippy::field_reassign_with_default)]
{
session_metadata_no_total.accumulated_total_tokens = Some(7500);
}
let result_with_no_total = check_compaction_needed(
&agent,
&messages,
Some(0.3), // 30% threshold
Some(&session_metadata_no_total),
)
.await
.unwrap();
// Should fall back to estimation since total_tokens is None
assert!(!result_with_no_total.needs_compaction);
assert!(result_with_no_total.current_tokens < 7500);
// Test with metadata that has no token counts - should fall back to estimation
let empty_metadata = Session::default();
let result_with_empty_metadata = check_compaction_needed(
&agent,
&messages,
Some(0.3), // 30% threshold
Some(&empty_metadata),
)
.await
.unwrap();
// Should fall back to estimation
assert!(!result_with_empty_metadata.needs_compaction);
assert!(result_with_empty_metadata.current_tokens < 7500);
}
#[tokio::test]
async fn test_auto_compact_end_to_end_with_metadata() {
use crate::session::Session;
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(10_000.into()),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
// Create some test messages
let messages = vec![
create_test_message("First message"),
create_test_message("Second message"),
create_test_message("Third message"),
create_test_message("Fourth message"),
create_test_message("Fifth message"),
];
// Create session metadata with high token count to trigger compaction
let mut session = Session::default();
#[allow(clippy::field_reassign_with_default)]
{
session.total_tokens = Some(9000); // High enough to trigger compaction
}
// Test full compaction flow with session metadata
let result = check_and_compact_messages(
&agent,
&messages,
Some(0.3), // 30% threshold
Some(&session),
)
.await
.unwrap();
// Should have triggered compaction
assert!(result.compacted);
assert!(result.summarization_usage.is_some());
// Verify the compacted messages are returned
assert!(!result.messages.is_empty());
// After visibility implementation, we keep all messages plus summary
// Original messages become user_visible only, summary becomes agent_visible only
assert!(result.messages.len() > messages.len());
}
#[tokio::test]
async fn test_auto_compact_with_comprehensive_session_metadata() {
let mock_provider = Arc::new(MockProvider {
model_config: ModelConfig::new("test-model")
.unwrap()
.with_context_limit(8_000.into()),
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let messages = vec![
create_test_message("Test message 1"),
create_test_message("Test message 2"),
create_test_message("Test message 3"),
];
// Use the helper function to create comprehensive non-null session metadata
let comprehensive_metadata = create_test_session_metadata(3, "/test/working/dir");
// Verify the helper created non-null metadata
assert_eq!(
comprehensive_metadata
.clone()
.conversation
.unwrap_or_default()
.len(),
3
);
assert_eq!(
comprehensive_metadata.working_dir.to_str().unwrap(),
"/test/working/dir"
);
assert_eq!(comprehensive_metadata.description, "Test session");
assert_eq!(
comprehensive_metadata.schedule_id,
Some("test_job".to_string())
);
assert_eq!(comprehensive_metadata.total_tokens, Some(100));
assert_eq!(comprehensive_metadata.input_tokens, Some(50));
assert_eq!(comprehensive_metadata.output_tokens, Some(50));
assert_eq!(comprehensive_metadata.accumulated_total_tokens, Some(100));
assert_eq!(comprehensive_metadata.accumulated_input_tokens, Some(50));
assert_eq!(comprehensive_metadata.accumulated_output_tokens, Some(50));
// Test compaction with the comprehensive metadata (low token count, shouldn't compact)
let result_low_tokens = check_compaction_needed(
&agent,
&messages,
Some(0.7), // 70% threshold
Some(&comprehensive_metadata),
)
.await
.unwrap();
assert!(!result_low_tokens.needs_compaction);
assert_eq!(result_low_tokens.current_tokens, 100); // Should use total_tokens from metadata
// Create a modified version with high token count to trigger compaction
let mut high_token_metadata = create_test_session_metadata(5, "/test/working/dir");
high_token_metadata.total_tokens = Some(6_000); // High enough to trigger compaction
high_token_metadata.input_tokens = Some(4_000);
high_token_metadata.output_tokens = Some(2_000);
high_token_metadata.accumulated_total_tokens = Some(12_000);
let result_high_tokens = check_compaction_needed(
&agent,
&messages,
Some(0.7), // 70% threshold
Some(&high_token_metadata),
)
.await
.unwrap();
assert!(result_high_tokens.needs_compaction);
assert_eq!(result_high_tokens.current_tokens, 6_000); // Should use total_tokens, not accumulated
// Test that metadata fields are preserved correctly in edge cases
let mut edge_case_metadata = create_test_session_metadata(10, "/edge/case/dir");
edge_case_metadata.total_tokens = None; // No total tokens
edge_case_metadata.accumulated_total_tokens = Some(7_000); // Has accumulated
let result_edge_case = check_compaction_needed(
&agent,
&messages,
Some(0.5), // 50% threshold
Some(&edge_case_metadata),
)
.await
.unwrap();
// Should fall back to estimation since total_tokens is None
assert!(result_edge_case.current_tokens < 7_000);
// With estimation, likely won't trigger compaction
assert!(!result_edge_case.needs_compaction);
}
}
crates/goose/src/context_mgmt/common.rs
-99
View File
@@ -1,99 +0,0 @@
use std::sync::Arc;
use rmcp::model::Tool;
use crate::conversation::message::Message;
use crate::{
providers::base::Provider,
token_counter::{AsyncTokenCounter, TokenCounter},
};
const ESTIMATE_FACTOR: f32 = 0.7;
pub const SYSTEM_PROMPT_TOKEN_OVERHEAD: usize = 3_000;
pub const TOOLS_TOKEN_OVERHEAD: usize = 5_000;
pub fn estimate_target_context_limit(provider: Arc<dyn Provider>) -> usize {
let model_context_limit = provider.get_model_config().context_limit();
// Our conservative estimate of the **target** context limit
// Our token count is an estimate since model providers often don't provide the tokenizer (eg. Claude)
let target_limit = (model_context_limit as f32 * ESTIMATE_FACTOR) as usize;
// subtract out overhead for system prompt and tools, but ensure we don't go negative
let overhead = SYSTEM_PROMPT_TOKEN_OVERHEAD + TOOLS_TOKEN_OVERHEAD;
if target_limit > overhead {
target_limit - overhead
} else {
// If overhead is larger than target limit, return a minimal usable limit
std::cmp::max(target_limit / 2, 1000)
}
}
pub fn get_messages_token_counts(token_counter: &TokenCounter, 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()
}
/// Async version of get_messages_token_counts for better performance
pub fn get_messages_token_counts_async(
token_counter: &AsyncTokenCounter,
messages: &[Message],
) -> Vec<usize> {
messages
.iter()
.filter(|m| m.is_agent_visible())
.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)]
pub struct ChatTokenCounts {
pub system: usize,
pub tools: usize,
pub messages: Vec<usize>,
}
#[allow(dead_code)]
pub fn get_token_counts(
token_counter: &TokenCounter,
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(token_counter, messages);
ChatTokenCounts {
system: system_prompt_token_count,
tools: tools_token_count,
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/mod.rs
+337 -5
View File
@@ -1,6 +1,338 @@
pub mod auto_compact;
mod common;
pub mod summarize;
pub mod truncate;
use crate::conversation::message::Message;
use crate::conversation::message::MessageMetadata;
use crate::conversation::Conversation;
use crate::prompt_template::render_global_file;
use crate::providers::base::{Provider, ProviderUsage};
use crate::{agents::Agent, config::Config, token_counter::create_token_counter};
use anyhow::Result;
use rmcp::model::Role;
use serde::Serialize;
use std::sync::Arc;
use tracing::{debug, info};
pub use common::*;
pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.8;
/// Result of auto-compaction check
#[derive(Debug)]
pub struct AutoCompactResult {
/// Whether compaction was performed
pub compacted: bool,
/// The messages after potential compaction
pub messages: Conversation,
/// Provider usage from summarization (if compaction occurred)
/// This contains the actual token counts after compaction
pub summarization_usage: Option<crate::providers::base::ProviderUsage>,
}
/// Result of checking if compaction is needed
#[derive(Debug)]
pub struct CompactionCheckResult {
/// Whether compaction is needed
pub needs_compaction: bool,
/// Current token count
pub current_tokens: usize,
/// Context limit being used
pub context_limit: usize,
/// Current usage ratio (0.0 to 1.0)
pub usage_ratio: f64,
/// Remaining tokens before compaction threshold
pub remaining_tokens: usize,
/// Percentage until compaction threshold (0.0 to 100.0)
pub percentage_until_compaction: f64,
}
#[derive(Serialize)]
struct SummarizeContext {
messages: String,
}
/// Check if messages need compaction and compact them if necessary
///
/// This function combines checking and compaction. It first checks if compaction
/// is needed based on the threshold, and if so, performs the compaction by
/// summarizing messages and updating their visibility metadata.
///
/// # Arguments
/// * `agent` - The agent to use for context management
/// * `messages` - The current message history
/// * `force_compact` - If true, skip the threshold check and force compaction
/// * `preserve_last_user_message` - If true and last message is not a user message, copy the most recent user message to the end
/// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config)
/// * `session_metadata` - Optional session metadata containing actual token counts
///
/// # Returns
/// * A tuple containing:
/// - `bool`: Whether compaction was performed
/// - `Conversation`: The potentially compacted messages
/// - `Vec<usize>`: Indices of removed messages (empty if no compaction)
/// - `Option<ProviderUsage>`: Provider usage from summarization (if compaction occurred)
pub async fn check_and_compact_messages(
agent: &Agent,
messages_with_user_message: &[Message],
force_compact: bool,
preserve_last_user_message: bool,
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::Session>,
) -> std::result::Result<(bool, Conversation, Vec<usize>, Option<ProviderUsage>), anyhow::Error> {
// Check if compaction is needed (unless forced)
if !force_compact {
let check_result = check_compaction_needed(
agent,
messages_with_user_message,
threshold_override,
session_metadata,
)
.await?;
// If no compaction is needed, return early
if !check_result.needs_compaction {
debug!(
"No compaction needed (usage: {:.1}% <= {:.1}% threshold)",
check_result.usage_ratio * 100.0,
check_result.percentage_until_compaction
);
return Ok((
false,
Conversation::new_unvalidated(messages_with_user_message.to_vec()),
Vec::new(),
None,
));
}
info!(
"Performing message compaction (usage: {:.1}%)",
check_result.usage_ratio * 100.0
);
} else {
info!("Forcing message compaction due to context limit exceeded");
}
// Perform the actual compaction
// Check if the most recent message is a user message
let (messages, preserved_user_message) =
if let Some(last_message) = messages_with_user_message.last() {
if matches!(last_message.role, rmcp::model::Role::User) {
// Remove the last user message before compaction
(
&messages_with_user_message[..messages_with_user_message.len() - 1],
Some(last_message.clone()),
)
} else if preserve_last_user_message {
// Last message is not a user message, but we want to preserve the most recent user message
// Find the most recent user message and copy it (don't remove from history)
let most_recent_user_message = messages_with_user_message
.iter()
.rev()
.find(|msg| matches!(msg.role, rmcp::model::Role::User))
.cloned();
(messages_with_user_message, most_recent_user_message)
} else {
(messages_with_user_message, None)
}
} else {
(messages_with_user_message, None)
};
let provider = agent.provider().await?;
let summary = summarize(provider.clone(), messages).await?;
let (summary_message, summarization_usage) = match summary {
Some((summary_message, provider_usage)) => (summary_message, Some(provider_usage)),
None => {
// No summary was generated (empty input)
tracing::warn!("Summarization failed. Returning empty messages.");
return Ok((false, Conversation::empty(), vec![], None));
}
};
// Create the final message list with updated visibility metadata:
// 1. Original messages become user_visible but not agent_visible
// 2. Summary message becomes agent_visible but not user_visible
// 3. Assistant messages to continue the conversation remain both user_visible and agent_visible
let mut final_messages = Vec::new();
let mut final_token_counts = Vec::new();
// Add all original messages with updated visibility (preserve user_visible, set agent_visible=false)
for msg in messages.iter().cloned() {
let updated_metadata = msg.metadata.with_agent_invisible();
let updated_msg = msg.with_metadata(updated_metadata);
final_messages.push(updated_msg);
// Token count doesn't matter for agent_visible=false messages, but we'll use 0
final_token_counts.push(0);
}
// Add the compaction marker (user_visible=true, agent_visible=false)
let compaction_marker = Message::assistant()
.with_conversation_compacted("Conversation compacted and summarized")
.with_metadata(MessageMetadata::user_only());
let compaction_marker_tokens: usize = 0; // Not counted since agent_visible=false
final_messages.push(compaction_marker);
final_token_counts.push(compaction_marker_tokens);
// Add the summary message (agent_visible=true, user_visible=false)
let summary_msg = summary_message.with_metadata(MessageMetadata::agent_only());
// For token counting purposes, we use the output tokens (the actual summary content)
// since that's what will be in the context going forward
let summary_tokens = summarization_usage
.as_ref()
.and_then(|usage| usage.usage.output_tokens)
.unwrap_or(0) as usize;
final_messages.push(summary_msg);
final_token_counts.push(summary_tokens);
// Add an assistant message to continue the conversation (agent_visible=true, user_visible=false)
let assistant_message = Message::assistant()
.with_text(
"The previous message contains a summary that was prepared because a context limit was reached.
Do not mention that you read a summary or that conversation summarization occurred
Just continue the conversation naturally based on the summarized context"
)
.with_metadata(MessageMetadata::agent_only());
let assistant_message_tokens: usize = 0; // Not counted since it's for agent context only
final_messages.push(assistant_message);
final_token_counts.push(assistant_message_tokens);
// Add back the preserved user message if it exists
if let Some(user_message) = preserved_user_message {
final_messages.push(user_message);
}
Ok((
true,
Conversation::new_unvalidated(final_messages),
final_token_counts,
summarization_usage,
))
}
/// Check if messages need compaction without performing the compaction
///
/// This function analyzes the current token usage and returns detailed information
/// about whether compaction is needed and how close we are to the threshold.
/// It prioritizes actual token counts from session metadata when available,
/// falling back to estimated counts if needed.
///
/// # Arguments
/// * `agent` - The agent to use for context management
/// * `messages` - The current message history
/// * `threshold_override` - Optional threshold override (defaults to GOOSE_AUTO_COMPACT_THRESHOLD config)
/// * `session_metadata` - Optional session metadata containing actual token counts
///
/// # Returns
/// * `CompactionCheckResult` containing detailed information about compaction needs
async fn check_compaction_needed(
agent: &Agent,
messages: &[Message],
threshold_override: Option<f64>,
session_metadata: Option<&crate::session::Session>,
) -> Result<CompactionCheckResult> {
// Get threshold from config or use override
let config = Config::global();
// TODO(Douwe): check the default here; it seems to reset to 0.3 sometimes
let threshold = threshold_override.unwrap_or_else(|| {
config
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(DEFAULT_COMPACTION_THRESHOLD)
});
let provider = agent.provider().await?;
let context_limit = provider.get_model_config().context_limit();
let (current_tokens, token_source) = match session_metadata.and_then(|m| m.total_tokens) {
Some(tokens) => (tokens as usize, "session metadata"),
None => {
let token_counter = create_token_counter()
.await
.map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?;
let token_counts: Vec<_> = messages
.iter()
.filter(|m| m.is_agent_visible())
.map(|msg| token_counter.count_chat_tokens("", std::slice::from_ref(msg), &[]))
.collect();
(token_counts.iter().sum(), "estimated")
}
};
let usage_ratio = current_tokens as f64 / context_limit as f64;
let threshold_tokens = (context_limit as f64 * threshold) as usize;
let remaining_tokens = threshold_tokens.saturating_sub(current_tokens);
let percentage_until_compaction = if usage_ratio < threshold {
(threshold - usage_ratio) * 100.0
} else {
0.0
};
let needs_compaction = if threshold <= 0.0 || threshold >= 1.0 {
usage_ratio > DEFAULT_COMPACTION_THRESHOLD
} else {
usage_ratio > threshold
};
debug!(
"Compaction check: {} / {} tokens ({:.1}%), threshold: {:.1}%, needs compaction: {}, source: {}",
current_tokens,
context_limit,
usage_ratio * 100.0,
threshold * 100.0,
needs_compaction,
token_source
);
Ok(CompactionCheckResult {
needs_compaction,
current_tokens,
context_limit,
usage_ratio,
remaining_tokens,
percentage_until_compaction,
})
}
async fn summarize(
provider: Arc<dyn Provider>,
messages: &[Message],
) -> anyhow::Result<Option<(Message, ProviderUsage)>, anyhow::Error> {
if messages.is_empty() {
return std::prelude::rust_2015::Ok(None);
}
// Format all messages as a single string for the summarization prompt
let messages_text = messages
.iter()
.map(|msg| format!("{:?}", msg))
.collect::<Vec<_>>()
.join("\n\n");
let context = SummarizeContext {
messages: messages_text,
};
// Render the one-shot summarization prompt
let system_prompt = render_global_file("summarize_oneshot.md", &context)?;
// Create a simple user message requesting summarization
let user_message = Message::user()
.with_text("Please summarize the conversation history provided in the system prompt.");
let summarization_request = vec![user_message];
// Send the request to the provider and fetch the response
let (mut response, mut provider_usage) = provider
.complete_fast(&system_prompt, &summarization_request, &[])
.await?;
// Set role to user as it will be used in following conversation as user content
response.role = Role::User;
// Ensure we have token counts, estimating if necessary
provider_usage
.ensure_tokens(&system_prompt, &summarization_request, &response, &[])
.await
.map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?;
std::prelude::rust_2015::Ok(Some((response, provider_usage)))
}
crates/goose/src/context_mgmt/summarize.rs
-187
View File
@@ -1,187 +0,0 @@
use crate::conversation::message::Message;
use crate::prompt_template::render_global_file;
use crate::providers::base::Provider;
use anyhow::Result;
use rmcp::model::Role;
use serde::Serialize;
use std::sync::Arc;
#[derive(Serialize)]
struct SummarizeContext {
messages: String,
}
use crate::providers::base::ProviderUsage;
/// Summarization function that uses the detailed prompt from the markdown template
pub async fn summarize_messages(
provider: Arc<dyn Provider>,
messages: &[Message],
) -> Result<Option<(Message, ProviderUsage)>, anyhow::Error> {
if messages.is_empty() {
return Ok(None);
}
// Format all messages as a single string for the summarization prompt
let messages_text = messages
.iter()
.map(|msg| format!("{:?}", msg))
.collect::<Vec<_>>()
.join("\n\n");
let context = SummarizeContext {
messages: messages_text,
};
// Render the one-shot summarization prompt
let system_prompt = render_global_file("summarize_oneshot.md", &context)?;
// Create a simple user message requesting summarization
let user_message = Message::user()
.with_text("Please summarize the conversation history provided in the system prompt.");
let summarization_request = vec![user_message];
// Send the request to the provider and fetch the response
let (mut response, mut provider_usage) = provider
.complete_fast(&system_prompt, &summarization_request, &[])
.await?;
// Set role to user as it will be used in following conversation as user content
response.role = Role::User;
// Ensure we have token counts, estimating if necessary
provider_usage
.ensure_tokens(&system_prompt, &summarization_request, &response, &[])
.await
.map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?;
Ok(Some((response, provider_usage)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
use crate::providers::base::{ProviderMetadata, ProviderUsage, Usage};
use crate::providers::errors::ProviderError;
use chrono::Utc;
use rmcp::model::Role;
use rmcp::model::Tool;
use rmcp::model::{AnnotateAble, RawTextContent};
use std::sync::Arc;
#[derive(Clone)]
struct MockProvider {
model_config: ModelConfig,
}
#[async_trait::async_trait]
impl Provider for MockProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete_with_model(
&self,
_model_config: &ModelConfig,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
Ok((
Message::new(
Role::Assistant,
Utc::now().timestamp(),
vec![MessageContent::Text(
RawTextContent {
text: "Summarized content".to_string(),
meta: None,
}
.no_annotation(),
)],
),
ProviderUsage::new(
"mock".to_string(),
Usage {
input_tokens: Some(100),
output_tokens: Some(50),
total_tokens: Some(150),
},
),
))
}
}
fn create_mock_provider() -> Result<Arc<dyn Provider>> {
let mock_model_config = ModelConfig::new("test-model")?.with_context_limit(Some(200_000));
Ok(Arc::new(MockProvider {
model_config: mock_model_config,
}))
}
fn create_test_messages() -> Vec<Message> {
vec![
set_up_text_message("Message 1", Role::User),
set_up_text_message("Message 2", Role::Assistant),
set_up_text_message("Message 3", Role::User),
]
}
fn set_up_text_message(text: &str, role: Role) -> Message {
Message::new(role, 0, vec![MessageContent::text(text.to_string())])
}
#[tokio::test]
async fn test_summarize_messages_basic() {
let provider = create_mock_provider().expect("failed to create mock provider");
let messages = create_test_messages();
let result = summarize_messages(Arc::clone(&provider), &messages).await;
assert!(result.is_ok(), "The function should return Ok.");
let summary_result = result.unwrap();
assert!(
summary_result.is_some(),
"The summary should contain a result."
);
let (summarized_message, provider_usage) = summary_result.unwrap();
assert_eq!(
summarized_message.role,
Role::User,
"The summarized message should be from the user."
);
assert!(
provider_usage.usage.input_tokens.unwrap_or(0) > 0,
"Should have input token count"
);
assert!(
provider_usage.usage.output_tokens.unwrap_or(0) > 0,
"Should have output token count"
);
}
#[tokio::test]
async fn test_summarize_messages_empty_input() {
let provider = create_mock_provider().expect("failed to create mock provider");
let messages: Vec<Message> = Vec::new();
let result = summarize_messages(Arc::clone(&provider), &messages).await;
assert!(result.is_ok(), "The function should return Ok.");
let summary_result = result.unwrap();
assert!(
summary_result.is_none(),
"The summary should be None for empty input."
);
}
}
crates/goose/src/context_mgmt/truncate.rs
-749
View File
@@ -1,749 +0,0 @@
use crate::conversation::message::{Message, MessageContent};
use crate::conversation::Conversation;
use crate::utils::safe_truncate;
use anyhow::{anyhow, Result};
use rmcp::model::{RawContent, ResourceContents, Role};
use std::collections::HashSet;
use std::ops::DerefMut;
use tracing::{debug, warn};
/// Maximum size for truncated content in characters
const MAX_TRUNCATED_CONTENT_SIZE: usize = 5000;
/// Handles messages that are individually larger than the context limit
/// by truncating their content rather than removing them entirely
fn handle_oversized_messages(
messages: &[Message],
token_counts: &[usize],
context_limit: usize,
strategy: &dyn TruncationStrategy,
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
let mut truncated_messages = Vec::new();
let mut truncated_token_counts = Vec::new();
let mut any_truncated = false;
// Create a basic token counter for re-estimating truncated content
// Note: This is a rough approximation since we don't have access to the actual tokenizer here
let estimate_tokens = |text: &str| -> usize {
// Rough approximation: 1 token per 4 characters for English text
(text.len() / 4).max(1)
};
for (i, (message, &original_tokens)) in messages.iter().zip(token_counts.iter()).enumerate() {
if original_tokens > context_limit {
warn!(
"Message {} has {} tokens, exceeding context limit of {}",
i, original_tokens, context_limit
);
// Try to truncate the message content
let truncated_message = truncate_message_content(message, MAX_TRUNCATED_CONTENT_SIZE)?;
let estimated_new_tokens =
estimate_message_tokens(&truncated_message, &estimate_tokens);
if estimated_new_tokens > context_limit {
// Even truncated message is too large, skip it entirely
warn!("Skipping message {} as even truncated version ({} tokens) exceeds context limit", i, estimated_new_tokens);
any_truncated = true;
continue;
}
truncated_messages.push(truncated_message);
truncated_token_counts.push(estimated_new_tokens);
any_truncated = true;
} else {
truncated_messages.push(message.clone());
truncated_token_counts.push(original_tokens);
}
}
if any_truncated {
debug!("Truncated large message content, now attempting normal truncation");
// After content truncation, try normal truncation if still needed
return truncate_messages(
&truncated_messages,
&truncated_token_counts,
context_limit,
strategy,
);
}
Ok((
Conversation::new_unvalidated(truncated_messages),
truncated_token_counts,
))
}
/// Truncates the content within a message while preserving its structure
fn truncate_message_content(message: &Message, max_content_size: usize) -> Result<Message> {
let mut new_message = message.clone();
for content in &mut new_message.content {
match content {
MessageContent::Text(text_content) => {
if text_content.text.chars().count() > max_content_size {
let truncated = format!(
"{}\n\n[... content truncated from {} to {} characters ...]",
safe_truncate(&text_content.text, max_content_size),
text_content.text.chars().count(),
max_content_size
);
text_content.text = truncated;
}
}
MessageContent::ToolResponse(tool_response) => {
if let Ok(ref mut result) = tool_response.tool_result {
for content_item in result {
if let RawContent::Text(ref mut text_content) = content_item.deref_mut() {
if text_content.text.chars().count() > max_content_size {
let truncated = format!(
"{}\n\n[... tool response truncated from {} to {} characters ...]",
safe_truncate(&text_content.text, max_content_size),
text_content.text.chars().count(),
max_content_size
);
text_content.text = truncated;
}
}
// Handle Resource content which might contain large text
else if let RawContent::Resource(ref mut resource_content) =
content_item.deref_mut()
{
if let ResourceContents::TextResourceContents { text, .. } =
&mut resource_content.resource
{
if text.chars().count() > max_content_size {
let truncated = format!(
"{}\n\n[... resource content truncated from {} to {} characters ...]",
safe_truncate(text, max_content_size),
text.chars().count(),
max_content_size
);
*text = truncated;
}
}
}
}
}
}
// Other content types are typically smaller, but we could extend this if needed
_ => {}
}
}
Ok(new_message)
}
/// Estimates token count for a message using a simple heuristic
fn estimate_message_tokens(message: &Message, estimate_fn: &dyn Fn(&str) -> usize) -> usize {
let mut total_tokens = 10; // Base overhead for message structure
for content in &message.content {
match content {
MessageContent::Text(text_content) => {
total_tokens += estimate_fn(&text_content.text);
}
MessageContent::ToolResponse(tool_response) => {
if let Ok(ref result) = tool_response.tool_result {
for content_item in result {
match &content_item.raw {
RawContent::Text(text_content) => {
total_tokens += estimate_fn(&text_content.text);
}
RawContent::Resource(resource) => {
match &resource.resource {
ResourceContents::TextResourceContents { text, .. } => {
total_tokens += estimate_fn(text);
}
_ => total_tokens += 5, // Small overhead for other resource types
}
}
_ => {
total_tokens += 5; // Small overhead for other content types
}
}
}
}
}
_ => total_tokens += 5, // Small overhead for other content types
}
}
total_tokens
}
/// Truncates the messages to fit within the model's context window.
/// Mutates the input messages and token counts in place.
/// Returns an error if it's impossible to truncate the messages within the context limit.
/// - messages: The vector of messages in the conversation.
/// - token_counts: A parallel vector containing the token count for each message.
/// - context_limit: The maximum allowed context length in tokens.
/// - strategy: The truncation strategy to use. Only option is OldestFirstTruncation.
pub fn truncate_messages(
messages: &[Message],
token_counts: &[usize],
context_limit: usize,
strategy: &dyn TruncationStrategy,
) -> Result<(Conversation, Vec<usize>), anyhow::Error> {
let mut messages = messages.to_owned();
let mut token_counts = token_counts.to_owned();
if messages.len() != token_counts.len() {
return Err(anyhow!(
"The vector for messages and token_counts must have same length"
));
}
// Step 1: Calculate total tokens
let mut total_tokens: usize = token_counts.iter().sum();
debug!("Total tokens before truncation: {}", total_tokens);
// Check if any individual message is larger than the context limit
// First, check for any message that's too large
let max_message_tokens = token_counts.iter().max().copied().unwrap_or(0);
if max_message_tokens > context_limit {
// Try to handle large messages by truncating their content
debug!(
"Found oversized message with {} tokens, attempting content truncation",
max_message_tokens
);
return handle_oversized_messages(&messages, &token_counts, context_limit, strategy);
}
let min_user_msg_tokens = messages
.iter()
.zip(token_counts.iter())
.filter(|(msg, _)| msg.role == Role::User && msg.has_only_text_content())
.map(|(_, &tokens)| tokens)
.min();
// If there are no valid user messages, or the smallest one is too big for the context
if min_user_msg_tokens.is_none() || min_user_msg_tokens.unwrap() > context_limit {
return Err(anyhow!(
"Not possible to truncate messages within context limit: no suitable user messages found"
));
}
if total_tokens <= context_limit {
return Ok((
Conversation::new_unvalidated(messages.to_vec()),
token_counts.to_vec(),
)); // No truncation needed
}
// Step 2: Determine indices to remove based on strategy
let indices_to_remove =
strategy.determine_indices_to_remove(&messages, &token_counts, context_limit)?;
// Circuit breaker: if we can't remove enough messages, fail gracefully
let tokens_to_remove: usize = indices_to_remove
.iter()
.map(|&i| token_counts.get(i).copied().unwrap_or(0))
.sum();
if total_tokens - tokens_to_remove > context_limit && !indices_to_remove.is_empty() {
debug!(
"Standard truncation insufficient: {} tokens remain after removing {} tokens",
total_tokens - tokens_to_remove,
tokens_to_remove
);
// Try more aggressive truncation or content truncation
return handle_oversized_messages(&messages, &token_counts, context_limit, strategy);
}
if indices_to_remove.is_empty() && total_tokens > context_limit {
return Err(anyhow!(
"Cannot truncate any messages: all messages may be essential or too large individually"
));
}
// Step 3: Remove the marked messages
// Vectorize the set and sort in reverse order to avoid shifting indices when removing
let mut indices_to_remove = indices_to_remove.iter().cloned().collect::<Vec<usize>>();
indices_to_remove.sort_unstable_by(|a, b| b.cmp(a));
for &index in &indices_to_remove {
if index < messages.len() {
let _ = messages.remove(index);
let removed_tokens = token_counts.remove(index);
total_tokens -= removed_tokens;
}
}
// Step 4: Ensure the last message is a user message with TextContent only
while let Some(last_msg) = messages.last() {
if last_msg.role != Role::User || !last_msg.has_only_text_content() {
let _ = messages.pop().ok_or(anyhow!("Failed to pop message"))?;
let removed_tokens = token_counts
.pop()
.ok_or(anyhow!("Failed to pop token count"))?;
total_tokens -= removed_tokens;
} else {
break;
}
}
// Step 5: Check first msg is a User message with TextContent only
while let Some(first_msg) = messages.first() {
if first_msg.role != Role::User || !first_msg.has_only_text_content() {
let _ = messages.remove(0);
let removed_tokens = token_counts.remove(0);
total_tokens -= removed_tokens;
} else {
break;
}
}
debug!("Total tokens after truncation: {}", total_tokens);
// Ensure we have at least one message remaining and it's within context limit
if messages.is_empty() {
return Err(anyhow!(
"Unable to preserve any messages within context limit"
));
}
if total_tokens > context_limit {
return Err(anyhow!(
"Unable to truncate messages within context window."
));
}
debug!("Truncation complete. Total tokens: {}", total_tokens);
Ok((
Conversation::new_unvalidated(messages.to_vec()),
token_counts.to_vec(),
))
}
/// Trait representing a truncation strategy
pub trait TruncationStrategy {
/// Determines the indices of messages to remove to fit within the context limit.
///
/// - `messages`: The list of messages in the conversation.
/// - `token_counts`: A parallel array containing the token count for each message.
/// - `context_limit`: The maximum allowed context length in tokens.
///
/// Returns a vector of indices to remove.
fn determine_indices_to_remove(
&self,
messages: &[Message],
token_counts: &[usize],
context_limit: usize,
) -> Result<HashSet<usize>>;
}
/// Strategy to truncate messages by removing the oldest first
pub struct OldestFirstTruncation;
impl TruncationStrategy for OldestFirstTruncation {
fn determine_indices_to_remove(
&self,
messages: &[Message],
token_counts: &[usize],
context_limit: usize,
) -> Result<HashSet<usize>> {
let mut indices_to_remove = HashSet::new();
let mut total_tokens: usize = token_counts.iter().sum();
let mut tool_ids_to_remove = HashSet::new();
for (i, message) in messages.iter().enumerate() {
if total_tokens <= context_limit {
break;
}
// Remove the message
indices_to_remove.insert(i);
total_tokens -= token_counts[i];
debug!(
"OldestFirst: Removing message at index {}. Tokens removed: {}",
i, token_counts[i]
);
// If it's a ToolRequest or ToolResponse, mark its pair for removal
if message.is_tool_call() || message.is_tool_response() {
message.get_tool_ids().iter().for_each(|id| {
tool_ids_to_remove.insert((i, id.to_string()));
});
}
}
// Now, find and remove paired ToolResponses or ToolRequests
for (i, message) in messages.iter().enumerate() {
let message_tool_ids = message.get_tool_ids();
// Find the other part of the pair - same tool_id but different message index
for (message_idx, tool_id) in &tool_ids_to_remove {
if message_idx != &i && message_tool_ids.contains(tool_id.as_str()) {
indices_to_remove.insert(i);
// No need to check other tool_ids for this message since it's already marked
break;
}
}
}
Ok(indices_to_remove)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::Message;
use anyhow::Result;
use rmcp::model::{CallToolRequestParam, Content};
use rmcp::object;
// Helper function to create a user text message with a specified token count
fn user_text(index: usize, tokens: usize) -> (Message, usize) {
let content = format!("User message {}", index);
(Message::user().with_text(content), tokens)
}
// Helper function to create an assistant text message with a specified token count
fn assistant_text(index: usize, tokens: usize) -> (Message, usize) {
let content = format!("Assistant message {}", index);
(Message::assistant().with_text(content), tokens)
}
// Helper function to create a tool request message with a specified token count
fn assistant_tool_request(
id: &str,
tool_call: CallToolRequestParam,
tokens: usize,
) -> (Message, usize) {
(
Message::assistant().with_tool_request(id, Ok(tool_call)),
tokens,
)
}
// Helper function to create a tool response message with a specified token count
fn user_tool_response(id: &str, result: Vec<Content>, tokens: usize) -> (Message, usize) {
(Message::user().with_tool_response(id, Ok(result)), tokens)
}
// Helper function to create a large tool response with massive content
fn large_tool_response(id: &str, large_text: String, tokens: usize) -> (Message, usize) {
(
Message::user().with_tool_response(id, Ok(vec![Content::text(large_text)])),
tokens,
)
}
// Helper function to create messages with alternating user and assistant
// text messages of a fixed token count
fn create_messages_with_counts(
num_pairs: usize,
tokens: usize,
remove_last: bool,
) -> (Conversation, Vec<usize>) {
let mut messages = Conversation::new_unvalidated((0..num_pairs).flat_map(|i| {
vec![
user_text(i * 2, tokens).0,
assistant_text((i * 2) + 1, tokens).0,
]
}));
if remove_last {
messages.pop();
}
let token_counts = vec![tokens; messages.len()];
(messages, token_counts)
}
#[test]
fn test_handle_oversized_single_message() -> Result<()> {
// Create a scenario similar to the real issue: one very large tool response
let large_content = "A".repeat(50000); // Very large content
let messages = vec![
user_text(1, 10).0,
assistant_tool_request(
"tool1",
CallToolRequestParam {
name: "read_file".into(),
arguments: Some(object!({"path": "large_file.txt"})),
},
20,
)
.0,
large_tool_response("tool1", large_content, 100000).0, // Massive tool response
user_text(2, 10).0,
];
let token_counts = vec![10, 20, 100000, 10]; // One message is huge
let context_limit = 5000; // Much smaller than the large message
let result = truncate_messages(
&messages,
&token_counts,
context_limit,
&OldestFirstTruncation,
);
// Should succeed by truncating the large content
assert!(
result.is_ok(),
"Should handle oversized message by content truncation"
);
let (truncated_messages, truncated_counts) = result.unwrap();
// Should have some messages remaining
assert!(
!truncated_messages.is_empty(),
"Should have some messages left"
);
// Total should be within limit
let total_tokens: usize = truncated_counts.iter().sum();
assert!(
total_tokens <= context_limit,
"Total tokens {} should be <= context limit {}",
total_tokens,
context_limit
);
Ok(())
}
#[test]
fn test_oldest_first_no_truncation() -> Result<()> {
let (messages, token_counts) = create_messages_with_counts(1, 10, false);
let context_limit = 25;
let result = truncate_messages(
messages.messages(),
&token_counts,
context_limit,
&OldestFirstTruncation,
)?;
assert_eq!(result.0.messages(), messages.messages());
assert_eq!(result.1, token_counts);
Ok(())
}
#[test]
fn test_complex_conversation_with_tools() -> Result<()> {
// Simulating a real conversation with multiple tool interactions
let tool_call1 = CallToolRequestParam {
name: "file_read".into(),
arguments: Some(object!({"path": "/tmp/test.txt"})),
};
let tool_call2 = CallToolRequestParam {
name: "database_query".into(),
arguments: Some(object!({"query": "SELECT * FROM users"})),
};
let messages = vec![
user_text(1, 15).0, // Initial user query
assistant_tool_request("tool1", tool_call1.clone(), 20).0,
user_tool_response(
"tool1",
vec![Content::text("File contents".to_string())],
10,
)
.0,
assistant_text(2, 25).0, // Assistant processes file contents
user_text(3, 10).0, // User follow-up
assistant_tool_request("tool2", tool_call2.clone(), 30).0,
user_tool_response(
"tool2",
vec![Content::text("Query results".to_string())],
20,
)
.0,
assistant_text(4, 35).0, // Assistant analyzes query results
user_text(5, 5).0, // Final user confirmation
];
let token_counts = vec![15, 20, 10, 25, 10, 30, 20, 35, 5];
let context_limit = 100; // Force truncation while preserving some tool interactions
let result = truncate_messages(
&messages,
&token_counts,
context_limit,
&OldestFirstTruncation,
)?;
let (truncated_messages, truncated_counts) = result;
// Verify that tool pairs are kept together and the conversation remains coherent
assert!(truncated_messages.len() >= 3); // At least one complete interaction should remain
assert!(truncated_messages.last().unwrap().role == Role::User); // Last message should be from user
// Verify tool pairs are either both present or both removed
let tool_ids: HashSet<_> = truncated_messages
.iter()
.flat_map(|m| m.get_tool_ids())
.collect();
// Each tool ID should appear 0 or 2 times (request + response)
for id in tool_ids {
let count = truncated_messages
.iter()
.flat_map(|m| m.get_tool_ids().into_iter())
.filter(|&tool_id| tool_id == id)
.count();
assert!(count == 0 || count == 2, "Tool pair was split: {}", id);
}
// Total should be within limit
let total_tokens: usize = truncated_counts.iter().sum();
assert!(total_tokens <= context_limit);
Ok(())
}
#[test]
fn test_edge_case_context_window() -> Result<()> {
// Test case where we're exactly at the context limit
let (messages, token_counts) = create_messages_with_counts(2, 25, false);
let context_limit = 100; // Exactly matches total tokens
let result = truncate_messages(
messages.messages(),
&token_counts,
context_limit,
&OldestFirstTruncation,
)?;
let (mut messages, mut token_counts) = result;
assert_eq!(messages.len(), 4); // No truncation needed
assert_eq!(token_counts.iter().sum::<usize>(), 100);
// Now add one more token to force truncation
messages.push(user_text(5, 1).0);
token_counts.push(1);
let result = truncate_messages(
messages.messages(),
&token_counts,
context_limit,
&OldestFirstTruncation,
)?;
let (messages, token_counts) = result;
assert!(token_counts.iter().sum::<usize>() <= context_limit);
assert!(messages.last().unwrap().role == Role::User);
Ok(())
}
#[test]
fn test_multi_tool_chain() -> Result<()> {
// Simulate a chain of dependent tool calls
let tool_calls = vec![
CallToolRequestParam {
name: "git_status".into(),
arguments: Some(object!({})),
},
CallToolRequestParam {
name: "git_diff".into(),
arguments: Some(object!({"file": "main.rs"})),
},
CallToolRequestParam {
name: "git_commit".into(),
arguments: Some(object!({"message": "Update"})),
},
];
let mut messages = Vec::new();
let mut token_counts = Vec::new();
// Build a chain of related tool calls
// 30 tokens each round
for (i, tool_call) in tool_calls.into_iter().enumerate() {
let id = format!("git_{}", i);
messages.push(user_text(i, 10).0);
token_counts.push(10);
messages.push(assistant_tool_request(&id, tool_call, 15).0);
token_counts.push(20);
}
let context_limit = 50; // Force partial truncation
let result = truncate_messages(
&messages,
&token_counts,
context_limit,
&OldestFirstTruncation,
)?;
let (truncated_messages, _) = result;
// Verify that remaining tool chains are complete
let remaining_tool_ids: HashSet<_> = truncated_messages
.iter()
.flat_map(|m| m.get_tool_ids())
.collect();
for _id in remaining_tool_ids {
// Count request/response pairs
let requests = truncated_messages
.iter()
.flat_map(|m| m.get_tool_request_ids().into_iter())
.count();
let responses = truncated_messages
.iter()
.flat_map(|m| m.get_tool_response_ids().into_iter())
.count();
assert_eq!(requests, 1, "Each remaining tool should have one request");
assert_eq!(responses, 1, "Each remaining tool should have one response");
}
Ok(())
}
#[test]
fn test_truncation_with_image_content() -> Result<()> {
// Create a conversation with image content mixed in
let messages = vec![
Message::user().with_image("base64_data", "image/png"), // 50 tokens
Message::assistant().with_text("I see the image"), // 10 tokens
Message::user().with_text("Can you describe it?"), // 10 tokens
Message::assistant().with_text("It shows..."), // 20 tokens
Message::user().with_text("Thanks!"), // 5 tokens
];
let token_counts = vec![50, 10, 10, 20, 5];
let context_limit = 45; // Force truncation
let result = truncate_messages(
&messages,
&token_counts,
context_limit,
&OldestFirstTruncation,
)?;
let (messages, token_counts) = result;
// Verify the conversation still makes sense
assert!(!messages.is_empty());
assert!(messages.last().unwrap().role == Role::User);
assert!(token_counts.iter().sum::<usize>() <= context_limit);
Ok(())
}
#[test]
fn test_error_cases() -> Result<()> {
// Test impossibly small context window
let (messages, token_counts) = create_messages_with_counts(1, 10, false);
let result = truncate_messages(
messages.messages(),
&token_counts,
5, // Impossibly small context
&OldestFirstTruncation,
);
assert!(result.is_err());
// Test unmatched token counts
let messages = vec![user_text(1, 10).0];
let token_counts = vec![10, 10]; // Mismatched length
let result = truncate_messages(&messages, &token_counts, 100, &OldestFirstTruncation);
assert!(result.is_err());
Ok(())
}
}
crates/goose/src/conversation/message.rs
+15 -57
View File
@@ -112,12 +112,7 @@ pub struct FrontendToolRequest {
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
pub struct ContextLengthExceeded {
pub msg: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
pub struct SummarizationRequested {
pub struct ConversationCompacted {
pub msg: String,
}
@@ -133,8 +128,7 @@ pub enum MessageContent {
FrontendToolRequest(FrontendToolRequest),
Thinking(ThinkingContent),
RedactedThinking(RedactedThinkingContent),
ContextLengthExceeded(ContextLengthExceeded),
SummarizationRequested(SummarizationRequested),
ConversationCompacted(ConversationCompacted),
}
impl fmt::Display for MessageContent {
@@ -162,10 +156,7 @@ impl fmt::Display for MessageContent {
},
MessageContent::Thinking(t) => write!(f, "[Thinking: {}]", t.thinking),
MessageContent::RedactedThinking(_r) => write!(f, "[RedactedThinking]"),
MessageContent::ContextLengthExceeded(r) => {
write!(f, "[ContextLengthExceeded: {}]", r.msg)
}
MessageContent::SummarizationRequested(r) => {
MessageContent::ConversationCompacted(r) => {
write!(f, "[SummarizationRequested: {}]", r.msg)
}
}
@@ -246,17 +237,13 @@ impl MessageContent {
})
}
pub fn context_length_exceeded<S: Into<String>>(msg: S) -> Self {
MessageContent::ContextLengthExceeded(ContextLengthExceeded { msg: msg.into() })
}
pub fn summarization_requested<S: Into<String>>(msg: S) -> Self {
MessageContent::SummarizationRequested(SummarizationRequested { msg: msg.into() })
pub fn conversation_compacted<S: Into<String>>(msg: S) -> Self {
MessageContent::ConversationCompacted(ConversationCompacted { msg: msg.into() })
}
// Add this new method to check for summarization requested content
pub fn as_summarization_requested(&self) -> Option<&SummarizationRequested> {
if let MessageContent::SummarizationRequested(ref summarization_requested) = self {
pub fn as_summarization_requested(&self) -> Option<&ConversationCompacted> {
if let MessageContent::ConversationCompacted(ref summarization_requested) = self {
Some(summarization_requested)
} else {
None
@@ -390,10 +377,8 @@ impl From<PromptMessage> for Message {
#[serde(rename_all = "camelCase")]
pub struct MessageMetadata {
/// Whether the message should be visible to the user in the UI
#[serde(default = "default_true")]
pub user_visible: bool,
/// Whether the message should be included in the agent's context window
#[serde(default = "default_true")]
pub agent_visible: bool,
}
@@ -464,28 +449,18 @@ impl MessageMetadata {
}
}
fn default_true() -> bool {
true
}
#[derive(ToSchema, Clone, PartialEq, Serialize, Deserialize, Debug)]
/// A message to or from an LLM
#[serde(rename_all = "camelCase")]
pub struct Message {
pub id: Option<String>,
pub role: Role,
#[serde(default = "default_created")]
pub created: i64,
#[serde(deserialize_with = "deserialize_sanitized_content")]
pub content: Vec<MessageContent>,
#[serde(default)]
pub metadata: MessageMetadata,
}
fn default_created() -> i64 {
0 // old messages do not have timestamps.
}
impl Message {
pub fn new(role: Role, created: i64, content: Vec<MessageContent>) -> Self {
Message {
@@ -605,11 +580,6 @@ impl Message {
self.with_content(MessageContent::redacted_thinking(data))
}
/// Add context length exceeded content to the message
pub fn with_context_length_exceeded<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::context_length_exceeded(msg))
}
/// Get the concatenated text content of the message, separated by newlines
pub fn as_concat_text(&self) -> String {
self.content
@@ -680,9 +650,8 @@ impl Message {
.all(|c| matches!(c, MessageContent::Text(_)))
}
/// Add summarization requested to the message
pub fn with_summarization_requested<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::summarization_requested(msg))
pub fn with_conversation_compacted<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::conversation_compacted(msg))
}
/// Set the visibility metadata for the message
@@ -837,7 +806,8 @@ mod tests {
}
}
}
]
],
"metadata": { "agentVisible": true, "userVisible": true }
}"#;
let message: Message = serde_json::from_str(json_str).unwrap();
@@ -1045,7 +1015,8 @@ mod tests {
"data": "base64data",
"mimeType": "image/png"
}}
]
],
"metadata": {{ "agentVisible": true, "userVisible": true }}
}}"#,
malicious_text
);
@@ -1073,7 +1044,8 @@ mod tests {
"content": [{
"type": "text",
"text": "Hello world 世界 🌍"
}]
}],
"metadata": { "agentVisible": true, "userVisible": true }
}"#;
let message: Message = serde_json::from_str(clean_json).unwrap();
@@ -1142,20 +1114,6 @@ mod tests {
let message: Message = serde_json::from_str(json_with_metadata).unwrap();
assert!(!message.is_user_visible());
assert!(message.is_agent_visible());
// Test without metadata (should use defaults)
let json_without_metadata = r#"{
"role": "user",
"created": 1640995200,
"content": [{
"type": "text",
"text": "Test"
}]
}"#;
let message: Message = serde_json::from_str(json_without_metadata).unwrap();
assert!(message.is_user_visible());
assert!(message.is_agent_visible());
}
#[test]
crates/goose/src/providers/formats/anthropic.rs
+1 -4
View File
@@ -90,10 +90,7 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
MessageContent::ToolConfirmationRequest(_tool_confirmation_request) => {
// Skip tool confirmation requests
}
MessageContent::ContextLengthExceeded(_) => {
// Skip
}
MessageContent::SummarizationRequested(_) => {
MessageContent::ConversationCompacted(_) => {
// Skip
}
MessageContent::Thinking(thinking) => {
crates/goose/src/providers/formats/bedrock.rs
+1 -4
View File
@@ -48,10 +48,7 @@ pub fn to_bedrock_message_content(content: &MessageContent) -> Result<bedrock::C
// Redacted thinking blocks are not supported in Bedrock - skip
bedrock::ContentBlock::Text("".to_string())
}
MessageContent::ContextLengthExceeded(_) => {
bail!("ContextLengthExceeded should not get passed to the provider")
}
MessageContent::SummarizationRequested(_) => {
MessageContent::ConversationCompacted(_) => {
bail!("SummarizationRequested should not get passed to the provider")
}
MessageContent::ToolRequest(tool_req) => {
crates/goose/src/providers/formats/databricks.rs
+1 -4
View File
@@ -127,10 +127,7 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
}
}
}
MessageContent::ContextLengthExceeded(_) => {
continue;
}
MessageContent::SummarizationRequested(_) => {
MessageContent::ConversationCompacted(_) => {
continue;
}
MessageContent::ToolResponse(response) => {
crates/goose/src/providers/formats/openai.rs
+1 -4
View File
@@ -95,10 +95,7 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
// Redacted thinking blocks are not directly used in OpenAI format
continue;
}
MessageContent::ContextLengthExceeded(_) => {
continue;
}
MessageContent::SummarizationRequested(_) => {
MessageContent::ConversationCompacted(_) => {
continue;
}
MessageContent::ToolRequest(request) => match &request.tool_call {
crates/goose/src/providers/formats/snowflake.rs
+1 -4
View File
@@ -53,10 +53,7 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
MessageContent::ToolConfirmationRequest(_) => {
// Skip tool confirmation requests
}
MessageContent::ContextLengthExceeded(_) => {
// Skip
}
MessageContent::SummarizationRequested(_) => {
MessageContent::ConversationCompacted(_) => {
// Skip
}
MessageContent::Thinking(_thinking) => {
crates/goose/src/providers/usage_estimator.rs
+2 -2
View File
@@ -1,6 +1,6 @@
use crate::conversation::message::Message;
use crate::providers::base::ProviderUsage;
use crate::token_counter::create_async_token_counter;
use crate::token_counter::create_token_counter;
use anyhow::Result;
use rmcp::model::Tool;
@@ -17,7 +17,7 @@ pub async fn ensure_usage_tokens(
return Ok(());
}
let token_counter = create_async_token_counter()
let token_counter = create_token_counter()
.await
.map_err(|e| anyhow::anyhow!("Failed to create token counter: {}", e))?;
crates/goose/src/providers/utils.rs
+28 -10
View File
@@ -1,6 +1,8 @@
use super::base::Usage;
use super::errors::GoogleErrorCode;
use crate::config::paths::Paths;
use crate::model::ModelConfig;
use crate::providers::errors::{OpenAIError, ProviderError};
use anyhow::Result;
use base64::Engine;
use regex::Regex;
@@ -12,8 +14,6 @@ use std::io::Read;
use std::path::Path;
use std::time::Duration;
use crate::providers::errors::{OpenAIError, ProviderError};
#[derive(serde::Deserialize)]
struct OpenAIErrorResponse {
error: OpenAIError,
@@ -461,14 +461,32 @@ pub fn emit_debug_trace<T1, T2>(
T1: ?Sized + Serialize,
T2: ?Sized + Serialize,
{
tracing::debug!(
model_config = %serde_json::to_string_pretty(model_config).unwrap_or_default(),
input = %serde_json::to_string_pretty(payload).unwrap_or_default(),
output = %serde_json::to_string_pretty(response).unwrap_or_default(),
input_tokens = ?usage.input_tokens.unwrap_or_default(),
output_tokens = ?usage.output_tokens.unwrap_or_default(),
total_tokens = ?usage.total_tokens.unwrap_or_default(),
);
let logs_dir = Paths::in_state_dir("logs");
if let Err(e) = std::fs::create_dir_all(&logs_dir) {
tracing::warn!("Failed to create logs directory: {}", e);
return;
}
let log_path = |i| logs_dir.join(format!("llm_request.{}.json", i));
for i in (0..4).rev() {
let _ = std::fs::rename(log_path(i), log_path(i + 1));
}
let data = serde_json::json!({
"model_config": model_config,
"input": payload,
"output": response,
"usage": usage,
});
if let Err(e) = std::fs::write(
log_path(0),
serde_json::to_string_pretty(&data).unwrap_or_default(),
) {
tracing::warn!("Failed to write log file: {}", e);
}
}
/// Safely parse a JSON string that may contain doubly-encoded or malformed JSON.
crates/goose/src/scheduler.rs
+6 -8
View File
@@ -1200,7 +1200,7 @@ async fn run_scheduled_job_internal(
}
if let Some(ref prompt_text) = recipe.prompt {
let mut all_session_messages =
let mut conversation =
Conversation::new_unvalidated(vec![Message::user().with_text(prompt_text.clone())]);
let session_config = SessionConfig {
@@ -1213,11 +1213,7 @@ async fn run_scheduled_job_internal(
};
match agent
.reply(
all_session_messages.clone(),
Some(session_config.clone()),
None,
)
.reply(conversation.clone(), Some(session_config.clone()), None)
.await
{
Ok(mut stream) => {
@@ -1231,11 +1227,13 @@ async fn run_scheduled_job_internal(
if msg.role == rmcp::model::Role::Assistant {
tracing::info!("[Job {}] Assistant: {:?}", job.id, msg.content);
}
all_session_messages.push(msg);
conversation.push(msg);
}
Ok(AgentEvent::McpNotification(_)) => {}
Ok(AgentEvent::ModelChange { .. }) => {}
Ok(AgentEvent::HistoryReplaced(_)) => {}
Ok(AgentEvent::HistoryReplaced(updated_conversation)) => {
conversation = updated_conversation;
}
Err(e) => {
tracing::error!(
"[Job {}] Error receiving message from agent: {}",
crates/goose/src/session/session_manager.rs
+31 -10
View File
@@ -18,7 +18,7 @@ use tokio::sync::OnceCell;
use tracing::{info, warn};
use utoipa::ToSchema;
const CURRENT_SCHEMA_VERSION: i32 = 2;
const CURRENT_SCHEMA_VERSION: i32 = 3;
static SESSION_STORAGE: OnceCell<Arc<SessionStorage>> = OnceCell::const_new();
@@ -425,7 +425,8 @@ impl SessionStorage {
content_json TEXT NOT NULL,
created_timestamp INTEGER NOT NULL,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
tokens INTEGER
tokens INTEGER,
metadata_json TEXT
)
"#,
)
@@ -610,6 +611,15 @@ impl SessionStorage {
.execute(&self.pool)
.await?;
}
3 => {
sqlx::query(
r#"
ALTER TABLE messages ADD COLUMN metadata_json TEXT
"#,
)
.execute(&self.pool)
.await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
@@ -768,15 +778,15 @@ impl SessionStorage {
}
async fn get_conversation(&self, session_id: &str) -> Result<Conversation> {
let rows = sqlx::query_as::<_, (String, String, i64)>(
"SELECT role, content_json, created_timestamp FROM messages WHERE session_id = ? ORDER BY timestamp",
let rows = sqlx::query_as::<_, (String, String, i64, Option<String>)>(
"SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp",
)
.bind(session_id)
.fetch_all(&self.pool)
.await?;
let mut messages = Vec::new();
for (role_str, content_json, created_timestamp) in rows {
for (role_str, content_json, created_timestamp, metadata_json) in rows {
let role = match role_str.as_str() {
"user" => Role::User,
"assistant" => Role::Assistant,
@@ -784,7 +794,12 @@ impl SessionStorage {
};
let content = serde_json::from_str(&content_json)?;
let message = Message::new(role, created_timestamp, content);
let metadata = metadata_json
.and_then(|json| serde_json::from_str(&json).ok())
.unwrap_or_default();
let mut message = Message::new(role, created_timestamp, content);
message.metadata = metadata;
messages.push(message);
}
@@ -792,16 +807,19 @@ impl SessionStorage {
}
async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> {
let metadata_json = serde_json::to_string(&message.metadata)?;
sqlx::query(
r#"
INSERT INTO messages (session_id, role, content_json, created_timestamp)
VALUES (?, ?, ?, ?)
INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json)
VALUES (?, ?, ?, ?, ?)
"#,
)
.bind(session_id)
.bind(role_to_string(&message.role))
.bind(serde_json::to_string(&message.content)?)
.bind(message.created)
.bind(metadata_json)
.execute(&self.pool)
.await?;
@@ -826,16 +844,19 @@ impl SessionStorage {
.await?;
for message in conversation.messages() {
let metadata_json = serde_json::to_string(&message.metadata)?;
sqlx::query(
r#"
INSERT INTO messages (session_id, role, content_json, created_timestamp)
VALUES (?, ?, ?, ?)
INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json)
VALUES (?, ?, ?, ?, ?)
"#,
)
.bind(session_id)
.bind(role_to_string(&message.role))
.bind(serde_json::to_string(&message.content)?)
.bind(message.created)
.bind(metadata_json)
.execute(&mut *tx)
.await?;
}
crates/goose/src/token_counter.rs
+14 -247
View File
@@ -20,16 +20,12 @@ const ENUM_INIT: isize = -3;
const ENUM_ITEM: usize = 3;
const FUNC_END: usize = 12;
pub struct AsyncTokenCounter {
pub struct TokenCounter {
tokenizer: Arc<CoreBPE>,
token_cache: Arc<DashMap<u64, usize>>,
}
pub struct TokenCounter {
tokenizer: Arc<CoreBPE>,
}
impl AsyncTokenCounter {
impl TokenCounter {
pub async fn new() -> Result<Self, String> {
let tokenizer = get_tokenizer().await?;
Ok(Self {
@@ -130,6 +126,9 @@ impl AsyncTokenCounter {
}
for message in messages {
if !message.metadata.agent_visible {
continue;
}
num_tokens += tokens_per_message;
for content in &message.content {
if let Some(content_text) = content.as_text() {
@@ -183,136 +182,6 @@ impl AsyncTokenCounter {
}
}
impl Default for TokenCounter {
fn default() -> Self {
Self::new()
}
}
impl TokenCounter {
pub fn new() -> Self {
let tokenizer = get_tokenizer_blocking().expect("Failed to initialize tokenizer");
Self { tokenizer }
}
pub fn count_tokens(&self, text: &str) -> usize {
let tokens = self.tokenizer.encode_with_special_tokens(text);
tokens.len()
}
pub fn count_tokens_for_tools(&self, tools: &[Tool]) -> usize {
let mut func_token_count = 0;
if !tools.is_empty() {
for tool in tools {
func_token_count += FUNC_INIT;
let name = &tool.name;
let description = &tool
.description
.as_ref()
.map(|d| d.as_ref())
.unwrap_or_default()
.trim_end_matches('.');
let line = format!("{}:{}", name, description);
func_token_count += self.count_tokens(&line);
if let Some(serde_json::Value::Object(properties)) =
tool.input_schema.get("properties")
{
if !properties.is_empty() {
func_token_count += PROP_INIT;
for (key, value) in properties {
func_token_count += PROP_KEY;
let p_name = key;
let p_type = value.get("type").and_then(|v| v.as_str()).unwrap_or("");
let p_desc = value
.get("description")
.and_then(|v| v.as_str())
.unwrap_or("")
.trim_end_matches('.');
let line = format!("{}:{}:{}", p_name, p_type, p_desc);
func_token_count += self.count_tokens(&line);
if let Some(enum_values) = value.get("enum").and_then(|v| v.as_array())
{
func_token_count =
func_token_count.saturating_add_signed(ENUM_INIT);
for item in enum_values {
if let Some(item_str) = item.as_str() {
func_token_count += ENUM_ITEM;
func_token_count += self.count_tokens(item_str);
}
}
}
}
}
}
}
func_token_count += FUNC_END;
}
func_token_count
}
pub fn count_chat_tokens(
&self,
system_prompt: &str,
messages: &[Message],
tools: &[Tool],
) -> usize {
let tokens_per_message = 4;
let mut num_tokens = 0;
if !system_prompt.is_empty() {
num_tokens += self.count_tokens(system_prompt) + tokens_per_message;
}
for message in messages {
num_tokens += tokens_per_message;
for content in &message.content {
if let Some(content_text) = content.as_text() {
num_tokens += self.count_tokens(content_text);
} else if let Some(tool_request) = content.as_tool_request() {
if let Ok(tool_call) = tool_request.tool_call.as_ref() {
let text = format!(
"{}:{}:{:?}",
tool_request.id, tool_call.name, tool_call.arguments
);
num_tokens += self.count_tokens(&text);
}
} else if let Some(tool_response_text) = content.as_tool_response_text() {
num_tokens += self.count_tokens(&tool_response_text);
} else {
continue;
}
}
}
if !tools.is_empty() {
num_tokens += self.count_tokens_for_tools(tools);
}
num_tokens += 3;
num_tokens
}
pub fn count_everything(
&self,
system_prompt: &str,
messages: &[Message],
tools: &[Tool],
resources: &[String],
) -> usize {
let mut num_tokens = self.count_chat_tokens(system_prompt, messages, tools);
if !resources.is_empty() {
for resource in resources {
num_tokens += self.count_tokens(resource);
}
}
num_tokens
}
}
async fn get_tokenizer() -> Result<Arc<CoreBPE>, String> {
let tokenizer = TOKENIZER
.get_or_init(|| async {
@@ -325,99 +194,17 @@ async fn get_tokenizer() -> Result<Arc<CoreBPE>, String> {
Ok(tokenizer.clone())
}
fn get_tokenizer_blocking() -> Result<Arc<CoreBPE>, String> {
if let Some(tokenizer) = TOKENIZER.get() {
return Ok(tokenizer.clone());
}
match tiktoken_rs::o200k_base() {
Ok(bpe) => {
let tokenizer = Arc::new(bpe);
let _ = TOKENIZER.set(tokenizer.clone());
Ok(tokenizer)
}
Err(e) => Err(format!("Failed to initialize o200k_base tokenizer: {}", e)),
}
}
pub async fn create_async_token_counter() -> Result<AsyncTokenCounter, String> {
AsyncTokenCounter::new().await
pub async fn create_token_counter() -> Result<TokenCounter, String> {
TokenCounter::new().await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conversation::message::{Message, MessageContent};
use rmcp::model::{Role, Tool};
use rmcp::object;
#[test]
fn test_count_chat_tokens() {
let counter = TokenCounter::new();
let system_prompt =
"You are a helpful assistant that can answer questions about the weather.";
let messages = vec![
Message::new(
Role::User,
0,
vec![MessageContent::text(
"What's the weather like in San Francisco?",
)],
),
Message::new(
Role::Assistant,
1,
vec![MessageContent::text(
"Looks like it's 60 degrees Fahrenheit in San Francisco.",
)],
),
Message::new(
Role::User,
2,
vec![MessageContent::text("How about New York?")],
),
];
let tools = vec![Tool::new(
"get_current_weather",
"Get the current weather in a given location",
object!({
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA"
},
"unit": {
"type": "string",
"description": "The unit of temperature to return",
"enum": ["celsius", "fahrenheit"]
}
},
"required": ["location"]
}),
)];
let token_count_without_tools = counter.count_chat_tokens(system_prompt, &messages, &[]);
println!("Total tokens without tools: {}", token_count_without_tools);
let token_count_with_tools = counter.count_chat_tokens(system_prompt, &messages, &tools);
println!("Total tokens with tools: {}", token_count_with_tools);
assert!(
token_count_without_tools > 0,
"Should have some tokens without tools"
);
assert!(
token_count_with_tools > token_count_without_tools,
"Should have more tokens with tools"
);
}
#[tokio::test]
async fn test_async_token_caching() {
let counter = create_async_token_counter().await.unwrap();
async fn test_token_caching() {
let counter = create_token_counter().await.unwrap();
let text = "This is a test for caching functionality";
@@ -434,8 +221,8 @@ mod tests {
}
#[tokio::test]
async fn test_async_cache_management() {
let counter = create_async_token_counter().await.unwrap();
async fn test_cache_management() {
let counter = create_token_counter().await.unwrap();
counter.count_tokens("First text");
counter.count_tokens("Second text");
@@ -454,7 +241,7 @@ mod tests {
#[tokio::test]
async fn test_concurrent_token_counter_creation() {
let handles: Vec<_> = (0..10)
.map(|_| tokio::spawn(async { create_async_token_counter().await.unwrap() }))
.map(|_| tokio::spawn(async { create_token_counter().await.unwrap() }))
.collect();
let counters: Vec<_> = futures::future::join_all(handles)
@@ -473,7 +260,7 @@ mod tests {
#[tokio::test]
async fn test_cache_eviction_behavior() {
let counter = create_async_token_counter().await.unwrap();
let counter = create_token_counter().await.unwrap();
let mut cached_texts = Vec::new();
for i in 0..50 {
@@ -493,7 +280,7 @@ mod tests {
#[tokio::test]
async fn test_concurrent_cache_operations() {
let counter = std::sync::Arc::new(create_async_token_counter().await.unwrap());
let counter = std::sync::Arc::new(create_token_counter().await.unwrap());
let handles: Vec<_> = (0..20)
.map(|i| {
@@ -518,24 +305,4 @@ mod tests {
assert!(counter.cache_size() > 0);
assert!(counter.cache_size() <= MAX_TOKEN_CACHE_SIZE);
}
#[test]
fn test_tokenizer_consistency() {
let sync_counter = TokenCounter::new();
let text = "This is a test for tokenizer consistency";
let sync_count = sync_counter.count_tokens(text);
assert!(sync_count > 0, "Sync tokenizer should produce tokens");
let short_text = "Hi";
let long_text = "This is a much longer text that should produce significantly more tokens than the short text";
let short_count = sync_counter.count_tokens(short_text);
let long_count = sync_counter.count_tokens(long_text);
assert!(
short_count < long_count,
"Longer text should have more tokens"
);
}
}
crates/goose/tests/agent.rs
+7 -13
View File
@@ -117,7 +117,7 @@ async fn run_truncate_test(
agent.update_provider(provider).await?;
let repeat_count = context_window + 10_000;
let large_message_content = "hello ".repeat(repeat_count);
let messages = Conversation::new(vec![
let conversation = Conversation::new(vec![
Message::user().with_text("hi there. what is 2 + 2?"),
Message::assistant().with_text("hey! I think it's 4."),
Message::user().with_text(&large_message_content),
@@ -130,7 +130,7 @@ async fn run_truncate_test(
])
.unwrap();
let reply_stream = agent.reply(messages, None, None).await?;
let reply_stream = agent.reply(conversation, None, None).await?;
tokio::pin!(reply_stream);
let mut responses = Vec::new();
@@ -143,8 +143,8 @@ async fn run_truncate_test(
Ok(AgentEvent::ModelChange { .. }) => {
// Model change events are informational, just continue
}
Ok(AgentEvent::HistoryReplaced(_)) => {
// Handle history replacement events if needed
Ok(AgentEvent::HistoryReplaced(_updated_conversation)) => {
// Should update the conversation here, but we're not reading it
}
Err(e) => {
println!("Error: {:?}", e);
@@ -175,14 +175,6 @@ async fn run_truncate_test(
assert!(text_content.text.to_lowercase().contains("no"));
assert!(!text_content.text.to_lowercase().contains("yes"));
}
goose::conversation::message::MessageContent::ContextLengthExceeded(_) => {
// This is an acceptable outcome for providers that don't truncate themselves
// and correctly report that the context length was exceeded.
println!(
"Received ContextLengthExceeded as expected for {:?}",
provider_type
);
}
_ => {
panic!(
"Unexpected message content type: {:?}",
@@ -1119,7 +1111,9 @@ mod max_turns_tests {
}
Ok(AgentEvent::McpNotification(_)) => {}
Ok(AgentEvent::ModelChange { .. }) => {}
Ok(AgentEvent::HistoryReplaced(_)) => {}
Ok(AgentEvent::HistoryReplaced(_updated_conversation)) => {
// We should update the conversation here, but we're not reading it
}
Err(e) => {
return Err(e);
}