From b6f3402d9b4ce137ee03c62d078536b31dbc4355 Mon Sep 17 00:00:00 2001 From: Alex Hancock Date: Mon, 10 Aug 2026 17:10:36 -0400 Subject: [PATCH] feat: compaction in the GDK (#11042) --- Cargo.lock | 18 ++ crates/goose-context-management/Cargo.toml | 26 ++ crates/goose-context-management/src/format.rs | 86 ++++++ crates/goose-context-management/src/lib.rs | 70 +++++ crates/goose-context-management/src/model.rs | 53 ++++ .../src/prompts/compaction.md | 0 .../src/prompts/compaction_summary.md | 0 .../goose-context-management/src/provider.rs | 106 +++++++ .../src}/structured.rs | 11 +- .../goose-context-management/src/summarize.rs | 175 +++++++++++ .../goose-context-management/src/templates.rs | 59 ++++ crates/goose-sdk/Cargo.toml | 2 + crates/goose-sdk/src/bindings.rs | 92 +++++- crates/goose/Cargo.toml | 1 + crates/goose/src/context_mgmt/mod.rs | 290 ++++-------------- crates/goose/src/prompt_template.rs | 38 ++- 16 files changed, 790 insertions(+), 237 deletions(-) create mode 100644 crates/goose-context-management/Cargo.toml create mode 100644 crates/goose-context-management/src/format.rs create mode 100644 crates/goose-context-management/src/lib.rs create mode 100644 crates/goose-context-management/src/model.rs rename crates/{goose => goose-context-management}/src/prompts/compaction.md (100%) rename crates/{goose => goose-context-management}/src/prompts/compaction_summary.md (100%) create mode 100644 crates/goose-context-management/src/provider.rs rename crates/{goose/src/context_mgmt => goose-context-management/src}/structured.rs (98%) create mode 100644 crates/goose-context-management/src/summarize.rs create mode 100644 crates/goose-context-management/src/templates.rs diff --git a/Cargo.lock b/Cargo.lock index ba9bc3b1f..bd122a6d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4950,6 +4950,7 @@ dependencies = [ "futures", "gethostname", "goose-acp-macros", + "goose-context-management", "goose-download-manager", "goose-mcp", "goose-providers", @@ -5111,6 +5112,22 @@ dependencies = [ "zip 8.6.0", ] +[[package]] +name = "goose-context-management" +version = "0.1.0-alpha.5" +dependencies = [ + "anyhow", + "async-trait", + "goose-providers", + "include_dir", + "minijinja", + "rmcp 3.0.0", + "serde", + "serde_json", + "tokio", + "tracing", +] + [[package]] name = "goose-download-manager" version = "0.1.0-alpha.5" @@ -5259,6 +5276,7 @@ dependencies = [ "anyhow", "base64 0.23.0", "futures", + "goose-context-management", "goose-providers", "goose-sdk-types", "rmcp 3.0.0", diff --git a/crates/goose-context-management/Cargo.toml b/crates/goose-context-management/Cargo.toml new file mode 100644 index 000000000..2725dddff --- /dev/null +++ b/crates/goose-context-management/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "goose-context-management" +version = "0.1.0-alpha.5" +edition.workspace = true +rust-version.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +description = "Conversation compaction for Goose" + +[lints] +workspace = true + +[dependencies] +goose-providers = { version = "0.1.0-alpha.5", path = "../goose-providers", features = ["rustls-tls"] } +anyhow = { workspace = true } +async-trait = { workspace = true } +include_dir = { workspace = true } +minijinja = { version = "2.18", default-features = false, features = ["loader", "multi_template", "serde"] } +rmcp = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true, features = ["preserve_order"] } +tracing = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } diff --git a/crates/goose-context-management/src/format.rs b/crates/goose-context-management/src/format.rs new file mode 100644 index 000000000..de333746f --- /dev/null +++ b/crates/goose-context-management/src/format.rs @@ -0,0 +1,86 @@ +use goose_providers::conversation::message::{ActionRequiredData, Message, MessageContent}; +use rmcp::model::Role; + +pub fn format_message_for_compacting(msg: &Message) -> String { + let content_parts: Vec = msg + .content + .iter() + .filter_map(|content| match content { + MessageContent::Text(text) => Some(text.text.clone()), + MessageContent::Image(img) => Some(format!("[image: {}]", img.mime_type)), + MessageContent::ToolRequest(req) => { + if let Ok(call) = &req.tool_call { + Some(format!( + "tool_request({}): {}", + call.name, + serde_json::to_string(&call.arguments) + .unwrap_or_else(|_| "<>".to_string()) + )) + } else { + Some("tool_request: [error]".to_string()) + } + } + MessageContent::ToolResponse(res) => { + if let Ok(result) = &res.tool_result { + let text_items: Vec = result + .content + .iter() + .filter_map(|content| { + content.as_text().map(|text_str| text_str.text.clone()) + }) + .collect(); + + if !text_items.is_empty() { + Some(format!("tool_response: {}", text_items.join("\n"))) + } else { + Some("tool_response: [non-text content]".to_string()) + } + } else { + Some("tool_response: [error]".to_string()) + } + } + MessageContent::ToolConfirmationRequest(req) => { + Some(format!("tool_confirmation_request: {}", req.tool_name)) + } + MessageContent::ActionRequired(action) => match &action.data { + ActionRequiredData::ToolConfirmation { tool_name, .. } => { + Some(format!("action_required(tool_confirmation): {}", tool_name)) + } + ActionRequiredData::Elicitation { message, .. } => { + Some(format!("action_required(elicitation): {}", message)) + } + ActionRequiredData::ElicitationResponse { id, .. } => { + Some(format!("action_required(elicitation_response): {}", id)) + } + ActionRequiredData::ToolConfirmationResponse { id, .. } => Some(format!( + "action_required(tool_confirmation_response): {}", + id + )), + }, + MessageContent::FrontendToolRequest(req) => { + if let Ok(call) = &req.tool_call { + Some(format!("frontend_tool_request: {}", call.name)) + } else { + Some("frontend_tool_request: [error]".to_string()) + } + } + MessageContent::Thinking(_) => None, + MessageContent::RedactedThinking(_) => None, + MessageContent::SystemNotification(notification) => { + Some(format!("system_notification: {}", notification.msg)) + } + MessageContent::Error(error) => Some(format!("error: {}", error.message)), + }) + .collect(); + + let role_str = match msg.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + + if content_parts.is_empty() { + format!("[{}]: ", role_str) + } else { + format!("[{}]: {}", role_str, content_parts.join("\n")) + } +} diff --git a/crates/goose-context-management/src/lib.rs b/crates/goose-context-management/src/lib.rs new file mode 100644 index 000000000..85fd951e6 --- /dev/null +++ b/crates/goose-context-management/src/lib.rs @@ -0,0 +1,70 @@ +//! Conversation compaction: summarizing a message history down to a single +//! message so a conversation can continue past a model's context window. +//! +//! Three layers, smallest first: +//! +//! * [`summarize`] - given a model and messages, produce one summary message. +//! * [`compact`] - the trait-based API ([`CompactionInput`] / +//! [`CompactionOutput`]) letting a caller read from and write back to its own +//! conversation representation. Rust only. +//! +//! Cross-language (Python/Kotlin) access is exposed by `goose-sdk`, which +//! wraps this crate in its uniffi bindings. + +pub mod format; +pub mod model; +pub mod provider; +pub mod structured; +pub mod summarize; +pub mod templates; + +use anyhow::Result; +use goose_providers::conversation::message::Message; +use goose_providers::conversation::token_usage::ProviderUsage; + +pub use format::format_message_for_compacting; +pub use model::{CompactionModel, ProviderModel, TokenEstimator}; +pub use provider::CompactingProvider; +pub use structured::{FileActivity, StructuredSummary}; +pub use summarize::{summarize, Summary}; +pub use templates::Templates; + +pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.8; + +/// Everything compaction reads from the caller's conversation. +pub trait CompactionInput { + fn messages(&self) -> Vec; + + fn templates(&self) -> Templates { + Templates::default() + } +} + +/// Where compaction writes its result back into the caller's conversation. +pub trait CompactionOutput { + fn set_summary(&mut self, summary: Message); + fn set_usage(&mut self, usage: ProviderUsage); +} + +pub async fn compact( + model: &dyn CompactionModel, + estimator: Option<&dyn TokenEstimator>, + input: &I, + output: &mut O, +) -> Result<()> +where + I: CompactionInput + ?Sized, + O: CompactionOutput + ?Sized, +{ + let templates = input.templates(); + let summary = summarize(model, estimator, &templates, &input.messages()).await?; + output.set_summary(summary.message); + output.set_usage(summary.usage); + Ok(()) +} + +impl CompactionInput for Vec { + fn messages(&self) -> Vec { + self.clone() + } +} diff --git a/crates/goose-context-management/src/model.rs b/crates/goose-context-management/src/model.rs new file mode 100644 index 000000000..81ffdd324 --- /dev/null +++ b/crates/goose-context-management/src/model.rs @@ -0,0 +1,53 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use goose_providers::base::Provider; +use goose_providers::conversation::message::Message; +use goose_providers::conversation::token_usage::ProviderUsage; +use goose_providers::errors::ProviderError; +use goose_providers::model::ModelConfig; + +/// The single completion call compaction needs. Implementations decide model +/// selection, fallbacks and session plumbing. +#[async_trait] +pub trait CompactionModel: Send + Sync { + async fn complete( + &self, + system: &str, + messages: &[Message], + ) -> Result<(Message, ProviderUsage), ProviderError>; +} + +/// Counts tokens for usage estimation and retained-context reporting. +#[async_trait] +pub trait TokenEstimator: Send + Sync { + async fn count_chat_tokens(&self, system: &str, messages: &[Message]) -> usize; + async fn count_text_tokens(&self, text: &str) -> usize; +} + +pub struct ProviderModel { + provider: Arc, + model_config: ModelConfig, +} + +impl ProviderModel { + pub fn new(provider: Arc, model_config: ModelConfig) -> Self { + Self { + provider, + model_config, + } + } +} + +#[async_trait] +impl CompactionModel for ProviderModel { + async fn complete( + &self, + system: &str, + messages: &[Message], + ) -> Result<(Message, ProviderUsage), ProviderError> { + self.provider + .complete(&self.model_config, system, messages, &[]) + .await + } +} diff --git a/crates/goose/src/prompts/compaction.md b/crates/goose-context-management/src/prompts/compaction.md similarity index 100% rename from crates/goose/src/prompts/compaction.md rename to crates/goose-context-management/src/prompts/compaction.md diff --git a/crates/goose/src/prompts/compaction_summary.md b/crates/goose-context-management/src/prompts/compaction_summary.md similarity index 100% rename from crates/goose/src/prompts/compaction_summary.md rename to crates/goose-context-management/src/prompts/compaction_summary.md diff --git a/crates/goose-context-management/src/provider.rs b/crates/goose-context-management/src/provider.rs new file mode 100644 index 000000000..a4301e5b8 --- /dev/null +++ b/crates/goose-context-management/src/provider.rs @@ -0,0 +1,106 @@ +//! Provider-level compaction: retry a completion that overflowed the context +//! window against a compacted history instead of surfacing the error. + +use std::sync::Arc; + +use async_trait::async_trait; +use goose_providers::base::{MessageStream, Provider}; +use goose_providers::conversation::message::Message; +use goose_providers::conversation::token_usage::ProviderUsage; +use goose_providers::errors::ProviderError; +use goose_providers::model::ModelConfig; + +use crate::model::ProviderModel; +use crate::templates::Templates; + +/// Wraps a provider so a `ContextLengthExceeded` response triggers compaction +/// and one retry with the summary standing in for the prior history. +pub struct CompactingProvider { + inner: Arc, + templates: Templates, +} + +impl CompactingProvider { + pub fn new(inner: Arc) -> Self { + Self { + inner, + templates: Templates::default(), + } + } + + pub fn with_templates(mut self, templates: Templates) -> Self { + self.templates = templates; + self + } + + async fn compacted_messages( + &self, + model_config: &ModelConfig, + messages: &[Message], + ) -> Result, ProviderError> { + let model = ProviderModel::new(self.inner.clone(), model_config.clone()); + let summary = crate::summarize(&model, None, &self.templates, messages) + .await + .map_err(|error| ProviderError::ExecutionError(error.to_string()))?; + Ok(vec![summary.message]) + } +} + +#[async_trait] +impl Provider for CompactingProvider { + fn get_name(&self) -> &str { + self.inner.get_name() + } + + async fn stream( + &self, + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[rmcp::model::Tool], + ) -> Result { + match self + .inner + .stream(model_config, system, messages, tools) + .await + { + Err(ProviderError::ContextLengthExceeded(_)) => { + let compacted = self.compacted_messages(model_config, messages).await?; + self.inner + .stream(model_config, system, &compacted, tools) + .await + } + other => other, + } + } + + async fn complete( + &self, + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[rmcp::model::Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + match self + .inner + .complete(model_config, system, messages, tools) + .await + { + Err(ProviderError::ContextLengthExceeded(_)) => { + let compacted = self.compacted_messages(model_config, messages).await?; + self.inner + .complete(model_config, system, &compacted, tools) + .await + } + other => other, + } + } + + async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { + self.inner.get_context_limit(model_config).await + } + + fn manages_own_context(&self) -> bool { + true + } +} diff --git a/crates/goose/src/context_mgmt/structured.rs b/crates/goose-context-management/src/structured.rs similarity index 98% rename from crates/goose/src/context_mgmt/structured.rs rename to crates/goose-context-management/src/structured.rs index f5e0dbca9..8507defc6 100644 --- a/crates/goose/src/context_mgmt/structured.rs +++ b/crates/goose-context-management/src/structured.rs @@ -1,4 +1,4 @@ -use crate::prompt_template::render_template; +use crate::templates::{self, COMPACTION_SUMMARY_TEMPLATE}; use goose_providers::json::safely_parse_json; use serde::{Deserialize, Serialize}; @@ -138,7 +138,14 @@ impl StructuredSummary { } pub fn render(&self) -> Result { - render_template("compaction_summary.md", self) + self.render_with( + &templates::builtin_template(COMPACTION_SUMMARY_TEMPLATE) + .expect("builtin compaction summary template"), + ) + } + + pub fn render_with(&self, template: &str) -> Result { + templates::render(template, self) } /// Drops blank entries so a response of blank strings counts as empty diff --git a/crates/goose-context-management/src/summarize.rs b/crates/goose-context-management/src/summarize.rs new file mode 100644 index 000000000..4cb060d5b --- /dev/null +++ b/crates/goose-context-management/src/summarize.rs @@ -0,0 +1,175 @@ +use anyhow::Result; +use goose_providers::conversation::message::{Message, MessageContent}; +use goose_providers::conversation::token_usage::ProviderUsage; +use goose_providers::errors::ProviderError; +use rmcp::model::Role; +use serde::Serialize; +use tracing::warn; + +use crate::format::format_message_for_compacting; +use crate::model::{CompactionModel, TokenEstimator}; +use crate::structured::StructuredSummary; +use crate::templates::{render, Templates}; + +const REMOVAL_PERCENTAGES: [u32; 5] = [0, 10, 20, 50, 100]; + +const SUMMARIZE_REQUEST_TEXT: &str = + "Please summarize the conversation history provided in the system prompt."; + +#[derive(Serialize)] +struct SummarizeContext { + messages: String, +} + +pub struct Summary { + pub message: Message, + pub usage: ProviderUsage, +} + +fn has_tool_response(msg: &Message) -> bool { + msg.content + .iter() + .any(|c| matches!(c, MessageContent::ToolResponse(_))) +} + +/// Drops tool responses from the middle outwards, where context is least +/// likely to matter, to fit an oversized history into the summarizer. +fn filter_tool_responses(messages: &[Message], remove_percent: u32) -> Vec<&Message> { + if remove_percent == 0 { + return messages.iter().collect(); + } + + let tool_indices: Vec = messages + .iter() + .enumerate() + .filter(|(_, msg)| has_tool_response(msg)) + .map(|(i, _)| i) + .collect(); + + if tool_indices.is_empty() { + return messages.iter().collect(); + } + + let num_to_remove = ((tool_indices.len() * remove_percent as usize) / 100).max(1); + let middle = tool_indices.len() / 2; + let mut indices_to_remove = Vec::new(); + + for i in 0..num_to_remove { + let offset = i / 2; + if i % 2 == 0 { + if middle > offset { + indices_to_remove.push(tool_indices[middle - offset - 1]); + } + } else if middle + offset < tool_indices.len() { + indices_to_remove.push(tool_indices[middle + offset]); + } + } + + messages + .iter() + .enumerate() + .filter(|(i, _)| !indices_to_remove.contains(i)) + .map(|(_, msg)| msg) + .collect() +} + +/// When the model didn't follow the structured output format (schema-ignoring +/// models, user-customized prompts), the raw response text is kept unchanged +/// as the summary. +fn apply_structured_summary(response: &mut Message, summary_template: &str) { + let Some(summary) = StructuredSummary::parse(&response.as_concat_text()) else { + return; + }; + match summary.render_with(summary_template) { + Ok(rendered) if !rendered.trim().is_empty() => { + response.content = vec![MessageContent::text(rendered)]; + } + Ok(_) => warn!( + "Structured compaction summary rendered empty (broken template override?), keeping raw output" + ), + Err(e) => warn!("Failed to render structured compaction summary, keeping raw output: {e}"), + } +} + +async fn ensure_usage_tokens( + usage: &mut ProviderUsage, + estimator: &dyn TokenEstimator, + system_prompt: &str, + request: &[Message], + response: &Message, +) { + if usage.usage.input_tokens.is_none() { + let count = estimator.count_chat_tokens(system_prompt, request).await; + usage.usage.input_tokens = Some(count as i32); + } + if usage.usage.output_tokens.is_none() { + let text = response + .content + .iter() + .map(|c| format!("{}", c)) + .collect::>() + .join(" "); + let count = estimator.count_text_tokens(&text).await; + usage.usage.output_tokens = Some(count as i32); + } + if let (Some(input), Some(output)) = (usage.usage.input_tokens, usage.usage.output_tokens) { + usage.usage.total_tokens = Some(input + output); + } +} + +/// Summarizes `messages` into a single user-role message, retrying with +/// progressively more tool responses removed when the summarizer itself +/// overflows its context window. +pub async fn summarize( + model: &dyn CompactionModel, + estimator: Option<&dyn TokenEstimator>, + templates: &Templates, + messages: &[Message], +) -> Result { + let request = vec![Message::user().with_text(SUMMARIZE_REQUEST_TEXT)]; + + for (attempt, &remove_percent) in REMOVAL_PERCENTAGES.iter().enumerate() { + let filtered = filter_tool_responses(messages, remove_percent); + let context = SummarizeContext { + messages: filtered + .iter() + .map(|&msg| format_message_for_compacting(msg)) + .collect::>() + .join("\n"), + }; + let system_prompt = render(&templates.compaction, &context)?; + + match model.complete(&system_prompt, &request).await { + Ok((mut response, mut usage)) => { + response.role = Role::User; + + // Usage must reflect the raw model output (billable tokens), + // so estimate before the response is rewritten to the smaller + // rendered summary. + if let Some(estimator) = estimator { + ensure_usage_tokens(&mut usage, estimator, &system_prompt, &request, &response) + .await; + } + + apply_structured_summary(&mut response, &templates.summary); + + return Ok(Summary { + message: response, + usage, + }); + } + Err(ProviderError::ContextLengthExceeded(_)) + if attempt < REMOVAL_PERCENTAGES.len() - 1 => {} + Err(ProviderError::ContextLengthExceeded(_)) => { + return Err(anyhow::anyhow!( + "Failed to compact: context limit exceeded even after removing all tool responses" + )); + } + Err(e) => return Err(e.into()), + } + } + + Err(anyhow::anyhow!( + "Unexpected: exhausted all attempts without returning" + )) +} diff --git a/crates/goose-context-management/src/templates.rs b/crates/goose-context-management/src/templates.rs new file mode 100644 index 000000000..4e7055d74 --- /dev/null +++ b/crates/goose-context-management/src/templates.rs @@ -0,0 +1,59 @@ +use include_dir::{include_dir, Dir}; +use minijinja::{Environment, Error as MiniJinjaError, Value as MJValue}; +use serde::Serialize; + +static PROMPTS: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/prompts"); + +pub const COMPACTION_TEMPLATE: &str = "compaction.md"; +pub const COMPACTION_SUMMARY_TEMPLATE: &str = "compaction_summary.md"; + +pub fn builtin_template(name: &str) -> Option { + PROMPTS + .get_file(name) + .map(|file| String::from_utf8_lossy(file.contents()).to_string()) +} + +/// Prompt sources for a compaction run, letting callers substitute +/// user-customized templates for the built-in ones. +#[derive(Debug, Clone)] +pub struct Templates { + pub compaction: String, + pub summary: String, +} + +impl Default for Templates { + fn default() -> Self { + Self { + compaction: builtin_template(COMPACTION_TEMPLATE).expect("builtin compaction template"), + summary: builtin_template(COMPACTION_SUMMARY_TEMPLATE) + .expect("builtin compaction summary template"), + } + } +} + +fn code_fence(code: String) -> String { + let longest_run = code + .chars() + .fold((0usize, 0usize), |(max, run), c| { + if c == '`' { + (max.max(run + 1), run + 1) + } else { + (max, 0) + } + }) + .0; + let fence = "`".repeat((longest_run + 1).max(3)); + format!("{fence}\n{}\n{fence}", code.trim_end_matches('\n')) +} + +pub fn render(template: &str, context: &T) -> Result { + let mut env = Environment::new(); + env.set_trim_blocks(true); + env.set_lstrip_blocks(true); + env.add_filter("code_fence", code_fence); + env.add_template("template", template)?; + let rendered = env + .get_template("template")? + .render(MJValue::from_serialize(context))?; + Ok(rendered.trim().to_string()) +} diff --git a/crates/goose-sdk/Cargo.toml b/crates/goose-sdk/Cargo.toml index d185719ca..817f47650 100644 --- a/crates/goose-sdk/Cargo.toml +++ b/crates/goose-sdk/Cargo.toml @@ -21,6 +21,7 @@ required-features = ["uniffi"] default = [] uniffi = [ "dep:uniffi", + "dep:goose-context-management", "dep:thiserror", "dep:anyhow", "dep:goose-providers", @@ -39,6 +40,7 @@ agent-client-protocol-schema = { workspace = true } uniffi = { version = "0.32", features = ["cli"], optional = true } thiserror = { version = "2", optional = true } goose-providers = { version = "0.1.0-alpha.5", path = "../goose-providers", features = ["rustls-tls"], optional = true } +goose-context-management = { version = "0.1.0-alpha.5", path = "../goose-context-management", optional = true } futures = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } tokio = { workspace = true, features = ["rt-multi-thread", "sync"], optional = true } diff --git a/crates/goose-sdk/src/bindings.rs b/crates/goose-sdk/src/bindings.rs index 4852a3263..79cbb1c65 100644 --- a/crates/goose-sdk/src/bindings.rs +++ b/crates/goose-sdk/src/bindings.rs @@ -25,7 +25,8 @@ use goose_providers::{ utils::sanitize_unicode_tags, }; use rmcp::model::{ - CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData, Role, Tool, + CallToolRequestParams, CallToolResult, ContentBlock as Content, ErrorCode, ErrorData, Role, + Tool, }; use serde_json::Value; @@ -755,6 +756,95 @@ impl Provider { ) -> Result { self.handle.complete(model, system, messages, tools).await } + + /// Summarizes a conversation down to a single message so it can continue + /// past this model's context window. + pub async fn compact( + &self, + model_name: String, + messages: Vec, + templates: Option, + ) -> Result { + let messages: Vec = messages + .iter() + .map(CompactionMessage::to_goose_message) + .collect(); + let templates = templates.map(Into::into).unwrap_or_default(); + let model = goose_context_management::ProviderModel::new( + self.handle.provider.clone(), + ModelConfig::new(&model_name), + ); + + let summary = run_on_runtime(async move { + goose_context_management::summarize(&model, None, &templates, &messages).await + }) + .await? + .map_err(GooseError::generic)?; + + Ok(CompactionSummary { + text: summary.message.as_concat_text(), + input_tokens: summary.usage.usage.input_tokens, + output_tokens: summary.usage.usage.output_tokens, + total_tokens: summary.usage.usage.total_tokens, + }) + } +} + +/// A text-only message. Compaction reads conversations as text, so this is the +/// whole input shape callers need across the language boundary. +#[derive(Debug, Clone, uniffi::Record)] +pub struct CompactionMessage { + pub role: MessageRole, + pub text: String, +} + +impl CompactionMessage { + fn to_goose_message(&self) -> Message { + let role = match self.role { + MessageRole::User | MessageRole::Tool => Role::User, + MessageRole::Assistant => Role::Assistant, + }; + let mut message = match role { + Role::User => Message::user(), + Role::Assistant => Message::assistant(), + } + .with_text(&self.text); + message.role = role; + message + } +} + +/// Overrides for the summarization and summary-rendering prompts. +#[derive(Debug, Clone, uniffi::Record)] +pub struct CompactionTemplates { + pub compaction: String, + pub summary: String, +} + +impl From for goose_context_management::Templates { + fn from(value: CompactionTemplates) -> Self { + Self { + compaction: value.compaction, + summary: value.summary, + } + } +} + +#[derive(Debug, Clone, uniffi::Record)] +pub struct CompactionSummary { + pub text: String, + pub input_tokens: Option, + pub output_tokens: Option, + pub total_tokens: Option, +} + +#[uniffi::export] +pub fn default_compaction_templates() -> CompactionTemplates { + let templates = goose_context_management::Templates::default(); + CompactionTemplates { + compaction: templates.compaction, + summary: templates.summary, + } } #[uniffi::export] diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index cef1da9d8..f826767f5 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -135,6 +135,7 @@ strum = { workspace = true } once_cell = { workspace = true } etcetera = { workspace = true } fs-err = { version = "3.1", default-features = false } +goose-context-management = { path = "../goose-context-management" } goose-providers = { path = "../goose-providers", default-features = false } goose-download-manager = { path = "../goose-download-manager" } goose-sdk-types = { path = "../goose-sdk-types" } diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index a451dcec2..58058bf13 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -1,10 +1,8 @@ -pub mod structured; +pub use goose_context_management::structured; -use crate::context_mgmt::structured::StructuredSummary; -use crate::conversation::message::{ActionRequiredData, MessageMetadata}; +use crate::conversation::message::MessageMetadata; use crate::conversation::message::{Message, MessageContent}; use crate::conversation::{merge_consecutive_messages, Conversation}; -use crate::prompt_template::render_template; use crate::providers::base::Provider; #[cfg(test)] use crate::providers::base::{stream_from_single_message, MessageStream}; @@ -17,13 +15,12 @@ use indoc::indoc; use rmcp::model::Role; #[cfg(test)] use rmcp::model::{Annotations, ContentBlock, TextContent}; -use serde::Serialize; use std::sync::Arc; use tokio::task::JoinHandle; use tracing::info; use tracing::log::warn; -pub const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.8; +pub use goose_context_management::DEFAULT_COMPACTION_THRESHOLD; pub(crate) const TOOLCALL_SUMMARIZATION_BATCH_SIZE: usize = 10; @@ -48,11 +45,6 @@ const MANUAL_COMPACT_CONTINUATION_TEXT: &str = Do not mention that you read a summary or that conversation summarization occurred. Just continue the conversation naturally based on the summarized context."; -#[derive(Serialize)] -struct SummarizeContext { - messages: String, -} - pub struct CompactionResult { pub conversation: Conversation, /// Billable usage of the summarization call, counting the raw model @@ -284,54 +276,61 @@ pub async fn check_if_compaction_needed( Ok(needs_compaction) } -fn filter_tool_responses(messages: &[Message], remove_percent: u32) -> Vec<&Message> { - fn has_tool_response(msg: &Message) -> bool { - msg.content - .iter() - .any(|c| matches!(c, MessageContent::ToolResponse(_))) +struct FastCompactionModel<'a> { + provider: &'a dyn Provider, + model_config: &'a ModelConfig, + session_id: &'a str, +} + +#[async_trait::async_trait] +impl goose_context_management::CompactionModel for FastCompactionModel<'_> { + async fn complete( + &self, + system: &str, + messages: &[Message], + ) -> Result<(Message, ProviderUsage), ProviderError> { + crate::model_config::complete_fast( + self.provider, + self.model_config, + self.session_id, + system, + messages, + &[], + ) + .await } +} - if remove_percent == 0 { - return messages.iter().collect(); - } +struct GooseTokenEstimator; - let tool_indices: Vec = messages - .iter() - .enumerate() - .filter(|(_, msg)| has_tool_response(msg)) - .map(|(i, _)| i) - .collect(); - - if tool_indices.is_empty() { - return messages.iter().collect(); - } - - let num_to_remove = ((tool_indices.len() * remove_percent as usize) / 100).max(1); - - let middle = tool_indices.len() / 2; - let mut indices_to_remove = Vec::new(); - - // Middle out - for i in 0..num_to_remove { - if i % 2 == 0 { - let offset = i / 2; - if middle > offset { - indices_to_remove.push(tool_indices[middle - offset - 1]); - } - } else { - let offset = i / 2; - if middle + offset < tool_indices.len() { - indices_to_remove.push(tool_indices[middle + offset]); +#[async_trait::async_trait] +impl goose_context_management::TokenEstimator for GooseTokenEstimator { + async fn count_chat_tokens(&self, system: &str, messages: &[Message]) -> usize { + match create_token_counter().await { + Ok(counter) => counter.count_chat_tokens(system, messages, &[]), + Err(error) => { + warn!("Failed to create token counter: {error}"); + 0 } } } - messages - .iter() - .enumerate() - .filter(|(i, _)| !indices_to_remove.contains(i)) - .map(|(_, msg)| msg) - .collect() + async fn count_text_tokens(&self, text: &str) -> usize { + match create_token_counter().await { + Ok(counter) => counter.count_tokens(text), + Err(error) => { + warn!("Failed to create token counter: {error}"); + 0 + } + } + } +} + +fn compaction_templates() -> Result { + Ok(goose_context_management::Templates { + compaction: crate::prompt_template::template_source("compaction.md")?, + summary: crate::prompt_template::template_source("compaction_summary.md")?, + }) } async fn do_compact( @@ -349,182 +348,23 @@ async fn do_compact( ) .agent_visible_messages(); - // Try progressively removing more tool response messages from the middle to reduce context length - let removal_percentages = [0, 10, 20, 50, 100]; - - for (attempt, &remove_percent) in removal_percentages.iter().enumerate() { - let filtered_messages = filter_tool_responses(&agent_visible_messages, remove_percent); - - let messages_text = filtered_messages - .iter() - .map(|&msg| format_message_for_compacting(msg)) - .collect::>() - .join("\n"); - - let context = SummarizeContext { - messages: messages_text, - }; - - let system_prompt = render_template("compaction.md", &context)?; - - let user_message = Message::user() - .with_text("Please summarize the conversation history provided in the system prompt."); - let summarization_request = vec![user_message]; - - match crate::model_config::complete_fast( - provider, - model_config, - session_id, - &system_prompt, - &summarization_request, - &[], - ) - .await - { - Ok((mut response, mut provider_usage)) => { - response.role = Role::User; - - // Usage must reflect the raw model output (billable tokens), - // so estimate before the response is rewritten to the smaller - // rendered summary. - crate::providers::usage_estimator::ensure_usage_tokens( - &mut provider_usage, - &system_prompt, - &summarization_request, - &response, - &[], - ) - .await - .map_err(|e| anyhow::anyhow!("Failed to ensure usage tokens: {}", e))?; - - apply_structured_summary(&mut response); - - return Ok((response, provider_usage)); - } - Err(e) => { - if matches!(e, ProviderError::ContextLengthExceeded(_)) { - if attempt < removal_percentages.len() - 1 { - continue; - } else { - return Err(anyhow::anyhow!( - "Failed to compact: context limit exceeded even after removing all tool responses" - )); - } - } - return Err(e.into()); - } - } - } - - Err(anyhow::anyhow!( - "Unexpected: exhausted all attempts without returning" - )) -} - -/// When the model didn't follow the structured output format (schema-ignoring -/// models, user-customized prompts), the raw response text is kept unchanged -/// as the summary. -fn apply_structured_summary(response: &mut Message) { - let Some(summary) = StructuredSummary::parse(&response.as_concat_text()) else { - return; + let model = FastCompactionModel { + provider, + model_config, + session_id, }; - match summary.render() { - Ok(rendered) if !rendered.trim().is_empty() => { - response.content = vec![MessageContent::text(rendered)]; - } - Ok(_) => warn!( - "Structured compaction summary rendered empty (broken template override?), keeping raw output" - ), - Err(e) => warn!( - "Failed to render structured compaction summary, keeping raw output: {}", - e - ), - } + let summary = goose_context_management::summarize( + &model, + Some(&GooseTokenEstimator), + &compaction_templates()?, + &agent_visible_messages, + ) + .await?; + + Ok((summary.message, summary.usage)) } -pub fn format_message_for_compacting(msg: &Message) -> String { - let content_parts: Vec = msg - .content - .iter() - .filter_map(|content| match content { - MessageContent::Text(text) => Some(text.text.clone()), - MessageContent::Image(img) => Some(format!("[image: {}]", img.mime_type)), - MessageContent::ToolRequest(req) => { - if let Ok(call) = &req.tool_call { - Some(format!( - "tool_request({}): {}", - call.name, - serde_json::to_string(&call.arguments) - .unwrap_or_else(|_| "<>".to_string()) - )) - } else { - Some("tool_request: [error]".to_string()) - } - } - MessageContent::ToolResponse(res) => { - if let Ok(result) = &res.tool_result { - let text_items: Vec = result - .content - .iter() - .filter_map(|content| { - content.as_text().map(|text_str| text_str.text.clone()) - }) - .collect(); - - if !text_items.is_empty() { - Some(format!("tool_response: {}", text_items.join("\n"))) - } else { - Some("tool_response: [non-text content]".to_string()) - } - } else { - Some("tool_response: [error]".to_string()) - } - } - MessageContent::ToolConfirmationRequest(req) => { - Some(format!("tool_confirmation_request: {}", req.tool_name)) - } - MessageContent::ActionRequired(action) => match &action.data { - ActionRequiredData::ToolConfirmation { tool_name, .. } => { - Some(format!("action_required(tool_confirmation): {}", tool_name)) - } - ActionRequiredData::Elicitation { message, .. } => { - Some(format!("action_required(elicitation): {}", message)) - } - ActionRequiredData::ElicitationResponse { id, .. } => { - Some(format!("action_required(elicitation_response): {}", id)) - } - ActionRequiredData::ToolConfirmationResponse { id, .. } => Some(format!( - "action_required(tool_confirmation_response): {}", - id - )), - }, - MessageContent::FrontendToolRequest(req) => { - if let Ok(call) = &req.tool_call { - Some(format!("frontend_tool_request: {}", call.name)) - } else { - Some("frontend_tool_request: [error]".to_string()) - } - } - MessageContent::Thinking(_) => None, - MessageContent::RedactedThinking(_) => None, - MessageContent::SystemNotification(notification) => { - Some(format!("system_notification: {}", notification.msg)) - } - MessageContent::Error(error) => Some(format!("error: {}", error.message)), - }) - .collect(); - - let role_str = match msg.role { - Role::User => "user", - Role::Assistant => "assistant", - }; - - if content_parts.is_empty() { - format!("[{}]: ", role_str) - } else { - format!("[{}]: {}", role_str, content_parts.join("\n")) - } -} +pub use goose_context_management::format_message_for_compacting; pub fn compute_tool_call_cutoff(context_limit: usize, compaction_threshold: f64) -> usize { let threshold = if compaction_threshold > 0.0 && compaction_threshold <= 1.0 { diff --git a/crates/goose/src/prompt_template.rs b/crates/goose/src/prompt_template.rs index 4d662b611..cd0055d81 100644 --- a/crates/goose/src/prompt_template.rs +++ b/crates/goose/src/prompt_template.rs @@ -64,6 +64,13 @@ pub struct Template { pub is_customized: bool, } +fn builtin_content(name: &str) -> Option { + CORE_PROMPTS_DIR + .get_file(name) + .map(|file| String::from_utf8_lossy(file.contents()).to_string()) + .or_else(|| goose_context_management::templates::builtin_template(name)) +} + fn user_prompts_dir() -> PathBuf { Paths::config_dir().join("prompts") } @@ -121,24 +128,39 @@ pub fn render_template(name: &str, context: &T) -> Result Result { + let user_path = user_prompts_dir().join(name); + if user_path.exists() { + return std::fs::read_to_string(&user_path).map_err(|e| { + MiniJinjaError::new( + minijinja::ErrorKind::InvalidOperation, + format!("Failed to read user template: {}", e), + ) + }); + } + builtin_content(name).ok_or_else(|| { + MiniJinjaError::new( + minijinja::ErrorKind::TemplateNotFound, + format!("Built-in template '{}' not found", name), + ) + }) +} + pub fn get_template(name: &str) -> Option