tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

perf(providers): keep turn-context out of the Anthropic prefix cache (#10030)

Browse Source
This commit is contained in:
filip
2026-06-29 22:09:37 -07:00
committed by GitHub
parent 37e62497d4
commit d43c692127
7 changed files with 431 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-providers/src/conversation.rs
+10
View File
@@ -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!("</{CURRENT_TIME_TAG}>\n<{WORKING_DIRECTORY_TAG}>"))
&& text.trim_end().ends_with(&format!("</{TURN_CONTEXT_TAG}>"))
}
pub fn effective_role(message: &Message) -> String {
if message.role == Role::User && has_tool_response(message) {
"tool".to_string()
crates/goose-providers/src/formats/anthropic.rs
+330 -16
View File
@@ -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<usize> {
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<JsonObject>) -> Arc<JsonObject> {
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!(
"<turn-context>\n\
<current-time>{time}</current-time>\n\
<working-directory>/Users/me/code/goose</working-directory>\n\
<turn-budget>{turn_budget}</turn-budget>\n\
</turn-context>"
)
}
fn sample_tools() -> Vec<Tool> {
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<Message> {
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<Value> = 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<Message> {
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:?}"
);
}
}
}
crates/goose/src/agents/moim.rs
+87 -5
View File
@@ -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", &timestamp.to_string()),
tag("working-directory", &working_dir.display().to_string()),
tag(CURRENT_TIME_TAG, &timestamp.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<i32>,
context_limit: Option<usize>,
turns_taken: u32,
max_turns: u32,
extension_parts: Vec<String>,
) -> 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!["<guardrail>stay on task</guardrail>".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(
"<turn-context>\n\
<current-time>2026-06-25 12:00:00</current-time>\n\
<session-id>abc</session-id>\n\
</turn-context>"
),
"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"
);
}
}
}
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__all_platform_extensions.snap
+1 -1
View File
@@ -8,7 +8,7 @@ goose is being developed as an open-source software project.
# Turn Context
Each turn may include a `<turn-context>` block prepended to the latest user message.
Each turn may include a `<turn-context>` block added to the request.
This block is generated by goose and contains current operational context such as:
- current time
- working directory
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__basic.snap
+1 -1
View File
@@ -8,7 +8,7 @@ goose is being developed as an open-source software project.
# Turn Context
Each turn may include a `<turn-context>` block prepended to the latest user message.
Each turn may include a `<turn-context>` block added to the request.
This block is generated by goose and contains current operational context such as:
- current time
- working directory
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__one_extension.snap
+1 -1
View File
@@ -8,7 +8,7 @@ goose is being developed as an open-source software project.
# Turn Context
Each turn may include a `<turn-context>` block prepended to the latest user message.
Each turn may include a `<turn-context>` block added to the request.
This block is generated by goose and contains current operational context such as:
- current time
- working directory
crates/goose/src/agents/snapshots/goose__agents__prompt_manager__tests__typical_setup.snap
+1 -1
View File
@@ -8,7 +8,7 @@ goose is being developed as an open-source software project.
# Turn Context
Each turn may include a `<turn-context>` block prepended to the latest user message.
Each turn may include a `<turn-context>` block added to the request.
This block is generated by goose and contains current operational context such as:
- current time
- working directory