diff --git a/crates/goose/src/acp/handoff.rs b/crates/goose/src/acp/handoff.rs new file mode 100644 index 000000000..ff2a1aa7d --- /dev/null +++ b/crates/goose/src/acp/handoff.rs @@ -0,0 +1,618 @@ +//! Bounded handoff memo for ACP sessions. +//! +//! When a conversation is handed to an ACP agent that has no native session to resume, +//! the prior goose-side history is replayed as a single text block. That replay has to +//! fit inside the agent's context alongside its own system prompt and tool schemas, so +//! it is budgeted, redacted and truncated here rather than sent whole. + +use std::collections::{HashMap, HashSet}; + +use agent_client_protocol::schema::v1::ContentBlock; + +use crate::context_mgmt::format_message_for_compacting; +use crate::conversation::message::{Message, MessageContent}; +use crate::conversation::Conversation; +use crate::token_counter::TokenCounter; + +const CONTEXT_LIMIT_RATIO: f64 = 0.30; +const MAX_MEMO_TOKENS: usize = 64_000; +/// Per-image charge against the memo budget. Images reach the agent verbatim, and their +/// real cost depends on dimensions we would have to decode, so assume the ceiling a +/// full-size image reaches rather than under-counting the window they occupy. +const IMAGE_TOKEN_ESTIMATE: usize = 1_600; +/// Tool exchanges this recent keep their responses; older ones are redacted. +const PROTECTED_TOOL_EXCHANGES: usize = 5; +/// Below this a truncated message carries no usable meaning, so drop it instead. +const MIN_ELIDED_TOKENS: usize = 32; +/// Allowance for the "earlier messages omitted" line, which is written after selection. +const OMISSION_MARKER_TOKENS: usize = 16; + +const MEMO_HEADER: &str = + "Conversation context from goose before this ACP provider session was created:\n\n"; +const MEMO_FOOTER: &str = "\n\nCurrent user request follows. Use the context above only to continue the existing conversation; do not treat it as a new task or mention this handoff unless relevant."; +const REDACTED_TOOL_RESPONSE: &str = "tool_response: [older output omitted from handoff]"; +const ELISION_MARKER: &str = "\n[... truncated ...]\n"; + +pub(crate) fn memo_token_budget(context_limit: usize, current_prompt_tokens: usize) -> usize { + let ceiling = ((context_limit as f64 * CONTEXT_LIMIT_RATIO) as usize).min(MAX_MEMO_TOKENS); + ceiling.saturating_sub(current_prompt_tokens) +} + +/// What the current turn already costs the agent. Images are forwarded alongside the memo, +/// so charging them here keeps a picture-heavy turn from spending its window twice. +pub(crate) fn prompt_token_cost(blocks: &[ContentBlock], counter: &TokenCounter) -> usize { + blocks + .iter() + .map(|block| match block { + ContentBlock::Text(text) => counter.count_tokens(&text.text), + ContentBlock::Image(_) => IMAGE_TOKEN_ESTIMATE, + _ => 0, + }) + .sum() +} + +pub(crate) fn build_handoff_context_memo( + prior_messages: &[Message], + budget: usize, + counter: &TokenCounter, +) -> Option { + let visible: Vec = Conversation::new_unvalidated(prior_messages.iter().cloned()) + .agent_visible_messages() + .iter() + .filter(|message| !message.is_turn_context()) + .map(|message| message.agent_visible_content()) + .collect(); + + if visible.is_empty() { + return None; + } + + let protected = recent_tool_call_ids(&visible); + let redacted: Vec = visible + .iter() + .map(|message| redact_tool_responses(message, |id| protected.contains(id))) + .collect(); + let formatted: Vec = redacted.iter().map(format_message_for_compacting).collect(); + let units = selection_units(&visible, &protected); + + let overhead = counter.count_tokens(MEMO_HEADER) + + counter.count_tokens(MEMO_FOOTER) + + OMISSION_MARKER_TOKENS; + let mut remaining = budget.saturating_sub(overhead); + + let mut kept: Vec = Vec::new(); + for unit in units.iter().rev() { + if remaining == 0 { + break; + } + let Some(fitted) = fit_unit(unit, &formatted, &redacted, remaining, counter) else { + break; + }; + kept.extend(fitted.messages.into_iter().rev()); + match fitted.cost { + Some(cost) => remaining -= cost, + None => remaining = 0, + } + } + + if kept.is_empty() { + return None; + } + + kept.reverse(); + let omitted = formatted.len() - kept.len(); + let mut body = String::new(); + if omitted > 0 { + body.push_str(&format!("[{omitted} earlier messages omitted]\n")); + } + body.push_str(&kept.join("\n")); + + Some(format!("{MEMO_HEADER}{body}{MEMO_FOOTER}")) +} + +/// Ids of the most recent tool exchanges, keyed by response so parallel and batched +/// calls are protected individually rather than by message position. +fn recent_tool_call_ids(messages: &[Message]) -> HashSet { + let mut ids: Vec<&str> = Vec::new(); + for message in messages { + for content in &message.content { + if let MessageContent::ToolResponse(response) = content { + ids.push(&response.id); + } + } + } + ids.into_iter() + .rev() + .take(PROTECTED_TOOL_EXCHANGES) + .map(str::to_string) + .collect() +} + +/// Contiguous message groups that are kept or dropped together. A protected tool response +/// travels with the request that produced it, so a tight budget can never leave one half of +/// an exchange orphaned. +fn selection_units(messages: &[Message], protected: &HashSet) -> Vec> { + let mut earliest: Vec = (0..messages.len()).collect(); + let mut request_at: HashMap<&str, usize> = HashMap::new(); + for (index, message) in messages.iter().enumerate() { + for content in &message.content { + match content { + MessageContent::ToolRequest(request) => { + request_at.insert(request.id.as_str(), index); + } + MessageContent::ToolResponse(response) if protected.contains(&response.id) => { + if let Some(&request_index) = request_at.get(response.id.as_str()) { + earliest[index] = earliest[index].min(request_index); + } + } + _ => {} + } + } + } + + let mut units: Vec> = Vec::new(); + let mut end = messages.len(); + while end > 0 { + let mut start = end - 1; + loop { + let extended = earliest[start..end].iter().copied().min().unwrap_or(start); + if extended == start { + break; + } + start = extended; + } + units.push((start..end).collect()); + end = start; + } + units.reverse(); + units +} + +struct FittedUnit { + messages: Vec, + /// `None` when the unit had to be degraded to fit, which ends selection. + cost: Option, +} + +/// Fit a whole unit into `budget`, degrading it only in ways that keep every exchange +/// it holds complete. +fn fit_unit( + unit: &[usize], + formatted: &[String], + redacted: &[Message], + budget: usize, + counter: &TokenCounter, +) -> Option { + let members: Vec = unit.iter().map(|&index| formatted[index].clone()).collect(); + let cost = unit_cost(&members, counter); + if cost <= budget { + return Some(FittedUnit { + messages: members, + cost: Some(cost), + }); + } + + // Eliding a message that carries protected responses would cut individual calls out of + // the middle of a batch. Degrade the exchange the way a stale one is degraded instead — + // requests intact, responses replaced whole — so nothing is left half-reported. + let degraded = if unit + .iter() + .any(|&index| holds_tool_response(&redacted[index])) + { + unit.iter() + .map(|&index| { + format_message_for_compacting(&redact_tool_responses(&redacted[index], |_| false)) + }) + .collect() + } else { + members + }; + + if unit_cost(°raded, counter) <= budget { + return Some(FittedUnit { + messages: degraded, + cost: None, + }); + } + + elide_unit_to_budget(degraded, budget, counter).map(|messages| FittedUnit { + messages, + cost: None, + }) +} + +/// Shrink the largest members until the whole unit fits. Nothing here carries a protected +/// response any more, so an oversized tool request is truncated rather than taking the +/// entire memo down with it. +fn elide_unit_to_budget( + mut members: Vec, + budget: usize, + counter: &TokenCounter, +) -> Option> { + for _ in 0..members.len() { + let costs: Vec = members + .iter() + .map(|message| counter.count_tokens(message) + 1) + .collect(); + let total: usize = costs.iter().sum(); + if total <= budget { + return Some(members); + } + let (index, largest) = costs + .iter() + .enumerate() + .max_by_key(|(_, &cost)| cost) + .map(|(index, &cost)| (index, cost))?; + let room = budget.checked_sub(total - largest + 1)?; + members[index] = elide_to_budget(&members[index], room, counter)?; + } + (unit_cost(&members, counter) <= budget).then_some(members) +} + +fn unit_cost(members: &[String], counter: &TokenCounter) -> usize { + members + .iter() + .map(|message| counter.count_tokens(message) + 1) + .sum() +} + +fn holds_tool_response(message: &Message) -> bool { + message + .content + .iter() + .any(|content| matches!(content, MessageContent::ToolResponse(_))) +} + +fn redact_tool_responses(message: &Message, keep: impl Fn(&str) -> bool) -> Message { + let should_redact = |content: &MessageContent| matches!(content, MessageContent::ToolResponse(response) if !keep(&response.id)); + if !message.content.iter().any(should_redact) { + return message.clone(); + } + + let content = message + .content + .iter() + .map(|content| { + if should_redact(content) { + MessageContent::text(REDACTED_TOOL_RESPONSE) + } else { + content.clone() + } + }) + .collect(); + + Message { + content, + ..message.clone() + } +} + +/// Middle-elide `text` so it fits in `budget` tokens, keeping its head and tail. +fn elide_to_budget(text: &str, budget: usize, counter: &TokenCounter) -> Option { + if budget < MIN_ELIDED_TOKENS { + return None; + } + + let total = counter.count_tokens(text).max(1); + let mut ratio = budget as f64 / total as f64; + for _ in 0..6 { + // Bytes to keep, so the floor below is a floor on the head and tail worth emitting + // rather than a token count. + let keep = ((text.len() as f64 * ratio * 0.9) as usize).min(text.len()); + if keep < 2 * MIN_ELIDED_TOKENS { + return None; + } + let head_end = floor_char_boundary(text, keep / 2); + let tail_start = ceil_char_boundary(text, text.len() - (keep - keep / 2)); + if tail_start <= head_end { + return None; + } + let candidate = format!( + "{}{ELISION_MARKER}{}", + text.get(..head_end)?, + text.get(tail_start..)? + ); + if counter.count_tokens(&candidate) <= budget { + return Some(candidate); + } + ratio *= 0.7; + } + None +} + +fn floor_char_boundary(text: &str, mut index: usize) -> usize { + while index > 0 && !text.is_char_boundary(index) { + index -= 1; + } + index +} + +fn ceil_char_boundary(text: &str, mut index: usize) -> usize { + while index < text.len() && !text.is_char_boundary(index) { + index += 1; + } + index +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::token_counter::create_token_counter; + use agent_client_protocol::schema::v1::{ImageContent, TextContent}; + use rmcp::model::{CallToolRequestParams, CallToolResult, ContentBlock as RmcpContent}; + + fn tool_exchange(id: &str, output: &str) -> Vec { + vec![ + Message::assistant().with_tool_request(id, Ok(CallToolRequestParams::new("read_file"))), + Message::user().with_tool_response( + id, + Ok(CallToolResult::success(vec![RmcpContent::text(output)])), + ), + ] + } + + fn memo_body(memo: &str) -> String { + memo.trim_start_matches(MEMO_HEADER) + .trim_end_matches(MEMO_FOOTER) + .to_string() + } + + #[test] + fn budget_is_capped_by_ratio_and_absolute_maximum() { + assert_eq!(memo_token_budget(100_000, 0), 30_000); + assert_eq!(memo_token_budget(1_000_000, 0), MAX_MEMO_TOKENS); + assert_eq!(memo_token_budget(100_000, 1_000), 29_000); + assert_eq!(memo_token_budget(1_000, 100_000), 0); + } + + #[tokio::test] + async fn images_in_the_current_turn_are_charged_against_the_budget() { + let counter = create_token_counter().await.unwrap(); + let text_only = vec![ContentBlock::Text(TextContent::new("current request"))]; + let with_image = vec![ + ContentBlock::Text(TextContent::new("current request")), + ContentBlock::Image(ImageContent::new("base64data", "image/png")), + ]; + + let text_cost = prompt_token_cost(&text_only, &counter); + let image_cost = prompt_token_cost(&with_image, &counter); + + assert_eq!(image_cost, text_cost + IMAGE_TOKEN_ESTIMATE); + assert!( + memo_token_budget(100_000, image_cost) < memo_token_budget(100_000, text_cost), + "an image has to shrink the memo's share of the window" + ); + } + + #[tokio::test] + async fn memo_stays_within_budget_and_drops_oldest_first() { + let counter = create_token_counter().await.unwrap(); + let messages: Vec = (0..200) + .map(|i| Message::user().with_text(format!("message {i} {}", "filler ".repeat(50)))) + .collect(); + + let memo = build_handoff_context_memo(&messages, 2_000, &counter).unwrap(); + + assert!(counter.count_tokens(&memo) <= 2_000); + assert!(memo.contains("message 199")); + assert!(!memo.contains("message 0 ")); + assert!(memo.contains("earlier messages omitted")); + } + + #[tokio::test] + async fn memo_keeps_chronological_order() { + let counter = create_token_counter().await.unwrap(); + let messages = vec![ + Message::user().with_text("first"), + Message::assistant().with_text("second"), + Message::user().with_text("third"), + ]; + + let memo = build_handoff_context_memo(&messages, 1_000, &counter).unwrap(); + let body = memo_body(&memo); + + let first = body.find("first").unwrap(); + let second = body.find("second").unwrap(); + let third = body.find("third").unwrap(); + assert!(first < second && second < third); + assert!(!body.contains("earlier messages omitted")); + } + + #[tokio::test] + async fn oversized_single_message_is_elided_not_dropped() { + let counter = create_token_counter().await.unwrap(); + let messages = vec![Message::user().with_text(format!( + "START {} END", + "an extremely long paragraph ".repeat(2_000) + ))]; + + let memo = build_handoff_context_memo(&messages, 500, &counter).unwrap(); + + assert!(counter.count_tokens(&memo) <= 500); + assert!(memo.contains("START")); + assert!(memo.contains("END")); + assert!(memo.contains("[... truncated ...]")); + } + + #[tokio::test] + async fn recent_tool_responses_are_kept_and_older_ones_redacted() { + let counter = create_token_counter().await.unwrap(); + let mut messages = vec![Message::user().with_text("start")]; + for i in 0..7 { + messages.extend(tool_exchange(&format!("call-{i}"), &format!("output-{i}"))); + } + + let memo = build_handoff_context_memo(&messages, 20_000, &counter).unwrap(); + + assert!(!memo.contains("output-0")); + assert!(!memo.contains("output-1")); + for i in 2..7 { + assert!(memo.contains(&format!("output-{i}")), "kept exchange {i}"); + } + assert!(memo.contains(REDACTED_TOOL_RESPONSE)); + assert!( + memo.contains("tool_request(read_file)"), + "tool requests survive redaction" + ); + } + + #[tokio::test] + async fn parallel_tool_responses_in_one_message_are_protected_individually() { + let counter = create_token_counter().await.unwrap(); + let batched = |ids: &[&str]| { + ids.iter().fold(Message::user(), |message, id| { + message.with_tool_response( + *id, + Ok(CallToolResult::success(vec![RmcpContent::text(format!( + "output-{id}" + ))])), + ) + }) + }; + let messages = vec![ + batched(&["a", "b", "c", "d"]), + batched(&["e", "f", "g", "h"]), + ]; + + let memo = build_handoff_context_memo(&messages, 20_000, &counter).unwrap(); + + for id in ["a", "b", "c"] { + assert!(!memo.contains(&format!("output-{id}")), "redacted {id}"); + } + for id in ["d", "e", "f", "g", "h"] { + assert!(memo.contains(&format!("output-{id}")), "kept {id}"); + } + } + + #[tokio::test] + async fn tight_budget_redacts_a_protected_exchange_instead_of_splitting_it() { + let counter = create_token_counter().await.unwrap(); + let messages = vec![ + Message::user().with_text("start"), + Message::assistant().with_tool_request( + "call-1", + Ok(CallToolRequestParams::new("read_file").with_arguments( + serde_json::json!({ "path": format!("src/{}.rs", "nested/".repeat(60)) }) + .as_object() + .unwrap() + .clone(), + )), + ), + Message::user().with_tool_response( + "call-1", + Ok(CallToolResult::success(vec![RmcpContent::text(format!( + "output-1 {}", + "filler ".repeat(200) + ))])), + ), + ]; + let request_tokens = counter.count_tokens(&format_message_for_compacting(&messages[1])); + let response_tokens = counter.count_tokens(&format_message_for_compacting(&messages[2])); + let overhead = counter.count_tokens(MEMO_HEADER) + + counter.count_tokens(MEMO_FOOTER) + + OMISSION_MARKER_TOKENS; + // Room for the request and a redacted response, but not for both in full. + let budget = overhead + request_tokens + response_tokens / 2; + + let memo = build_handoff_context_memo(&messages, budget, &counter).unwrap(); + + assert!(counter.count_tokens(&memo) <= budget); + assert!( + memo.contains("tool_request(read_file)"), + "the request survives with its response" + ); + assert!(memo.contains(REDACTED_TOOL_RESPONSE)); + assert!( + !memo.contains("output-1"), + "a protected response is replaced whole, never elided mid-call" + ); + } + + #[tokio::test] + async fn an_oversized_tool_request_is_elided_rather_than_losing_the_memo() { + let counter = create_token_counter().await.unwrap(); + let messages = vec![ + Message::user().with_text("start"), + Message::assistant().with_tool_request( + "call-1", + Ok(CallToolRequestParams::new("write_file").with_arguments( + serde_json::json!({ "contents": format!("HEAD {} TAIL", "payload ".repeat(4_000)) }) + .as_object() + .unwrap() + .clone(), + )), + ), + Message::user().with_tool_response( + "call-1", + Ok(CallToolResult::success(vec![RmcpContent::text("written")])), + ), + ]; + + let memo = build_handoff_context_memo(&messages, 500, &counter).unwrap(); + + assert!(counter.count_tokens(&memo) <= 500); + assert!( + memo.contains(ELISION_MARKER.trim()), + "the request is truncated, not dropped" + ); + assert!( + memo.contains("write_file"), + "the exchange is still readable" + ); + } + + #[tokio::test] + async fn a_protected_response_is_never_kept_without_its_request() { + let counter = create_token_counter().await.unwrap(); + let messages = vec![ + Message::user().with_text("start"), + Message::assistant().with_tool_request( + "call-1", + Ok(CallToolRequestParams::new("read_file").with_arguments( + serde_json::json!({ "path": format!("src/{}.rs", "nested/".repeat(400)) }) + .as_object() + .unwrap() + .clone(), + )), + ), + Message::user().with_tool_response( + "call-1", + Ok(CallToolResult::success(vec![RmcpContent::text("output-1")])), + ), + ]; + let response_tokens = counter.count_tokens(&format_message_for_compacting(&messages[2])); + let overhead = counter.count_tokens(MEMO_HEADER) + + counter.count_tokens(MEMO_FOOTER) + + OMISSION_MARKER_TOKENS; + // Room for the response and a truncated request, but not for the request in full. + let budget = overhead + response_tokens + 120; + + // Newest-first selection over bare messages would keep the small response here and + // drop the request that explains it. + let memo = build_handoff_context_memo(&messages, budget, &counter).unwrap(); + + assert!( + memo.contains("tool_response"), + "the exchange is in the memo" + ); + assert!( + memo.contains("tool_request(read_file)"), + "an orphaned response tells the agent a call happened but not what was asked" + ); + } + + #[tokio::test] + async fn zero_budget_produces_no_memo() { + let counter = create_token_counter().await.unwrap(); + let messages = vec![Message::user().with_text("prior context")]; + + assert!(build_handoff_context_memo(&messages, 0, &counter).is_none()); + } + + #[tokio::test] + async fn empty_history_produces_no_memo() { + let counter = create_token_counter().await.unwrap(); + + assert!(build_handoff_context_memo(&[], 10_000, &counter).is_none()); + } +} diff --git a/crates/goose/src/acp/mod.rs b/crates/goose/src/acp/mod.rs index 731a419b5..3e73d59c6 100644 --- a/crates/goose/src/acp/mod.rs +++ b/crates/goose/src/acp/mod.rs @@ -1,5 +1,6 @@ mod common; pub(crate) mod fs; +mod handoff; mod mcp_app_proxy; mod provider; mod response_builder; @@ -15,6 +16,11 @@ pub use provider::{ extension_configs_to_mcp_servers, AcpProvider, AcpProviderConfig, ACP_CURRENT_MODEL, }; +/// `data.reason` on a prompt error raised because the agent's account is out of credits. +/// Set by the ACP server, read by the provider to tell a spent account apart from a +/// prompt the agent could not accept. +pub(crate) const CREDITS_EXHAUSTED_REASON: &str = "credits_exhausted"; + pub(crate) fn configured_model_for_provider( config: &crate::config::Config, provider_name: &str, diff --git a/crates/goose/src/acp/provider.rs b/crates/goose/src/acp/provider.rs index 3eb3cd64b..f1505a4b7 100644 --- a/crates/goose/src/acp/provider.rs +++ b/crates/goose/src/acp/provider.rs @@ -32,15 +32,15 @@ use tokio::process::{Child, Command}; use tokio::sync::{mpsc, oneshot, Mutex as TokioMutex}; use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _}; +use crate::acp::handoff::{build_handoff_context_memo, memo_token_budget, prompt_token_cost}; use crate::acp::{map_permission_response, PermissionDecision}; use crate::config::{ExtensionConfig, GooseMode}; -use crate::context_mgmt::format_message_for_compacting; use crate::conversation::message::{Message, MessageContent, TOOL_META_EXTERNAL_DISPATCH_KEY}; -use crate::conversation::Conversation; use crate::permission::permission_confirmation::PrincipalType; use crate::permission::{Permission, PermissionConfirmation}; use crate::providers::base::{MessageStream, PermissionRouting, Provider}; use crate::subprocess::configure_subprocess; +use crate::token_counter::create_token_counter; use crate::utils::sanitize_unicode_tags; use goose_providers::errors::ProviderError; use goose_providers::model::ModelConfig; @@ -152,6 +152,23 @@ enum AcpUpdate { Error(agent_client_protocol::Error), } +/// Whether dropping the handoff memo could plausibly change the outcome. An agent that +/// rejected the very first update has told us nothing except that it disliked the prompt, +/// and the memo is the only part we added — but a spent account or a missing credential +/// says nothing about the prompt at all, so retrying would burn the single fallback the +/// session gets and consume a memo the agent never actually refused. +fn retry_without_memo_could_help(error: &agent_client_protocol::Error) -> bool { + if error.code == agent_client_protocol::schema::v1::ErrorCode::AuthRequired { + return false; + } + error + .data + .as_ref() + .and_then(|data| data.get("reason")) + .and_then(serde_json::Value::as_str) + != Some(crate::acp::CREDITS_EXHAUSTED_REASON) +} + fn provider_error_from_acp(error: agent_client_protocol::Error) -> ProviderError { if error.code == agent_client_protocol::schema::v1::ErrorCode::AuthRequired { ProviderError::Authentication(error.to_string()) @@ -517,6 +534,29 @@ impl AcpProvider { include_context: first_prompt && has_handoff_context(messages), } } + + /// Prior conversation, bounded against the agent's context window. The agent's own + /// system prompt and tool schemas are invisible to us, hence the conservative share. + async fn bounded_handoff_memo( + &self, + model_config: &ModelConfig, + messages: &[Message], + current_prompt: &[ContentBlock], + ) -> Option { + let last_user_index = last_user_message_index(messages)?; + let counter = match create_token_counter().await { + Ok(counter) => counter, + Err(error) => { + tracing::error!(%error, "no token counter, dropping ACP handoff context"); + return None; + } + }; + + let context_limit = self.get_context_limit(model_config).await.ok()?; + let budget = memo_token_budget(context_limit, prompt_token_cost(current_prompt, &counter)); + + build_handoff_context_memo(&messages[..last_user_index], budget, &counter) + } } fn fresh_text_run() -> (String, i64) { @@ -634,7 +674,7 @@ impl Provider for AcpProvider { ProviderError::RequestFailed(format!("Failed to set ACP model option: {e}")) })?; - let current_prompt_blocks = messages_to_prompt(messages, false); + let current_prompt_blocks = messages_to_prompt(messages, None); if current_prompt_blocks.is_empty() { return Ok(Box::pin(futures::stream::empty())); } @@ -642,24 +682,59 @@ impl Provider for AcpProvider { let claim = self.claim_handoff_context(messages); let mut handoff_claim_guard = HandoffContextClaimGuard::new(self.handoff_context_sent.clone(), claim.first_prompt); - let prompt_blocks = if claim.include_context { - messages_to_prompt(messages, true) + let memo = if claim.include_context { + self.bounded_handoff_memo(model_config, messages, ¤t_prompt_blocks) + .await } else { - current_prompt_blocks + None + }; + if claim.include_context && memo.is_none() { + // Nothing fit beside this turn's prompt, so the context never left goose. Give + // it back rather than marking a handoff that never happened as done — a single + // oversized turn would otherwise cost the session its whole history. + handoff_claim_guard.rollback(); + } + // A memo is only ever an estimate of what the agent will accept, so keep the bare + // prompt to retry with. Without it a bad estimate leaves the session unresumable. + let (prompt_blocks, mut bare_retry_blocks) = match memo { + Some(memo) => ( + messages_to_prompt(messages, Some(memo)), + Some(current_prompt_blocks), + ), + None => (current_prompt_blocks, None), }; // Drop any tool-call buffer state left over from a prior prompt // (e.g. cancelled or interrupted before its terminal status arrived). if let Ok(mut buffer) = self.pending_tool_updates.lock() { buffer.clear(); } - let mut rx = match self.prompt(session_id, prompt_blocks).await { + let mut rx = match self.prompt(session_id.clone(), prompt_blocks).await { Ok(rx) => rx, - Err(e) => { - return Err(ProviderError::RequestFailed(format!( - "Failed to send ACP prompt: {e}" - ))); - } + Err(e) => match bare_retry_blocks.take() { + Some(blocks) => { + // Consume the handoff before retrying. The memo is the only thing this + // attempt added, so rebuilding it next time would reproduce the same + // rejection and leave the session permanently unresumable. + handoff_claim_guard.commit(); + self.prompt(session_id.clone(), blocks) + .await + .map_err(|retry_error| { + ProviderError::RequestFailed(format!( + "Failed to send ACP prompt: {retry_error}" + )) + })? + } + // Nothing was added to this prompt, so the guard rolls the claim back as + // it drops and the next attempt can still carry the context. + None => { + return Err(ProviderError::RequestFailed(format!( + "Failed to send ACP prompt: {e}" + ))); + } + }, }; + let bare_retry = + bare_retry_blocks.map(|blocks| (self.tx.as_ref().unwrap().clone(), session_id, blocks)); let pending_confirmations = self.pending_confirmations.clone(); let goose_mode = *self @@ -672,12 +747,15 @@ impl Provider for AcpProvider { Ok(Box::pin(try_stream! { let mut suppress_text = false; + let mut bare_retry = bare_retry; + let mut updates_seen = 0usize; let mut rejected_tool_calls: HashSet = HashSet::new(); // Stable id+timestamp per contiguous run so Desktop coalesces chunks into one bubble. let mut text_run: Option<(String, i64)> = None; let mut thought_run: Option<(String, i64)> = None; while let Some(update) = rx.recv().await { + updates_seen += 1; match update { AcpUpdate::Text(text) => { if !suppress_text { @@ -822,10 +900,35 @@ impl Provider for AcpProvider { break; } AcpUpdate::Error(e) => { + let retry_could_help = retry_without_memo_could_help(&e); + let error = provider_error_from_acp(e); + if updates_seen == 1 && retry_could_help { + if let Some((tx, session_id, blocks)) = bare_retry.take() { + // Consume the handoff before retrying. The agent has already + // seen and rejected this memo, so rebuilding it on a later + // turn would reproduce the rejection forever. + handoff_claim_guard.commit(); + let (response_tx, response_rx) = mpsc::channel(64); + let request = ClientRequest::Prompt { + session_id, + content: blocks, + response_tx, + }; + if tx.send(request).await.is_ok() { + tracing::error!( + %error, + "ACP prompt with handoff context rejected, retrying without it" + ); + rx = response_rx; + updates_seen = 0; + continue; + } + } + } // Reset before yielding so an immediate retry can include the handoff even // while the failed stream value is still alive. handoff_claim_guard.rollback(); - Err(provider_error_from_acp(e))?; + Err(error)?; } } } @@ -1584,7 +1687,7 @@ fn filter_supported_servers( .collect() } -fn messages_to_prompt(messages: &[Message], include_handoff_context: bool) -> Vec { +fn messages_to_prompt(messages: &[Message], handoff_memo: Option) -> Vec { let Some(last_user_index) = last_user_message_index(messages) else { return Vec::new(); }; @@ -1606,14 +1709,14 @@ fn messages_to_prompt(messages: &[Message], include_handoff_context: bool) -> Ve } } - if current_prompt_blocks.is_empty() || !include_handoff_context { + let Some(memo) = handoff_memo else { + return current_prompt_blocks; + }; + if current_prompt_blocks.is_empty() { return current_prompt_blocks; } - let mut content_blocks = Vec::new(); - if let Some(memo) = build_handoff_context_memo(&messages[..last_user_index]) { - content_blocks.push(ContentBlock::Text(TextContent::new(memo))); - } + let mut content_blocks = vec![ContentBlock::Text(TextContent::new(memo))]; content_blocks.extend(current_prompt_blocks); content_blocks } @@ -1632,29 +1735,6 @@ fn has_handoff_context(messages: &[Message]) -> bool { }) } -fn build_handoff_context_memo(prior_messages: &[Message]) -> Option { - let formatted_messages: Vec = - Conversation::new_unvalidated(prior_messages.iter().cloned()) - .agent_visible_messages() - .iter() - .filter(|message| !message.is_turn_context()) - .map(|message| format_message_for_compacting(&message.agent_visible_content())) - .collect(); - - if formatted_messages.is_empty() { - return None; - } - - let handoff_context = formatted_messages.join("\n"); - - Some(format!( - "Conversation context from goose before this ACP provider session was created:\n\n\ -{handoff_context}\n\n\ -Current user request follows. Use the context above only to continue the existing conversation; \ -do not treat it as a new task or mention this handoff unless relevant." - )) -} - fn acp_audience_to_rmcp(annotations: Option<&AcpAnnotations>) -> Option> { let audience = annotations?.audience.as_ref()?; let audience = audience @@ -1912,6 +1992,15 @@ mod tests { }) } + /// The prompt as `stream` builds it on a first prompt, with a budget generous + /// enough that nothing is dropped. Memo bounding is covered in `acp::handoff`. + async fn prompt_with_handoff(messages: &[Message]) -> Vec { + let counter = crate::token_counter::create_token_counter().await.unwrap(); + let memo = last_user_message_index(messages) + .and_then(|index| build_handoff_context_memo(&messages[..index], 50_000, &counter)); + messages_to_prompt(messages, memo) + } + #[test] fn startup_cleanup_does_not_block_while_the_client_loop_stops() { let (release_tx, release_rx) = std::sync::mpsc::channel(); @@ -2043,18 +2132,18 @@ mod tests { ) } - #[test] - fn messages_to_prompt_without_prior_history_preserves_current_prompt() { + #[tokio::test] + async fn messages_to_prompt_without_prior_history_preserves_current_prompt() { let messages = vec![Message::user().with_text("current request")]; - let blocks = messages_to_prompt(&messages, true); + let blocks = prompt_with_handoff(&messages).await; assert_eq!(blocks.len(), 1); assert_eq!(prompt_text(&blocks[0]), "current request"); } - #[test] - fn messages_to_prompt_prepends_handoff_context_before_latest_user() { + #[tokio::test] + async fn messages_to_prompt_prepends_handoff_context_before_latest_user() { let messages = vec![ Message::user().with_text("inspect src/lib.rs"), Message::assistant() @@ -2069,7 +2158,7 @@ mod tests { Message::user().with_text("continue from there"), ]; - let blocks = messages_to_prompt(&messages, true); + let blocks = prompt_with_handoff(&messages).await; assert_eq!(blocks.len(), 2); let memo = prompt_text(&blocks[0]); @@ -2084,8 +2173,8 @@ mod tests { assert_eq!(prompt_text(&blocks[1]), "continue from there"); } - #[test] - fn messages_to_prompt_skips_turn_context_events() { + #[tokio::test] + async fn messages_to_prompt_skips_turn_context_events() { use crate::conversation::message::MessageMetadata; let turn_context = |text: &str| { @@ -2101,7 +2190,7 @@ mod tests { turn_context("new cwd /repo"), ]; - let blocks = messages_to_prompt(&messages, true); + let blocks = prompt_with_handoff(&messages).await; assert_eq!(blocks.len(), 2); assert!( @@ -2115,8 +2204,8 @@ mod tests { ); } - #[test] - fn messages_to_prompt_drops_user_only_acp_rows_from_handoff() { + #[tokio::test] + async fn messages_to_prompt_drops_user_only_acp_rows_from_handoff() { let user_only = TextContent::new("SECRET_USER_ONLY") .annotations(AcpAnnotations::new().audience(vec![AcpRole::User])); let messages = vec![ @@ -2125,7 +2214,7 @@ mod tests { Message::user().with_text("current request"), ]; - let blocks = messages_to_prompt(&messages, true); + let blocks = prompt_with_handoff(&messages).await; assert_eq!(blocks.len(), 2); let memo = prompt_text(&blocks[0]); @@ -2135,8 +2224,8 @@ mod tests { assert_eq!(prompt_text(&blocks[1]), "current request"); } - #[test] - fn messages_to_prompt_keeps_latest_user_images_after_handoff_memo() { + #[tokio::test] + async fn messages_to_prompt_keeps_latest_user_images_after_handoff_memo() { let messages = vec![ Message::assistant().with_text("prior answer"), Message::user() @@ -2144,7 +2233,7 @@ mod tests { .with_text("describe this"), ]; - let blocks = messages_to_prompt(&messages, true); + let blocks = prompt_with_handoff(&messages).await; assert_eq!(blocks.len(), 3); assert!(prompt_text(&blocks[0]).contains("[assistant]: prior answer")); @@ -2158,8 +2247,8 @@ mod tests { assert_eq!(prompt_text(&blocks[2]), "describe this"); } - #[test] - fn messages_to_prompt_excludes_user_only_current_and_handoff_content() { + #[tokio::test] + async fn messages_to_prompt_excludes_user_only_current_and_handoff_content() { use rmcp::model::{Annotations, TextContent}; fn user_only_text(text: &str) -> MessageContent { @@ -2178,7 +2267,8 @@ mod tests { .with_content(user_only_text("SECRET_CURRENT")), ]; - let rendered = messages_to_prompt(&messages, true) + let rendered = prompt_with_handoff(&messages) + .await .iter() .filter_map(|block| match block { ContentBlock::Text(text) => Some(text.text.as_str()), @@ -2193,8 +2283,8 @@ mod tests { assert!(!rendered.contains("SECRET_CURRENT")); } - #[test] - fn messages_to_prompt_drops_handoff_when_current_content_is_user_only() { + #[tokio::test] + async fn messages_to_prompt_drops_handoff_when_current_content_is_user_only() { use rmcp::model::{Annotations, TextContent}; let current = MessageContent::Text( @@ -2206,7 +2296,7 @@ mod tests { Message::user().with_content(current), ]; - assert!(messages_to_prompt(&messages, true).is_empty()); + assert!(prompt_with_handoff(&messages).await.is_empty()); } #[tokio::test] @@ -2411,7 +2501,7 @@ mod tests { } #[tokio::test] - async fn streamed_error_on_first_prompt_resends_handoff_context() { + async fn streamed_error_after_bare_retry_consumes_handoff_context() { use futures::StreamExt; let (tx, mut rx) = mpsc::channel(1); @@ -2420,18 +2510,19 @@ mod tests { Message::assistant().with_text("prior answer"), Message::user().with_text("current request"), ]; - let (retry_content_tx, retry_content_rx) = oneshot::channel(); + let (next_content_tx, next_content_rx) = oneshot::channel(); - // Serve the first prompt like a harness that accepts the request but - // fails while processing it (e.g. because the prompt is too large), - // then capture the retry. + // Serve the first prompt and its memo-free retry like a harness that accepts the + // request but fails while processing it, then capture the following turn. let server = tokio::spawn(async move { - if let Some(ClientRequest::Prompt { response_tx, .. }) = rx.recv().await { - let _ = response_tx - .send(AcpUpdate::Error( - agent_client_protocol::Error::internal_error().data("prompt too large"), - )) - .await; + for _ in 0..2 { + if let Some(ClientRequest::Prompt { response_tx, .. }) = rx.recv().await { + let _ = response_tx + .send(AcpUpdate::Error( + agent_client_protocol::Error::internal_error().data("prompt too large"), + )) + .await; + } } if let Some(ClientRequest::Prompt { content, @@ -2439,7 +2530,7 @@ mod tests { .. }) = rx.recv().await { - let _ = retry_content_tx.send(content); + let _ = next_content_tx.send(content); let _ = response_tx .send(AcpUpdate::Complete(StopReason::EndTurn, None)) .await; @@ -2453,12 +2544,15 @@ mod tests { "expected streamed error, got {first:?}" ); - let mut retry_stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); - let retry_content = retry_content_rx.await.unwrap(); - assert_eq!(retry_content.len(), 2); - assert!(prompt_text(&retry_content[0]).contains("[assistant]: prior answer")); - assert_eq!(prompt_text(&retry_content[1]), "current request"); - assert!(retry_stream.next().await.is_none()); + let mut next_stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); + let next_content = next_content_rx.await.unwrap(); + assert_eq!( + next_content.len(), + 1, + "a rejected memo must not be rebuilt on the next turn" + ); + assert_eq!(prompt_text(&next_content[0]), "current request"); + assert!(next_stream.next().await.is_none()); server.await.unwrap(); } @@ -2559,7 +2653,21 @@ mod tests { } #[tokio::test] - async fn failed_first_prompt_send_rolls_back_handoff_context_claim() { + async fn failed_first_prompt_send_without_handoff_rolls_back_claim() { + let (tx, rx) = mpsc::channel(1); + drop(rx); + let (provider, model) = test_provider_with_tx(Some(tx)); + let messages = vec![Message::user().with_text("current request")]; + + let result = provider.stream(&model, "", &messages, &[]).await; + + assert!(matches!(result, Err(ProviderError::RequestFailed(_)))); + let next_claim = provider.claim_handoff_context(&messages); + assert!(next_claim.first_prompt); + } + + #[tokio::test] + async fn failed_handoff_send_consumes_the_claim() { let (tx, rx) = mpsc::channel(1); drop(rx); let (provider, model) = test_provider_with_tx(Some(tx)); @@ -2572,8 +2680,231 @@ mod tests { assert!(matches!(result, Err(ProviderError::RequestFailed(_)))); let next_claim = provider.claim_handoff_context(&messages); - assert!(next_claim.first_prompt); - assert!(next_claim.include_context); + assert!( + !next_claim.include_context, + "a memo that already failed to send must not be rebuilt" + ); + } + + fn expect_prompt(request: ClientRequest) -> (Vec, mpsc::Sender) { + match request { + ClientRequest::Prompt { + content, + response_tx, + .. + } => (content, response_tx), + _ => panic!("expected ACP prompt request"), + } + } + + #[tokio::test] + async fn rejected_handoff_prompt_retries_once_without_the_memo() { + use futures::StreamExt; + + let (tx, mut rx) = mpsc::channel(2); + let (provider, model) = test_provider_with_tx(Some(tx)); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + + let handle = tokio::spawn(async move { + let mut stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); + let mut results = Vec::new(); + while let Some(item) = stream.next().await { + results.push(item); + } + results + }); + + let (content, response_tx) = expect_prompt(rx.recv().await.unwrap()); + assert_eq!(content.len(), 2, "memo precedes the current prompt"); + response_tx + .send(AcpUpdate::Error( + agent_client_protocol::Error::invalid_request().data("Prompt is too long"), + )) + .await + .unwrap(); + + let (content, response_tx) = expect_prompt(rx.recv().await.unwrap()); + assert_eq!(content.len(), 1, "retry carries the current prompt only"); + assert_eq!(prompt_text(&content[0]), "current request"); + response_tx + .send(AcpUpdate::Complete(StopReason::EndTurn, None)) + .await + .unwrap(); + + let results = handle.await.unwrap(); + assert!(results.iter().all(|item| item.is_ok())); + } + + #[tokio::test] + async fn rejected_retry_surfaces_the_error() { + use futures::StreamExt; + + let (tx, mut rx) = mpsc::channel(2); + let (provider, model) = test_provider_with_tx(Some(tx)); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + + let handle = tokio::spawn(async move { + let mut stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); + let mut results = Vec::new(); + while let Some(item) = stream.next().await { + results.push(item); + } + results + }); + + for _ in 0..2 { + let (_, response_tx) = expect_prompt(rx.recv().await.unwrap()); + response_tx + .send(AcpUpdate::Error( + agent_client_protocol::Error::invalid_request().data("Prompt is too long"), + )) + .await + .unwrap(); + } + + let results = handle.await.unwrap(); + assert!(matches!( + results.as_slice(), + [Err(ProviderError::RequestFailed(_))] + )); + assert!(rx.try_recv().is_err(), "exactly one retry"); + } + + #[tokio::test] + async fn auth_failure_surfaces_instead_of_retrying_without_the_memo() { + use futures::StreamExt; + + let (tx, mut rx) = mpsc::channel(2); + let (provider, model) = test_provider_with_tx(Some(tx)); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + + let handle = tokio::spawn(async move { + let mut stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); + let mut results = Vec::new(); + while let Some(item) = stream.next().await { + results.push(item); + } + results + }); + + let (_, response_tx) = expect_prompt(rx.recv().await.unwrap()); + response_tx + .send(AcpUpdate::Error( + agent_client_protocol::Error::auth_required(), + )) + .await + .unwrap(); + + // As above: a retry would leave the stream waiting on a prompt nothing serves. + let results = tokio::time::timeout(std::time::Duration::from_secs(10), handle) + .await + .expect("stream ended without retrying") + .unwrap(); + assert!(matches!( + results.as_slice(), + [Err(ProviderError::Authentication(_))] + )); + assert!(rx.try_recv().is_err(), "no retry on an auth failure"); + } + + #[tokio::test] + async fn a_budget_too_small_for_a_memo_keeps_the_claim() { + use futures::StreamExt; + + let (tx, mut rx) = mpsc::channel(1); + let (provider, model) = test_provider_with_tx(Some(tx)); + // A window this small leaves no room for a memo beside the current prompt. + provider.context_size.store(64, Ordering::Relaxed); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + let server = tokio::spawn(async move { + let Some(ClientRequest::Prompt { + content, + response_tx, + .. + }) = rx.recv().await + else { + return Vec::new(); + }; + let _ = response_tx + .send(AcpUpdate::Complete(StopReason::EndTurn, None)) + .await; + content + }); + + let mut stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); + assert!(stream.next().await.is_none()); + let content = server.await.unwrap(); + assert_eq!(content.len(), 1, "no memo fit beside the prompt"); + + assert!( + provider.claim_handoff_context(&messages).include_context, + "context that never left goose must still be handed off later" + ); + } + + #[tokio::test] + async fn exhausted_credits_surface_without_spending_the_handoff() { + use futures::StreamExt; + + let (tx, mut rx) = mpsc::channel(2); + let (provider, model) = test_provider_with_tx(Some(tx)); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + + let handle = tokio::spawn(async move { + let mut stream = provider.stream(&model, "", &messages, &[]).await.unwrap(); + let mut results = Vec::new(); + while let Some(item) = stream.next().await { + results.push(item); + } + (provider, results) + }); + + let (_, response_tx) = expect_prompt(rx.recv().await.unwrap()); + response_tx + .send(AcpUpdate::Error( + agent_client_protocol::Error::internal_error() + .data(serde_json::json!({ "reason": crate::acp::CREDITS_EXHAUSTED_REASON })), + )) + .await + .unwrap(); + + // A retry here would leave the stream waiting on a prompt nothing serves, so bound + // the wait rather than hanging the suite on a regression. + let (provider, results) = tokio::time::timeout(std::time::Duration::from_secs(10), handle) + .await + .expect("stream ended without retrying") + .unwrap(); + assert!(matches!( + results.as_slice(), + [Err(ProviderError::RequestFailed(_))] + )); + assert!( + rx.try_recv().is_err(), + "a spent account is not a prompt the agent refused" + ); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + assert!( + provider.claim_handoff_context(&messages).include_context, + "the memo survives so a topped-up account still resumes the conversation" + ); } fn test_provider_with_model_option( @@ -2863,8 +3194,8 @@ mod tests { assert_eq!(*provider.goose_mode.lock().unwrap(), GooseMode::Auto); } - #[test] - fn messages_to_prompt_includes_all_prior_handoff_context() { + #[tokio::test] + async fn messages_to_prompt_includes_all_prior_handoff_context() { let messages = vec![ Message::user().with_text("older context that should be retained"), Message::assistant().with_text("middle context"), @@ -2872,7 +3203,7 @@ mod tests { Message::user().with_text("current request"), ]; - let blocks = messages_to_prompt(&messages, true); + let blocks = prompt_with_handoff(&messages).await; assert_eq!(blocks.len(), 2); let memo = prompt_text(&blocks[0]); diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 2bda88ed1..fdbe5ea68 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -1373,7 +1373,7 @@ fn prompt_error_from_message_content( let mut data = serde_json::Map::new(); data.insert( "reason".to_string(), - serde_json::Value::String("credits_exhausted".to_string()), + serde_json::Value::String(crate::acp::CREDITS_EXHAUSTED_REASON.to_string()), ); Some( agent_client_protocol::Error::new(-32603, error.message.clone()) @@ -1390,7 +1390,7 @@ fn credits_exhausted_prompt_error( let mut data = serde_json::Map::new(); data.insert( "reason".to_string(), - serde_json::Value::String("credits_exhausted".to_string()), + serde_json::Value::String(crate::acp::CREDITS_EXHAUSTED_REASON.to_string()), ); if let Some(url) = notification