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

fix: insert tool pair summaries at chronological position in conversation (#9087)

Browse Source 
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2026-05-11 14:54:16 -04:00
committed by GitHub
parent a38922d95e
commit 4dd8ef3b88
3 changed files with 85 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+39 -33
View File
@@ -1414,6 +1414,7 @@ impl Agent {
}); });
let mut compaction_attempts = 0; let mut compaction_attempts = 0;
let mut last_assistant_text = String::new(); let mut last_assistant_text = String::new();
let mut tool_pair_summarization_done = false;
loop { loop {
if is_token_cancelled(&cancel_token) { if is_token_cancelled(&cancel_token) {
@@ -1460,13 +1461,17 @@ impl Agent {
.count() .count()
.saturating_sub(pre_turn_tool_count); .saturating_sub(pre_turn_tool_count);
let tool_pair_summarization_task = crate::context_mgmt::maybe_summarize_tool_pairs( let tool_pair_summarization_task = if tool_pair_summarization_done {
self.provider().await?, None
session_config.id.clone(), } else {
conversation.clone(), crate::context_mgmt::maybe_summarize_tool_pairs(
tool_call_cut_off, self.provider().await?,
current_turn_tool_count, session_config.id.clone(),
); conversation.clone(),
tool_call_cut_off,
current_turn_tool_count,
)
};
let mut no_tools_called = true; let mut no_tools_called = true;
let mut messages_to_add = Conversation::default(); let mut messages_to_add = Conversation::default();
@@ -1936,39 +1941,40 @@ impl Agent {
} }
if is_token_cancelled(&cancel_token) { if is_token_cancelled(&cancel_token) {
tool_pair_summarization_task.abort(); if let Some(ref task) = tool_pair_summarization_task {
task.abort();
}
} }
if let Ok(summaries) = tool_pair_summarization_task.await { if let Some(task) = tool_pair_summarization_task {
let mut updated_messages = conversation.messages().clone(); tool_pair_summarization_done = true;
if let Ok(summaries) = task.await {
for (summary_msg, tool_id) in summaries { for (summary_msg, tool_id) in summaries {
let matching: Vec<&mut Message> = updated_messages let matching_ids: Vec<String> = conversation.messages()
.iter_mut() .iter()
.filter(|msg| { .filter(|msg| {
msg.id.is_some() && msg.content.iter().any(|c| match c { msg.id.is_some() && msg.content.iter().any(|c| match c {
MessageContent::ToolRequest(req) => req.id == tool_id, MessageContent::ToolRequest(req) => req.id == tool_id,
MessageContent::ToolResponse(resp) => resp.id == tool_id, MessageContent::ToolResponse(resp) => resp.id == tool_id,
_ => false, _ => false,
})
}) })
}) .filter_map(|msg| msg.id.clone())
.collect(); .collect();
if matching.len() == 2 { if matching_ids.len() == 2 {
for msg in matching { for id in &matching_ids {
let id = msg.id.as_ref().unwrap(); SessionManager::update_message_metadata(&session_config.id, id, |metadata| {
msg.metadata = msg.metadata.with_agent_invisible(); metadata.with_agent_invisible()
SessionManager::update_message_metadata(&session_config.id, id, |metadata| { }).await?;
metadata.with_agent_invisible() }
}).await?; session_manager.add_message(&session_config.id, &summary_msg).await?;
} else {
warn!("Expected a tool request/reply pair, but found {} matching messages",
matching_ids.len());
} }
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 { for msg in &messages_to_add {
crates/goose/src/context_mgmt/mod.rs
+11 -7
View File
@@ -533,13 +533,17 @@ pub fn maybe_summarize_tool_pairs(
conversation: Conversation, conversation: Conversation,
cutoff: usize, cutoff: usize,
protect_last_n: usize, protect_last_n: usize,
) -> JoinHandle<Vec<(Message, String)>> { ) -> Option<JoinHandle<Vec<(Message, String)>>> {
tokio::spawn(async move { if !tool_pair_summarization_enabled() || provider.manages_own_context() {
if !tool_pair_summarization_enabled() || provider.manages_own_context() { return None;
return Vec::new(); }
}
let tool_ids = tool_ids_to_summarize(&conversation, cutoff, protect_last_n); let tool_ids = tool_ids_to_summarize(&conversation, cutoff, protect_last_n);
if tool_ids.is_empty() {
return None;
}
Some(tokio::spawn(async move {
let mut results = Vec::new(); let mut results = Vec::new();
for tool_id in tool_ids { for tool_id in tool_ids {
match summarize_tool_call(provider.as_ref(), &session_id, &conversation, &tool_id).await match summarize_tool_call(provider.as_ref(), &session_id, &conversation, &tool_id).await
@@ -551,7 +555,7 @@ pub fn maybe_summarize_tool_pairs(
} }
} }
results results
}) }))
} }
#[cfg(test)] #[cfg(test)]
crates/goose/tests/agent.rs
+35 -8
View File
@@ -561,11 +561,11 @@ mod tests {
&self, &self,
_model_config: &ModelConfig, _model_config: &ModelConfig,
_session_id: &str, _session_id: &str,
_system_prompt: &str, system_prompt: &str,
_messages: &[Message], _messages: &[Message],
tools: &[Tool], _tools: &[Tool],
) -> Result<MessageStream, ProviderError> { ) -> Result<MessageStream, ProviderError> {
let message = if tools.is_empty() { let message = if system_prompt.contains("summarize a tool call") {
// Summarization call — return a unique summary // Summarization call — return a unique summary
let n = self.summary_count.fetch_add(1, Ordering::SeqCst); let n = self.summary_count.fetch_add(1, Ordering::SeqCst);
Message::assistant().with_text(format!("Summary of tool call #{}", n)) Message::assistant().with_text(format!("Summary of tool call #{}", n))
@@ -619,21 +619,25 @@ mod tests {
agent.update_provider(provider, &session.id).await?; agent.update_provider(provider, &session.id).await?;
// Pre-populate: start with a user message, then 13 tool call/response pairs // Pre-populate 13 tool pairs (need > cutoff + batch_size = 12 to trigger).
// (need > cutoff + 10 = 12 to trigger batch summarization) // Timestamps in the past so DB ordering places summaries before current turn.
let initial_msg = Message::user().with_text("help me read some files"); let base_ts = chrono::Utc::now().timestamp() - 100;
let mut initial_msg = Message::user().with_text("help me read some files");
initial_msg.created = base_ts;
session_manager session_manager
.add_message(&session.id, &initial_msg) .add_message(&session.id, &initial_msg)
.await?; .await?;
for i in 0..13 { for i in 0..13 {
let call_id = format!("precall_{}", i); let call_id = format!("precall_{}", i);
let req_msg = Message::assistant() let mut req_msg = Message::assistant()
.with_tool_request(&call_id, Ok(CallToolRequestParams::new("read_file"))) .with_tool_request(&call_id, Ok(CallToolRequestParams::new("read_file")))
.with_generated_id(); .with_generated_id();
req_msg.created = base_ts + i as i64 + 1;
session_manager.add_message(&session.id, &req_msg).await?; session_manager.add_message(&session.id, &req_msg).await?;
let resp_msg = Message::user() let mut resp_msg = Message::user()
.with_tool_response( .with_tool_response(
&call_id, &call_id,
Ok(CallToolResult::success(vec![RawContent::text(format!( Ok(CallToolResult::success(vec![RawContent::text(format!(
@@ -643,6 +647,7 @@ mod tests {
.no_annotation()])), .no_annotation()])),
) )
.with_generated_id(); .with_generated_id();
resp_msg.created = base_ts + i as i64 + 1;
session_manager.add_message(&session.id, &resp_msg).await?; session_manager.add_message(&session.id, &resp_msg).await?;
} }
@@ -717,6 +722,28 @@ mod tests {
invisible_tool_msgs.len() invisible_tool_msgs.len()
); );
// Summaries must appear before the current turn's reply, not after it
let agent_visible: Vec<&Message> = messages
.iter()
.filter(|m| m.metadata.agent_visible)
.collect();
let last_summary_pos = agent_visible
.iter()
.rposition(|m| m.as_concat_text().starts_with("Summary of tool call #"))
.expect("Should have at least one summary");
let agent_reply_pos = agent_visible
.iter()
.position(|m| m.as_concat_text().contains("Done processing."))
.expect("Should have the agent reply");
assert!(
last_summary_pos < agent_reply_pos,
"Summaries appeared after the current turn's reply: last_summary={}, reply={}",
last_summary_pos,
agent_reply_pos,
);
// Clean up the config override // Clean up the config override
Config::global().delete("GOOSE_TOOL_CALL_CUTOFF").unwrap(); Config::global().delete("GOOSE_TOOL_CALL_CUTOFF").unwrap();