diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index d682423761..045a788449 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -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 = 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 { diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 1cde777b04..7f3a1f9733 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -533,13 +533,17 @@ pub fn maybe_summarize_tool_pairs( conversation: Conversation, cutoff: usize, protect_last_n: usize, -) -> JoinHandle> { - tokio::spawn(async move { - if !tool_pair_summarization_enabled() || provider.manages_own_context() { - return Vec::new(); - } +) -> Option>> { + 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)] diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 96aac0ec71..81fbe3ecc8 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -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 { - 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();