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

perf(providers): keep turn-context out of the OpenAI-compatible prefix cache (#10937)

Browse Source
This commit is contained in:
filip
2026-08-04 14:09:45 -07:00
committed by GitHub
parent 06f218cd14
commit affd1cea10
2 changed files with 212 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-provider-types/src/formats/openai.rs
+205
View File
@@ -209,6 +209,11 @@ pub fn format_messages_with_options(
image_format: &ImageFormat,
options: OpenAiFormatOptions,
) -> Vec<Value> {
let extracted = 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.
@@ -518,9 +523,60 @@ 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<Message>, 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<Value>, 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.
///
@@ -4879,4 +4935,153 @@ data: [DONE]"#;
json!("<think>\nreasoning\n</think>")
);
}
mod cache_prefix_stability {
use super::*;
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 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![ContentBlock::text(
"fn main() { run(); }",
)])),
),
]
}
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() {
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
.iter()
.filter(|m| {
m["content"]
.as_str()
.is_some_and(|c| c.contains("<turn-context>"))
})
.count();
assert_eq!(occurrences, 1, "turn-context must appear exactly once");
assert_eq!(
spec[2]["content"],
json!("Now add error handling to it."),
"the source user message must keep its genuine text untouched"
);
}
#[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:?}"
);
}
}
}
}
crates/goose/tests/acp_common_tests/mod.rs
+7 -7
View File
@@ -22,7 +22,7 @@ use std::sync::Arc;
use std::time::Duration;
const SHELL_TEST_CONTENT: &str = "test-shell-content-98765";
const TURN_CONTEXT_CLOSE: &str = r#"</turn-context>\n"#;
const TURN_CONTEXT_OPEN: &str = r#"\n<turn-context>"#;
const OPENAI_SESSION_NAME_RESPONSE: &str = r#"data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1766229303,"model":"gpt-5-nano","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1766229303,"model":"gpt-5-nano","choices":[{"index":0,"delta":{"content":"Generated Test Title"},"finish_reason":null}]}
@@ -42,7 +42,7 @@ async fn new_basic_session<C: Connection>(config: TestConnectionConfig) -> Basic
let expected_session_id = C::expected_session_id();
let openai = OpenAiFixture::new(
vec![(
format!("{TURN_CONTEXT_CLOSE}what is 1+1"),
format!("what is 1+1{TURN_CONTEXT_OPEN}"),
include_str!("../acp_test_data/openai_basic.txt"),
)],
expected_session_id.clone(),
@@ -106,7 +106,7 @@ pub async fn run_session_name_update_notification<C: Connection>() {
let openai = OpenAiFixture::new(
vec![
(
format!("{TURN_CONTEXT_CLOSE}what should we call this conversation?"),
format!("what should we call this conversation?{TURN_CONTEXT_OPEN}"),
include_str!("../acp_test_data/openai_basic.txt"),
),
(
@@ -1147,7 +1147,7 @@ pub async fn run_prompt_basic<C: Connection>() {
let expected_session_id = C::expected_session_id();
let openai = OpenAiFixture::new(
vec![(
format!("{TURN_CONTEXT_CLOSE}what is 1+1"),
format!("what is 1+1{TURN_CONTEXT_OPEN}"),
include_str!("../acp_test_data/openai_basic.txt"),
)],
expected_session_id.clone(),
@@ -1198,7 +1198,7 @@ pub async fn run_prompt_codemode<C: Connection>() {
let openai = OpenAiFixture::new(
vec![
(
format!("{TURN_CONTEXT_CLOSE}{prompt}"),
format!("{prompt}{TURN_CONTEXT_OPEN}"),
include_str!("../acp_test_data/openai_builtin_search.txt"),
),
(
@@ -1246,7 +1246,7 @@ pub async fn run_prompt_image<C: Connection>() {
vec![
(
format!(
"{TURN_CONTEXT_CLOSE}Use the get_image tool and describe what you see in its result."
"Use the get_image tool and describe what you see in its result.{TURN_CONTEXT_OPEN}"
),
include_str!("../acp_test_data/openai_image_tool_call.txt"),
),
@@ -1322,7 +1322,7 @@ pub async fn run_prompt_mcp<C: Connection>() {
let openai = OpenAiFixture::new(
vec![
(
format!("{TURN_CONTEXT_CLOSE}Use the get_code tool and output only its result."),
format!("Use the get_code tool and output only its result.{TURN_CONTEXT_OPEN}"),
include_str!("../acp_test_data/openai_tool_call.txt"),
),
(