ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues 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 last_assistant_text = String::new();
let mut tool_pair_summarization_done = false;
loop {
if is_token_cancelled(&cancel_token) {
@@ -1460,13 +1461,17 @@ impl Agent {
.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 tool_pair_summarization_task = if tool_pair_summarization_done {
None
} else {
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();
@@ -1936,39 +1941,40 @@ impl Agent {
}
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 {
let mut updated_messages = conversation.messages().clone();
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,
if let Some(task) = tool_pair_summarization_task {
tool_pair_summarization_done = true;
if let Ok(summaries) = task.await {
for (summary_msg, tool_id) in summaries {
let matching_ids: Vec<String> = conversation.messages()
.iter()
.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();
.filter_map(|msg| msg.id.clone())
.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_ids.len() == 2 {
for id in &matching_ids {
SessionManager::update_message_metadata(&session_config.id, id, |metadata| {
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 {
crates/goose/src/context_mgmt/mod.rs
+11 -7
View File
@@ -533,13 +533,17 @@ pub fn maybe_summarize_tool_pairs(
conversation: Conversation,
cutoff: usize,
protect_last_n: usize,
) -> JoinHandle<Vec<(Message, String)>> {
tokio::spawn(async move {
if !tool_pair_summarization_enabled() || provider.manages_own_context() {
return Vec::new();
}
) -> Option<JoinHandle<Vec<(Message, String)>>> {
if !tool_pair_summarization_enabled() || provider.manages_own_context() {
return None;
}
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();
for tool_id in tool_ids {
match summarize_tool_call(provider.as_ref(), &session_id, &conversation, &tool_id).await
@@ -551,7 +555,7 @@ pub fn maybe_summarize_tool_pairs(
}
}
results
})
}))
}
#[cfg(test)]
crates/goose/tests/agent.rs
+35 -8
View File
@@ -561,11 +561,11 @@ mod tests {
&self,
_model_config: &ModelConfig,
_session_id: &str,
_system_prompt: &str,
system_prompt: &str,
_messages: &[Message],
tools: &[Tool],
_tools: &[Tool],
) -> 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
let n = self.summary_count.fetch_add(1, Ordering::SeqCst);
Message::assistant().with_text(format!("Summary of tool call #{}", n))
@@ -619,21 +619,25 @@ mod tests {
agent.update_provider(provider, &session.id).await?;
// Pre-populate: start with a user message, then 13 tool call/response pairs
// (need > cutoff + 10 = 12 to trigger batch summarization)
let initial_msg = Message::user().with_text("help me read some files");
// Pre-populate 13 tool pairs (need > cutoff + batch_size = 12 to trigger).
// Timestamps in the past so DB ordering places summaries before current turn.
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
.add_message(&session.id, &initial_msg)
.await?;
for i in 0..13 {
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_generated_id();
req_msg.created = base_ts + i as i64 + 1;
session_manager.add_message(&session.id, &req_msg).await?;
let resp_msg = Message::user()
let mut resp_msg = Message::user()
.with_tool_response(
&call_id,
Ok(CallToolResult::success(vec![RawContent::text(format!(
@@ -643,6 +647,7 @@ mod tests {
.no_annotation()])),
)
.with_generated_id();
resp_msg.created = base_ts + i as i64 + 1;
session_manager.add_message(&session.id, &resp_msg).await?;
}
@@ -717,6 +722,28 @@ mod tests {
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
Config::global().delete("GOOSE_TOOL_CALL_CUTOFF").unwrap();