tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 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)
})
}