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

Optimize tool summarization (#7938)

Browse Source 
Co-authored-by: Michael Neale <michael.neale@gmail.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2026-03-18 09:09:18 -04:00
committed by GitHub
parent 493566dff2
commit 475968db64
7 changed files with 458 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+68 -32
View File
@@ -363,15 +363,29 @@ impl Agent {
self.tool_inspection_manager.apply_tool_annotations(&tools);
}
let tool_call_cut_off = match Config::global().get_param::<usize>("GOOSE_TOOL_CALL_CUTOFF")
{
Ok(v) => v,
Err(_) => {
let context_limit = self
.provider()
.await
.map(|p| p.get_model_config().context_limit())
.unwrap_or(crate::model::DEFAULT_CONTEXT_LIMIT);
let compaction_threshold = Config::global()
.get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
.unwrap_or(crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD);
crate::context_mgmt::compute_tool_call_cutoff(context_limit, compaction_threshold)
}
};
Ok(ReplyContext {
conversation,
tools,
toolshim_tools,
system_prompt,
goose_mode,
tool_call_cut_off: Config::global()
.get_param::<usize>("GOOSE_TOOL_CALL_CUTOFF")
.unwrap_or(10),
tool_call_cut_off,
initial_messages,
})
}
@@ -1127,6 +1141,15 @@ impl Agent {
});
}
// Count tool calls present before this reply — everything added during
// the reply loop is part of the current turn and should not be summarized.
let pre_turn_tool_count = conversation
.messages()
.iter()
.flat_map(|m| m.content.iter())
.filter(|c| matches!(c, MessageContent::ToolRequest(_)))
.count();
let working_dir = session.working_dir.clone();
let reply_stream_span = tracing::info_span!(target: "goose::agents::agent", "reply_stream", session.id = %session_config.id);
let inner = Box::pin(async_stream::try_stream! {
@@ -1162,13 +1185,6 @@ impl Agent {
break;
}
let tool_pair_summarization_task = crate::context_mgmt::maybe_summarize_tool_pair(
self.provider().await?,
session_config.id.clone(),
conversation.clone(),
tool_call_cut_off,
);
let conversation_with_moim = super::moim::inject_moim(
&session_config.id,
conversation.clone(),
@@ -1185,6 +1201,20 @@ impl Agent {
&toolshim_tools,
).await?;
let current_turn_tool_count = conversation.messages().iter()
.flat_map(|m| m.content.iter())
.filter(|c| matches!(c, MessageContent::ToolRequest(_)))
.count()
.saturating_sub(pre_turn_tool_count);
let tool_pair_summarization_task = crate::context_mgmt::maybe_summarize_tool_pairs(
self.provider().await?,
session_config.id.clone(),
conversation.clone(),
tool_call_cut_off,
current_turn_tool_count,
);
let mut no_tools_called = true;
let mut messages_to_add = Conversation::default();
let mut tools_updated = false;
@@ -1641,34 +1671,40 @@ impl Agent {
}
}
if let Ok(Some((summary_msg, tool_id))) = tool_pair_summarization_task.await {
if is_token_cancelled(&cancel_token) {
tool_pair_summarization_task.abort();
}
if let Ok(summaries) = tool_pair_summarization_task.await {
let mut updated_messages = conversation.messages().clone();
let matching: Vec<&mut Message> = updated_messages
.iter_mut()
.filter(|msg| {
msg.id.is_some() && msg.content.iter().any(|c| match c {
MessageContent::ToolRequest(req) => req.id == tool_id,
MessageContent::ToolResponse(resp) => resp.id == tool_id,
_ => false,
for (summary_msg, tool_id) in summaries {
let matching: Vec<&mut Message> = updated_messages
.iter_mut()
.filter(|msg| {
msg.id.is_some() && msg.content.iter().any(|c| match c {
MessageContent::ToolRequest(req) => req.id == tool_id,
MessageContent::ToolResponse(resp) => resp.id == tool_id,
_ => false,
})
})
})
.collect();
.collect();
if matching.len() == 2 {
for msg in matching {
let id = msg.id.as_ref().unwrap();
msg.metadata = msg.metadata.with_agent_invisible();
SessionManager::update_message_metadata(&session_config.id, id, |metadata| {
metadata.with_agent_invisible()
}).await?;
if matching.len() == 2 {
for msg in matching {
let id = msg.id.as_ref().unwrap();
msg.metadata = msg.metadata.with_agent_invisible();
SessionManager::update_message_metadata(&session_config.id, id, |metadata| {
metadata.with_agent_invisible()
}).await?;
}
messages_to_add.push(summary_msg);
} else {
warn!("Expected a tool request/reply pair, but found {} matching messages",
matching.len());
}
conversation = Conversation::new_unvalidated(updated_messages);
messages_to_add.push(summary_msg);
} else {
warn!("Expected a tool request/reply pair, but found {} matching messages",
matching.len());
}
conversation = Conversation::new_unvalidated(updated_messages);
}
for msg in &messages_to_add {