From 9e31cfd7f86b8d7abf9a890b9bf45cce2a16da55 Mon Sep 17 00:00:00 2001 From: Douwe Osinga Date: Wed, 28 Jan 2026 11:35:08 -0500 Subject: [PATCH] Summarize old tool calls (#6119) Co-authored-by: Douwe Osinga Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.5 --- crates/goose-cli/src/commands/term.rs | 3 +- crates/goose/src/agents/agent.rs | 51 +++- crates/goose/src/context_mgmt/mod.rs | 245 +++++++++++++++++++- crates/goose/src/conversation/message.rs | 10 +- crates/goose/src/session/session_manager.rs | 104 ++++++++- 5 files changed, 396 insertions(+), 17 deletions(-) diff --git a/crates/goose-cli/src/commands/term.rs b/crates/goose-cli/src/commands/term.rs index 92dee120..22fade45 100644 --- a/crates/goose-cli/src/commands/term.rs +++ b/crates/goose-cli/src/commands/term.rs @@ -186,7 +186,8 @@ pub async fn handle_term_log(command: String) -> Result<()> { chrono::Utc::now().timestamp_millis(), vec![MessageContent::text(command)], ) - .with_metadata(MessageMetadata::user_only()); + .with_metadata(MessageMetadata::user_only()) + .with_generated_id(); let session_manager = SessionManager::instance(); session_manager.add_message(&session_id, &message).await?; diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index f4099437..4f4a8083 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -70,6 +70,7 @@ pub struct ReplyContext { pub toolshim_tools: Vec, pub system_prompt: String, pub goose_mode: GooseMode, + pub tool_call_cut_off: usize, pub initial_messages: Vec, } @@ -282,7 +283,10 @@ impl Agent { let mut messages = Vec::new(); let manager = self.config.session_manager.clone(); let mut elicitation_rx = ActionRequiredManager::global().request_rx.lock().await; - while let Ok(elicitation_message) = elicitation_rx.try_recv() { + while let Ok(mut elicitation_message) = elicitation_rx.try_recv() { + if elicitation_message.id.is_none() { + elicitation_message = elicitation_message.with_generated_id(); + } if let Err(e) = manager.add_message(session_id, &elicitation_message).await { warn!("Failed to save elicitation message to session: {}", e); } @@ -321,6 +325,9 @@ impl Agent { toolshim_tools, system_prompt, goose_mode: self.config.goose_mode, + tool_call_cut_off: Config::global() + .get_param::("GOOSE_TOOL_CALL_CUTOFF") + .unwrap_or(10), initial_messages, }) } @@ -1079,6 +1086,7 @@ impl Agent { mut tools, mut toolshim_tools, mut system_prompt, + tool_call_cut_off, goose_mode, initial_messages, } = context; @@ -1130,6 +1138,13 @@ impl Agent { break; } + let tool_pair_summarization_task = crate::context_mgmt::maybe_summarize_tool_pair( + self.provider().await?, + session_config.id.clone(), + conversation.clone(), + tool_call_cut_off, + ); + let conversation_with_moim = super::moim::inject_moim( &session_config.id, conversation.clone(), @@ -1202,9 +1217,7 @@ impl Agent { } let tool_response_messages: Vec>> = (0..num_tool_requests) - .map(|_| Arc::new(Mutex::new(Message::user().with_id( - format!("msg_{}", Uuid::new_v4()) - )))) + .map(|_| Arc::new(Mutex::new(Message::user().with_generated_id()))) .collect(); let mut request_to_response_map = HashMap::new(); @@ -1529,6 +1542,36 @@ impl Agent { } } + if let Ok(Some((summary_msg, tool_id))) = tool_pair_summarization_task.await { + let mut updated_messages = conversation.messages().clone(); + + 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, + }) + }) + .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?; + } + conversation = Conversation::new_unvalidated(updated_messages); + messages_to_add.push(summary_msg); + } else { + warn!("Expected a tool request/reply pair, but found {} matching messages", + matching.len()); + } + } + for msg in &messages_to_add { session_manager.add_message(&session_config.id, msg).await?; } diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index e06eb37f..fb45061b 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -6,8 +6,12 @@ use crate::providers::base::{Provider, ProviderUsage}; use crate::providers::errors::ProviderError; use crate::{config::Config, token_counter::create_token_counter}; use anyhow::Result; +use indoc::indoc; use rmcp::model::Role; use serde::Serialize; +use std::sync::Arc; +use tokio::task::JoinHandle; +use tracing::log::warn; use tracing::{debug, info}; pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.8; @@ -418,6 +422,116 @@ 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 { + let messages = conversation.messages(); + + let mut tool_call_count = 0; + let mut first_tool_call_id = None; + + for msg in messages.iter() { + if !msg.is_agent_visible() { + continue; + } + + 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; + } + } + } + } + None +} + +pub async fn summarize_tool_call( + provider: &dyn Provider, + session_id: &str, + conversation: &Conversation, + tool_id: &str, +) -> Result { + let messages = conversation.messages(); + + let matching_messages: Vec<&Message> = messages + .iter() + .filter(|m| { + m.content.iter().any(|c| match c { + MessageContent::ToolRequest(req) => req.id == tool_id, + MessageContent::ToolResponse(resp) => resp.id == tool_id, + _ => false, + }) + }) + .collect(); + + if matching_messages.is_empty() { + return Err(anyhow::anyhow!( + "No messages found for tool id: {}", + tool_id + )); + } + + let formatted = matching_messages + .iter() + .map(|msg| format_message_for_compacting(msg)) + .collect::>() + .join("\n"); + + let user_message = Message::user().with_text(formatted); + let summarization_request = vec![user_message]; + + let system_prompt = indoc! {r#" + 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 + 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 + .complete_fast(session_id, system_prompt, &summarization_request, &[]) + .await?; + + response.role = Role::User; + response.created = matching_messages.last().unwrap().created; + response.metadata = MessageMetadata::agent_only(); + + Ok(response.with_generated_id()) +} + +pub fn maybe_summarize_tool_pair( + provider: Arc, + session_id: String, + conversation: Conversation, + cutoff: usize, +) -> JoinHandle> { + tokio::spawn(async move { + if let Some(tool_id) = tool_id_to_summarize(&conversation, cutoff) { + match summarize_tool_call(provider.as_ref(), &session_id, &conversation, &tool_id).await + { + Ok(summary) => Some((summary, tool_id)), + Err(e) => { + warn!("Failed to summarize tool pair: {}", e); + None + } + } + } else { + None + } + }) +} + #[cfg(test)] mod tests { use super::*; @@ -579,11 +693,140 @@ mod tests { let conversation = Conversation::new_unvalidated(messages); let result = compact_messages(&provider, "test-session-id", &conversation, false).await; - // Should succeed after progressive removal assert!( result.is_ok(), "Should succeed with progressive removal: {:?}", result.err() ); } + + #[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 { + vec![ + Message::assistant() + .with_tool_request( + call_id, + Ok(CallToolRequestParams { + task: None, + name: tool_name.to_string().into(), + arguments: None, + meta: None, + }), + ) + .with_id(call_id), + Message::user() + .with_tool_response( + call_id, + Ok(rmcp::model::CallToolResult { + content: vec![RawContent::text(response_text).no_annotation()], + structured_content: None, + is_error: Some(false), + meta: None, + }), + ) + .with_id(response_id), + ] + } + + 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", + )); + + let conversation = Conversation::new_unvalidated(messages); + + let result = tool_id_to_summarize(&conversation, 2); + assert!( + result.is_some(), + "Should return a pair to summarize when tool calls exceed cutoff" + ); + + 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"); + } } diff --git a/crates/goose/src/conversation/message.rs b/crates/goose/src/conversation/message.rs index 419f2fb9..acde27ad 100644 --- a/crates/goose/src/conversation/message.rs +++ b/crates/goose/src/conversation/message.rs @@ -1,4 +1,6 @@ +use crate::conversation::tool_result_serde; use crate::mcp_utils::ToolResult; +use crate::utils::sanitize_unicode_tags; use chrono::Utc; use rmcp::model::{ AnnotateAble, CallToolRequestParams, CallToolResult, Content, ImageContent, JsonObject, @@ -9,9 +11,7 @@ use serde::{Deserialize, Deserializer, Serialize}; use std::collections::HashSet; use std::fmt; use utoipa::ToSchema; - -use crate::conversation::tool_result_serde; -use crate::utils::sanitize_unicode_tags; +use uuid::Uuid; #[derive(ToSchema)] pub enum ToolCallResult { @@ -719,6 +719,10 @@ impl Message { self } + pub fn with_generated_id(self) -> Self { + self.with_id(format!("msg_{}", Uuid::new_v4())) + } + /// Add any MessageContent to the message pub fn with_content(mut self, content: MessageContent) -> Self { self.content.push(content); diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 4b59c0b1..12af879b 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -18,7 +18,7 @@ use std::sync::{Arc, LazyLock}; use tracing::{info, warn}; use utoipa::ToSchema; -pub const CURRENT_SCHEMA_VERSION: i32 = 6; +pub const CURRENT_SCHEMA_VERSION: i32 = 7; pub const SESSIONS_FOLDER: &str = "sessions"; pub const DB_NAME: &str = "sessions.db"; @@ -363,6 +363,18 @@ impl SessionManager { .search_chat_history(query, limit, after_date, before_date, exclude_session_id) .await } + + pub async fn update_message_metadata(id: &str, message_id: &str, f: F) -> Result<()> + where + F: FnOnce( + crate::conversation::message::MessageMetadata, + ) -> crate::conversation::message::MessageMetadata, + { + Self::instance() + .storage + .update_message_metadata(id, message_id, f) + .await + } } pub struct SessionStorage { @@ -575,6 +587,7 @@ impl SessionStorage { r#" CREATE TABLE messages ( id INTEGER PRIMARY KEY AUTOINCREMENT, + message_id TEXT, session_id TEXT NOT NULL REFERENCES sessions(id), role TEXT NOT NULL, content_json TEXT NOT NULL, @@ -594,6 +607,9 @@ impl SessionStorage { sqlx::query("CREATE INDEX idx_messages_timestamp ON messages(timestamp)") .execute(pool) .await?; + sqlx::query("CREATE INDEX idx_messages_message_id ON messages(message_id)") + .execute(pool) + .await?; sqlx::query("CREATE INDEX idx_sessions_updated ON sessions(updated_at DESC)") .execute(pool) .await?; @@ -760,6 +776,7 @@ impl SessionStorage { Ok(()) } + #[allow(clippy::too_many_lines)] async fn apply_migration(pool: &Pool, version: i32) -> Result<()> { match version { 1 => { @@ -839,6 +856,28 @@ impl SessionStorage { .execute(pool) .await?; } + 7 => { + sqlx::query( + r#" + ALTER TABLE messages ADD COLUMN message_id TEXT + "#, + ) + .execute(pool) + .await?; + + sqlx::query( + r#" + UPDATE messages + SET message_id = 'msg_' || session_id || '_' || id + "#, + ) + .execute(pool) + .await?; + + sqlx::query("CREATE INDEX idx_messages_message_id ON messages(message_id)") + .execute(pool) + .await?; + } _ => { anyhow::bail!("Unknown migration version: {}", version); } @@ -1036,16 +1075,16 @@ impl SessionStorage { async fn get_conversation(&self, session_id: &str) -> Result { let pool = self.pool().await?; - let rows = sqlx::query_as::<_, (String, String, i64, Option)>( - "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp", + let rows = sqlx::query_as::<_, (String, String, i64, Option, Option)>( + "SELECT role, content_json, created_timestamp, metadata_json, message_id FROM messages WHERE session_id = ? ORDER BY timestamp", ) .bind(session_id) .fetch_all(pool) .await?; let mut messages = Vec::new(); - for (idx, (role_str, content_json, created_timestamp, metadata_json)) in - rows.into_iter().enumerate() + for (role_str, content_json, created_timestamp, metadata_json, message_id) in + rows.into_iter() { let role = match role_str.as_str() { "user" => Role::User, @@ -1060,7 +1099,9 @@ impl SessionStorage { let mut message = Message::new(role, created_timestamp, content); message.metadata = metadata; - message = message.with_id(format!("msg_{}_{}", session_id, idx)); + if let Some(id) = message_id { + message = message.with_id(id); + } messages.push(message); } @@ -1073,12 +1114,18 @@ impl SessionStorage { let metadata_json = serde_json::to_string(&message.metadata)?; + let message_id = message + .id + .clone() + .unwrap_or_else(|| format!("msg_{}_{}", session_id, uuid::Uuid::new_v4())); + sqlx::query( r#" - INSERT INTO messages (session_id, role, content_json, created_timestamp, metadata_json) - VALUES (?, ?, ?, ?, ?) + INSERT INTO messages (message_id, session_id, role, content_json, created_timestamp, metadata_json) + VALUES (?, ?, ?, ?, ?, ?) "#, ) + .bind(message_id) .bind(session_id) .bind(role_to_string(&message.role)) .bind(serde_json::to_string(&message.content)?) @@ -1344,6 +1391,47 @@ impl SessionStorage { .execute() .await } + + async fn update_message_metadata( + &self, + session_id: &str, + message_id: &str, + f: F, + ) -> Result<()> + where + F: FnOnce( + crate::conversation::message::MessageMetadata, + ) -> crate::conversation::message::MessageMetadata, + { + let mut tx = self.pool.begin().await?; + + let current_metadata_json = sqlx::query_scalar::<_, String>( + "SELECT metadata_json FROM messages WHERE message_id = ? AND session_id = ?", + ) + .bind(message_id) + .bind(session_id) + .fetch_one(&mut *tx) + .await?; + + let current_metadata: crate::conversation::message::MessageMetadata = + serde_json::from_str(¤t_metadata_json)?; + + let new_metadata = f(current_metadata); + let metadata_json = serde_json::to_string(&new_metadata)?; + + sqlx::query( + "UPDATE messages SET metadata_json = ? WHERE message_id = ? AND session_id = ?", + ) + .bind(metadata_json) + .bind(message_id) + .bind(session_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + Ok(()) + } } #[cfg(test)]