tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 31 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-cli/src/session/export.rs
+1 -1
View File
@@ -369,7 +369,7 @@ pub fn message_to_markdown(message: &Message, export_all_content: bool) -> Strin
md.push_str("**Thinking:**\n");
md.push_str("> *Thinking was redacted*\n\n");
}
MessageContent::SummarizationRequested(summarization) => {
MessageContent::ConversationCompacted(summarization) => {
md.push_str(&format!("*{}*\n\n", summarization.msg));
}
_ => {
crates/goose-cli/src/session/mod.rs
+25 -170
View File
@@ -47,6 +47,7 @@ use std::sync::Arc;
use std::time::Instant;
use tokio;
use tokio_util::sync::CancellationToken;
use tracing::warn;
pub enum RunMode {
Normal,
@@ -158,19 +159,6 @@ impl CliSession {
self.session_id.as_ref()
}
async fn summarize_context_messages(
messages: &mut Conversation,
agent: &Agent,
message_suffix: &str,
) -> Result<()> {
let (summarized_messages, _, _) = agent.summarize_context(messages.messages()).await?;
let msg = format!("Context maxed out\n{}\n{}", "-".repeat(50), message_suffix);
output::render_text(&msg, Some(Color::Yellow), true);
*messages = summarized_messages;
Ok(())
}
/// Add a stdio extension to the session
///
/// # Arguments
@@ -674,10 +662,24 @@ impl CliSession {
// Get the provider for summarization
let _provider = self.agent.provider().await?;
// Get session metadata if available
let session_metadata_for_compact =
if let Some(ref session_id) = self.session_id {
SessionManager::get_session(session_id, false).await.ok()
} else {
None
};
// Call the summarize_context method
let (summarized_messages, _token_counts, summarization_usage) = self
.agent
.summarize_context(self.messages.messages())
let (_, summarized_messages, _token_counts, summarization_usage) =
goose::context_mgmt::check_and_compact_messages(
&self.agent,
self.messages.messages(),
true,
false,
None,
session_metadata_for_compact.as_ref(),
)
.await?;
// Update the session messages with the summarized ones
@@ -937,84 +939,7 @@ impl CliSession {
permission,
},).await;
}
} else if let Some(MessageContent::ContextLengthExceeded(_)) = message.content.first() {
output::hide_thinking();
// Check for user-configured default context strategy
let config = Config::global();
let context_strategy = config.get_param::<String>("GOOSE_CONTEXT_STRATEGY")
.unwrap_or_else(|_| if interactive { "prompt".to_string() } else { "summarize".to_string() });
let selected = match context_strategy.as_str() {
"clear" => "clear",
"truncate" => "truncate",
"summarize" => "summarize",
_ => {
if interactive {
// In interactive mode with no default, ask the user what to do
let prompt = "The model's context length is maxed out. You will need to reduce the # msgs. Do you want to?".to_string();
cliclack::select(prompt)
.item("clear", "Clear Session", "Removes all messages from Goose's memory")
.item("truncate", "Truncate Messages", "Removes old messages till context is within limits")
.item("summarize", "Summarize Session", "Summarize the session to reduce context length")
.interact()?
} else {
// In headless mode, default to summarize
"summarize"
}
}
};
match selected {
"clear" => {
self.messages.clear();
let msg = if context_strategy == "clear" {
format!("Context maxed out - automatically cleared session.\n{}", "-".repeat(50))
} else {
format!("Session cleared.\n{}", "-".repeat(50))
};
output::render_text(&msg, Some(Color::Yellow), true);
break; // exit the loop to hand back control to the user
}
"truncate" => {
// Truncate messages to fit within context length
let (truncated_messages, _) = self.agent.truncate_context(self.messages.messages()).await?;
let msg = if context_strategy == "truncate" {
format!("Context maxed out - automatically truncated messages.\n{}\nGoose tried its best to truncate messages for you.", "-".repeat(50))
} else {
format!("Context maxed out\n{}\nGoose tried its best to truncate messages for you.", "-".repeat(50))
};
output::render_text("", Some(Color::Yellow), true);
output::render_text(&msg, Some(Color::Yellow), true);
self.messages = truncated_messages;
}
"summarize" => {
// Use the helper function to summarize context
let message_suffix = if context_strategy == "summarize" {
"Goose automatically summarized messages for you."
} else if interactive {
"Goose summarized messages for you."
} else {
"Goose automatically summarized messages to continue processing."
};
Self::summarize_context_messages(&mut self.messages, &self.agent, message_suffix).await?;
}
_ => {
unreachable!()
}
}
// Restart the stream after handling ContextLengthExceeded
stream = self
.agent
.reply(
self.messages.clone(),
session_config.clone(),
None
)
.await?;
}
// otherwise we have a model/tool to render
else {
for content in &message.content {
if let MessageContent::ToolRequest(tool_request) = content {
@@ -1171,10 +1096,10 @@ impl CliSession {
_ => (),
}
}
Some(Ok(AgentEvent::HistoryReplaced(new_messages))) => {
self.messages = Conversation::new_unvalidated(new_messages.clone());
}
Some(Ok(AgentEvent::ModelChange { model, mode })) => {
Some(Ok(AgentEvent::HistoryReplaced(updated_conversation))) => {
self.messages = updated_conversation;
}
Some(Ok(AgentEvent::ModelChange { model, mode })) => {
// Log model change if in debug mode
if self.debug {
eprintln!("Model changed to {} in {} mode", model, mode);
@@ -1182,94 +1107,24 @@ impl CliSession {
}
Some(Err(e)) => {
// TODO(Douwe): Delete this
// Check if it's a ProviderError::ContextLengthExceeded
if e.downcast_ref::<goose::providers::errors::ProviderError>()
.map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_)))
.unwrap_or(false) {
output::render_text(
"Context limit reached. Performing auto-compaction...",
"Compaction requested. Should have happened in the agent!",
Some(Color::Yellow),
true
);
// Try auto-compaction first - keep the stream alive!
if let Ok(compact_result) = goose::context_mgmt::auto_compact::perform_compaction(&self.agent, self.messages.messages()).await {
self.messages = compact_result.messages;
if let Some(session_id) = &self.session_id {
SessionManager::replace_conversation(session_id, &self.messages).await?;
}
output::render_text(
"Compaction complete. Conversation has been automatically compacted to continue.",
Some(Color::Yellow),
true
);
// Restart the stream after successful compaction - keep the stream alive!
stream = self
.agent
.reply(
self.messages.clone(),
session_config.clone(),
Some(cancel_token.clone())
)
.await?;
continue;
}
// Auto-compaction failed, fall through to common error handling below
warn!("Compaction requested. Should have happened in the agent!");
}
eprintln!("Error: {}", e);
cancel_token_clone.cancel();
drop(stream);
if let Err(e) = self.handle_interrupted_messages(false).await {
eprintln!("Error handling interruption: {}", e);
}
// Check if it's a ProviderError::ContextLengthExceeded
if e.downcast_ref::<goose::providers::errors::ProviderError>()
.map(|provider_error| matches!(provider_error, goose::providers::errors::ProviderError::ContextLengthExceeded(_)))
.unwrap_or(false) {
output::render_error(&format!("Error: Context length exceeded: {}", e));
let prompt = "The tool calling loop was interrupted. How would you like to proceed?";
let selected = match cliclack::select(prompt.to_string())
.item("clear", "Clear Session", "Removes all messages from Goose's memory")
.item("summarize", "Summarize Session", "Summarize the session to reduce context length")
.interact()
{
Ok(choice) => Some(choice),
Err(e) => {
if e.kind() == std::io::ErrorKind::Interrupted {
// If interrupted, do nothing and let user handle it manually
output::render_text("Operation cancelled. You can use /clear or /summarize to continue.", Some(Color::Yellow), true);
None
} else {
return Err(e.into());
}
}
};
if let Some(choice) = selected {
match choice {
"clear" => {
self.messages.clear();
let msg = format!("Session cleared.\n{}", "-".repeat(50));
output::render_text(&msg, Some(Color::Yellow), true);
}
"summarize" => {
// Use the helper function to summarize context
let message_suffix = "Goose summarized messages for you.";
if let Err(e) = Self::summarize_context_messages(&mut self.messages, &self.agent, message_suffix).await {
output::render_error(&format!("Failed to summarize: {}", e));
output::render_text("Consider using /clear to start fresh.", Some(Color::Yellow), true);
}
}
_ => {
unreachable!()
}
}
}
} else {
output::render_error(
"The error above was an exception we were not able to handle.\n\
crates/goose-cli/src/session/output.rs
+1 -1
View File
@@ -185,7 +185,7 @@ pub fn render_message(message: &Message, debug: bool) {
println!("\n{}", style("Thinking:").dim().italic());
print_markdown("Thinking was redacted", theme);
}
MessageContent::SummarizationRequested(summarization) => {
MessageContent::ConversationCompacted(summarization) => {
println!("\n{}", style(&summarization.msg).yellow());
}
_ => {