diff --git a/Cargo.lock b/Cargo.lock index d553e17d0ba..624d83bf136 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1202,8 +1202,10 @@ dependencies = [ "getrandom 0.4.3", "goose", "goose-agent", + "goose-context-management", "goose-provider-types", "goose-providers", + "goose-sdk-types", "hex", "idna_adapter", "nix 0.31.3", diff --git a/crates/buzz-agent/Cargo.toml b/crates/buzz-agent/Cargo.toml index 434f8c475dd..f816bd1fcc0 100644 --- a/crates/buzz-agent/Cargo.toml +++ b/crates/buzz-agent/Cargo.toml @@ -34,6 +34,8 @@ goose = { git = "https://github.com/aaif-goose/goose", rev = "0a9749b1cbf38d1820 # exact same Goose revision as `goose` so shared conversation types stay unified. goose-agent = { git = "https://github.com/aaif-goose/goose", rev = "0a9749b1cbf38d182080a53286c4942629f66565" } goose-provider-types = { git = "https://github.com/aaif-goose/goose", rev = "0a9749b1cbf38d182080a53286c4942629f66565" } +goose-context-management = { git = "https://github.com/aaif-goose/goose", rev = "0a9749b1cbf38d182080a53286c4942629f66565" } +goose-sdk-types = { git = "https://github.com/aaif-goose/goose", rev = "0a9749b1cbf38d182080a53286c4942629f66565" } # ModelConfig: goose re-exports it from goose-providers; goose::model_config # keeps its own copy private. goose-providers = { git = "https://github.com/aaif-goose/goose", rev = "0a9749b1cbf38d182080a53286c4942629f66565" } diff --git a/crates/buzz-agent/src/lib.rs b/crates/buzz-agent/src/lib.rs index 8aff009334f..eab24b3f7ef 100644 --- a/crates/buzz-agent/src/lib.rs +++ b/crates/buzz-agent/src/lib.rs @@ -902,7 +902,7 @@ async fn cancel_session(app: &Arc, params: Value) { /// Construct the goose provider. See [`crate::provider`]. async fn build_provider( provider_name: &str, -) -> Result, AgentError> { +) -> Result, AgentError> { crate::provider::build(provider_name).await } diff --git a/crates/buzz-agent/src/loop_drive.rs b/crates/buzz-agent/src/loop_drive.rs index 17f1fc6e4d4..721dae97087 100644 --- a/crates/buzz-agent/src/loop_drive.rs +++ b/crates/buzz-agent/src/loop_drive.rs @@ -13,7 +13,7 @@ //! | tool surface | `Agent::list_tools` | //! | tool execution | `Agent::dispatch_tool_call` | //! | system prompt | `Agent::build_turn_system_prompt` → `PromptManager` | -//! | compaction | `goose::context_mgmt::{check_if_compaction_needed, compact_messages}` | +//! | compaction | `goose_context_management` with buzz threshold policy | //! | conversation store | buzz's own [`crate::turn_state::TurnState`], in memory | //! //! Everything that decides *turn shape* stays here: round structure, the @@ -158,13 +158,8 @@ pub async fn run_turn( // The turn's conversation lives here, not in a database. See // `crate::turn_state` for why goose never needed one. - let (_provider, model_config, _model_id) = ctx.model.snapshot().await; - let model_config = Some(model_config); - let mut state = crate::turn_state::TurnState::new( - ctx.session_id.to_string(), - ctx.working_dir.clone(), - model_config, - ); + let mut state = + crate::turn_state::TurnState::new(ctx.session_id.to_string(), ctx.working_dir.clone()); for message in ctx.history.iter().cloned() { state.push(message); } @@ -203,7 +198,6 @@ pub async fn run_turn( ctx.model.clone(), Arc::clone(ctx.mcp), ctx.hook_extension.map(str::to_string), - ctx.session_id.to_string(), ), ctx.cancel.clone(), ); @@ -621,7 +615,7 @@ fn warn_if_silent_turn(published: bool, text_is_empty: bool, output_tokens: Opti fn accumulate_usage( tokens: &mut super::agent::TurnTokens, - usage: &goose::providers::base::ProviderUsage, + usage: &goose_provider_types::conversation::token_usage::ProviderUsage, ) { let u = &usage.usage; if let Some(i) = u.input_tokens { @@ -716,7 +710,7 @@ mod tests { use super::*; fn turn_state() -> crate::turn_state::TurnState { - crate::turn_state::TurnState::new("s".to_string(), std::path::PathBuf::from("/tmp"), None) + crate::turn_state::TurnState::new("s".to_string(), std::path::PathBuf::from("/tmp")) } fn usage( @@ -726,7 +720,7 @@ mod tests { total: Option, read: Option, write: Option, - ) -> goose::providers::base::ProviderUsage { + ) -> goose_provider_types::conversation::token_usage::ProviderUsage { let mut provider_usage = goose_provider_types::conversation::token_usage::Usage::new( Some(input), Some(output), @@ -736,7 +730,10 @@ mod tests { // `Usage::new` synthesizes a total when absent; tests need to exercise // a provider response that genuinely omitted it. provider_usage.total_tokens = total; - goose::providers::base::ProviderUsage::new(model.to_string(), provider_usage) + goose_provider_types::conversation::token_usage::ProviderUsage::new( + model.to_string(), + provider_usage, + ) } #[test] diff --git a/crates/buzz-agent/src/mcp.rs b/crates/buzz-agent/src/mcp.rs index 63738a668d2..495fd0a067f 100644 --- a/crates/buzz-agent/src/mcp.rs +++ b/crates/buzz-agent/src/mcp.rs @@ -166,7 +166,7 @@ pub struct McpRegistry { init_timeout: Duration, tool_timeout: Duration, hook_timeout: Duration, - skills: Vec, + skills: Vec, } impl McpRegistry { @@ -219,7 +219,7 @@ impl McpRegistry { skills: goose::skills::discover_skills(Some(std::path::Path::new(cwd))) .into_iter() .filter(|skill| { - skill.source_type != goose::custom_requests::SourceType::BuiltinSkill + skill.source_type != goose_sdk_types::custom_requests::SourceType::BuiltinSkill }) .collect(), }; @@ -370,7 +370,7 @@ impl McpRegistry { tools } - pub fn skills(&self) -> &[goose::custom_requests::SourceEntry] { + pub fn skills(&self) -> &[goose_sdk_types::custom_requests::SourceEntry] { &self.skills } diff --git a/crates/buzz-agent/src/model.rs b/crates/buzz-agent/src/model.rs index 15792edf6e9..fd4d084f3bd 100644 --- a/crates/buzz-agent/src/model.rs +++ b/crates/buzz-agent/src/model.rs @@ -18,7 +18,7 @@ use std::sync::Arc; -use goose::providers::base::Provider; +use goose_providers::base::Provider; use goose_providers::model::ModelConfig; use tokio::sync::RwLock; @@ -97,7 +97,7 @@ mod tests { _system: &str, _messages: &[Message], _tools: &[rmcp::model::Tool], - ) -> Result { + ) -> Result { unreachable!("tests never call the model") } } diff --git a/crates/buzz-agent/src/ops.rs b/crates/buzz-agent/src/ops.rs index ac0eb28aabc..971a03f90f0 100644 --- a/crates/buzz-agent/src/ops.rs +++ b/crates/buzz-agent/src/ops.rs @@ -29,8 +29,10 @@ use goose_agent::operation::{ applied, assistant_turn_count, messages_since_kickoff, not_applicable, yielded, ConversationEffect, Emitter, Operation, OperationResult, }; -use goose_provider_types::conversation::message::{Message, ToolRequest}; -use goose_provider_types::conversation::Conversation; +use goose_provider_types::conversation::message::{ + Message, MessageContent, MessageMetadata, ToolRequest, +}; +use goose_provider_types::conversation::{merge_consecutive_messages, Conversation}; use crate::types::StopReason; @@ -413,12 +415,131 @@ impl Operation for BuzzSteerOperation { } } +fn auto_compact_threshold() -> f64 { + std::env::var("GOOSE_AUTO_COMPACT_THRESHOLD") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(goose_context_management::DEFAULT_COMPACTION_THRESHOLD) +} + +fn context_limit_override() -> Result> { + std::env::var("GOOSE_CONTEXT_LIMIT") + .ok() + .map(|value| { + value + .parse::() + .map_err(|error| anyhow::anyhow!("invalid GOOSE_CONTEXT_LIMIT {value:?}: {error}")) + }) + .transpose() +} + +fn needs_compaction( + provider_manages_context: bool, + total_tokens: Option, + context_limit: usize, + threshold: f64, +) -> bool { + if provider_manages_context || threshold <= 0.0 || threshold >= 1.0 { + return false; + } + context_limit > 0 + && total_tokens.is_some_and(|tokens| { + let current_tokens = tokens.max(0) as f64; + current_tokens / context_limit as f64 > threshold + }) +} + +fn compacted_conversation(conversation: &Conversation, summary: Message) -> Conversation { + const CONTINUATION: &str = "Your context was compacted. The previous message contains a summary of the conversation so far.\nDo not mention that you read a summary or that conversation summarization occurred.\nJust continue the conversation naturally based on the summarized context."; + const TOOL_CONTINUATION: &str = "Your context was compacted. The previous message contains a summary of the conversation so far.\nDo not mention that you read a summary or that conversation summarization occurred.\nContinue calling tools as necessary to complete the task."; + + let messages = conversation.messages(); + let preserved = messages + .iter() + .enumerate() + .rev() + .find_map(|(index, message)| { + if !message.is_agent_visible() + || message.is_turn_context() + || message.role != rmcp::model::Role::User + { + return None; + } + let projected = message.agent_visible_content(); + let has_text = projected + .content + .iter() + .any(|content| matches!(content, MessageContent::Text(_))); + let has_tool_content = projected.content.iter().any(|content| { + matches!( + content, + MessageContent::ToolRequest(_) | MessageContent::ToolResponse(_) + ) + }); + if !has_text || has_tool_content { + return None; + } + let message = projected + .content + .into_iter() + .filter(|content| matches!(content, MessageContent::Text(_))) + .fold( + Message::user().with_metadata(MessageMetadata::agent_only()), + Message::with_content, + ); + Some((index, message)) + }); + + let is_most_recent = preserved + .as_ref() + .is_some_and(|(index, _)| messages[*index + 1..].iter().all(Message::is_turn_context)); + let continuation_text = if is_most_recent { + CONTINUATION + } else { + TOOL_CONTINUATION + }; + + let mut compacted = messages + .iter() + .cloned() + .map(|message| { + let metadata = message.metadata.clone().with_agent_invisible(); + message.with_metadata(metadata) + }) + .collect::>(); + let summary = summary.with_metadata(MessageMetadata::agent_only()); + let continuation = Message::assistant() + .with_text(continuation_text) + .with_metadata(MessageMetadata::agent_only()); + let continuation_created = continuation.created; + let (continuation, _) = merge_consecutive_messages(vec![summary, continuation]); + compacted.extend(continuation); + + if let Some((index, mut message)) = preserved { + message.created = continuation_created; + compacted.push(message); + if let Some(turn_context) = messages[index + 1..] + .iter() + .rev() + .find(|message| message.is_turn_context() && message.is_agent_visible()) + { + let mut carried = turn_context.clone(); + carried.id = None; + if let Some(latest) = compacted.iter().map(|message| message.created).max() { + carried.created = carried.created.max(latest); + } + compacted.push(carried); + } + } + + Conversation::new_unvalidated(compacted) +} + /// Compacts the conversation when it approaches the context limit. /// -/// The mechanism is entirely goose's (`check_if_compaction_needed` / -/// `compact_messages`); this operation is the buzz-specific part around it: -/// the `_PostCompact` hook that re-injects buzz-dev-mcp's todo state, which is -/// what buzz-agent's old context-handoff existed to preserve. +/// Summarization comes from the standalone Goose context-management crate. +/// This operation owns the threshold policy and the Buzz-specific +/// `_PostCompact` hook that re-injects buzz-dev-mcp's todo state. /// /// Not goose's `CompactionOperation`: that one yields to the client and emits /// its own user-facing notification. In buzz a yield ends the turn and the @@ -428,7 +549,6 @@ pub struct BuzzCompactionOperation { model: crate::model::SessionModel, mcp: Arc, hook_extension: Option, - session_id: String, } impl BuzzCompactionOperation { @@ -436,13 +556,11 @@ impl BuzzCompactionOperation { model: crate::model::SessionModel, mcp: Arc, hook_extension: Option, - session_id: String, ) -> Self { Self { model, mcp, hook_extension, - session_id, } } } @@ -460,39 +578,35 @@ impl Operation for BuzzCompactionOperation { _emit: &Emitter, ) -> Result> { let (provider, model_config, _model_id) = self.model.snapshot().await; - // Goose's legacy threshold facade still takes its application Session. - // Keep that coupling at this single adapter until the granular context - // crate exposes thresholding over model config + token occupancy. - let legacy_session = goose::session::Session { - id: session.id.clone(), - working_dir: session.working_dir.clone(), - model_config: session.model_config.clone(), - usage: goose_provider_types::conversation::token_usage::Usage { - total_tokens: session.total_tokens, - ..Default::default() - }, - conversation: Some(conversation.clone()), - ..Default::default() - }; - if !goose::context_mgmt::check_if_compaction_needed( - provider.as_ref(), - conversation, - None, - &legacy_session, - ) - .await? - { + let threshold = auto_compact_threshold(); + let context_limit_override = context_limit_override()?; + let context_limit = provider + .get_context_limit(&model_config.model_name, context_limit_override) + .await; + if !needs_compaction( + provider.manages_own_context(), + session.total_tokens, + context_limit, + threshold, + ) { return not_applicable(); } - let result = goose::context_mgmt::compact_messages( - provider.as_ref(), - &model_config, - &self.session_id, - conversation, - false, + let visible_messages = conversation + .messages() + .iter() + .filter(|message| message.is_agent_visible() && !message.is_turn_context()) + .cloned() + .collect::>(); + let model = goose_context_management::ProviderModel::new(provider, model_config); + let summary = goose_context_management::summarize( + &model, + None, + &goose_context_management::Templates::default(), + &visible_messages, ) .await?; + let compacted = compacted_conversation(conversation, summary.message); tracing::info!(target: "buzz_agent::compaction", "history compacted"); @@ -500,7 +614,7 @@ impl Operation for BuzzCompactionOperation { // the running total is reset by the driving loop's `apply_effects` // instead: the old total described a conversation that no longer // exists, and carrying it forward would re-trigger compaction at once. - let mut effects = vec![ConversationEffect::ReplaceConversation(result.conversation)]; + let mut effects = vec![ConversationEffect::ReplaceConversation(compacted)]; if let Some(extension) = &self.hook_extension { if let Some(reported) = crate::hooks::post_compact_state(&self.mcp, extension).await { @@ -626,6 +740,47 @@ mod tests { .with_visibility(false, true) } + #[test] + fn compaction_threshold_requires_reported_occupancy_above_the_boundary() { + assert!(!needs_compaction(false, None, 100_000, 0.8)); + assert!(!needs_compaction(false, Some(80_000), 100_000, 0.8)); + assert!(needs_compaction(false, Some(80_001), 100_000, 0.8)); + assert!(!needs_compaction(true, Some(90_000), 100_000, 0.8)); + assert!(!needs_compaction(false, Some(90_000), 100_000, 0.0)); + assert!(!needs_compaction(false, Some(90_000), 100_000, 1.0)); + assert!(!needs_compaction(false, Some(1), 0, 0.8)); + } + + #[test] + fn compaction_preserves_the_latest_prompt_and_turn_context() { + let first = Message::user().with_text("old prompt"); + let answer = Message::assistant().with_text("old answer"); + let latest = Message::user().with_text("latest prompt"); + let mut turn_context = Message::user() + .with_text("current context") + .with_visibility(false, true); + turn_context.metadata.turn_context = true; + let compacted = compacted_conversation( + &Conversation::new_unvalidated(vec![first, answer, latest, turn_context]), + Message::assistant().with_text("summary"), + ); + let visible = compacted + .messages() + .iter() + .filter(|message| message.is_agent_visible()) + .collect::>(); + + assert!(visible + .iter() + .any(|message| message.as_concat_text().contains("summary"))); + assert!(visible + .iter() + .any(|message| message.as_concat_text() == "latest prompt")); + assert!(visible.iter().any(|message| message.is_turn_context())); + assert!(!compacted.messages()[0].is_agent_visible()); + assert!(!compacted.messages()[1].is_agent_visible()); + } + #[tokio::test] async fn under_budget_does_not_apply() { let outcome = Outcome::new(); diff --git a/crates/buzz-agent/src/prompt.rs b/crates/buzz-agent/src/prompt.rs index 514be7d41fc..6d80d661c79 100644 --- a/crates/buzz-agent/src/prompt.rs +++ b/crates/buzz-agent/src/prompt.rs @@ -15,7 +15,7 @@ use std::path::Path; use std::sync::Arc; use goose::agents::PromptManager; -use goose::config::GooseMode; +use goose_provider_types::goose_mode::GooseMode; use tokio::sync::Mutex; /// The system prompt for one session. diff --git a/crates/buzz-agent/src/provider.rs b/crates/buzz-agent/src/provider.rs index 0371e3607f3..a6c9816dee9 100644 --- a/crates/buzz-agent/src/provider.rs +++ b/crates/buzz-agent/src/provider.rs @@ -32,7 +32,7 @@ use std::sync::Arc; -use goose::providers::base::Provider; +use goose_providers::base::Provider; use goose_providers::databricks_auth::{ DatabricksAuth, DatabricksOauthTokenProvider, DatabricksRefreshHook, }; @@ -112,9 +112,7 @@ fn databricks_v2_with_buzz_oauth(host: String) -> Result, Agen .ok() .filter(|token| !token.trim().is_empty()) })), - // Keeps the `agent-session-id` request header goose attaches for - // provider-side attribution; dropping it would silently lose that. - Some(goose::session_context::session_id_request_builder()), + None, refresh_hook, ) .map_err(|e| AgentError::Llm(format!("databricks provider: {e}")))?; diff --git a/crates/buzz-agent/src/skills.rs b/crates/buzz-agent/src/skills.rs index b936d35f35d..25e23d3bc99 100644 --- a/crates/buzz-agent/src/skills.rs +++ b/crates/buzz-agent/src/skills.rs @@ -12,7 +12,7 @@ //! model the tool is worth calling. The wording matches goose's own so the two //! do not drift into describing the same tool differently. -use goose::custom_requests::{SourceEntry, SourceType}; +use goose_sdk_types::custom_requests::{SourceEntry, SourceType}; /// Provider-native definition for the in-process skill loader. pub fn load_skill_tool() -> rmcp::model::Tool { diff --git a/crates/buzz-agent/src/turn_state.rs b/crates/buzz-agent/src/turn_state.rs index cce3ed8ef42..db074cbca83 100644 --- a/crates/buzz-agent/src/turn_state.rs +++ b/crates/buzz-agent/src/turn_state.rs @@ -15,8 +15,6 @@ pub struct TurnSession { pub id: String, /// Working directory used by prompt discovery and tool hints. pub working_dir: std::path::PathBuf, - /// Model configuration used by the temporary legacy compaction adapter. - pub model_config: Option, /// Latest provider-reported occupancy for the current conversation. pub total_tokens: Option, conversation: Conversation, @@ -38,16 +36,11 @@ pub struct TurnState { } impl TurnState { - pub fn new( - id: String, - working_dir: std::path::PathBuf, - model_config: Option, - ) -> Self { + pub fn new(id: String, working_dir: std::path::PathBuf) -> Self { Self { session: TurnSession { id, working_dir, - model_config, total_tokens: None, conversation: Conversation::new_unvalidated(Vec::new()), }, @@ -82,7 +75,7 @@ mod tests { use super::*; fn state() -> TurnState { - TurnState::new("s1".to_string(), std::path::PathBuf::from("/tmp"), None) + TurnState::new("s1".to_string(), std::path::PathBuf::from("/tmp")) } #[test] @@ -110,7 +103,7 @@ mod tests { #[test] fn runtime_metadata_is_preserved() { - let mut state = TurnState::new("abc".to_string(), std::path::PathBuf::from("/work"), None); + let mut state = TurnState::new("abc".to_string(), std::path::PathBuf::from("/work")); state.set_total_tokens(Some(1234)); assert_eq!(state.session().id, "abc"); assert_eq!( diff --git a/crates/buzz-agent/src/types.rs b/crates/buzz-agent/src/types.rs index 33119e5abae..2cf6dd2bce1 100644 --- a/crates/buzz-agent/src/types.rs +++ b/crates/buzz-agent/src/types.rs @@ -7,7 +7,7 @@ //! `context_pressure_bytes`) that fed the bespoke handoff heuristic. //! //! Goose owns all of that now: conversation state is `goose::conversation`, -//! tool plumbing is `rmcp`, and compaction is `goose::context_mgmt`. What +//! tool plumbing is `rmcp`, and compaction is `goose_context_management`. What //! survives here is only what crosses the ACP wire. use serde::Deserialize;