tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 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/context_mgmt/mod.rs
+145 -144
View File
@@ -18,10 +18,13 @@ use tracing::log::warn;
pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.8;
/// Feature flag to enable/disable tool pair summarization.
/// Set to `false` to disable summarizing old tool call/response pairs.
/// TODO: Re-enable once tool summarization stability issues are resolved.
const ENABLE_TOOL_PAIR_SUMMARIZATION: bool = false;
const TOOLCALL_SUMMARIZATION_BATCH_SIZE: usize = 10;
fn tool_pair_summarization_enabled() -> bool {
Config::global()
.get_param::<bool>("GOOSE_TOOL_PAIR_SUMMARIZATION")
.unwrap_or(true)
}
const CONVERSATION_CONTINUATION_TEXT: &str =
"Your context was compacted. The previous message contains a summary of the conversation so far.
@@ -418,13 +421,24 @@ fn format_message_for_compacting(msg: &Message) -> String {
}
}
/// Find the id of a tool call to summarize. We only do this if we have more than
/// cutoff tool calls that aren't summarized yet
pub fn tool_id_to_summarize(conversation: &Conversation, cutoff: usize) -> Option<String> {
pub fn compute_tool_call_cutoff(context_limit: usize, compaction_threshold: f64) -> usize {
let threshold = if compaction_threshold > 0.0 && compaction_threshold <= 1.0 {
compaction_threshold
} else {
DEFAULT_COMPACTION_THRESHOLD
};
let effective_limit = (context_limit as f64 * threshold) as usize;
(3 * effective_limit / 20_000).clamp(10, 500)
}
pub fn tool_ids_to_summarize(
conversation: &Conversation,
cutoff: usize,
protect_last_n: usize,
) -> Vec<String> {
let messages = conversation.messages();
let mut tool_call_count = 0;
let mut first_tool_call_id = None;
let mut tool_call_ids: Vec<String> = Vec::new();
for msg in messages.iter() {
if !msg.is_agent_visible() {
@@ -433,17 +447,21 @@ pub fn tool_id_to_summarize(conversation: &Conversation, cutoff: usize) -> Optio
for content in &msg.content {
if let MessageContent::ToolRequest(req) = content {
if first_tool_call_id.is_none() {
first_tool_call_id = Some(req.id.clone());
}
tool_call_count += 1;
if tool_call_count > cutoff {
return first_tool_call_id;
}
tool_call_ids.push(req.id.clone());
}
}
}
None
// Never summarize the last N tool calls (current turn)
let eligible = tool_call_ids.len().saturating_sub(protect_last_n);
if eligible <= cutoff + TOOLCALL_SUMMARIZATION_BATCH_SIZE {
return Vec::new();
}
tool_call_ids
.into_iter()
.take(TOOLCALL_SUMMARIZATION_BATCH_SIZE)
.collect()
}
pub async fn summarize_tool_call(
@@ -482,17 +500,16 @@ pub async fn summarize_tool_call(
let summarization_request = vec![user_message];
let system_prompt = indoc! {r#"
Your task is to summarize a tool call & response pair to save tokens
Your task is to summarize a tool call & response pair to save tokens.
reply with a single message that describe what happened. Typically a toolcall
is asks for something using a bunch of parameters and then the result is also some
Reply with a single message that describes what happened. Typically a tool call
asks for something using a bunch of parameters and then the result is also some
structured output. So the tool might ask to look up something on github and the
reply might be a json document. So you could reply with something like:
"A call to github was made to get the project status"
if that is what it was.
"#};
let (mut response, _) = provider
@@ -506,31 +523,30 @@ pub async fn summarize_tool_call(
Ok(response.with_generated_id())
}
pub fn maybe_summarize_tool_pair(
pub fn maybe_summarize_tool_pairs(
provider: Arc<dyn Provider>,
session_id: String,
conversation: Conversation,
cutoff: usize,
) -> JoinHandle<Option<(Message, String)>> {
protect_last_n: usize,
) -> JoinHandle<Vec<(Message, String)>> {
tokio::spawn(async move {
// Tool pair summarization is currently disabled via feature flag.
// See ENABLE_TOOL_PAIR_SUMMARIZATION constant above.
if !ENABLE_TOOL_PAIR_SUMMARIZATION {
return None;
if !tool_pair_summarization_enabled() || provider.manages_own_context() {
return Vec::new();
}
if let Some(tool_id) = tool_id_to_summarize(&conversation, cutoff) {
let tool_ids = tool_ids_to_summarize(&conversation, cutoff, protect_last_n);
let mut results = Vec::new();
for tool_id in tool_ids {
match summarize_tool_call(provider.as_ref(), &session_id, &conversation, &tool_id).await
{
Ok(summary) => Some((summary, tool_id)),
Ok(summary) => results.push((summary, tool_id)),
Err(e) => {
warn!("Failed to summarize tool pair: {}", e);
None
}
}
} else {
None
}
results
})
}
@@ -544,6 +560,30 @@ mod tests {
use async_trait::async_trait;
use rmcp::model::{AnnotateAble, CallToolRequestParams, RawContent, Tool};
fn create_tool_pair(
call_id: &str,
response_id: &str,
tool_name: &str,
response_text: &str,
) -> Vec<Message> {
vec![
Message::assistant()
.with_tool_request(
call_id,
Ok(CallToolRequestParams::new(tool_name.to_string())),
)
.with_id(call_id),
Message::user()
.with_tool_response(
call_id,
Ok(rmcp::model::CallToolResult::success(vec![
RawContent::text(response_text).no_annotation(),
])),
)
.with_id(response_id),
]
}
struct MockProvider {
message: Message,
config: ModelConfig,
@@ -677,125 +717,86 @@ mod tests {
);
}
#[tokio::test]
async fn test_tool_pair_summarization_workflow() {
fn create_tool_pair(
call_id: &str,
response_id: &str,
tool_name: &str,
response_text: &str,
) -> Vec<Message> {
vec![
Message::assistant()
.with_tool_request(
call_id,
Ok(CallToolRequestParams::new(tool_name.to_string())),
)
.with_id(call_id),
Message::user()
.with_tool_response(
call_id,
Ok(rmcp::model::CallToolResult::success(vec![
RawContent::text(response_text).no_annotation(),
])),
)
.with_id(response_id),
]
#[test]
fn test_compute_tool_call_cutoff_scales_with_context() {
// Default threshold (0.8)
assert_eq!(compute_tool_call_cutoff(128_000, 0.8), 15); // 102K effective
assert_eq!(compute_tool_call_cutoff(200_000, 0.8), 24); // 160K effective
assert_eq!(compute_tool_call_cutoff(1_000_000, 0.8), 120); // 800K effective
// Clamp at minimum
assert_eq!(compute_tool_call_cutoff(50_000, 0.8), 10);
assert_eq!(compute_tool_call_cutoff(10_000, 0.8), 10);
// Clamp at maximum (500)
assert_eq!(compute_tool_call_cutoff(10_000_000, 0.8), 500);
// Lower compaction threshold means earlier summarization
assert_eq!(compute_tool_call_cutoff(200_000, 0.3), 10); // 60K effective
assert_eq!(compute_tool_call_cutoff(1_000_000, 0.5), 75); // 500K effective
// Invalid threshold falls back to default 0.8
assert_eq!(compute_tool_call_cutoff(200_000, 0.0), 24); // falls back to 0.8
assert_eq!(compute_tool_call_cutoff(200_000, -1.0), 24); // falls back to 0.8
}
#[test]
fn test_tool_ids_to_summarize_triggers_at_cutoff_plus_batch() {
// cutoff=5, so we need >5+10=15 to trigger. 15 exactly should NOT trigger.
let mut messages = vec![Message::user().with_text("hello")];
for i in 0..15 {
messages.extend(create_tool_pair(
&format!("call{}", i),
&format!("resp{}", i),
"read_file",
"content",
));
}
let conversation = Conversation::new_unvalidated(messages);
let result = tool_ids_to_summarize(&conversation, 5, 0);
assert!(result.is_empty(), "Exactly cutoff+batch should not trigger");
let summary_response = Message::assistant()
.with_text("Tool call to list files and response with file listing");
let provider = MockProvider::new(summary_response, 1000);
let mut messages = vec![Message::user().with_text("list files").with_id("msg_1")];
messages.extend(create_tool_pair(
"call1",
"response1",
"shell",
"file1.txt\nfile2.txt",
));
messages.extend(create_tool_pair(
"call2",
"response2",
"read_file",
"content of file1",
));
messages.extend(create_tool_pair(
"call3",
"response3",
"read_file",
"content of file2",
));
// 16 tool calls: now exceeds cutoff+10, should return a batch of 10
let mut messages = vec![Message::user().with_text("hello")];
for i in 0..16 {
messages.extend(create_tool_pair(
&format!("call{}", i),
&format!("resp{}", i),
"read_file",
"content",
));
}
let conversation = Conversation::new_unvalidated(messages);
let result = tool_ids_to_summarize(&conversation, 5, 0);
assert_eq!(result.len(), TOOLCALL_SUMMARIZATION_BATCH_SIZE);
assert_eq!(result[0], "call0");
assert_eq!(result[9], "call9");
}
#[test]
fn test_tool_ids_to_summarize_protects_current_turn() {
// 20 tool pairs, cutoff=2 → 20 > 12, would normally trigger
let mut messages = vec![Message::user().with_text("hello")];
for i in 0..20 {
messages.extend(create_tool_pair(
&format!("call{}", i),
&format!("resp{}", i),
"read_file",
"content",
));
}
let conversation = Conversation::new_unvalidated(messages);
let result = tool_id_to_summarize(&conversation, 2);
// No protection: 20 eligible, 20 > 12 → batch of 10
let result = tool_ids_to_summarize(&conversation, 2, 0);
assert_eq!(result.len(), TOOLCALL_SUMMARIZATION_BATCH_SIZE);
// Protect last 8: 12 eligible, 12 <= 12 → nothing
let result = tool_ids_to_summarize(&conversation, 2, 8);
assert!(
result.is_some(),
"Should return a pair to summarize when tool calls exceed cutoff"
result.is_empty(),
"Should not summarize when protected count leaves eligible <= cutoff + batch"
);
let tool_call_id = result.unwrap();
assert_eq!(tool_call_id, "call1");
let summary = summarize_tool_call(&provider, "test-session", &conversation, &tool_call_id)
.await
.unwrap();
assert_eq!(summary.role, Role::User);
assert!(summary.metadata.agent_visible);
assert!(!summary.metadata.user_visible);
let mut updated_messages = conversation.messages().clone();
for msg in updated_messages.iter_mut() {
let has_matching_content = msg.content.iter().any(|c| match c {
MessageContent::ToolRequest(req) => req.id == tool_call_id,
MessageContent::ToolResponse(resp) => resp.id == tool_call_id,
_ => false,
});
if has_matching_content {
msg.metadata = msg.metadata.with_agent_invisible();
}
}
updated_messages.push(summary);
let updated_conversation = Conversation::new_unvalidated(updated_messages);
let messages = updated_conversation.messages();
let call1_msg = messages
.iter()
.find(|m| m.id.as_deref() == Some("call1"))
.unwrap();
assert!(
!call1_msg.is_agent_visible(),
"Original call should not be agent visible"
);
let response1_msg = messages
.iter()
.find(|m| m.id.as_deref() == Some("response1"))
.unwrap();
assert!(
!response1_msg.is_agent_visible(),
"Original response should not be agent visible"
);
let summary_msg = messages
.iter()
.find(|m| {
m.metadata.agent_visible
&& !m.metadata.user_visible
&& m.as_concat_text().contains("Tool call")
})
.unwrap();
assert!(
!summary_msg.is_user_visible(),
"Summary should not be user visible"
);
let result = tool_id_to_summarize(&updated_conversation, 3);
assert!(result.is_none(), "Nothing left to summarize");
// Protect last 7: 13 eligible, 13 > 12 → batch of 10
let result = tool_ids_to_summarize(&conversation, 2, 7);
assert_eq!(result.len(), TOOLCALL_SUMMARIZATION_BATCH_SIZE);
assert_eq!(result[0], "call0");
}
}