Optimize tool summarization (#7938)

Co-authored-by: Michael Neale <michael.neale@gmail.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>

crates/goose/tests/agent.rs

@@ -495,6 +495,234 @@ mod tests {
         }
     }
 
+    #[cfg(test)]
+    mod tool_pair_summarization_tests {
+        use super::*;
+        use async_trait::async_trait;
+        use goose::agents::SessionConfig;
+        use goose::config::base::Config;
+        use goose::config::GooseMode;
+        use goose::conversation::message::Message;
+        use goose::model::ModelConfig;
+        use goose::providers::base::{
+            stream_from_single_message, MessageStream, Provider, ProviderDef, ProviderMetadata,
+            ProviderUsage, Usage,
+        };
+        use goose::providers::errors::ProviderError;
+        use goose::session::session_manager::SessionType;
+        use rmcp::model::{AnnotateAble, CallToolRequestParams, CallToolResult, RawContent, Tool};
+        use std::path::PathBuf;
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        use std::sync::Arc;
+
+        /// Mock provider that returns text for the main reply and summaries for
+        /// summarization calls. Distinguishes by checking if tools are empty
+        /// (summarization calls pass no tools).
+        struct SummarizationTestProvider {
+            summary_count: AtomicUsize,
+        }
+
+        impl SummarizationTestProvider {
+            fn new() -> Self {
+                Self {
+                    summary_count: AtomicUsize::new(0),
+                }
+            }
+        }
+
+        impl ProviderDef for SummarizationTestProvider {
+            type Provider = Self;
+
+            fn metadata() -> ProviderMetadata {
+                ProviderMetadata {
+                    name: "mock-summarization".to_string(),
+                    display_name: "Mock Summarization Provider".to_string(),
+                    description: "Mock provider for summarization tests".to_string(),
+                    default_model: "mock-model".to_string(),
+                    known_models: vec![],
+                    model_doc_link: "".to_string(),
+                    config_keys: vec![],
+                    setup_steps: vec![],
+                }
+            }
+
+            fn from_env(
+                _model: ModelConfig,
+                _extensions: Vec<goose::config::ExtensionConfig>,
+            ) -> futures::future::BoxFuture<'static, anyhow::Result<Self>> {
+                Box::pin(async { Ok(Self::new()) })
+            }
+        }
+
+        #[async_trait]
+        impl Provider for SummarizationTestProvider {
+            async fn stream(
+                &self,
+                _model_config: &ModelConfig,
+                _session_id: &str,
+                _system_prompt: &str,
+                _messages: &[Message],
+                tools: &[Tool],
+            ) -> Result<MessageStream, ProviderError> {
+                let message = if tools.is_empty() {
+                    // 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))
+                } else {
+                    // Main agent reply — return plain text so the loop exits
+                    Message::assistant().with_text("Done processing.")
+                };
+
+                let usage = ProviderUsage::new(
+                    "mock-model".to_string(),
+                    Usage::new(Some(10), Some(5), Some(15)),
+                );
+                Ok(stream_from_single_message(message, usage))
+            }
+
+            fn get_model_config(&self) -> ModelConfig {
+                ModelConfig::new("mock-model").unwrap()
+            }
+
+            fn get_name(&self) -> &str {
+                "mock-summarization"
+            }
+        }
+
+        /// Test that batch tool pair summarization preserves all summaries.
+        ///
+        /// Pre-populates a session with enough tool call/response pairs to trigger
+        /// batch summarization, runs agent.reply(), then verifies:
+        /// - All 10 summaries are present in the final conversation
+        /// - The original tool pairs are marked invisible
+        #[tokio::test]
+        async fn test_batch_summarization_preserves_all_summaries() -> Result<()> {
+            // Set a low cutoff so we don't need hundreds of tool pairs.
+            // cutoff=2 means we need >2+10=12 visible tool pairs to trigger.
+            Config::global()
+                .set_param("GOOSE_TOOL_CALL_CUTOFF", 2)
+                .unwrap();
+
+            let agent = Agent::new();
+            let session_manager = agent.config.session_manager.clone();
+            let provider = Arc::new(SummarizationTestProvider::new());
+
+            let session = session_manager
+                .create_session(
+                    PathBuf::from("."),
+                    "summarization-test".to_string(),
+                    SessionType::Hidden,
+                    GooseMode::Auto,
+                )
+                .await?;
+
+            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");
+            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()
+                    .with_tool_request(&call_id, Ok(CallToolRequestParams::new("read_file")))
+                    .with_generated_id();
+                session_manager.add_message(&session.id, &req_msg).await?;
+
+                let resp_msg = Message::user()
+                    .with_tool_response(
+                        &call_id,
+                        Ok(CallToolResult::success(vec![RawContent::text(format!(
+                            "content of file {}",
+                            i
+                        ))
+                        .no_annotation()])),
+                    )
+                    .with_generated_id();
+                session_manager.add_message(&session.id, &resp_msg).await?;
+            }
+
+            // Send a user message to trigger the reply loop
+            let user_message = Message::user().with_text("summarize what you found");
+
+            let session_config = SessionConfig {
+                id: session.id.clone(),
+                schedule_id: None,
+                max_turns: Some(1),
+                retry_config: None,
+            };
+
+            let reply_stream = agent.reply(user_message, session_config, None).await?;
+            tokio::pin!(reply_stream);
+
+            // Drain the stream
+            while let Some(event) = reply_stream.next().await {
+                match event {
+                    Ok(AgentEvent::Message(_)) => {}
+                    Ok(_) => {}
+                    Err(e) => return Err(e),
+                }
+            }
+
+            // Load the final session and inspect the conversation
+            let final_session = session_manager.get_session(&session.id, true).await?;
+            let conversation = final_session
+                .conversation
+                .expect("Session should have a conversation");
+            let messages = conversation.messages();
+
+            // Count summaries: messages that are agent-visible, not user-visible,
+            // and contain our summary text pattern
+            let summaries: Vec<&Message> = messages
+                .iter()
+                .filter(|m| {
+                    m.metadata.agent_visible
+                        && !m.metadata.user_visible
+                        && m.as_concat_text().starts_with("Summary of tool call #")
+                })
+                .collect();
+
+            assert_eq!(
+                summaries.len(),
+                10,
+                "Expected 10 summaries (one full batch), got {}. Summary texts: {:?}",
+                summaries.len(),
+                summaries
+                    .iter()
+                    .map(|m| m.as_concat_text())
+                    .collect::<Vec<_>>()
+            );
+
+            // Verify each summary is unique
+            let summary_texts: std::collections::HashSet<String> =
+                summaries.iter().map(|m| m.as_concat_text()).collect();
+            assert_eq!(summary_texts.len(), 10, "All 10 summaries should be unique");
+
+            // Count invisible tool pairs: original pairs that were summarized
+            // should have agent_visible=false
+            let invisible_tool_msgs: Vec<&Message> = messages
+                .iter()
+                .filter(|m| !m.metadata.agent_visible && (m.is_tool_call() || m.is_tool_response()))
+                .collect();
+
+            // Each summarized pair = 2 invisible messages (request + response)
+            assert_eq!(
+                invisible_tool_msgs.len(),
+                20, // 10 pairs × 2 messages
+                "Expected 20 invisible tool messages (10 summarized pairs), got {}",
+                invisible_tool_msgs.len()
+            );
+
+            // Clean up the config override
+            Config::global().delete("GOOSE_TOOL_CALL_CUTOFF").unwrap();
+
+            Ok(())
+        }
+    }
+
     #[cfg(test)]
     mod extension_manager_tests {
         use super::*;