From 66051ec7d2ea4d5c2701fca2aa762fcecf82a487 Mon Sep 17 00:00:00 2001 From: filip <44206832+filipkujawa@users.noreply.github.com> Date: Thu, 6 Aug 2026 20:08:49 -0700 Subject: [PATCH] Cache-safe request assembly: append-only turn context and declared cache semantics (#11022) --- crates/goose-cli/src/commands/term.rs | 50 +- crates/goose-cli/src/session/mod.rs | 81 +++- .../src/cache_semantics.rs | 203 +++++++++ .../goose-provider-types/src/conversation.rs | 87 +++- .../src/conversation/message.rs | 13 + .../src/formats/anthropic.rs | 315 ++----------- .../src/formats/databricks.rs | 220 +-------- .../src/formats/openai.rs | 245 +--------- crates/goose-provider-types/src/lib.rs | 1 + .../tests/prefix_invariance.rs | 428 ++++++++++++++++++ crates/goose-providers/src/lib.rs | 4 +- crates/goose-providers/src/openai.rs | 1 - crates/goose/src/acp/provider.rs | 36 +- crates/goose/src/agents/agent.rs | 77 +++- crates/goose/src/agents/extension_manager.rs | 5 +- crates/goose/src/agents/moim.rs | 334 ++++---------- .../agents/platform_extensions/chatrecall.rs | 6 +- .../platform_extensions/orchestrator.rs | 19 +- crates/goose/src/agents/reply_parts.rs | 5 +- ...nager__tests__all_platform_extensions.snap | 1 + ..._agents__prompt_manager__tests__basic.snap | 2 +- ..._prompt_manager__tests__one_extension.snap | 2 +- ..._prompt_manager__tests__typical_setup.snap | 2 +- .../goose/src/agents/state_machine/ops_llm.rs | 33 +- .../agents/state_machine/ops_toolcalling.rs | 2 + .../src/agents/state_machine/tests/mod.rs | 29 ++ .../agents/state_machine/tests/pipeline.rs | 1 + .../tests/recipe_scheduling_lifecycle.rs | 7 +- crates/goose/src/context_mgmt/mod.rs | 242 +++++++++- crates/goose/src/providers/litellm.rs | 70 +-- crates/goose/src/providers/openrouter.rs | 87 +--- crates/goose/src/scheduler.rs | 8 +- .../goose/src/security/adversary_inspector.rs | 39 +- crates/goose/src/security/scanner.rs | 30 ++ .../goose/src/session/chat_history_search.rs | 57 ++- crates/goose/src/session/session_manager.rs | 58 ++- crates/goose/src/session/session_naming.rs | 2 +- crates/goose/tests/agent.rs | 11 +- crates/goose/tests/compaction.rs | 7 +- 39 files changed, 1589 insertions(+), 1231 deletions(-) create mode 100644 crates/goose-provider-types/src/cache_semantics.rs create mode 100644 crates/goose-provider-types/tests/prefix_invariance.rs diff --git a/crates/goose-cli/src/commands/term.rs b/crates/goose-cli/src/commands/term.rs index 9c80fc1e0e..41a39ee22c 100644 --- a/crates/goose-cli/src/commands/term.rs +++ b/crates/goose-cli/src/commands/term.rs @@ -248,6 +248,20 @@ pub async fn handle_term_log(command: String) -> Result<()> { Ok(()) } +fn shell_history_text(newest_first_tail: &[&Message]) -> Option { + let history: Vec = newest_first_tail + .iter() + .filter(|m| !m.is_turn_context()) + .rev() + .map(|m| m.as_concat_text()) + .collect(); + if history.is_empty() { + None + } else { + Some(history.join("\n")) + } +} + pub async fn handle_term_run(prompt: Vec) -> Result<()> { let prompt = prompt.join(" "); let session_id = std::env::var("AGENT_SESSION_ID").map_err(|_| { @@ -291,20 +305,12 @@ pub async fn handle_term_run(prompt: Vec) -> Result<()> { } } - let prompt_with_context = if user_messages_after_last_assistant.is_empty() { - prompt - } else { - let history = user_messages_after_last_assistant - .iter() - .rev() // back to chronological order - .map(|m| m.as_concat_text()) - .collect::>() - .join("\n"); - - format!( + let prompt_with_context = match shell_history_text(&user_messages_after_last_assistant) { + Some(history) => format!( "\n{}\n\n\n{}", history, prompt - ) + ), + None => prompt, }; let config = SessionBuilderConfig { @@ -407,4 +413,24 @@ mod tests { assert!(!script.contains("command_not_found")); } + + #[test] + fn shell_history_skips_turn_context_events() { + use goose::conversation::message::MessageMetadata; + + let older = Message::user().with_text("git status"); + let block = Message::user() + .with_text("cwd /repo") + .with_metadata(MessageMetadata::agent_only().with_turn_context()); + let newer = Message::user().with_text("cargo build"); + let newest_first_tail = vec![&newer, &block, &older]; + + assert_eq!( + shell_history_text(&newest_first_tail).unwrap(), + "git status\ncargo build" + ); + + let only_block = vec![&block]; + assert_eq!(shell_history_text(&only_block), None); + } } diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 6ad4fe4861..87b1cd569c 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -12,7 +12,7 @@ mod thinking; use crate::session::task_execution_display::{ format_task_execution_notification, TASK_EXECUTION_NOTIFICATION_TYPE, }; -use goose::conversation::{fix_conversation, Conversation}; +use goose::conversation::{fix_conversation, merge_consecutive_messages_for_request, Conversation}; use std::env; use std::io::Write; use std::str::FromStr; @@ -69,8 +69,16 @@ const SHELL_STATUS_MAX_LINES: usize = 3; const SHELL_STATUS_RESERVED_WIDTH: usize = 2; fn planner_provider_messages(plan_messages: &Conversation) -> Conversation { - let projected_messages = plan_messages.agent_visible_messages(); - fix_conversation(Conversation::new_unvalidated(projected_messages)).0 + // The planner prompt has no turn-context instructions; drop the blocks. + let projected_messages: Vec = plan_messages + .agent_visible_messages() + .into_iter() + .filter(|message| !message.is_turn_context()) + .collect(); + let fixed = fix_conversation(Conversation::new_unvalidated(projected_messages)).0; + Conversation::new_unvalidated(merge_consecutive_messages_for_request( + fixed.messages().clone(), + )) } #[derive(Serialize, Deserialize, Debug)] @@ -1645,24 +1653,29 @@ impl CliSession { &Message::assistant().with_text(interrupt_prompt), self.debug, ); - } else if let Some(last_msg) = self.messages.last() { - if last_msg.role == rmcp::model::Role::User { - match last_msg.content.first() { - Some(MessageContent::ToolResponse(_)) => { - self.push_message(Message::assistant().with_text(interrupt_prompt)); - output::render_message( - &Message::assistant().with_text(interrupt_prompt), - self.debug, - ); - } - Some(_) => { - self.messages.pop(); - let assistant_msg = Message::assistant().with_text(interrupt_prompt); - self.push_message(assistant_msg.clone()); - output::render_message(&assistant_msg, self.debug); - } - None => { - // Empty message content — nothing to do, just continue gracefully + } else { + while self.messages.last().is_some_and(Message::is_turn_context) { + self.messages.pop(); + } + if let Some(last_msg) = self.messages.last() { + if last_msg.role == rmcp::model::Role::User { + match last_msg.content.first() { + Some(MessageContent::ToolResponse(_)) => { + self.push_message(Message::assistant().with_text(interrupt_prompt)); + output::render_message( + &Message::assistant().with_text(interrupt_prompt), + self.debug, + ); + } + Some(_) => { + self.messages.pop(); + let assistant_msg = Message::assistant().with_text(interrupt_prompt); + self.push_message(assistant_msg.clone()); + output::render_message(&assistant_msg, self.debug); + } + None => { + // Empty message content — nothing to do, just continue gracefully + } } } } @@ -2646,6 +2659,32 @@ mod tests { .contains("hidden separator")); } + #[test] + fn planner_history_excludes_turn_context_events() { + use goose::conversation::message::MessageMetadata; + + let history = Conversation::new_unvalidated([ + Message::user().with_text("plan the refactor"), + Message::user() + .with_text("cwd /repo, todo: ship v2") + .with_metadata(MessageMetadata::agent_only().with_turn_context()), + Message::assistant().with_text("on it"), + ]); + + let provider_text = planner_provider_messages(&history) + .agent_visible_messages() + .iter() + .map(|message| message.as_concat_text()) + .collect::>() + .join("\n"); + + assert!(provider_text.contains("plan the refactor")); + assert!( + !provider_text.contains("turn-context"), + "the planner prompt has no turn-context instructions, so blocks must not reach it" + ); + } + #[test] fn test_format_elapsed_time_under_60_seconds() { // Test sub-second duration diff --git a/crates/goose-provider-types/src/cache_semantics.rs b/crates/goose-provider-types/src/cache_semantics.rs new file mode 100644 index 0000000000..edbbda2c70 --- /dev/null +++ b/crates/goose-provider-types/src/cache_semantics.rs @@ -0,0 +1,203 @@ +//! Prompt-cache semantics declared per (provider, model), instead of implied +//! by whichever format module a request flows through. + +use serde_json::{json, Value}; + +use crate::formats::openai::is_openai_responses_model; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CacheSemantics { + /// Caller places markers; reuse needs an exact match of the marked bytes. + ExplicitBreakpoints { max_breakpoints: usize }, + /// The longest matching stored prefix is reused implicitly. + ImplicitTolerant, + /// Only extends a prompt reproduced byte-for-byte from the start. + ImplicitStrict, + /// No known prompt cache. + Uncached, +} + +impl CacheSemantics { + /// Unknown pairs default to `ImplicitStrict`, which is safe for every + /// cache. + pub fn for_model(provider_name: &str, model_name: &str) -> Self { + match provider_name { + "anthropic" | "minimax" | "zai" | "kimi_code" => { + CacheSemantics::ExplicitBreakpoints { max_breakpoints: 4 } + } + "aws_bedrock" | "databricks" | "gcp_vertex_ai" if model_name.contains("claude") => { + CacheSemantics::ExplicitBreakpoints { max_breakpoints: 4 } + } + "openrouter" | "litellm" if model_name.starts_with("anthropic/") => { + CacheSemantics::ExplicitBreakpoints { max_breakpoints: 4 } + } + "openai" | "azure_openai" | "github_copilot" => { + if is_openai_responses_model(model_name) { + CacheSemantics::ImplicitStrict + } else { + CacheSemantics::ImplicitTolerant + } + } + "moonshot" | "custom_deepseek" | "groq" | "together" | "fireworks-ai" | "mistral" + | "zhipu" | "alibaba" => CacheSemantics::ImplicitTolerant, + "snowflake" | "sagemaker_tgi" => CacheSemantics::Uncached, + _ => CacheSemantics::ImplicitStrict, + } + } + + pub fn uses_explicit_breakpoints(self) -> bool { + matches!(self, CacheSemantics::ExplicitBreakpoints { .. }) + } +} + +/// Anthropic-dialect breakpoints for OpenAI-style chat payloads (OpenRouter, +/// LiteLLM, Databricks). +pub fn apply_chat_payload_breakpoints(payload: &mut Value) { + if let Some(messages) = payload + .get_mut("messages") + .and_then(|messages| messages.as_array_mut()) + { + let mut user_count = 0; + for message in messages.iter_mut().rev() { + if message.get("role") != Some(&json!("user")) { + continue; + } + if mark_last_content_block(message) { + user_count += 1; + if user_count >= 2 { + break; + } + } + } + + if let Some(system_message) = messages + .iter_mut() + .find(|message| message.get("role") == Some(&json!("system"))) + { + mark_last_content_block(system_message); + } + } + + if let Some(last_tool) = payload + .get_mut("tools") + .and_then(|tools| tools.as_array_mut()) + .and_then(|tools| tools.last_mut()) + { + if let Some(function) = last_tool.get_mut("function").and_then(Value::as_object_mut) { + function.insert("cache_control".to_string(), json!({ "type": "ephemeral" })); + } + } +} + +fn mark_last_content_block(message: &mut Value) -> bool { + let Some(content) = message.get_mut("content") else { + return false; + }; + if let Some(text) = content.as_str() { + *content = json!([{ + "type": "text", + "text": text, + "cache_control": { "type": "ephemeral" } + }]); + return true; + } + if let Some(block) = content + .as_array_mut() + .and_then(|blocks| blocks.last_mut()) + .and_then(Value::as_object_mut) + { + block.insert("cache_control".to_string(), json!({ "type": "ephemeral" })); + return true; + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn for_model_resolves_registered_identifiers_and_defaults_to_strict() { + let explicit = CacheSemantics::ExplicitBreakpoints { max_breakpoints: 4 }; + for (provider, model, expected) in [ + ("anthropic", "claude-sonnet-4-5", explicit), + ("kimi_code", "kimi-for-coding", explicit), + ("aws_bedrock", "anthropic.claude-sonnet-4-5", explicit), + ("minimax", "MiniMax-M2.5", explicit), + ("openrouter", "anthropic/claude-sonnet-4-5", explicit), + ( + "openrouter", + "openai/gpt-5.2", + CacheSemantics::ImplicitStrict, + ), + ("openai", "gpt-5.2", CacheSemantics::ImplicitStrict), + ("openai", "gpt-4.1-mini", CacheSemantics::ImplicitTolerant), + ("snowflake", "llama-3.3-70b", CacheSemantics::Uncached), + ( + "some_new_vendor", + "some-model", + CacheSemantics::ImplicitStrict, + ), + ] { + assert_eq!( + CacheSemantics::for_model(provider, model), + expected, + "{provider}/{model}" + ); + } + } + + #[test] + fn breakpoints_cover_system_tools_and_last_two_user_messages() { + let mut payload = json!({ + "model": "anthropic/claude-sonnet-4-5", + "messages": [ + {"role": "system", "content": "be careful"}, + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "second"}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": "third"} + ], + "tools": [ + {"type": "function", "function": {"name": "a"}}, + {"type": "function", "function": {"name": "b"}} + ] + }); + apply_chat_payload_breakpoints(&mut payload); + + let messages = payload["messages"].as_array().unwrap(); + assert_eq!( + messages[0]["content"][0]["cache_control"]["type"], + "ephemeral" + ); + assert!(messages[1]["content"].is_string(), "only last two users"); + assert_eq!( + messages[3]["content"][0]["cache_control"]["type"], + "ephemeral" + ); + assert_eq!( + messages[5]["content"][0]["cache_control"]["type"], + "ephemeral" + ); + let tools = payload["tools"].as_array().unwrap(); + assert!(tools[0]["function"].get("cache_control").is_none()); + assert_eq!(tools[1]["function"]["cache_control"]["type"], "ephemeral"); + } + + #[test] + fn array_content_user_messages_get_a_breakpoint() { + let mut payload = json!({ + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "look at this"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,xyz"}} + ]} + ] + }); + apply_chat_payload_breakpoints(&mut payload); + let blocks = payload["messages"][0]["content"].as_array().unwrap(); + assert!(blocks[0].get("cache_control").is_none()); + assert_eq!(blocks[1]["cache_control"]["type"], "ephemeral"); + } +} diff --git a/crates/goose-provider-types/src/conversation.rs b/crates/goose-provider-types/src/conversation.rs index 56e3cadbc4..53572b6ee9 100644 --- a/crates/goose-provider-types/src/conversation.rs +++ b/crates/goose-provider-types/src/conversation.rs @@ -518,14 +518,33 @@ fn fix_tool_calling(mut messages: Vec) -> (Vec, Vec) { (messages, issues) } +/// Never merges across visibility or turn-context boundaries, so the result +/// is safe to persist. pub fn merge_consecutive_messages(messages: Vec) -> (Vec, Vec) { + merge_consecutive(messages, false) +} + +/// Merges regardless of visibility, for providers that require strict role +/// alternation. Never persist the result. +pub fn merge_consecutive_messages_for_request(messages: Vec) -> Vec { + merge_consecutive(messages, true).0 +} + +fn merge_consecutive( + messages: Vec, + across_visibility: bool, +) -> (Vec, Vec) { let mut issues = Vec::new(); let mut merged_messages: Vec = Vec::new(); for message in messages { if let Some(last) = merged_messages.last_mut() { let effective = effective_role(&message); - if effective_role(last) == effective { + if effective_role(last) == effective + && (across_visibility + || (last.metadata.user_visible == message.metadata.user_visible + && last.metadata.turn_context == message.metadata.turn_context)) + { last.content.extend(message.content); issues.push(format!("Merged consecutive {} messages", effective)); continue; @@ -613,12 +632,6 @@ pub const TURN_CONTEXT_TAG: &str = "turn-context"; pub const CURRENT_TIME_TAG: &str = "current-time"; pub const WORKING_DIRECTORY_TAG: &str = "working-directory"; -pub fn is_turn_context_text(text: &str) -> bool { - text.starts_with(&format!("<{TURN_CONTEXT_TAG}>\n<{CURRENT_TIME_TAG}>")) - && text.contains(&format!("\n<{WORKING_DIRECTORY_TAG}>")) - && text.trim_end().ends_with(&format!("")) -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EffectiveRole { User, @@ -713,7 +726,9 @@ pub fn debug_conversation_fix( #[cfg(test)] mod tests { - use crate::conversation::message::{InferenceMetadata, Message, MessageContentBlock}; + use crate::conversation::message::{ + InferenceMetadata, Message, MessageContentBlock, MessageMetadata, + }; use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation}; use rmcp::model::{CallToolRequestParams, Role}; use rmcp::object; @@ -846,6 +861,62 @@ mod tests { assert_eq!(fixed[0].content.len(), 2); } + #[test] + fn merge_does_not_cross_user_visibility() { + let visible = Message::user().with_text("what is in main.rs?"); + let hidden = Message::user() + .with_text("frozen") + .with_visibility(false, true); + + let (fixed, _) = fix_conversation(Conversation::new_unvalidated(vec![visible, hidden])); + assert_eq!( + fixed.messages().len(), + 2, + "the persistable form keeps the agent-only event separate" + ); + assert!(fixed.messages()[0].is_user_visible()); + assert!(!fixed.messages()[1].is_user_visible()); + + let merged = + crate::conversation::merge_consecutive_messages_for_request(fixed.messages().clone()); + assert_eq!( + merged.len(), + 1, + "the request form merges to satisfy role alternation" + ); + assert_eq!(merged[0].content.len(), 2); + } + + #[test] + fn merge_does_not_erase_turn_context_marker() { + let preserved_prompt = Message::user() + .with_text("the preserved prompt") + .with_metadata(MessageMetadata::agent_only()); + let carried_event = Message::user() + .with_text("frozen") + .with_metadata(MessageMetadata::agent_only().with_turn_context()); + + let (fixed, _) = fix_conversation(Conversation::new_unvalidated(vec![ + preserved_prompt, + carried_event, + ])); + assert_eq!( + fixed.messages().len(), + 2, + "the persistable form must not fold a turn-context event into the prompt" + ); + assert!(!fixed.messages()[0].is_turn_context()); + assert!(fixed.messages()[1].is_turn_context()); + + let merged = + crate::conversation::merge_consecutive_messages_for_request(fixed.messages().clone()); + assert_eq!( + merged.len(), + 1, + "the request form still merges to satisfy role alternation" + ); + } + #[test] fn test_orphaned_tools_and_empty_messages() { use rmcp::model::ContentBlock; diff --git a/crates/goose-provider-types/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs index f158b9c605..65b0628109 100644 --- a/crates/goose-provider-types/src/conversation/message.rs +++ b/crates/goose-provider-types/src/conversation/message.rs @@ -742,6 +742,9 @@ pub struct MessageMetadata { /// without matching user-visible text. Never sent to providers. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub steer: bool, + /// Whether this message is a per-turn context event appended by the agent. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub turn_context: bool, #[serde(default, skip_serializing_if = "Option::is_none")] pub usage: Option>, /// What an operation did to this message, keyed by operation name. Read back @@ -759,6 +762,7 @@ impl Default for MessageMetadata { inference: None, output_token_limit_reached: false, steer: false, + turn_context: false, usage: None, operations: None, } @@ -846,6 +850,11 @@ impl MessageMetadata { self.steer = true; self } + + pub fn with_turn_context(mut self) -> Self { + self.turn_context = true; + self + } } #[derive(Clone, PartialEq, Serialize, Deserialize, Debug)] @@ -1228,6 +1237,10 @@ impl Message { pub fn is_agent_visible(&self) -> bool { self.metadata.agent_visible } + + pub fn is_turn_context(&self) -> bool { + self.metadata.turn_context + } } #[derive(Debug, Clone, Default, Serialize, Deserialize)] diff --git a/crates/goose-provider-types/src/formats/anthropic.rs b/crates/goose-provider-types/src/formats/anthropic.rs index c14b21db9e..0472dffff0 100644 --- a/crates/goose-provider-types/src/formats/anthropic.rs +++ b/crates/goose-provider-types/src/formats/anthropic.rs @@ -404,27 +404,16 @@ fn format_messages_with_options( })); } - // The volatile turn-context must sit after every cache breakpoint, or it invalidates the - // message-level cached prefix (Anthropic hashes tools -> system -> messages). Move it to the - // tail and place cache_control on the last non-turn-context block. - relocate_turn_context_to_tail(&mut anthropic_messages); - + // The last two user messages extend the cached prefix each turn. let mut user_count = 0; for message in anthropic_messages.iter_mut().rev() { if message.get(ROLE_FIELD) != Some(&json!(USER_ROLE)) { continue; } - let Some(content_array) = message + if let Some(block) = message .get_mut(CONTENT_FIELD) .and_then(|content| content.as_array_mut()) - else { - continue; - }; - let Some(target) = cache_control_target_index(content_array) else { - continue; - }; - if let Some(block) = content_array - .get_mut(target) + .and_then(|content_array| content_array.last_mut()) .and_then(|b| b.as_object_mut()) { block.insert( @@ -441,52 +430,6 @@ fn format_messages_with_options( anthropic_messages } -fn relocate_turn_context_to_tail(messages: &mut [Value]) { - let Some(last) = messages.len().checked_sub(1) else { - return; - }; - let source = messages.iter().enumerate().rev().find_map(|(mi, m)| { - m.get(CONTENT_FIELD) - .and_then(|c| c.as_array()) - .and_then(|a| a.iter().position(is_turn_context_block)) - .map(|bi| (mi, bi)) - }); - let Some((mi, bi)) = source else { - return; - }; - if mi != last - && messages[mi] - .get(CONTENT_FIELD) - .and_then(|c| c.as_array()) - .map_or(0, |a| a.len()) - <= 1 - { - return; - } - let block = messages[mi][CONTENT_FIELD] - .as_array_mut() - .unwrap() - .remove(bi); - messages[last][CONTENT_FIELD] - .as_array_mut() - .unwrap() - .push(block); -} - -fn cache_control_target_index(content_array: &[Value]) -> Option { - content_array - .iter() - .rposition(|block| !is_turn_context_block(block)) -} - -fn is_turn_context_block(block: &Value) -> bool { - block.get(TYPE_FIELD).and_then(Value::as_str) == Some(TEXT_TYPE) - && block - .get(TEXT_TYPE) - .and_then(Value::as_str) - .is_some_and(crate::conversation::is_turn_context_text) -} - fn anthropic_flavored_input_schema(input_schema: Arc) -> Arc { if input_schema.is_empty() { return Arc::new(json_object!({ @@ -2543,29 +2486,12 @@ mod tests { assert!(parts.tool_errors.is_empty()); } - /// Anthropic prefix caching only pays off when the bytes up to a cache - /// breakpoint are identical turn over turn. The per-turn turn-context block - /// (timestamp, turn budget, compaction state) changes on every call, so if - /// it ever lands inside a cached prefix every request becomes a cache write - /// instead of a read. These tests pin the property that keeps caching alive - /// so a future refactor of the formatter or the turn-context format can't - /// silently regress it. - mod cache_prefix_stability { + /// Placement shape only; cross-request prefix stability is covered by + /// `tests/prefix_invariance.rs`. + mod cache_breakpoint_placement { use super::*; use rmcp::model::CallToolResult; - /// A turn-context block whose shape matches what `is_turn_context_text` - /// recognizes, varying only the volatile fields. - fn turn_context(time: &str, turn_budget: &str) -> String { - format!( - "\n\ - {time}\n\ - /Users/me/code/goose\n\ - {turn_budget}\n\ - " - ) - } - fn sample_tools() -> Vec { vec![ Tool::new( @@ -2587,11 +2513,21 @@ mod tests { ] } - /// A realistic multi-turn conversation. `inject_moim` prepends the - /// turn-context block to the latest genuine user message, so it sits as - /// the first text block of the final user message here. - fn conversation(turn_context_block: &str) -> Vec { - vec![ + fn breakpoints(messages: &[Value]) -> Vec<(usize, usize)> { + let mut found = Vec::new(); + for (mi, message) in messages.iter().enumerate() { + for (bi, block) in message["content"].as_array().unwrap().iter().enumerate() { + if block.get(CACHE_CONTROL_FIELD).is_some() { + found.push((mi, bi)); + } + } + } + found + } + + #[test] + fn breakpoints_cover_tools_system_and_last_two_user_messages() { + let messages = vec![ Message::user().with_text("What does the main entrypoint do?"), Message::assistant().with_tool_request( "tool_1", @@ -2605,200 +2541,45 @@ mod tests { ])), ), Message::assistant().with_text("It calls `run()`."), - Message::user() - .with_text(turn_context_block) - .with_text("Now add error handling to it."), - ] - } - - /// The (message index, block index) of the last block carrying a - /// `cache_control` marker, scanning in canonical order. This is the far - /// edge of the furthest cached prefix. - fn last_breakpoint(messages: &[Value]) -> Option<(usize, usize)> { - let mut found = None; - for (mi, message) in messages.iter().enumerate() { - for (bi, block) in message["content"].as_array().unwrap().iter().enumerate() { - if block.get(CACHE_CONTROL_FIELD).is_some() { - found = Some((mi, bi)); - } - } - } - found - } - - fn find_turn_context(messages: &[Value]) -> Option<(usize, usize)> { - messages.iter().enumerate().find_map(|(mi, message)| { - message["content"] - .as_array() - .unwrap() - .iter() - .position(is_turn_context_block) - .map(|bi| (mi, bi)) - }) - } - - /// The exact bytes Anthropic hashes for its furthest cache breakpoint: - /// tools, then system, then messages truncated at the last - /// `cache_control` marker. Everything after that point is outside every - /// cached prefix and may change freely turn to turn. - fn cached_prefix(payload: &Value) -> String { - let messages = payload["messages"].as_array().unwrap(); - let (last_mi, last_bi) = last_breakpoint(messages) - .expect("request must carry at least one cache_control breakpoint"); - - let prefix_messages: Vec = messages - .iter() - .take(last_mi + 1) - .enumerate() - .map(|(mi, message)| { - let mut message = message.clone(); - if mi == last_mi { - message["content"] - .as_array_mut() - .unwrap() - .truncate(last_bi + 1); - } - message - }) - .collect(); - - json!({ - "tools": payload.get("tools"), - "system": payload.get("system"), - "messages": prefix_messages, - }) - .to_string() - } - - /// The production tool-loop case: `inject_moim` prepends turn-context to - /// the latest *genuine* user message, but the request then ends with a - /// later `tool_result` message. The block must be relocated *across* - /// messages to land after the trailing breakpoint, not merely reordered - /// within its own message. - fn tool_loop_conversation(turn_context_block: &str) -> Vec { - vec![ - Message::user().with_text("What does the main entrypoint do?"), - Message::assistant().with_text("Let me read it."), - Message::user() - .with_text(turn_context_block) - .with_text("Now add error handling to it."), - Message::assistant().with_tool_request( - "tool_1", - Ok(CallToolRequestParams::new("read_file") - .with_arguments(object!({"path": "src/main.rs"}))), - ), - Message::user().with_tool_response( - "tool_1", - Ok(CallToolResult::success(vec![ - rmcp::model::ContentBlock::text("fn main() { run(); }"), - ])), - ), - ] - } - - fn request_with(messages: &[Message]) -> Value { - create_request_with_default_options( + Message::user().with_text("Now add error handling to it."), + ]; + let req = create_request_with_default_options( &cfg("claude-sonnet-4-5"), "You are a careful coding assistant.", - messages, + &messages, &sample_tools(), ) - .unwrap() - } + .unwrap(); - fn request(turn_context_block: &str) -> Value { - request_with(&conversation(turn_context_block)) - } - - #[test] - fn cached_prefix_is_invariant_to_turn_context_changes() { - let req_a = request(&turn_context("2026-06-25 12:00:00", "14/40 used")); - let req_b = request(&turn_context("2026-06-25 13:47:00", "31/40 used")); - - assert_ne!( - req_a.to_string(), - req_b.to_string(), - "test setup is vacuous: the two requests are byte-identical, so the \ - turn-context never reached the request body" - ); + let tools = req["tools"].as_array().unwrap(); + assert!(tools[0].get(CACHE_CONTROL_FIELD).is_none()); + assert!(tools[1].get(CACHE_CONTROL_FIELD).is_some()); + assert!(req["system"][0].get(CACHE_CONTROL_FIELD).is_some()); + let request_messages = req["messages"].as_array().unwrap(); + let marked = breakpoints(request_messages); assert_eq!( - cached_prefix(&req_a), - cached_prefix(&req_b), - "the cached prefix changed when only the volatile turn-context changed; \ - prefix caching will collapse into a per-turn cache write" - ); - - assert!( - !cached_prefix(&req_a).contains("12:00:00"), - "the volatile turn-context timestamp leaked into the cached prefix" + marked, + vec![(2, 0), (4, 0)], + "message breakpoints should sit on the last block of the last two user messages" ); } #[test] - fn turn_context_sits_after_every_cache_breakpoint() { - let req = request(&turn_context("2026-06-25 12:00:00", "14/40 used")); - let messages = req["messages"].as_array().unwrap(); + fn breakpoints_land_on_the_last_content_block() { + let messages = vec![Message::user() + .with_text("Here is the context.") + .with_text("And the question.")]; + let req = create_request_with_default_options( + &cfg("claude-sonnet-4-5"), + "You are a careful coding assistant.", + &messages, + &sample_tools(), + ) + .unwrap(); - for message in messages { - for block in message["content"].as_array().unwrap() { - if block.get(CACHE_CONTROL_FIELD).is_some() { - assert!( - !is_turn_context_block(block), - "a cache_control breakpoint landed on the volatile turn-context block" - ); - } - } - } - - let breakpoint = last_breakpoint(messages).expect("a breakpoint should exist"); - let turn_context = find_turn_context(messages) - .expect("the turn-context block should survive into the formatted request"); - assert!( - turn_context > breakpoint, - "turn-context at {turn_context:?} is not after the last cache breakpoint at \ - {breakpoint:?}, so it sits inside a cached prefix" - ); - } - - /// Guards the tool-loop path: turn-context is injected onto an earlier - /// genuine user message while the request ends with a `tool_result`, so - /// keeping it out of the cached prefix requires relocating it across - /// messages. A regression that only reorders within a message would - /// pass the tests above but fail here. - #[test] - fn cached_prefix_is_invariant_in_tool_loop() { - let req_a = request_with(&tool_loop_conversation(&turn_context( - "2026-06-25 12:00:00", - "14/40", - ))); - let req_b = request_with(&tool_loop_conversation(&turn_context( - "2026-06-25 13:47:00", - "31/40", - ))); - - assert_ne!( - req_a.to_string(), - req_b.to_string(), - "test setup is vacuous: turn-context never reached the request body" - ); - - assert_eq!( - cached_prefix(&req_a), - cached_prefix(&req_b), - "the cached prefix changed when only the volatile turn-context changed during a \ - tool loop; the block was not relocated past the trailing tool_result breakpoint" - ); - - let messages = req_a["messages"].as_array().unwrap(); - let breakpoint = last_breakpoint(messages).expect("a breakpoint should exist"); - let turn_context = find_turn_context(messages) - .expect("the turn-context block should survive into the formatted request"); - assert!( - turn_context > breakpoint, - "turn-context at {turn_context:?} was not relocated across messages to after the \ - last breakpoint at {breakpoint:?}" - ); + let request_messages = req["messages"].as_array().unwrap(); + assert_eq!(breakpoints(request_messages), vec![(0, 1)]); } } } diff --git a/crates/goose-provider-types/src/formats/databricks.rs b/crates/goose-provider-types/src/formats/databricks.rs index 3eccc400c9..236ed1d2a8 100644 --- a/crates/goose-provider-types/src/formats/databricks.rs +++ b/crates/goose-provider-types/src/formats/databricks.rs @@ -1,3 +1,4 @@ +use crate::cache_semantics::{apply_chat_payload_breakpoints, CacheSemantics}; use crate::conversation::message::{Message, MessageContentBlock}; use crate::formats::anthropic::{ adaptive_output_effort, model_supports_temperature, thinking_budget_tokens, @@ -475,87 +476,6 @@ fn is_claude_model(model_name: &str) -> bool { model_name.contains("claude") } -/// Add Anthropic-style cache_control fields to the request payload for Claude models. -/// This enables prompt caching to reduce costs when using Claude via Databricks. -/// -/// Cache control is added to: -/// - The system message -/// - The last two user messages (for incremental caching across turns) -/// - The last tool definition (so all tools are cached as a single prefix) -pub fn apply_cache_control_for_claude(payload: &mut Value) { - if let Some(messages_spec) = payload - .as_object_mut() - .and_then(|obj| obj.get_mut("messages")) - .and_then(|messages| messages.as_array_mut()) - { - // Add cache_control to the last two user messages for incremental caching. - // The last message gets cached so future turns can read from it. - // The second-to-last user message is also cached to read from the previous cache. - let mut user_count = 0; - for message in messages_spec.iter_mut().rev() { - if message.get("role") == Some(&json!("user")) { - if let Some(content) = message.get_mut("content") { - if let Some(content_str) = content.as_str() { - *content = json!([{ - "type": "text", - "text": content_str, - "cache_control": { "type": "ephemeral" } - }]); - } else if let Some(content_array) = content.as_array_mut() { - // Content is already an array, add cache_control to the last element - if let Some(last_content) = content_array.last_mut() { - if let Some(obj) = last_content.as_object_mut() { - obj.insert( - "cache_control".to_string(), - json!({ "type": "ephemeral" }), - ); - } - } - } - } - user_count += 1; - if user_count >= 2 { - break; - } - } - } - - // Add cache_control to the system message - if let Some(system_message) = messages_spec - .iter_mut() - .find(|msg| msg.get("role") == Some(&json!("system"))) - { - if let Some(content) = system_message.get_mut("content") { - if let Some(content_str) = content.as_str() { - *system_message = json!({ - "role": "system", - "content": [{ - "type": "text", - "text": content_str, - "cache_control": { "type": "ephemeral" } - }] - }); - } - } - } - } - - // Add cache_control to the last tool definition - if let Some(tools_spec) = payload - .as_object_mut() - .and_then(|obj| obj.get_mut("tools")) - .and_then(|tools| tools.as_array_mut()) - { - if let Some(last_tool) = tools_spec.last_mut() { - if let Some(function) = last_tool.get_mut("function") { - if let Some(obj) = function.as_object_mut() { - obj.insert("cache_control".to_string(), json!({ "type": "ephemeral" })); - } - } - } - } -} - #[allow(clippy::too_many_lines)] pub fn create_request( model_config: &ModelConfig, @@ -658,9 +578,9 @@ pub fn create_request_for_provider( ); } - // Apply cache control for Claude models to enable prompt caching - if is_claude_model(&model_config.model_name) { - apply_cache_control_for_claude(&mut payload); + if CacheSemantics::for_model("databricks", &model_config.model_name).uses_explicit_breakpoints() + { + apply_chat_payload_breakpoints(&mut payload); } // Add request_params to the payload (e.g., anthropic_beta for extended context) @@ -1534,138 +1454,6 @@ mod tests { assert!(!is_claude_model("databricks-meta-llama-3-3-70b")); } - #[test] - fn test_apply_cache_control_for_claude_system_message() -> anyhow::Result<()> { - let mut payload = json!({ - "model": "databricks-claude-sonnet-4", - "messages": [ - { - "role": "system", - "content": "You are a helpful assistant." - }, - { - "role": "user", - "content": "Hello" - } - ] - }); - - apply_cache_control_for_claude(&mut payload); - - let messages = payload["messages"].as_array().unwrap(); - let system_msg = &messages[0]; - - // System message content should be converted to array with cache_control - assert!(system_msg["content"].is_array()); - let content = system_msg["content"].as_array().unwrap(); - assert_eq!(content.len(), 1); - assert_eq!(content[0]["type"], "text"); - assert_eq!(content[0]["text"], "You are a helpful assistant."); - assert_eq!(content[0]["cache_control"]["type"], "ephemeral"); - - Ok(()) - } - - #[test] - fn test_apply_cache_control_for_claude_user_messages() -> anyhow::Result<()> { - let mut payload = json!({ - "model": "databricks-claude-sonnet-4", - "messages": [ - { - "role": "system", - "content": "You are helpful" - }, - { - "role": "user", - "content": "First question" - }, - { - "role": "assistant", - "content": "First answer" - }, - { - "role": "user", - "content": "Second question" - }, - { - "role": "assistant", - "content": "Second answer" - }, - { - "role": "user", - "content": "Third question" - } - ] - }); - - apply_cache_control_for_claude(&mut payload); - - let messages = payload["messages"].as_array().unwrap(); - - // First user message should NOT have cache_control (only last 2) - let first_user = &messages[1]; - assert_eq!(first_user["content"], "First question"); - - // Second-to-last user message should have cache_control - let second_user = &messages[3]; - assert!(second_user["content"].is_array()); - assert_eq!( - second_user["content"][0]["cache_control"]["type"], - "ephemeral" - ); - - // Last user message should have cache_control - let last_user = &messages[5]; - assert!(last_user["content"].is_array()); - assert_eq!( - last_user["content"][0]["cache_control"]["type"], - "ephemeral" - ); - - Ok(()) - } - - #[test] - fn test_apply_cache_control_for_claude_tools() -> anyhow::Result<()> { - let mut payload = json!({ - "model": "databricks-claude-sonnet-4", - "messages": [ - { - "role": "system", - "content": "You are helpful" - } - ], - "tools": [ - { - "type": "function", - "function": { - "name": "tool1", - "description": "First tool" - } - }, - { - "type": "function", - "function": { - "name": "tool2", - "description": "Second tool" - } - } - ] - }); - - apply_cache_control_for_claude(&mut payload); - - let tools = payload["tools"].as_array().unwrap(); - - // First tool should NOT have cache_control - assert!(tools[0]["function"].get("cache_control").is_none()); - - // Last tool should have cache_control - assert_eq!(tools[1]["function"]["cache_control"]["type"], "ephemeral"); - - Ok(()) - } - #[test] fn test_format_messages_with_thought_signature_metadata() -> anyhow::Result<()> { let mut metadata = serde_json::Map::new(); diff --git a/crates/goose-provider-types/src/formats/openai.rs b/crates/goose-provider-types/src/formats/openai.rs index 52760a043b..2f24a0d2a4 100644 --- a/crates/goose-provider-types/src/formats/openai.rs +++ b/crates/goose-provider-types/src/formats/openai.rs @@ -67,10 +67,6 @@ pub fn is_reserved_request_param_key(key: &str) -> bool { pub struct OpenAiFormatOptions { pub preserve_thinking_context: bool, pub thinking_preservation_format: Option, - /// Keep the turn-context block on its source user message. The implicit - /// cache behind OpenAI's Responses-stack models only extends append-only - /// prompts, so the relocated tail would cap hits at the static prefix. - pub turn_context_in_place: bool, } fn merge_reasoning_text(prefix: &str, suffix: &str) -> String { @@ -213,15 +209,6 @@ pub fn format_messages_with_options( image_format: &ImageFormat, options: OpenAiFormatOptions, ) -> Vec { - let extracted = if options.turn_context_in_place { - None - } else { - extract_turn_context(messages) - }; - let (messages, turn_context) = match &extracted { - Some((stripped, text)) => (stripped.as_slice(), Some(text.as_str())), - None => (messages, None), - }; let mut messages_spec = Vec::new(); let mut pending_assistant_reasoning = String::new(); // Reasoning to propagate across consecutive tool-call messages in the same turn. @@ -531,60 +518,9 @@ pub fn format_messages_with_options( inline_reasoning_content(&mut messages_spec, format); } - if let Some(text) = turn_context { - append_turn_context_tail(&mut messages_spec, text); - } - messages_spec } -/// The volatile turn-context block busts implicit prefix caching from -/// mid-history; re-emitted at the request tail instead. Mirrors formats/anthropic.rs. -fn extract_turn_context(messages: &[Message]) -> Option<(Vec, String)> { - let (mi, bi) = messages.iter().enumerate().rev().find_map(|(mi, m)| { - if m.role != Role::User { - return None; - } - m.content - .iter() - .position(|block| { - matches!(block, MessageContentBlock::Text(t) - if crate::conversation::is_turn_context_text(&t.text)) - }) - .map(|bi| (mi, bi)) - })?; - if mi + 1 != messages.len() && messages[mi].content.len() <= 1 { - return None; - } - let mut messages = messages.to_vec(); - let MessageContentBlock::Text(text) = messages[mi].content.remove(bi) else { - return None; - }; - Some((messages, text.text)) -} - -/// Merges into a trailing user message when one exists; strict chat templates -/// reject consecutive user messages. -fn append_turn_context_tail(messages_spec: &mut Vec, text: &str) { - if let Some(last) = messages_spec.last_mut() { - if last["role"] == json!("user") { - match last.get_mut("content") { - Some(Value::String(existing)) => { - existing.push('\n'); - existing.push_str(text); - return; - } - Some(Value::Array(blocks)) => { - blocks.push(json!({"type": "text", "text": text})); - return; - } - _ => {} - } - } - } - messages_spec.push(json!({"role": "user", "content": text})); -} - /// Rewrites `reasoning_content` into the message `content` for models that reject a /// separate reasoning field on replay. /// @@ -1628,8 +1564,6 @@ pub fn create_request_for_model_with_options( "content": system }); - let mut format_options = format_options; - format_options.turn_context_in_place |= is_reasoning_model; let messages_spec = format_messages_with_options(messages, image_format, format_options); let mut tools_spec = format_tools(tools)?; @@ -4849,7 +4783,6 @@ data: [DONE]"#; OpenAiFormatOptions { preserve_thinking_context: true, thinking_preservation_format: Some(format), - ..Default::default() }, ) } @@ -4965,8 +4898,8 @@ data: [DONE]"#; Message::user().with_text("What does the main entrypoint do?"), Message::assistant().with_text("Let me read it."), Message::user() - .with_text(turn_context_block) - .with_text("Now add error handling to it."), + .with_text("Now add error handling to it.") + .with_text(turn_context_block), Message::assistant().with_tool_request( "tool_1", Ok(CallToolRequestParams::new("read_file") @@ -4981,182 +4914,26 @@ data: [DONE]"#; ] } - fn prefix(spec: &[Value]) -> String { - json!(spec[..spec.len() - 1]).to_string() - } - #[test] - fn formatted_prefix_is_invariant_to_turn_context_changes() { - let spec_a = format_messages( - &tool_loop_conversation(&turn_context("2026-08-03 12:00:00", "14/40 used")), - &ImageFormat::OpenAi, - ); - let spec_b = format_messages( - &tool_loop_conversation(&turn_context("2026-08-03 13:47:00", "31/40 used")), - &ImageFormat::OpenAi, - ); - - assert_ne!( - json!(spec_a).to_string(), - json!(spec_b).to_string(), - "test setup is vacuous: the turn-context never reached the request body" - ); - assert_eq!( - prefix(&spec_a), - prefix(&spec_b), - "the formatted prefix changed when only the volatile turn-context changed; \ - implicit prefix caching will re-prefill the conversation tail every request" - ); - } - - #[test] - fn turn_context_is_relocated_to_a_trailing_user_message() { + fn turn_context_stays_on_its_source_message() { let block = turn_context("2026-08-03 12:00:00", "14/40 used"); let spec = format_messages(&tool_loop_conversation(&block), &ImageFormat::OpenAi); - let last = spec.last().unwrap(); - assert_eq!(last["role"], "user"); - assert_eq!(last["content"], json!(block)); - - let occurrences = spec + let occurrences: Vec = spec .iter() - .filter(|m| { + .enumerate() + .filter_map(|(i, m)| { m["content"] .as_str() .is_some_and(|c| c.contains("")) + .then_some(i) }) - .count(); - assert_eq!(occurrences, 1, "turn-context must appear exactly once"); - + .collect(); assert_eq!( - spec[2]["content"], - json!("Now add error handling to it."), - "the source user message must keep its genuine text untouched" + occurrences, + vec![2], + "turn-context must appear exactly once, on its source message" ); } - - #[test] - fn first_request_of_a_turn_keeps_turn_context_inside_the_user_message() { - let block = turn_context("2026-08-03 12:00:00", "1/40 used"); - let messages = vec![Message::user() - .with_text(&block) - .with_text("What does the main entrypoint do?")]; - let spec = format_messages(&messages, &ImageFormat::OpenAi); - - assert_eq!(spec.len(), 1); - assert_eq!( - spec[0]["content"], - json!(format!("What does the main entrypoint do?\n{block}")) - ); - } - - #[test] - fn mid_history_sole_turn_context_is_left_in_place() { - let block = turn_context("2026-08-03 12:00:00", "14/40 used"); - let messages = vec![ - Message::user().with_text(&block), - Message::assistant().with_text("Understood."), - Message::user().with_text("Please refactor the argument parser."), - ]; - let spec = format_messages(&messages, &ImageFormat::OpenAi); - - assert_eq!(spec.len(), 3); - assert_eq!(spec[0]["content"], json!(block)); - assert_eq!( - spec[2]["content"], - json!("Please refactor the argument parser.") - ); - } - - #[test] - fn turn_context_merges_into_the_synthetic_image_message() { - let block = turn_context("2026-08-03 12:00:00", "14/40 used"); - let mut messages = tool_loop_conversation(&block); - messages.pop(); - messages.push(Message::user().with_tool_response( - "tool_1", - Ok(CallToolResult::success(vec![ContentBlock::image( - "aGVsbG8=", - "image/png", - )])), - )); - let spec = format_messages(&messages, &ImageFormat::OpenAi); - - let last = spec.last().unwrap(); - assert_eq!(last["role"], "user"); - let blocks = last["content"].as_array().unwrap(); - assert!(blocks.iter().any(|b| b["type"] == json!("image_url"))); - assert_eq!(blocks.last().unwrap()["text"], json!(block)); - for pair in spec.windows(2) { - assert!( - !(pair[0]["role"] == json!("user") && pair[1]["role"] == json!("user")), - "consecutive user messages reached the wire: {pair:?}" - ); - } - } - - fn responses_stack_request(messages: &[Message]) -> Vec { - let request = create_request( - &test_model_config("openai/gpt-5.6-terra"), - "system", - messages, - &[], - &ImageFormat::OpenAi, - false, - ) - .unwrap(); - request["messages"].as_array().unwrap().clone() - } - - #[test] - fn responses_stack_tool_loop_requests_are_append_only() { - let block = turn_context("2026-08-03 12:00:00", "14/40 used"); - let earlier = tool_loop_conversation(&block); - let mut later = earlier.clone(); - later.push( - Message::assistant().with_tool_request( - "tool_2", - Ok(CallToolRequestParams::new("read_file") - .with_arguments(object!({"path": "src/lib.rs"}))), - ), - ); - later.push(Message::user().with_tool_response( - "tool_2", - Ok(CallToolResult::success(vec![ContentBlock::text( - "pub fn run() {}", - )])), - )); - - let spec_earlier = responses_stack_request(&earlier); - let spec_later = responses_stack_request(&later); - - assert!(spec_later.len() > spec_earlier.len()); - for (i, earlier_msg) in spec_earlier.iter().enumerate() { - assert_eq!( - json!(earlier_msg).to_string(), - json!(spec_later[i]).to_string(), - "request N must be a byte prefix of request N+1 at index {i}" - ); - } - } - - #[test] - fn chat_stack_models_still_relocate_turn_context_to_the_tail() { - let block = turn_context("2026-08-03 12:00:00", "14/40 used"); - let request = create_request( - &test_model_config("openai/gpt-4.1-mini"), - "system", - &tool_loop_conversation(&block), - &[], - &ImageFormat::OpenAi, - false, - ) - .unwrap(); - let spec = request["messages"].as_array().unwrap(); - - let last = spec.last().unwrap(); - assert_eq!(last["role"], "user"); - assert_eq!(last["content"], json!(block)); - } } } diff --git a/crates/goose-provider-types/src/lib.rs b/crates/goose-provider-types/src/lib.rs index 3ab49ef87e..dd7a87dda7 100644 --- a/crates/goose-provider-types/src/lib.rs +++ b/crates/goose-provider-types/src/lib.rs @@ -1,4 +1,5 @@ pub mod base; +pub mod cache_semantics; pub mod canonical; pub mod conversation; pub mod errors; diff --git a/crates/goose-provider-types/tests/prefix_invariance.rs b/crates/goose-provider-types/tests/prefix_invariance.rs new file mode 100644 index 0000000000..5ec61e0eb6 --- /dev/null +++ b/crates/goose-provider-types/tests/prefix_invariance.rs @@ -0,0 +1,428 @@ +//! Across the consecutive requests of a session, the cache-relevant bytes a +//! provider has already seen must never change. +//! +//! Implicit caches require request N to be a verbatim item-prefix of request +//! N+1; explicit-breakpoint caches require the bytes up to the last +//! breakpoint to recur at the same positions, modulo the markers themselves. +//! Session states are shaped as the agent persists them and run through the +//! same projection and `fix_conversation` pass as the real request path. + +use goose_provider_types::cache_semantics::{apply_chat_payload_breakpoints, CacheSemantics}; +use goose_provider_types::conversation::message::{Message, MessageMetadata}; +use goose_provider_types::conversation::{ + fix_conversation, merge_consecutive_messages_for_request, Conversation, +}; +use goose_provider_types::formats::anthropic::{self, AnthropicFormatOptions}; +use goose_provider_types::formats::{databricks, openai, openai_responses}; +use goose_provider_types::images::ImageFormat; +use goose_provider_types::model::ModelConfig; +use rmcp::model::{CallToolRequestParams, CallToolResult, ContentBlock, Tool}; +use rmcp::object; +use serde_json::{json, Value}; + +const SYSTEM: &str = "You are a careful coding assistant."; + +fn tools() -> Vec { + vec![Tool::new( + "read_file", + "Read a file from disk", + object!({ + "type": "object", + "properties": { "path": { "type": "string" } } + }), + )] +} + +fn turn_block(turn: usize) -> String { + format!( + "\n\ + 2026-08-06 12:0{turn}:00 +02:00\n\ + /Users/me/code/goose\n\ + step {turn} of the plan\n\ + " + ) +} + +fn turn_context_message(turn: usize) -> Message { + Message::user() + .with_text(turn_block(turn)) + .with_metadata(MessageMetadata::agent_only().with_turn_context()) +} + +fn tool_exchange(id: &str, path: &str) -> [Message; 2] { + [ + Message::assistant().with_tool_request( + id, + Ok(CallToolRequestParams::new("read_file").with_arguments(object!({ "path": path }))), + ), + Message::user().with_tool_response( + id, + Ok(CallToolResult::success(vec![ContentBlock::text( + "fn main() { run(); }", + )])), + ), + ] +} + +/// Two turns with one tool-loop iteration each; every state appends to the +/// previous one. +fn session_states() -> Vec> { + let u1 = Message::user().with_text("What does the main entrypoint do?"); + let u2 = Message::user().with_text("Now add error handling to it."); + let (tc1, tc2) = (turn_context_message(1), turn_context_message(2)); + let [a1, t1] = tool_exchange("call_1", "src/main.rs"); + let [a2, t2] = tool_exchange("call_2", "src/lib.rs"); + let a_text = Message::assistant().with_text("It calls `run()`."); + + vec![ + vec![u1.clone(), tc1.clone()], + vec![u1.clone(), tc1.clone(), a1.clone(), t1.clone()], + vec![ + u1.clone(), + tc1.clone(), + a1.clone(), + t1.clone(), + a_text.clone(), + u2.clone(), + tc2.clone(), + ], + vec![u1, tc1, a1, t1, a_text, u2, tc2, a2, t2], + ] +} + +/// The projection every request goes through before any formatter +/// (`stream_response_from_provider`). +fn provider_view(messages: &[Message]) -> Vec { + let projected = + Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages(); + let (fixed, _) = fix_conversation(Conversation::new_unvalidated(projected)); + merge_consecutive_messages_for_request(fixed.messages().clone()) +} + +fn anthropic_request(messages: &[Message]) -> Value { + anthropic::create_request( + "anthropic", + &ModelConfig::new("claude-sonnet-4-5"), + SYSTEM, + &provider_view(messages), + &tools(), + AnthropicFormatOptions::default(), + ) + .unwrap() +} + +fn openai_chat_messages(model: &str, messages: &[Message]) -> Vec { + openai::create_request( + &ModelConfig::new(model), + SYSTEM, + &provider_view(messages), + &tools(), + &ImageFormat::OpenAi, + false, + ) + .unwrap()["messages"] + .as_array() + .unwrap() + .clone() +} + +fn responses_input(messages: &[Message]) -> Vec { + openai_responses::create_responses_request( + &ModelConfig::new("gpt-5.2"), + SYSTEM, + &provider_view(messages), + &tools(), + ) + .unwrap()["input"] + .as_array() + .unwrap() + .clone() +} + +fn openrouter_style_request(messages: &[Message]) -> Value { + let mut payload = openai::create_request( + &ModelConfig::new("anthropic/claude-sonnet-4-5"), + SYSTEM, + &provider_view(messages), + &tools(), + &ImageFormat::OpenAi, + false, + ) + .unwrap(); + assert!( + CacheSemantics::for_model("openrouter", "anthropic/claude-sonnet-4-5") + .uses_explicit_breakpoints() + ); + apply_chat_payload_breakpoints(&mut payload); + payload +} + +fn databricks_request(messages: &[Message]) -> Value { + databricks::create_request( + &ModelConfig::new("databricks-claude-sonnet-4"), + SYSTEM, + &provider_view(messages), + &tools(), + &ImageFormat::OpenAi, + ) + .unwrap() +} + +/// Strip breakpoint markers and normalize string content to block form: +/// upstream caches hash the normalized bytes, not the marker encoding. +fn canonical(value: &Value) -> Value { + match value { + Value::Object(map) => Value::Object( + map.iter() + .filter(|(k, _)| k.as_str() != "cache_control") + .map(|(k, v)| { + if k == "content" && v.is_string() { + ( + k.clone(), + json!([{ "type": "text", "text": v.as_str().unwrap() }]), + ) + } else { + (k.clone(), canonical(v)) + } + }) + .collect(), + ), + Value::Array(items) => Value::Array(items.iter().map(canonical).collect()), + other => other.clone(), + } +} + +fn strict_violations(prev: &[Value], next: &[Value]) -> Vec { + if next.len() < prev.len() { + return vec![format!( + "request shrank from {} to {} items", + prev.len(), + next.len() + )]; + } + prev.iter() + .zip(next) + .enumerate() + .filter(|(_, (p, n))| p != n) + .map(|(i, (p, n))| format!("items diverge at index {i}:\n sent: {p}\n next: {n}")) + .collect() +} + +fn last_breakpoint(messages: &[Value]) -> Option<(usize, usize)> { + let mut found = None; + for (mi, message) in messages.iter().enumerate() { + let blocks = match message["content"].as_array() { + Some(blocks) => blocks.len(), + None => 1, + }; + for bi in 0..blocks { + let block = match message["content"].as_array() { + Some(array) => &array[bi], + None => message, + }; + if block.get("cache_control").is_some() { + found = Some((mi, bi)); + } + } + } + found +} + +/// Everything the previous request cached (tools, system, messages up to its +/// last breakpoint) must recur verbatim in the next request. +fn explicit_violations(prev_request: &Value, next_request: &Value) -> Vec { + let mut violations = Vec::new(); + + let prev_messages = prev_request["messages"].as_array().unwrap(); + let Some((last_mi, last_bi)) = last_breakpoint(prev_messages) else { + violations.push("previous request carries no message breakpoint".into()); + return violations; + }; + + let prev = canonical(prev_request); + let next = canonical(next_request); + for section in ["tools", "system"] { + if prev[section] != next[section] { + violations.push(format!( + "cached section `{section}` changed between requests" + )); + } + } + + let prev_messages = prev["messages"].as_array().unwrap(); + let next_messages = next["messages"].as_array().unwrap(); + if next_messages.len() <= last_mi { + violations.push("messages shrank below the previous breakpoint".into()); + return violations; + } + for mi in 0..last_mi { + if prev_messages[mi] != next_messages[mi] { + violations.push(format!( + "cached bytes changed at message {mi}:\n sent: {}\n next: {}", + prev_messages[mi], next_messages[mi] + )); + } + } + + let prev_blocks = prev_messages[last_mi]["content"].as_array().unwrap(); + match next_messages[last_mi]["content"].as_array() { + Some(next_blocks) + if next_blocks.len() > last_bi + && prev_messages[last_mi]["role"] == next_messages[last_mi]["role"] => + { + for bi in 0..=last_bi { + if prev_blocks[bi] != next_blocks[bi] { + violations.push(format!( + "cached bytes changed at message {last_mi} block {bi}:\n sent: {}\n next: {}", + prev_blocks[bi], next_blocks[bi] + )); + } + } + } + _ => violations.push(format!("cached message {last_mi} changed shape")), + } + violations +} + +fn consecutive(items: &[T]) -> impl Iterator { + items.windows(2).enumerate().map(|(i, w)| (i, &w[0], &w[1])) +} + +fn assert_clean(label: &str, violations: Vec) { + assert!( + violations.is_empty(), + "{label}: cache-relevant bytes changed between consecutive requests:\n{}", + violations.join("\n") + ); +} + +#[test] +fn anthropic_requests_stay_breakpoint_compatible() { + let requests: Vec = session_states() + .iter() + .map(|m| anthropic_request(m)) + .collect(); + assert_ne!(requests[0], requests[1], "vacuous: requests are identical"); + for (i, prev, next) in consecutive(&requests) { + assert_clean( + &format!("anthropic pair {i}"), + explicit_violations(prev, next), + ); + } +} + +/// The chat fleet's caches are merely tolerant, but with no relocation left +/// the strict relation holds anyway. +#[test] +fn openai_chat_requests_are_append_only() { + let requests: Vec> = session_states() + .iter() + .map(|m| openai_chat_messages("gpt-4.1-mini", m)) + .collect(); + for (i, prev, next) in consecutive(&requests) { + assert_clean( + &format!("openai chat pair {i}"), + strict_violations(prev, next), + ); + } +} + +#[test] +fn responses_requests_are_append_only() { + let requests: Vec> = session_states() + .iter() + .map(|m| responses_input(m)) + .collect(); + for (i, prev, next) in consecutive(&requests) { + assert_clean( + &format!("responses pair {i}"), + strict_violations(prev, next), + ); + } +} + +#[test] +fn responses_stack_chat_requests_are_append_only() { + let requests: Vec> = session_states() + .iter() + .map(|m| openai_chat_messages("gpt-5.6-terra", m)) + .collect(); + for (i, prev, next) in consecutive(&requests) { + assert_clean( + &format!("responses-stack chat pair {i}"), + strict_violations(prev, next), + ); + } +} + +/// The OpenRouter/LiteLLM path: an OpenAI chat payload post-processed with +/// Anthropic-dialect breakpoints by the shared writer. +#[test] +fn openrouter_style_requests_stay_breakpoint_compatible() { + let requests: Vec = session_states() + .iter() + .map(|m| openrouter_style_request(m)) + .collect(); + for (i, prev, next) in consecutive(&requests) { + assert_clean( + &format!("openrouter pair {i}"), + explicit_violations(prev, next), + ); + } +} + +#[test] +fn databricks_requests_stay_breakpoint_compatible() { + let requests: Vec = session_states() + .iter() + .map(|m| databricks_request(m)) + .collect(); + for (i, prev, next) in consecutive(&requests) { + assert_clean( + &format!("databricks pair {i}"), + explicit_violations(prev, next), + ); + } +} + +#[test] +fn seeded_regression_volatile_bytes_in_cached_prefix_are_caught() { + let request_for = |stamp: &str| { + let block = + format!("\n{stamp}\n"); + let [a1, t1] = tool_exchange("call_1", "src/main.rs"); + anthropic_request(&[ + Message::user() + .with_text(block) + .with_visibility(false, true), + Message::user().with_text("Now add error handling to it."), + a1, + t1, + ]) + }; + let violations = explicit_violations( + &request_for("2026-08-06 12:00:00"), + &request_for("2026-08-06 12:01:00"), + ); + assert!( + !violations.is_empty(), + "expected per-request volatile bytes inside the cached prefix to be caught" + ); +} + +#[test] +fn seeded_regression_relocated_tail_is_caught() { + let states = session_states(); + let prev = responses_input(&states[0]); + + let mut relocated_state = states[1].clone(); + let block = relocated_state.remove(1); + relocated_state.push(block); + let next = responses_input(&relocated_state); + + let violations = strict_violations(&prev, &next); + assert!( + !violations.is_empty(), + "expected a relocated turn-context block to violate the append-only relation" + ); +} diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index 7454264586..fb75fac7b9 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -6,8 +6,8 @@ pub mod databricks_auth; pub mod databricks_v2; pub mod google; pub use goose_provider_types::{ - base, canonical, conversation, errors, formats, goose_mode, images, json, model, permission, - request_log, retry, thinking, utils, + base, cache_semantics, canonical, conversation, errors, formats, goose_mode, images, json, + model, permission, request_log, retry, thinking, utils, }; pub mod declarative; pub mod http_status; diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index e8d03ccb83..b256ca989d 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -782,7 +782,6 @@ impl Provider for OpenAiProvider { preserve_thinking_context: self.preserve_thinking_context || thinking_preservation_format.is_some(), thinking_preservation_format, - ..Default::default() }, )?; diff --git a/crates/goose/src/acp/provider.rs b/crates/goose/src/acp/provider.rs index ec3d4e193d..f2b0b4e0dd 100644 --- a/crates/goose/src/acp/provider.rs +++ b/crates/goose/src/acp/provider.rs @@ -1442,14 +1442,14 @@ fn messages_to_prompt(messages: &[Message], include_handoff_context: bool) -> Ve fn last_user_message_index(messages: &[Message]) -> Option { messages .iter() - .rposition(|m| m.role == Role::User && m.is_agent_visible()) + .rposition(|m| m.role == Role::User && m.is_agent_visible() && !m.is_turn_context()) } fn has_handoff_context(messages: &[Message]) -> bool { last_user_message_index(messages).is_some_and(|last_user_index| { messages[..last_user_index] .iter() - .any(Message::is_agent_visible) + .any(|m| m.is_agent_visible() && !m.is_turn_context()) }) } @@ -1458,6 +1458,7 @@ fn build_handoff_context_memo(prior_messages: &[Message]) -> Option { 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(); @@ -1810,6 +1811,37 @@ mod tests { assert_eq!(prompt_text(&blocks[1]), "continue from there"); } + #[test] + fn messages_to_prompt_skips_turn_context_events() { + use crate::conversation::message::MessageMetadata; + + let turn_context = |text: &str| { + Message::user() + .with_text(text) + .with_metadata(MessageMetadata::agent_only().with_turn_context()) + }; + let messages = vec![ + Message::user().with_text("inspect src/lib.rs"), + turn_context("old cwd /repo"), + Message::assistant().with_text("I found the file"), + Message::user().with_text("continue from there"), + turn_context("new cwd /repo"), + ]; + + let blocks = messages_to_prompt(&messages, true); + + assert_eq!(blocks.len(), 2); + assert!( + !prompt_text(&blocks[0]).contains("turn-context"), + "handoff memo must not include turn-context events" + ); + assert_eq!( + prompt_text(&blocks[1]), + "continue from there", + "the current prompt must be the user's request, not a trailing turn-context event" + ); + } + #[test] fn messages_to_prompt_drops_user_only_acp_rows_from_handoff() { let user_only = TextContent::new("SECRET_USER_ONLY") diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 0d877809bf..69591818d9 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -46,7 +46,9 @@ use crate::conversation::message::{ ActionRequiredData, InferenceMetadata, Message, MessageContent, MessageUsage, ProviderMetadata, SystemNotificationType, ToolRequest, }; -use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation}; +use crate::conversation::{ + debug_conversation_fix, fix_conversation, merge_consecutive_messages_for_request, Conversation, +}; use crate::mcp_utils::ToolResult; use crate::permission::permission_inspector::PermissionInspector; use crate::permission::permission_judge::PermissionCheckResult; @@ -150,7 +152,6 @@ pub struct ReplyContext { pub system_prompt: String, pub goose_mode: GooseMode, pub tool_call_cut_off: usize, - pub initial_messages: Vec, pub model_config: goose_providers::model::ModelConfig, } @@ -787,8 +788,6 @@ impl Agent { ) ); } - let initial_messages = conversation.messages().clone(); - let (tools, toolshim_tools, system_prompt, model_config) = self .prepare_tools_and_prompt(session_id, working_dir) .await?; @@ -820,7 +819,6 @@ impl Agent { system_prompt, goose_mode, tool_call_cut_off, - initial_messages, model_config, }) } @@ -2154,7 +2152,6 @@ impl Agent { mut system_prompt, tool_call_cut_off, goose_mode, - initial_messages, model_config, } = context; @@ -2264,7 +2261,27 @@ impl Agent { let turn_start_compaction_info = super::moim::compute_compaction_info(&session_config.id, &self.extension_manager) .await; - let turn_start_turns_taken = turns_taken; + + if let Some(turn_context) = super::moim::turn_context_message( + &session_config.id, + &self.extension_manager, + turns_taken, + max_turns, + turn_start, + turn_start_compaction_info, + ) + .await + { + persist_and_push_message_with_id( + &session_manager, + &session_config.id, + &mut conversation, + turn_context, + ) + .await?; + } + // Snapshot after the turn-context append so a retry keeps the sent prefix. + let initial_messages = conversation.messages().clone(); loop { if is_token_cancelled(&cancel_token) { @@ -2359,23 +2376,12 @@ impl Agent { break; } - let conversation_with_moim = super::moim::inject_moim( - &session_config.id, - conversation.clone(), - &self.extension_manager, - turn_start_turns_taken, - max_turns, - turn_start, - turn_start_compaction_info.clone(), - ) - .await; - let mut stream = crate::agents::reply_parts::stream_response_from_provider( self.provider().await?, model_config.clone(), &session_config.id, &system_prompt, - conversation_with_moim.messages(), + conversation.messages(), &tools, &toolshim_tools, ).await?; @@ -3734,7 +3740,7 @@ impl Agent { .filter(super::reply_parts::is_tool_visible_to_model) .collect(); - messages = Conversation::new_unvalidated(messages.agent_visible_messages()); + messages = Conversation::new_unvalidated(recipe_conversation_history(&messages)); messages.push(Message::user().with_text(recipe_prompt)); let (messages, issues) = fix_conversation(messages); @@ -3743,6 +3749,9 @@ impl Agent { .iter() .for_each(|issue| tracing::warn!(recipe.conversation.issue = issue)); } + let messages = Conversation::new_unvalidated(merge_consecutive_messages_for_request( + messages.messages().clone(), + )); tracing::debug!( "Added recipe prompt to messages, total messages: {}", @@ -3908,6 +3917,15 @@ impl Agent { } } +fn recipe_conversation_history(messages: &Conversation) -> Vec { + // The recipe prompt has no turn-context instructions; drop the blocks. + messages + .agent_visible_messages() + .into_iter() + .filter(|message| !message.is_turn_context()) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -3923,6 +3941,25 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use tempfile::TempDir; + #[test] + fn recipe_history_excludes_turn_context_events() { + use crate::conversation::message::MessageMetadata; + + let history = Conversation::new_unvalidated([ + Message::user().with_text("build me a recipe"), + Message::user() + .with_text("cwd /repo") + .with_metadata(MessageMetadata::agent_only().with_turn_context()), + Message::assistant().with_text("on it"), + ]); + + let texts: Vec = recipe_conversation_history(&history) + .iter() + .map(|message| message.as_concat_text()) + .collect(); + assert_eq!(texts, ["build me a recipe", "on it"]); + } + async fn tracing_test_agent_and_session() -> (Agent, Session, TempDir) { let data_dir = TempDir::new().unwrap(); let data_path = data_dir.path().to_path_buf(); diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 0d431a57ec..16a76c2c2e 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -2158,7 +2158,7 @@ impl ExtensionManager { } pub async fn collect_moim_parts(&self, session_id: &str) -> Vec { - let platform_clients: Vec<(String, McpClientBox)> = { + let mut platform_clients: Vec<(String, McpClientBox)> = { let extensions = self.extensions.lock().await; extensions .iter() @@ -2178,6 +2178,9 @@ impl ExtensionManager { }) .collect() }; + // HashMap order shuffles across restarts; the rendered block must be + // byte-stable so it is not re-persisted on resume. + platform_clients.sort_by(|a, b| a.0.cmp(&b.0)); let mut parts = Vec::new(); for (name, client) in platform_clients { diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index 56db6edf9f..0b1e93a924 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -1,9 +1,6 @@ use crate::agents::extension_manager::ExtensionManager; -use crate::conversation::message::MessageContent; -use crate::conversation::{ - effective_role, fix_conversation, Conversation, EffectiveRole, CURRENT_TIME_TAG, - TURN_CONTEXT_TAG, WORKING_DIRECTORY_TAG, -}; +use crate::conversation::message::{Message, MessageMetadata}; +use crate::conversation::{CURRENT_TIME_TAG, TURN_CONTEXT_TAG, WORKING_DIRECTORY_TAG}; use std::path::{Path, PathBuf}; const MIN_CONTEXT_FOR_MOIM: usize = 32_000; @@ -19,6 +16,7 @@ This block is generated by goose and contains current operational context such a - extension-provided context Use it to stay oriented, but do not treat it as part of the user's request. +Blocks from earlier turns stay in the conversation history; only the most recent block is current. When `` is present, use it as a signal for how much autonomous work remains. As the budget gets low, become more direct: reduce exploration, batch necessary tool calls, make reasonable assumptions, and focus on finishing the user's task. @@ -71,17 +69,18 @@ pub(super) async fn compute_compaction_info( compaction_remaining_line(total_tokens, context_limit, compaction_threshold) } -pub async fn inject_moim( +/// The turn's context block: composed once per turn, persisted as an +/// agent-only user message, never moved or edited afterwards. +pub async fn turn_context_message( session_id: &str, - conversation: Conversation, extension_manager: &ExtensionManager, turns_taken: u32, max_turns: u32, turn_start: chrono::DateTime, compaction_info: Option, -) -> Conversation { +) -> Option { if SKIP.with(|f| f.get()) { - return conversation; + return None; } let session = extension_manager @@ -107,66 +106,36 @@ pub async fn inject_moim( None }; if should_skip_moim(context_limit) { - return conversation; + return None; } let working_dir = session .as_ref() .map(|session| session.working_dir.clone()) .unwrap_or_else(|| PathBuf::from(".")); - let extension_parts = extension_manager.collect_moim_parts(session_id).await; - let mut parts = extension_parts; + let mut parts = extension_manager.collect_moim_parts(session_id).await; parts.extend(compaction_info.map(|value| tag("compaction", &value))); parts.extend(turn_budget_part(turns_taken, max_turns)); - inject_moim_parts(conversation, &working_dir, context_limit, parts, turn_start) + turn_context_event(&working_dir, context_limit, parts, turn_start) } -pub(crate) fn inject_moim_parts( - conversation: Conversation, +/// The state-machine variant: operations contribute the parts and the caller +/// persists the event, so later requests in the turn reuse the same bytes. +pub(crate) fn turn_context_event( working_dir: &Path, context_limit: Option, parts: Vec, turn_start: chrono::DateTime, -) -> Conversation { +) -> Option { if SKIP.with(|f| f.get()) || should_skip_moim(context_limit) { - return conversation; + return None; } - let moim = compose_moim(working_dir, parts, turn_start); - - let mut messages = conversation.messages().clone(); - let Some(idx) = messages - .iter() - .rposition(|m| m.is_agent_visible() && effective_role(m) == EffectiveRole::User) - else { - return conversation; - }; - let insert_idx = messages[idx] - .content - .iter() - .take_while(|content| matches!(content, MessageContent::ToolResponse(_))) - .count(); - messages[idx] - .content - .insert(insert_idx, MessageContent::text(moim)); - - let (fixed, issues) = fix_conversation(Conversation::new_unvalidated(messages)); - - let has_unexpected_issues = issues.iter().any(|issue| { - !issue.contains("Merged consecutive user messages") - && !issue.contains("Merged consecutive assistant messages") - && !issue.contains("Added placeholder to empty tool result") - && !issue.contains("Trimmed trailing whitespace from assistant message") - && !issue.contains("Removed trailing assistant message") - && !issue.contains("Merged text content") - }); - - if has_unexpected_issues { - tracing::warn!("MOIM injection caused unexpected issues: {:?}", issues); - return conversation; - } - - fixed + Some( + Message::user() + .with_text(compose_moim(working_dir, parts, turn_start)) + .with_metadata(MessageMetadata::agent_only().with_turn_context()), + ) } fn should_skip_moim(context_limit: Option) -> bool { @@ -252,21 +221,8 @@ fn turn_budget_part(turns_taken: u32, max_turns: u32) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::conversation::message::Message; - use rmcp::model::CallToolRequestParams; - fn text_at(message: &crate::conversation::message::Message, index: usize) -> &str { - message.content[index].as_text().unwrap() - } - - fn is_moim(content: &MessageContent) -> bool { - content - .as_text() - .is_some_and(|text| text.starts_with(&format!("<{}>\n", TURN_CONTEXT_TAG))) - } - - #[tokio::test] - async fn test_moim_prepended_to_latest_user_message() { + async fn session_and_manager() -> (String, ExtensionManager, tempfile::TempDir) { let temp_dir = tempfile::tempdir().unwrap(); let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); let session = em @@ -280,209 +236,83 @@ mod tests { ) .await .unwrap(); - - let conv = Conversation::new_unvalidated(vec![ - Message::user().with_text("Hello"), - Message::assistant().with_text("Hi"), - Message::user().with_text("Bye"), - ]); - let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; - let msgs = result.messages(); - - assert_eq!(msgs.len(), 3); - assert_eq!(text_at(&msgs[0], 0), "Hello"); - assert_eq!(text_at(&msgs[1], 0), "Hi"); - assert!(is_moim(&msgs[2].content[0])); - assert_eq!(text_at(&msgs[2], 1), "Bye"); + (session.id, em, temp_dir) } #[tokio::test] - async fn test_moim_injection_no_assistant() { - let temp_dir = tempfile::tempdir().unwrap(); - let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); - let session = em - .get_context() - .session_manager - .create_session( - PathBuf::from("/test/dir"), - "test".to_string(), - crate::session::SessionType::User, - crate::config::GooseMode::Auto, - ) + async fn turn_context_message_is_an_agent_only_user_message() { + let (session_id, em, _tmp) = session_and_manager().await; + + let message = turn_context_message(&session_id, &em, 0, 100, chrono::Local::now(), None) .await - .unwrap(); + .expect("turn context should be produced"); - let conv = Conversation::new_unvalidated(vec![Message::user().with_text("Hello")]); - let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; - - assert_eq!(result.messages().len(), 1); - assert!(is_moim(&result.messages()[0].content[0])); - assert_eq!(text_at(&result.messages()[0], 1), "Hello"); + assert_eq!(message.role, rmcp::model::Role::User); + assert!(message.is_agent_visible()); + assert!(!message.is_user_visible()); + assert!(message.is_turn_context()); + let text = message.content[0].as_text().unwrap(); + assert!(text.starts_with(&format!("<{TURN_CONTEXT_TAG}>\n"))); + assert!(text.trim_end().ends_with(&format!(""))); } #[tokio::test] - async fn test_moim_skips_user_messages_not_visible_to_agent() { - let temp_dir = tempfile::tempdir().unwrap(); - let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); - let session = em - .get_context() - .session_manager - .create_session( - PathBuf::from("/test/dir"), - "test".to_string(), - crate::session::SessionType::User, - crate::config::GooseMode::Auto, - ) + async fn turn_context_bytes_are_stable_for_a_turn() { + let (session_id, em, _tmp) = session_and_manager().await; + let turn_start = chrono::Local::now(); + + let first = turn_context_message(&session_id, &em, 0, 100, turn_start, None) + .await + .unwrap(); + let second = turn_context_message(&session_id, &em, 0, 100, turn_start, None) .await .unwrap(); - let conv = Conversation::new_unvalidated(vec![ - Message::user().with_text("agent visible"), - Message::assistant().with_text("reply"), - Message::user().with_text("user only").user_only(), - ]); - let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; - let msgs = result.messages(); - - assert_eq!(msgs.len(), 2); - assert!(is_moim(&msgs[0].content[0])); - assert_eq!(text_at(&msgs[0], 1), "agent visible"); - assert_eq!(text_at(&msgs[1], 0), "user only"); - assert!(!msgs[1].is_agent_visible()); + assert_eq!( + first.content[0].as_text(), + second.content[0].as_text(), + "the same turn inputs must render byte-identical blocks" + ); } - #[tokio::test] - async fn test_moim_with_tool_calls() { - let temp_dir = tempfile::tempdir().unwrap(); - let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); - let session = em - .get_context() - .session_manager - .create_session( - PathBuf::from("/test/dir"), - "test".to_string(), - crate::session::SessionType::User, - crate::config::GooseMode::Auto, - ) - .await - .unwrap(); + #[test] + fn turn_context_event_is_marked_and_skips_small_contexts() { + let event = turn_context_event( + Path::new("/test/dir"), + Some(200_000), + vec![], + chrono::Local::now(), + ) + .expect("large-context models should get a turn-context event"); + assert!(event.is_turn_context()); + assert!(!event.is_user_visible()); - let conv = Conversation::new_unvalidated(vec![ - Message::user().with_text("Search for something"), - Message::assistant() - .with_text("I'll search for you") - .with_tool_request("search_1", Ok(CallToolRequestParams::new("search"))), - Message::user() - .with_tool_response("search_1", Ok(rmcp::model::CallToolResult::success(vec![]))), - ]); - - let result = inject_moim(&session.id, conv, &em, 0, 100, chrono::Local::now(), None).await; - let msgs = result.messages(); - - assert_eq!(msgs.len(), 3); - assert!(is_moim(&msgs[0].content[0])); - assert_eq!(text_at(&msgs[0], 1), "Search for something"); - assert!(matches!( - &msgs[2].content[0], - MessageContent::ToolResponse(_) - )); - assert_eq!(msgs[2].content.len(), 1); + assert!(turn_context_event( + Path::new("/test/dir"), + Some(16_000), + vec![], + chrono::Local::now(), + ) + .is_none()); } - /// The turn-context block is produced here (`compose_moim`, in goose) but - /// recognized in a different crate (`is_turn_context_text`, in - /// goose-providers) by matching its textual shape. That recognition is the - /// only thing that keeps the volatile block out of the Anthropic cache - /// prefix. The two are coupled by string layout alone, with no shared type - /// or compiler check, so a reorder or rename on either side would silently - /// collapse prompt caching. These tests fail loudly if they drift apart. - mod turn_context_detector_coupling { - use super::*; - use crate::conversation::is_turn_context_text; - use std::path::Path; + #[test] + fn current_time_is_timezone_aware() { + let block = compose_moim( + Path::new("/Users/me/code/goose"), + vec![], + chrono::Local::now(), + ); + let time_line = block + .lines() + .find(|line| line.starts_with(&format!("<{CURRENT_TIME_TAG}>"))) + .expect("turn-context block should contain a current-time line"); - fn moim( - total_tokens: Option, - context_limit: Option, - turns_taken: u32, - max_turns: u32, - mut parts: Vec, - ) -> String { - parts.extend( - compaction_remaining_line(total_tokens, context_limit, 0.8) - .map(|value| tag("compaction", &value)), - ); - parts.extend(turn_budget_part(turns_taken, max_turns)); - compose_moim( - Path::new("/Users/me/code/goose"), - parts, - chrono::Local::now(), - ) - } + let value = time_line + .trim_start_matches(&format!("<{CURRENT_TIME_TAG}>")) + .trim_end_matches(&format!("")); - #[test] - fn compose_moim_output_is_recognized_by_the_anthropic_detector() { - let guardrail = || vec!["stay on task".to_string()]; - let cases = [ - ("minimal", moim(None, None, 0, 0, vec![])), - ( - "with compaction line", - moim(Some(100_000), Some(200_000), 0, 0, vec![]), - ), - ("with turn-budget line", moim(None, None, 20, 40, vec![])), - ( - "with extension context", - moim(None, None, 0, 0, guardrail()), - ), - ( - "fully populated", - moim(Some(100_000), Some(200_000), 20, 40, guardrail()), - ), - ]; - - for (label, block) in &cases { - assert!( - is_turn_context_text(block), - "is_turn_context_text rejected real compose_moim output ({label}); the \ - producer in goose and the detector in goose-providers have drifted, so the \ - volatile block will no longer be kept out of the Anthropic cache prefix:\n\ - {block}" - ); - } - } - - #[test] - fn detector_rejects_text_that_is_not_a_turn_context_block() { - assert!( - !is_turn_context_text("Please refactor the argument parser."), - "ordinary user text must not be treated as a turn-context block" - ); - assert!( - !is_turn_context_text( - "\n\ - 2026-06-25 12:00:00\n\ - abc\n\ - " - ), - "a block whose first field after current-time is not working-directory must not \ - match: the detector's shape checks are what discriminate the injected block" - ); - } - - #[test] - fn current_time_is_timezone_aware() { - let block = moim(None, None, 0, 0, vec![]); - let time_line = block - .lines() - .find(|line| line.starts_with(&format!("<{CURRENT_TIME_TAG}>"))) - .expect("turn-context block should contain a current-time line"); - - let value = time_line - .trim_start_matches(&format!("<{CURRENT_TIME_TAG}>")) - .trim_end_matches(&format!("")); - - chrono::DateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S %:z") - .expect("current-time should include a numeric UTC offset"); - } + chrono::DateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S %:z") + .expect("current-time should include a numeric UTC offset"); } } diff --git a/crates/goose/src/agents/platform_extensions/chatrecall.rs b/crates/goose/src/agents/platform_extensions/chatrecall.rs index 7f76f0868e..0e437bfdbf 100644 --- a/crates/goose/src/agents/platform_extensions/chatrecall.rs +++ b/crates/goose/src/agents/platform_extensions/chatrecall.rs @@ -41,7 +41,11 @@ pub struct ChatRecallClient { } fn format_agent_visible_excerpt(conversation: &Conversation) -> Option<(usize, String)> { - let messages = conversation.agent_visible_messages(); + let messages: Vec<_> = conversation + .agent_visible_messages() + .into_iter() + .filter(|message| !message.is_turn_context()) + .collect(); let total = messages.len(); if total == 0 { return None; diff --git a/crates/goose/src/agents/platform_extensions/orchestrator.rs b/crates/goose/src/agents/platform_extensions/orchestrator.rs index 1da47dda79..b8e0852d00 100644 --- a/crates/goose/src/agents/platform_extensions/orchestrator.rs +++ b/crates/goose/src/agents/platform_extensions/orchestrator.rs @@ -336,12 +336,13 @@ impl OrchestratorClient { ) -> Result { let provider = self.get_provider().await?; - let conversation_text = Conversation::new_unvalidated(messages.iter().cloned()) - .agent_visible_messages() - .iter() - .map(format_message_for_compacting) - .collect::>() - .join("\n"); + let conversation_text = agent_visible_session_messages(&Conversation::new_unvalidated( + messages.iter().cloned(), + )) + .iter() + .map(format_message_for_compacting) + .collect::>() + .join("\n"); let system = "You are a helpful assistant. Summarize the following conversation concisely, \ @@ -578,7 +579,11 @@ impl OrchestratorClient { } fn agent_visible_session_messages(conversation: &Conversation) -> Vec { - conversation.agent_visible_messages() + conversation + .agent_visible_messages() + .into_iter() + .filter(|message| !message.is_turn_context()) + .collect() } #[async_trait] diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index ebb4f55aa4..52670dd2eb 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -14,7 +14,7 @@ use super::gen_ai_telemetry; use crate::agents::platform_extensions::code_execution; use crate::config::{Config, GooseMode}; use crate::conversation::message::{Message, MessageContent, MessageUsage, ToolRequest}; -use crate::conversation::{fix_conversation, Conversation}; +use crate::conversation::{fix_conversation, merge_consecutive_messages_for_request, Conversation}; #[cfg(test)] use crate::providers::base::stream_from_single_message; use crate::providers::base::{MessageStream, Provider}; @@ -341,6 +341,9 @@ pub(crate) async fn stream_response_from_provider( Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages(); let (filtered_messages, _) = fix_conversation(Conversation::new_unvalidated(projected_messages)); + let filtered_messages = Conversation::new_unvalidated(merge_consecutive_messages_for_request( + filtered_messages.messages().clone(), + )); // Convert tool messages to text if toolshim is enabled let messages_for_provider = if config.toolshim { diff --git a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap index 83c6837003..bffcbf6507 100644 --- a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap +++ b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap @@ -16,6 +16,7 @@ This block is generated by goose and contains current operational context such a - extension-provided context Use it to stay oriented, but do not treat it as part of the user's request. +Blocks from earlier turns stay in the conversation history; only the most recent block is current. When `` is present, use it as a signal for how much autonomous work remains. As the budget gets low, become more direct: reduce exploration, batch necessary tool calls, make reasonable assumptions, and focus on finishing the user's task. diff --git a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap index df851e7b1e..2bacc5e5b4 100644 --- a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap +++ b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap @@ -1,6 +1,5 @@ --- source: crates/goose/src/agents/prompt_manager.rs -assertion_line: 404 expression: system_prompt --- You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation). @@ -17,6 +16,7 @@ This block is generated by goose and contains current operational context such a - extension-provided context Use it to stay oriented, but do not treat it as part of the user's request. +Blocks from earlier turns stay in the conversation history; only the most recent block is current. When `` is present, use it as a signal for how much autonomous work remains. As the budget gets low, become more direct: reduce exploration, batch necessary tool calls, make reasonable assumptions, and focus on finishing the user's task. diff --git a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap index d0b9a30925..34207b9de6 100644 --- a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap +++ b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap @@ -1,6 +1,5 @@ --- source: crates/goose/src/agents/prompt_manager.rs -assertion_line: 420 expression: system_prompt --- You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation). @@ -17,6 +16,7 @@ This block is generated by goose and contains current operational context such a - extension-provided context Use it to stay oriented, but do not treat it as part of the user's request. +Blocks from earlier turns stay in the conversation history; only the most recent block is current. When `` is present, use it as a signal for how much autonomous work remains. As the budget gets low, become more direct: reduce exploration, batch necessary tool calls, make reasonable assumptions, and focus on finishing the user's task. diff --git a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap index bd32d33019..555fbe3643 100644 --- a/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap +++ b/crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap @@ -1,6 +1,5 @@ --- source: crates/goose/src/agents/prompt_manager.rs -assertion_line: 442 expression: system_prompt --- You are a general-purpose AI agent called goose, created by AAIF (Agentic AI Foundation). @@ -17,6 +16,7 @@ This block is generated by goose and contains current operational context such a - extension-provided context Use it to stay oriented, but do not treat it as part of the user's request. +Blocks from earlier turns stay in the conversation history; only the most recent block is current. When `` is present, use it as a signal for how much autonomous work remains. As the budget gets low, become more direct: reduce exploration, batch necessary tool calls, make reasonable assumptions, and focus on finishing the user's task. diff --git a/crates/goose/src/agents/state_machine/ops_llm.rs b/crates/goose/src/agents/state_machine/ops_llm.rs index c693838b70..b1e2d29049 100644 --- a/crates/goose/src/agents/state_machine/ops_llm.rs +++ b/crates/goose/src/agents/state_machine/ops_llm.rs @@ -434,18 +434,32 @@ impl Inference for InferenceRunner<'_> { .get_context_limit(&self.model_config) .await .unwrap_or_else(|_| self.model_config.context_limit()); - let turn_start = messages_since_kickoff(conversation)? + let turn = messages_since_kickoff(conversation)?; + let turn_start = turn .first() .and_then(|message| chrono::DateTime::from_timestamp(message.created, 0)) .map(|timestamp| timestamp.with_timezone(&chrono::Local)) .unwrap_or_else(chrono::Local::now); - let conversation_for_provider = crate::agents::moim::inject_moim_parts( - Conversation::new_unvalidated(messages_for_provider), + // Persist a fresh turn-context event only when its bytes changed + // (e.g. the turn budget ticked); earlier events stay in place. + let last_turn_context = turn + .iter() + .rev() + .find(|message| message.is_turn_context()) + .map(Message::as_concat_text); + let turn_context = crate::agents::moim::turn_context_event( &session.working_dir, Some(context_limit), input.moim_parts, turn_start, - ); + ) + .filter(|event| Some(event.as_concat_text()) != last_turn_context); + if let Some(event) = &turn_context { + messages_for_provider.push(event.clone()); + } + let conversation_for_provider = Conversation::new_unvalidated(messages_for_provider); + let mut usage_effects: Vec = + turn_context.into_iter().map(StateEffect::from).collect(); let stream = crate::agents::reply_parts::stream_response_from_provider( self.provider.clone(), @@ -460,7 +474,10 @@ impl Inference for InferenceRunner<'_> { let mut stream = match stream { Ok(stream) => stream, - Err(err) => return applied(self.error_outcome(&err, emit).await), + Err(err) => { + usage_effects.extend(self.error_outcome(&err, emit).await); + return applied(usage_effects); + } }; let requested_model = self.model_config.model_name.clone(); @@ -477,7 +494,6 @@ impl Inference for InferenceRunner<'_> { }); let mut accumulator = Conversation::empty(); - let mut usage_effects = Vec::new(); let mut tool_request_ids = std::collections::HashSet::new(); loop { tokio::select! { @@ -540,7 +556,10 @@ impl Inference for InferenceRunner<'_> { return yielded_with(usage_effects); } - if usage_effects.is_empty() { + let has_recorded_usage = usage_effects + .iter() + .any(|effect| matches!(effect, StateEffect::RecordUsage(_))); + if !has_recorded_usage { let mut usage = ProviderUsage::new( self.model_config.model_name.clone(), goose_providers::conversation::token_usage::Usage::default(), diff --git a/crates/goose/src/agents/state_machine/ops_toolcalling.rs b/crates/goose/src/agents/state_machine/ops_toolcalling.rs index 175c94d071..519faf1861 100644 --- a/crates/goose/src/agents/state_machine/ops_toolcalling.rs +++ b/crates/goose/src/agents/state_machine/ops_toolcalling.rs @@ -648,6 +648,8 @@ impl Operation for ToolExecutionOperation<'_> { if extensions.is_empty() { return Ok(prompt_parts); } + // HashMap order shuffles across restarts and would bust the prompt cache. + extensions.sort_by(|a, b| a.name.cmp(&b.name)); let mut lines = vec![ "# Extensions".to_string(), diff --git a/crates/goose/src/agents/state_machine/tests/mod.rs b/crates/goose/src/agents/state_machine/tests/mod.rs index 753a939f98..f5f2ad7ee6 100644 --- a/crates/goose/src/agents/state_machine/tests/mod.rs +++ b/crates/goose/src/agents/state_machine/tests/mod.rs @@ -49,6 +49,35 @@ async fn max_turns_counts_inference_calls_and_injects_budget() -> Result<()> { Ok(()) } +#[tokio::test] +async fn turn_context_is_persisted_once_per_turn_and_reused_across_inferences() -> Result<()> { + let (pipeline, api) = test_pipeline().await?; + api.on("add one").call(ADD, value(1)); + api.on("result: 1").reply("The total is 1"); + api.on("hello").reply("hi there!"); + + let result = pipeline.run(["add one", "hello"]).await?; + + let conversation = result.conversation(); + let events: Vec<_> = conversation + .messages() + .iter() + .filter(|message| message.is_turn_context()) + .collect(); + assert_eq!( + events.len(), + 2, + "one turn-context event per turn; the turn's second inference reuses it" + ); + assert!(events.iter().all(|event| !event.is_user_visible())); + assert!(api + .calls() + .iter() + .all(|call| call.input_contains(""))); + + Ok(()) +} + #[tokio::test] async fn goal_starts_nudges_and_clears_when_met() -> Result<()> { let (pipeline, api) = test_pipeline().await?; diff --git a/crates/goose/src/agents/state_machine/tests/pipeline.rs b/crates/goose/src/agents/state_machine/tests/pipeline.rs index f661ca788e..34d8bb3f90 100644 --- a/crates/goose/src/agents/state_machine/tests/pipeline.rs +++ b/crates/goose/src/agents/state_machine/tests/pipeline.rs @@ -970,6 +970,7 @@ impl TestRun { .conversation() .messages() .iter() + .filter(|message| !message.is_turn_context()) .flat_map(|message| { message .content diff --git a/crates/goose/src/agents/state_machine/tests/recipe_scheduling_lifecycle.rs b/crates/goose/src/agents/state_machine/tests/recipe_scheduling_lifecycle.rs index ed10eef866..50747d8fd3 100644 --- a/crates/goose/src/agents/state_machine/tests/recipe_scheduling_lifecycle.rs +++ b/crates/goose/src/agents/state_machine/tests/recipe_scheduling_lifecycle.rs @@ -266,7 +266,12 @@ async fn scheduler_is_advertised_only_when_configured_and_manages_jobs() -> Resu assert_eq!(scheduled_sessions.len(), 1); assert_eq!( scheduled_sessions[0].1.message_count, - listed.conversation().messages().len() + listed + .conversation() + .messages() + .iter() + .filter(|message| message.is_user_visible()) + .count() ); api.on("list scheduled runs").call( MANAGE_SCHEDULE_TOOL_NAME_COMPLETE, diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index aa66359939..a451dcec27 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -100,9 +100,13 @@ pub async fn compact_messages( has_text && !has_tool_content }; - let (preserved_user_message, is_most_recent) = if !manual_compact { + // Turn-context events are agent-appended, never the message to preserve. + let (preserved_user_message, preserved_idx, is_most_recent) = if !manual_compact { let found_msg = messages.iter().enumerate().rev().find_map(|(idx, msg)| { - if !msg.is_agent_visible() || !matches!(msg.role, rmcp::model::Role::User) { + if !msg.is_agent_visible() + || msg.is_turn_context() + || !matches!(msg.role, rmcp::model::Role::User) + { return None; } @@ -123,13 +127,13 @@ pub async fn compact_messages( }); if let Some((idx, msg)) = found_msg { - let is_last = idx == messages.len() - 1; - (Some(msg), is_last) + let is_last = messages[idx + 1..].iter().all(Message::is_turn_context); + (Some(msg), Some(idx), is_last) } else { - (None, false) + (None, None, false) } } else { - (None, false) + (None, None, false) }; let messages_to_compact = messages.as_slice(); @@ -175,6 +179,25 @@ pub async fn compact_messages( final_messages.push(user_msg); } + // Carry the turn's own context event (it follows the preserved prompt) so + // a mid-turn retry keeps it; anything earlier belongs to a previous turn. + if let Some(carry_from) = preserved_idx.map(|idx| idx + 1) { + if let Some(turn_context) = messages_to_compact[carry_from..] + .iter() + .rev() + .find(|msg| msg.is_turn_context() && msg.is_agent_visible()) + { + let mut carried = turn_context.clone(); + carried.id = None; + // Storage reloads order by created_timestamp; the copy must keep + // its appended position, not resurface at the original event's time. + if let Some(latest) = final_messages.iter().map(|msg| msg.created).max() { + carried.created = carried.created.max(latest); + } + final_messages.push(carried); + } + } + let conversation = Conversation::new_unvalidated(final_messages); let retained_context_tokens = match count_context_tokens(&conversation).await { Ok(tokens) => tokens, @@ -317,8 +340,14 @@ async fn do_compact( session_id: &str, messages: &[Message], ) -> Result<(Message, ProviderUsage), anyhow::Error> { - let agent_visible_messages = - Conversation::new_unvalidated(messages.iter().cloned()).agent_visible_messages(); + // Keep stale per-turn state out of the summary. + let agent_visible_messages = Conversation::new_unvalidated( + messages + .iter() + .filter(|msg| !msg.is_turn_context()) + .cloned(), + ) + .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]; @@ -697,6 +726,7 @@ mod tests { message: Message, config: ModelConfig, max_tool_responses: Option, + captured_system: std::sync::Mutex>, } impl MockProvider { @@ -715,6 +745,7 @@ mod tests { request_headers: None, }, max_tool_responses: None, + captured_system: std::sync::Mutex::new(None), } } @@ -733,10 +764,11 @@ mod tests { async fn stream( &self, _model_config: &ModelConfig, - _system: &str, + system: &str, messages: &[Message], _tools: &[Tool], ) -> Result { + *self.captured_system.lock().unwrap() = Some(system.to_string()); // If max_tool_responses is set, fail if we have too many if let Some(max) = self.max_tool_responses { let tool_response_count = messages @@ -956,6 +988,198 @@ mod tests { assert!(!user_text.contains("assistant-only preprompt")); } + #[tokio::test] + async fn preserved_user_message_skips_turn_context_events() { + let conversation = Conversation::new_unvalidated([ + Message::user().with_text("earlier request"), + Message::assistant().with_text("earlier response"), + Message::user().with_text("the real current request"), + Message::user() + .with_text("frozen block") + .with_metadata(MessageMetadata::agent_only().with_turn_context()), + ]); + let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000); + + let compacted = compact_messages( + &provider, + &provider.config, + "test-session-id", + &conversation, + false, + ) + .await + .unwrap() + .conversation; + + let preserved: Vec<_> = compacted + .messages() + .iter() + .filter(|message| message.is_agent_visible() && message.role == Role::User) + .collect(); + let turn_context_events: Vec<_> = preserved + .iter() + .filter(|message| message.is_turn_context()) + .collect(); + let [carried_block] = turn_context_events.as_slice() else { + panic!("expected exactly one carried turn-context event"); + }; + assert!( + carried_block.as_concat_text().contains("frozen block"), + "the turn's context event must be carried forward for the mid-turn retry" + ); + let last = compacted.messages().last().unwrap(); + assert!( + last.is_turn_context() && last.is_agent_visible(), + "the carried event must trail the preserved user message" + ); + let user_prompt = preserved[preserved.len() - 2]; + assert!( + user_prompt + .as_concat_text() + .contains("the real current request"), + "the user's prompt must survive compaction verbatim" + ); + assert!(!user_prompt.is_turn_context()); + + let continuation = compacted + .messages() + .iter() + .find(|message| message.role == Role::Assistant && message.is_agent_visible()) + .unwrap() + .as_concat_text(); + assert!( + continuation.contains(CONVERSATION_CONTINUATION_TEXT), + "a trailing turn-context event must not demote the compaction to a tool-loop continuation" + ); + } + + #[tokio::test] + async fn stale_turn_context_from_an_earlier_turn_is_not_carried() { + let conversation = Conversation::new_unvalidated([ + Message::user().with_text("earlier request"), + Message::user() + .with_text("stale block") + .with_metadata(MessageMetadata::agent_only().with_turn_context()), + Message::assistant().with_text("earlier response"), + Message::user().with_text("the new prompt"), + ]); + let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000); + + let compacted = compact_messages( + &provider, + &provider.config, + "test-session-id", + &conversation, + false, + ) + .await + .unwrap() + .conversation; + + assert!( + !compacted + .messages() + .iter() + .any(|message| message.is_agent_visible() && message.is_turn_context()), + "pre-turn compaction must not resurrect a previous turn's context event" + ); + let last = compacted.messages().last().unwrap(); + assert_eq!(last.as_concat_text(), "the new prompt"); + } + + #[tokio::test] + async fn carried_turn_context_stays_last_after_persist_and_reload() { + let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000); + let mut prompt = Message::user().with_text("the real current request"); + prompt.created -= 3600; + let mut block = Message::user() + .with_text("frozen block") + .with_metadata(MessageMetadata::agent_only().with_turn_context()); + block.created -= 3600; + let conversation = Conversation::new_unvalidated([prompt, block]); + + let compacted = compact_messages( + &provider, + &provider.config, + "test-session-id", + &conversation, + false, + ) + .await + .unwrap() + .conversation; + assert!(compacted.messages().last().unwrap().is_turn_context()); + + let temp_dir = tempfile::tempdir().unwrap(); + let manager = crate::session::SessionManager::new(temp_dir.path().to_path_buf()); + let session = manager + .create_session( + std::path::PathBuf::from("/tmp/test"), + "carry order".to_string(), + crate::session::session_manager::SessionType::User, + crate::config::GooseMode::default(), + ) + .await + .unwrap(); + manager + .replace_conversation(&session.id, &compacted) + .await + .unwrap(); + + let reloaded = manager + .get_session(&session.id, true) + .await + .unwrap() + .conversation + .unwrap(); + assert_eq!(reloaded.messages().len(), compacted.messages().len()); + let last = reloaded.messages().last().unwrap(); + assert!( + last.is_turn_context(), + "the carried event must not resurface at its original timestamp on reload" + ); + } + + #[tokio::test] + async fn summarizer_input_excludes_turn_context_events() { + let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000); + let turn_context = |text: &str| { + Message::user() + .with_text(text) + .with_metadata(MessageMetadata::agent_only().with_turn_context()) + }; + let conversation = Conversation::new_unvalidated([ + Message::user().with_text("please refactor the parser"), + turn_context("cwd /old/dir"), + Message::assistant().with_text("working on it"), + turn_context("cwd /new/dir"), + ]); + + let compacted = compact_messages( + &provider, + &provider.config, + "test-session-id", + &conversation, + false, + ) + .await + .unwrap() + .conversation; + + let system = provider.captured_system.lock().unwrap().clone().unwrap(); + assert!(system.contains("please refactor the parser")); + assert!( + !system.contains("/old/dir") && !system.contains("/new/dir"), + "turn-context events must not reach the summarizer as dialogue" + ); + + let carried = compacted.messages().last().unwrap(); + assert!( + carried.is_turn_context() && carried.as_concat_text().contains("/new/dir"), + "the newest turn-context event must still be carried forward" + ); + } + #[tokio::test] async fn tool_pair_summary_projects_nested_audiences_before_provider_input() { let provider = MockProvider::new(Message::assistant().with_text("summary"), 1000); diff --git a/crates/goose/src/providers/litellm.rs b/crates/goose/src/providers/litellm.rs index e9caf70f57..73ed45005b 100644 --- a/crates/goose/src/providers/litellm.rs +++ b/crates/goose/src/providers/litellm.rs @@ -1,10 +1,11 @@ use anyhow::Result; use async_trait::async_trait; use futures::future::BoxFuture; +use goose_providers::cache_semantics::apply_chat_payload_breakpoints; use goose_providers::conversation::token_usage::ProviderUsage; use goose_providers::errors::ProviderError; use goose_providers::images::ImageFormat; -use serde_json::{json, Value}; +use serde_json::Value; use std::collections::HashMap; use super::api_client::{ApiClient, AuthMethod}; @@ -248,7 +249,7 @@ impl Provider for LiteLLMProvider { )?; if self.supports_cache_control(model_config).await { - payload = update_request_for_cache_control(&payload); + apply_chat_payload_breakpoints(&mut payload); } let response = self @@ -280,71 +281,6 @@ impl Provider for LiteLLMProvider { } } -/// Updates the request payload to include cache control headers for automatic prompt caching -/// Adds ephemeral cache control to the last 2 user messages, system message, and last tool -pub fn update_request_for_cache_control(original_payload: &Value) -> Value { - let mut payload = original_payload.clone(); - - if let Some(messages_spec) = payload - .as_object_mut() - .and_then(|obj| obj.get_mut("messages")) - .and_then(|messages| messages.as_array_mut()) - { - let mut user_count = 0; - for message in messages_spec.iter_mut().rev() { - if message.get("role") == Some(&json!("user")) { - if let Some(content) = message.get_mut("content") { - if let Some(content_str) = content.as_str() { - *content = json!([{ - "type": "text", - "text": content_str, - "cache_control": { "type": "ephemeral" } - }]); - } - } - user_count += 1; - if user_count >= 2 { - break; - } - } - } - - if let Some(system_message) = messages_spec - .iter_mut() - .find(|msg| msg.get("role") == Some(&json!("system"))) - { - if let Some(content) = system_message.get_mut("content") { - if let Some(content_str) = content.as_str() { - *system_message = json!({ - "role": "system", - "content": [{ - "type": "text", - "text": content_str, - "cache_control": { "type": "ephemeral" } - }] - }); - } - } - } - } - - if let Some(tools_spec) = payload - .as_object_mut() - .and_then(|obj| obj.get_mut("tools")) - .and_then(|tools| tools.as_array_mut()) - { - if let Some(last_tool) = tools_spec.last_mut() { - if let Some(function) = last_tool.get_mut("function") { - function - .as_object_mut() - .unwrap() - .insert("cache_control".to_string(), json!({ "type": "ephemeral" })); - } - } - } - payload -} - fn parse_custom_headers(headers_str: String) -> HashMap { let mut headers = HashMap::new(); for line in headers_str.lines() { diff --git a/crates/goose/src/providers/openrouter.rs b/crates/goose/src/providers/openrouter.rs index b3a679f423..01e1c9bc96 100644 --- a/crates/goose/src/providers/openrouter.rs +++ b/crates/goose/src/providers/openrouter.rs @@ -11,6 +11,7 @@ use super::openai_compatible::{handle_status, stream_openai_compat}; use super::retry::ProviderRetry; use crate::conversation::message::Message; use crate::providers::formats::openrouter as openrouter_format; +use goose_providers::cache_semantics::{apply_chat_payload_breakpoints, CacheSemantics}; use goose_providers::errors::ProviderError; use goose_providers::formats::openai::create_request; use goose_providers::model::ModelConfig; @@ -21,7 +22,6 @@ pub const OPENROUTER_PROVIDER_NAME: &str = "openrouter"; const OPENROUTER_PARAMETERS_CONFIG_KEY: &str = "OPENROUTER_PARAMETERS"; pub const OPENROUTER_DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4"; pub const OPENROUTER_DEFAULT_FAST_MODEL: &str = "google/gemini-2.5-flash"; -pub const OPENROUTER_MODEL_PREFIX_ANTHROPIC: &str = "anthropic"; // OpenRouter can run many models, we suggest the default pub const OPENROUTER_KNOWN_MODELS: &[&str] = &[ @@ -98,79 +98,6 @@ fn is_mandatory_reasoning_error(error: &ProviderError) -> bool { matches!(error, ProviderError::RequestFailed(message) if message.contains("Reasoning is mandatory")) } -/// Update the request when using anthropic model. -/// For anthropic model, we can enable prompt caching to save cost. Since openrouter is the OpenAI compatible -/// endpoint, we need to modify the open ai request to have anthropic cache control field. -fn update_request_for_anthropic(original_payload: &Value) -> Value { - let mut payload = original_payload.clone(); - - if let Some(messages_spec) = payload - .as_object_mut() - .and_then(|obj| obj.get_mut("messages")) - .and_then(|messages| messages.as_array_mut()) - { - // Add "cache_control" to the last and second-to-last "user" messages. - // During each turn, we mark the final message with cache_control so the conversation can be - // incrementally cached. The second-to-last user message is also marked for caching with the - // cache_control parameter, so that this checkpoint can read from the previous cache. - let mut user_count = 0; - for message in messages_spec.iter_mut().rev() { - if message.get("role") == Some(&json!("user")) { - if let Some(content) = message.get_mut("content") { - if let Some(content_str) = content.as_str() { - *content = json!([{ - "type": "text", - "text": content_str, - "cache_control": { "type": "ephemeral" } - }]); - } - } - user_count += 1; - if user_count >= 2 { - break; - } - } - } - - // Update the system message to have cache_control field. - if let Some(system_message) = messages_spec - .iter_mut() - .find(|msg| msg.get("role") == Some(&json!("system"))) - { - if let Some(content) = system_message.get_mut("content") { - if let Some(content_str) = content.as_str() { - *system_message = json!({ - "role": "system", - "content": [{ - "type": "text", - "text": content_str, - "cache_control": { "type": "ephemeral" } - }] - }); - } - } - } - } - - if let Some(tools_spec) = payload - .as_object_mut() - .and_then(|obj| obj.get_mut("tools")) - .and_then(|tools| tools.as_array_mut()) - { - // Add "cache_control" to the last tool spec, if any. This means that all tool definitions, - // will be cached as a single prefix. - if let Some(last_tool) = tools_spec.last_mut() { - if let Some(function) = last_tool.get_mut("function") { - function - .as_object_mut() - .unwrap() - .insert("cache_control".to_string(), json!({ "type": "ephemeral" })); - } - } - } - payload -} - fn is_gemini_model(model_name: &str) -> bool { model_name.starts_with("google/") } @@ -350,8 +277,10 @@ impl Provider for OpenRouterProvider { } } - if supports_cache_control(model_config) { - payload = update_request_for_anthropic(&payload); + if CacheSemantics::for_model(OPENROUTER_PROVIDER_NAME, &model_config.model_name) + .uses_explicit_breakpoints() + { + apply_chat_payload_breakpoints(&mut payload); } if is_gemini_model(&model_config.model_name) { @@ -386,12 +315,6 @@ impl Provider for OpenRouterProvider { } } -fn supports_cache_control(model: &ModelConfig) -> bool { - model - .model_name - .starts_with(OPENROUTER_MODEL_PREFIX_ANTHROPIC) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose/src/scheduler.rs b/crates/goose/src/scheduler.rs index 23914b98c7..ed9dad2f3c 100644 --- a/crates/goose/src/scheduler.rs +++ b/crates/goose/src/scheduler.rs @@ -1493,7 +1493,13 @@ mod tests { .unwrap(); assert_eq!( sessions[0].message_count, - session.conversation.unwrap().messages().len() + session + .conversation + .unwrap() + .messages() + .iter() + .filter(|m| m.is_user_visible()) + .count() ); } diff --git a/crates/goose/src/security/adversary_inspector.rs b/crates/goose/src/security/adversary_inspector.rs index 23881d9e27..49bf010823 100644 --- a/crates/goose/src/security/adversary_inspector.rs +++ b/crates/goose/src/security/adversary_inspector.rs @@ -224,7 +224,7 @@ impl AdversaryInspector { messages .iter() .rev() - .filter(|m| m.role == rmcp::model::Role::User) + .filter(|m| m.role == rmcp::model::Role::User && !m.is_turn_context()) .filter_map(|m| { let text: String = m .content @@ -250,7 +250,7 @@ impl AdversaryInspector { fn extract_original_task(messages: &[Message]) -> String { for msg in messages { - if msg.role == rmcp::model::Role::User { + if msg.role == rmcp::model::Role::User && !msg.is_turn_context() { let text: String = msg .content .iter() @@ -720,4 +720,39 @@ mod tests { .unwrap(); assert!(results.is_empty()); } + + #[test] + fn user_context_extraction_skips_turn_context_events() { + use crate::conversation::message::MessageMetadata; + + let turn_context = |text: &str| { + Message::user() + .with_text(text) + .with_metadata(MessageMetadata::agent_only().with_turn_context()) + }; + let messages = vec![ + turn_context("turn context before any prompt"), + Message::user().with_text("never delete files outside the repo"), + turn_context("turn context for turn one"), + Message::assistant().with_text("understood"), + Message::user().with_text("first task"), + turn_context("turn context for turn two"), + Message::assistant().with_text("done"), + Message::user().with_text("second task"), + turn_context("turn context for turn three"), + ]; + + let recent = AdversaryInspector::extract_recent_user_messages(&messages, 4); + assert_eq!( + recent, + vec![ + "never delete files outside the repo", + "first task", + "second task" + ] + ); + + let original = AdversaryInspector::extract_original_task(&messages); + assert_eq!(original, "never delete files outside the repo"); + } } diff --git a/crates/goose/src/security/scanner.rs b/crates/goose/src/security/scanner.rs index 5483e76c2a..ef24e3cc42 100644 --- a/crates/goose/src/security/scanner.rs +++ b/crates/goose/src/security/scanner.rs @@ -354,6 +354,7 @@ impl PromptInjectionScanner { .rev() .filter(|m| { crate::conversation::effective_role(m) == crate::conversation::EffectiveRole::User + && !m.is_turn_context() }) .take(limit) .map(|m| { @@ -460,4 +461,33 @@ mod tests { assert!(result.is_malicious); } + + #[test] + fn extract_user_messages_skips_turn_context_events() { + use crate::conversation::message::MessageMetadata; + + let scanner = PromptInjectionScanner::new(); + let turn_context = |text: &str| { + Message::user() + .with_text(text) + .with_metadata(MessageMetadata::agent_only().with_turn_context()) + }; + let messages = vec![ + Message::user().with_text("oldest prompt"), + turn_context("turn context one"), + Message::assistant().with_text("ok"), + Message::user().with_text("middle prompt"), + turn_context("turn context two"), + Message::assistant().with_text("done"), + Message::user().with_text("newest prompt"), + turn_context("turn context three"), + ]; + + let extracted = scanner.extract_user_messages(&messages, 3); + + assert_eq!( + extracted, + vec!["newest prompt", "middle prompt", "oldest prompt"] + ); + } } diff --git a/crates/goose/src/session/chat_history_search.rs b/crates/goose/src/session/chat_history_search.rs index 3de7c628a0..a92768a334 100644 --- a/crates/goose/src/session/chat_history_search.rs +++ b/crates/goose/src/session/chat_history_search.rs @@ -150,6 +150,13 @@ impl<'a> ChatHistorySearch<'a> { END, 1 ) = 1 + AND COALESCE( + CASE + WHEN json_valid(m.metadata_json) + THEN json_extract(m.metadata_json, '$.turnContext') + END, + 0 + ) = 0 AND EXISTS ( SELECT 1 FROM json_each(m.content_json) AS content WHERE json_extract(content.value, '$.type') = 'text' @@ -263,12 +270,29 @@ impl<'a> ChatHistorySearch<'a> { ) -> Result> { let mut session_totals: HashMap = HashMap::new(); for session_id in session_messages.keys() { - let count: i64 = - sqlx::query_scalar("SELECT COUNT(*) FROM messages WHERE session_id = ?") - .bind(session_id) - .fetch_one(self.pool) - .await - .unwrap_or(0); + let count: i64 = sqlx::query_scalar( + r#" + SELECT COUNT(*) FROM messages m WHERE m.session_id = ? + AND COALESCE( + CASE + WHEN json_valid(m.metadata_json) + THEN json_extract(m.metadata_json, '$.agentVisible') + END, + 1 + ) = 1 + AND COALESCE( + CASE + WHEN json_valid(m.metadata_json) + THEN json_extract(m.metadata_json, '$.turnContext') + END, + 0 + ) = 0 + "#, + ) + .bind(session_id) + .fetch_one(self.pool) + .await + .unwrap_or(0); session_totals.insert(session_id.clone(), count as usize); } Ok(session_totals) @@ -443,5 +467,26 @@ mod tests { .await .unwrap(); assert_eq!(hidden_only.total_matches, 0); + + insert_message( + &pool, + &Message::user() + .with_text("needle in operational context") + .with_metadata(MessageMetadata::agent_only().with_turn_context()), + now, + ) + .await; + let needle = ChatHistorySearch::new(&pool, "needle", Some(10), None, None, None, vec![]) + .execute() + .await + .unwrap(); + assert_eq!( + needle.total_matches, 1, + "turn-context events are operational context, not searchable conversation" + ); + assert_eq!( + needle.results[0].total_messages_in_session, 2, + "session totals must not count turn-context or user-only rows" + ); } } diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 6cf6f285eb..77cf36e3cc 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -365,12 +365,14 @@ fn message_keyword_clause(keyword_count: usize) -> String { .collect::>() .join(" OR "); + let visible = user_visible_message_sql("mq.metadata_json"); format!( r#" EXISTS ( SELECT 1 FROM messages mq WHERE mq.session_id = s.id + AND {visible} AND EXISTS ( SELECT 1 FROM json_each(mq.content_json) @@ -619,7 +621,7 @@ impl SessionManager { let user_message_count = conversation .messages() .iter() - .filter(|m| matches!(m.role, Role::User)) + .filter(|m| matches!(m.role, Role::User) && m.is_user_visible()) .count(); let should_generate_name = if provider.manages_own_context() { @@ -714,6 +716,10 @@ fn normalized_message_timestamp_sql(column: &str) -> String { ) } +fn user_visible_message_sql(column: &str) -> String { + format!("COALESCE(json_extract({column}, '$.userVisible'), 1) != 0") +} + fn session_sort_at(session: &Session) -> DateTime { session.last_message_at.unwrap_or(session.updated_at) } @@ -1629,7 +1635,11 @@ impl SessionStorage { if include_messages { let conv = self.get_conversation(&session.id).await?; - session.message_count = conv.messages().len(); + session.message_count = conv + .messages() + .iter() + .filter(|m| m.is_user_visible()) + .count(); session.last_message_at = conv .messages() .iter() @@ -1638,7 +1648,8 @@ impl SessionStorage { session.conversation = Some(conv); } else { let sql = format!( - "SELECT COUNT(*), MAX({}) FROM messages WHERE session_id = ?", + "SELECT COUNT(*) FILTER (WHERE {}), MAX({}) FROM messages WHERE session_id = ?", + user_visible_message_sql("metadata_json"), normalized_message_timestamp_sql("created_timestamp") ); let (count, last_message_timestamp): (i64, Option) = sqlx::query_as(&sql) @@ -1979,7 +1990,7 @@ impl SessionStorage { s.schedule_id, s.recipe_json, s.user_recipe_values_json, s.provider_name, s.model_config_json, s.goose_mode, s.archived_at, s.project_id, s.parent_session_id, - COUNT(m.id) as message_count, + COUNT(m.id) FILTER (WHERE {}) as message_count, MAX({}) as last_message_timestamp, {} as sort_timestamp FROM sessions s @@ -1990,6 +2001,7 @@ impl SessionStorage { {} {} "#, + user_visible_message_sql("m.metadata_json"), normalized_message_timestamp, sort_timestamp_sql, message_join, @@ -2985,6 +2997,44 @@ mod tests { assert_eq!(with_messages.last_message_at, Some(expected)); } + #[tokio::test] + async fn test_message_count_excludes_agent_only_messages() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + let session = sm + .create_session( + PathBuf::from("/tmp/test"), + "Hidden events".to_string(), + SessionType::User, + GooseMode::default(), + ) + .await + .unwrap(); + + add_user_message(&sm, &session.id).await; + sm.add_message( + &session.id, + &Message::user() + .with_text("frozen") + .with_visibility(false, true), + ) + .await + .unwrap(); + sm.add_message(&session.id, &Message::assistant().with_text("hi")) + .await + .unwrap(); + + let counted = sm.get_session(&session.id, false).await.unwrap(); + assert_eq!(counted.message_count, 2); + + let with_messages = sm.get_session(&session.id, true).await.unwrap(); + assert_eq!(with_messages.message_count, 2); + + let listed = sm.list_sessions().await.unwrap(); + let listed_session = listed.iter().find(|s| s.id == session.id).unwrap(); + assert_eq!(listed_session.message_count, 2); + } + #[tokio::test] async fn test_truncate_conversation_from_message_keeps_same_second_previous_rows() { let temp_dir = TempDir::new().unwrap(); diff --git a/crates/goose/src/session/session_naming.rs b/crates/goose/src/session/session_naming.rs index f5130bb5de..4e0f34d3e1 100644 --- a/crates/goose/src/session/session_naming.rs +++ b/crates/goose/src/session/session_naming.rs @@ -75,7 +75,7 @@ fn extract_short_title(text: &str) -> String { fn get_initial_user_messages(messages: &Conversation) -> Vec { messages .iter() - .filter(|m| m.role == rmcp::model::Role::User) + .filter(|m| m.role == rmcp::model::Role::User && m.is_user_visible()) .take(MSG_COUNT_FOR_SESSION_NAME_GENERATION) .map(|m| { m.content diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index 0ab078db8f..5f52e9e3aa 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -1404,7 +1404,11 @@ mod tests { .messages() .to_vec(); - let user_count = messages.iter().filter(|m| m.role == Role::User).count(); + // Turn-context events are excluded; this test is about streaming deltas. + let user_count = messages + .iter() + .filter(|m| m.role == Role::User && m.is_user_visible()) + .count(); let asst_count = messages .iter() .filter(|m| m.role == Role::Assistant) @@ -1453,7 +1457,10 @@ mod tests { .messages() .to_vec(); - let user_count2 = messages2.iter().filter(|m| m.role == Role::User).count(); + let user_count2 = messages2 + .iter() + .filter(|m| m.role == Role::User && m.is_user_visible()) + .count(); let asst_count2 = messages2 .iter() .filter(|m| m.role == Role::Assistant) diff --git a/crates/goose/tests/compaction.rs b/crates/goose/tests/compaction.rs index 5f92de7dcc..4d1b9132dc 100644 --- a/crates/goose/tests/compaction.rs +++ b/crates/goose/tests/compaction.rs @@ -320,7 +320,12 @@ fn assert_conversation_compacted(conversation: &Conversation) { "Message after compaction at index {} should be agent visible", idx ); - if idx == continuation_end && matches!(msg.role, rmcp::model::Role::User) { + if msg.is_turn_context() { + assert!( + !msg.is_user_visible(), + "Carried turn-context event should be user-invisible" + ); + } else if idx == continuation_end && matches!(msg.role, rmcp::model::Role::User) { assert!( !msg.is_user_visible(), "Projected preserved user message should be user-invisible"