From d43c692127aceeb19fd8b0bff278da52fdeccc61 Mon Sep 17 00:00:00 2001 From: filip <44206832+filipkujawa@users.noreply.github.com> Date: Mon, 29 Jun 2026 22:09:37 -0700 Subject: [PATCH] perf(providers): keep turn-context out of the Anthropic prefix cache (#10030) --- crates/goose-providers/src/conversation.rs | 10 + .../goose-providers/src/formats/anthropic.rs | 346 +++++++++++++++++- crates/goose/src/agents/moim.rs | 92 ++++- ...nager__tests__all_platform_extensions.snap | 2 +- ..._agents__prompt_manager__tests__basic.snap | 2 +- ..._prompt_manager__tests__one_extension.snap | 2 +- ..._prompt_manager__tests__typical_setup.snap | 2 +- 7 files changed, 431 insertions(+), 25 deletions(-) diff --git a/crates/goose-providers/src/conversation.rs b/crates/goose-providers/src/conversation.rs index 4f8c92197..457a6b329 100644 --- a/crates/goose-providers/src/conversation.rs +++ b/crates/goose-providers/src/conversation.rs @@ -505,6 +505,16 @@ fn has_tool_response(message: &Message) -> bool { .any(|content| matches!(content, MessageContent::ToolResponse(_))) } +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!("")) +} + pub fn effective_role(message: &Message) -> String { if message.role == Role::User && has_tool_response(message) { "tool".to_string() diff --git a/crates/goose-providers/src/formats/anthropic.rs b/crates/goose-providers/src/formats/anthropic.rs index 163c7f318..acdce55cf 100644 --- a/crates/goose-providers/src/formats/anthropic.rs +++ b/crates/goose-providers/src/formats/anthropic.rs @@ -310,7 +310,6 @@ fn format_messages_with_options( } } - // If no messages, add a default one if anthropic_messages.is_empty() { anthropic_messages.push(json!({ ROLE_FIELD: USER_ROLE, @@ -321,23 +320,33 @@ fn format_messages_with_options( })); } - // 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. + // 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); + let mut user_count = 0; for message in anthropic_messages.iter_mut().rev() { - if message.get(ROLE_FIELD) == Some(&json!(USER_ROLE)) { - if let Some(content) = message.get_mut(CONTENT_FIELD) { - if let Some(content_array) = content.as_array_mut() { - if let Some(last_content) = content_array.last_mut() { - last_content.as_object_mut().unwrap().insert( - CACHE_CONTROL_FIELD.to_string(), - json!({ TYPE_FIELD: "ephemeral" }), - ); - } - } - } + if message.get(ROLE_FIELD) != Some(&json!(USER_ROLE)) { + continue; + } + let Some(content_array) = 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(|b| b.as_object_mut()) + { + block.insert( + CACHE_CONTROL_FIELD.to_string(), + json!({ TYPE_FIELD: "ephemeral" }), + ); user_count += 1; if user_count >= 2 { break; @@ -348,6 +357,52 @@ 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!({ @@ -2179,4 +2234,263 @@ mod tests { assert_eq!(parts.tool_calls, vec!["write"]); 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 { + 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( + "read_file", + "Read a file from disk", + object!({ + "type": "object", + "properties": { "path": { "type": "string" } } + }), + ), + Tool::new( + "write_file", + "Write a file to disk", + object!({ + "type": "object", + "properties": { "path": { "type": "string" }, "content": { "type": "string" } } + }), + ), + ] + } + + /// 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![ + Message::user().with_text("What does the main entrypoint do?"), + 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::Content::text( + "fn main() { run(); }", + )])), + ), + 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::Content::text( + "fn main() { run(); }", + )])), + ), + ] + } + + fn request_with(messages: &[Message]) -> Value { + create_request_with_default_options( + &cfg("claude-sonnet-4-5"), + "You are a careful coding assistant.", + messages, + &sample_tools(), + ) + .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" + ); + + 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" + ); + } + + #[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(); + + 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:?}" + ); + } + } } diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index a6b8407a3..c6ef1d99f 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -1,14 +1,16 @@ use crate::agents::extension_manager::ExtensionManager; use crate::conversation::message::MessageContent; -use crate::conversation::{effective_role, fix_conversation, Conversation}; +use crate::conversation::{ + effective_role, fix_conversation, Conversation, CURRENT_TIME_TAG, TURN_CONTEXT_TAG, + WORKING_DIRECTORY_TAG, +}; use std::path::{Path, PathBuf}; const MIN_CONTEXT_FOR_MOIM: usize = 32_000; -const TURN_CONTEXT_TAG: &str = "turn-context"; const SYSTEM_PROMPT_BLOCK_TEMPLATE: &str = r#"# Turn Context -Each turn may include a `<{turn_context_tag}>` block prepended to the latest user message. +Each turn may include a `<{turn_context_tag}>` block added to the request. This block is generated by goose and contains current operational context such as: - current time - working directory @@ -143,8 +145,8 @@ fn compose_moim( let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:00"); let mut lines = vec![ open_tag(TURN_CONTEXT_TAG), - tag("current-time", ×tamp.to_string()), - tag("working-directory", &working_dir.display().to_string()), + tag(CURRENT_TIME_TAG, ×tamp.to_string()), + tag(WORKING_DIRECTORY_TAG, &working_dir.display().to_string()), ]; if let Some(value) = @@ -356,4 +358,84 @@ mod tests { )); assert_eq!(msgs[2].content.len(), 1); } + + /// 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; + + fn moim( + total_tokens: Option, + context_limit: Option, + turns_taken: u32, + max_turns: u32, + extension_parts: Vec, + ) -> String { + compose_moim( + Path::new("/Users/me/code/goose"), + total_tokens, + context_limit, + 0.8, + turns_taken, + max_turns, + extension_parts, + ) + } + + #[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" + ); + } + } } 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 7f3621190..93fe2496a 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 @@ -8,7 +8,7 @@ goose is being developed as an open-source software project. # Turn Context -Each turn may include a `` block prepended to the latest user message. +Each turn may include a `` block added to the request. This block is generated by goose and contains current operational context such as: - current time - working directory 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 6b3dca14b..df851e7b1 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 @@ -8,7 +8,7 @@ goose is being developed as an open-source software project. # Turn Context -Each turn may include a `` block prepended to the latest user message. +Each turn may include a `` block added to the request. This block is generated by goose and contains current operational context such as: - current time - working directory 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 379772033..d0b9a3092 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 @@ -8,7 +8,7 @@ goose is being developed as an open-source software project. # Turn Context -Each turn may include a `` block prepended to the latest user message. +Each turn may include a `` block added to the request. This block is generated by goose and contains current operational context such as: - current time - working directory 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 fba2cb26e..bd32d3301 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 @@ -8,7 +8,7 @@ goose is being developed as an open-source software project. # Turn Context -Each turn may include a `` block prepended to the latest user message. +Each turn may include a `` block added to the request. This block is generated by goose and contains current operational context such as: - current time - working directory