Merge branch 'fix/proper-round-pairs'
CI / All (ubuntu-latest) (push) Failing after 30s
CI / All (macos-latest) (push) Canceled after 0s
CI / All (windows-latest) (push) Canceled after 0s

This commit is contained in:
Dark-Alex-17 committed 2026-09-17 16:33:31 -06:00
commit 56ed4664cd
6 files changed
+297 -15

No files matched your search

+100 -6
View File
@@ -423,6 +423,7 @@ fn build_chat_completions_body(
tool_results, tool_results,
text, text,
sequence, sequence,
round_starts,
}) => { }) => {
if !sequence { if !sequence {
let mut assistant_parts = vec![]; let mut assistant_parts = vec![];
@@ -472,11 +473,15 @@ fn build_chat_completions_body(
} }
} else { } else {
// One pair per round: Claude can reuse tool_use IDs across API calls. // One pair per round: Claude can reuse tool_use IDs across API calls.
// A round boundary is detected by the presence of round text, but // Round boundaries are recorded in `round_starts` at merge time.
// rounds where the model emitted only tool calls (no narration) // Splitting per round also keeps the serialized history append-only
// carry no text marker. As a backstop, also split whenever a // across tool-loop iterations, so prompt caching keeps hitting on
// toolUseId would repeat within the current assistant message. // earlier rounds instead of rewriting them every request.
// Converse rejects duplicate toolUseIds in a single message. // Sessions recorded before round tracking carry no `round_starts`;
// for those, fall back to the presence of round text, plus an ID
// backstop: split whenever a toolUseId would repeat within the
// current assistant message. Converse rejects duplicate toolUseIds
// in a single message.
let mut messages = vec![]; let mut messages = vec![];
let mut assistant_parts: Vec<Value> = vec![]; let mut assistant_parts: Vec<Value> = vec![];
let mut user_parts: Vec<Value> = vec![]; let mut user_parts: Vec<Value> = vec![];
@@ -487,7 +492,11 @@ fn build_chat_completions_body(
.id .id
.as_deref() .as_deref()
.is_some_and(|id| chunk_ids.contains(id)); .is_some_and(|id| chunk_ids.contains(id));
if index > 0 && (tool_result.text.is_some() || id_collision) { if index > 0
&& (round_starts.contains(&index)
|| tool_result.text.is_some()
|| id_collision)
{
messages.push( messages.push(
json!({ "role": "assistant", "content": assistant_parts }), json!({ "role": "assistant", "content": assistant_parts }),
); );
@@ -896,6 +905,7 @@ mod tests {
tool_results, tool_results,
text: text.to_string(), text: text.to_string(),
sequence, sequence,
round_starts: vec![],
}), }),
), ),
], ],
@@ -912,6 +922,29 @@ mod tests {
build_body_with(tool_results, "", true) build_body_with(tool_results, "", true)
} }
fn build_body_rounds(tool_results: Vec<ToolResult>, round_starts: Vec<usize>) -> Value {
let data = ChatCompletionsData {
messages: vec![
Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
Message::new(
MessageRole::Assistant,
MessageContent::ToolCalls(MessageContentToolCalls {
tool_results,
text: String::new(),
sequence: true,
round_starts,
}),
),
],
temperature: None,
top_p: None,
reasoning_effort: None,
functions: None,
stream: false,
};
build_chat_completions_body(data, &Model::new("bedrock", "test"), false).unwrap()
}
fn tool_result_content_block(output: Value) -> Value { fn tool_result_content_block(output: Value) -> Value {
let body = build_body(vec![tool_result(output)]); let body = build_body(vec![tool_result(output)]);
body["messages"][2]["content"][0]["toolResult"]["content"][0].clone() body["messages"][2]["content"][0]["toolResult"]["content"][0].clone()
@@ -1067,6 +1100,9 @@ mod tests {
assert_unique_tool_use_ids_per_message(&body); assert_unique_tool_use_ids_per_message(&body);
} }
/// Legacy sessions recorded before round tracking carry no
/// `round_starts`; without a text marker or an ID collision their
/// textless rounds still serialize merged (the pre-round-tracking shape).
#[test] #[test]
fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() { fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() {
let body = build_body(vec![ let body = build_body(vec![
@@ -1081,6 +1117,64 @@ mod tests {
assert_unique_tool_use_ids_per_message(&body); assert_unique_tool_use_ids_per_message(&body);
} }
/// Regression: textless tool-loop rounds must split into one
/// assistant/user pair per recorded round so the serialized history stays
/// append-only across iterations and prompt caching keeps hitting on
/// earlier rounds (mirrors the claude client's round_starts handling).
#[test]
fn sequence_splits_on_round_starts_without_text() {
let body = build_body_rounds(
vec![
sequence_tool_result("toolu_A", None),
sequence_tool_result("toolu_B", None),
sequence_tool_result("toolu_C", None),
],
vec![1, 2],
);
let messages = body["messages"].as_array().unwrap();
// user + 3 × (assistant, user) — one pair per round.
assert_eq!(messages.len(), 7, "body: {body}");
for pair in messages[1..].chunks(2) {
assert_eq!(pair[0]["role"], "assistant", "body: {body}");
assert_eq!(pair[0]["content"].as_array().unwrap().len(), 1);
assert_eq!(pair[1]["role"], "user", "body: {body}");
}
assert_unique_tool_use_ids_per_message(&body);
}
/// The serialized history must be append-only across tool-loop
/// iterations: request N+1's messages must start with exactly request N's
/// messages, so the prompt-cache prefix keeps matching.
#[test]
fn round_history_is_append_only_across_iterations() {
let two = build_body_rounds(
vec![
sequence_tool_result("toolu_A", None),
sequence_tool_result("toolu_B", None),
],
vec![1],
);
let three = build_body_rounds(
vec![
sequence_tool_result("toolu_A", None),
sequence_tool_result("toolu_B", None),
sequence_tool_result("toolu_C", None),
],
vec![1, 2],
);
let two_messages = two["messages"].as_array().unwrap();
let three_messages = three["messages"].as_array().unwrap();
assert_eq!(
&three_messages[..two_messages.len()],
&two_messages[..],
"two: {two}, three: {three}"
);
}
#[test] #[test]
fn non_sequence_emits_assistant_user_pair() { fn non_sequence_emits_assistant_user_pair() {
let body = build_body_with( let body = build_body_with(
+102 -8
View File
@@ -401,6 +401,7 @@ pub fn claude_build_chat_completions_body(
tool_results, tool_results,
text, text,
sequence, sequence,
round_starts,
}) => { }) => {
if !sequence { if !sequence {
let mut assistant_parts = vec![]; let mut assistant_parts = vec![];
@@ -456,14 +457,18 @@ pub fn claude_build_chat_completions_body(
} }
} else { } else {
// One pair per round: Claude can reuse tool_use IDs across API calls. // One pair per round: Claude can reuse tool_use IDs across API calls.
// A round boundary is detected by the presence of round text, but // Round boundaries are recorded in `round_starts` at merge time.
// rounds where the model emitted only tool calls (no narration) // Splitting per round also keeps the serialized history append-only
// carry no text marker. As a backstop, also split whenever a // across tool-loop iterations, so prefix-based prompt caching keeps
// tool_use ID would repeat within the current assistant message — // hitting on earlier rounds instead of rewriting them every request.
// the API rejects duplicate tool_use IDs in a single message. // Sessions recorded before round tracking carry no `round_starts`;
// for those, fall back to the presence of round text, plus an ID
// backstop: split whenever a tool_use ID would repeat within the
// current assistant message — the API rejects duplicate tool_use
// IDs in a single message.
let mut messages = vec![]; let mut messages = vec![];
let mut assistant_parts: Vec<serde_json::Value> = vec![]; let mut assistant_parts: Vec<Value> = vec![];
let mut user_parts: Vec<serde_json::Value> = vec![]; let mut user_parts: Vec<Value> = vec![];
let mut chunk_ids: HashSet<&str> = HashSet::new(); let mut chunk_ids: HashSet<&str> = HashSet::new();
for (index, tool_result) in tool_results.iter().enumerate() { for (index, tool_result) in tool_results.iter().enumerate() {
let id_collision = tool_result let id_collision = tool_result
@@ -471,7 +476,11 @@ pub fn claude_build_chat_completions_body(
.id .id
.as_deref() .as_deref()
.is_some_and(|id| chunk_ids.contains(id)); .is_some_and(|id| chunk_ids.contains(id));
if index > 0 && (tool_result.text.is_some() || id_collision) { if index > 0
&& (round_starts.contains(&index)
|| tool_result.text.is_some()
|| id_collision)
{
messages.push( messages.push(
json!({ "role": "assistant", "content": assistant_parts }), json!({ "role": "assistant", "content": assistant_parts }),
); );
@@ -735,6 +744,7 @@ mod tests {
tool_results, tool_results,
text: text.to_string(), text: text.to_string(),
sequence, sequence,
round_starts: vec![],
}), }),
), ),
], ],
@@ -752,6 +762,30 @@ mod tests {
build_body_with(tool_results, "", true, prompt_cache) build_body_with(tool_results, "", true, prompt_cache)
} }
fn build_body_rounds(tool_results: Vec<ToolResult>, round_starts: Vec<usize>) -> Value {
let data = ChatCompletionsData {
messages: vec![
Message::new(MessageRole::User, MessageContent::Text("hello".to_string())),
Message::new(
MessageRole::Assistant,
MessageContent::ToolCalls(MessageContentToolCalls {
tool_results,
text: String::new(),
sequence: true,
round_starts,
}),
),
],
temperature: None,
top_p: None,
reasoning_effort: None,
functions: None,
stream: false,
};
claude_build_chat_completions_body(data, &Model::new("claude", "claude-test"), false)
.unwrap()
}
fn multi_turn_data(functions: Option<Vec<FunctionDeclaration>>) -> ChatCompletionsData { fn multi_turn_data(functions: Option<Vec<FunctionDeclaration>>) -> ChatCompletionsData {
ChatCompletionsData { ChatCompletionsData {
messages: vec![ messages: vec![
@@ -823,6 +857,63 @@ mod tests {
assert_unique_tool_use_ids_per_message(&body); assert_unique_tool_use_ids_per_message(&body);
} }
/// Regression: tool-loop rounds where the model emitted no narration text
/// (the norm with thinking models) must still serialize as one
/// assistant/user pair per round. Collapsing them into one growing pair
/// rewrites the whole accumulated history on every request, which defeats
/// prefix-based prompt caching (cache reads stall at the system prompt).
#[test]
fn sequence_splits_on_round_starts_without_text() {
let body = build_body_rounds(
vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", None),
],
vec![1, 2],
);
let messages = body["messages"].as_array().unwrap();
// user + 3 × (assistant, user) — one pair per round.
assert_eq!(messages.len(), 7, "body: {body}");
for pair in messages[1..].chunks(2) {
assert_eq!(pair[0]["role"], "assistant", "body: {body}");
assert_eq!(pair[0]["content"].as_array().unwrap().len(), 1);
assert_eq!(pair[1]["role"], "user", "body: {body}");
}
assert_unique_tool_use_ids_per_message(&body);
}
/// The serialized history must be append-only across tool-loop
/// iterations: request N+1's messages must start with exactly request N's
/// messages, so the prompt-cache prefix keeps matching and the cache read
/// grows monotonically instead of re-writing every round each request.
#[test]
fn round_history_is_append_only_across_iterations() {
let two = build_body_rounds(
vec![tool_result("toolu_A", None), tool_result("toolu_B", None)],
vec![1],
);
let three = build_body_rounds(
vec![
tool_result("toolu_A", None),
tool_result("toolu_B", None),
tool_result("toolu_C", None),
],
vec![1, 2],
);
let two_messages = two["messages"].as_array().unwrap();
let three_messages = three["messages"].as_array().unwrap();
assert_eq!(
&three_messages[..two_messages.len()],
&two_messages[..],
"two: {two}, three: {three}"
);
}
#[test] #[test]
fn sequence_splits_on_reused_id_in_textless_round() { fn sequence_splits_on_reused_id_in_textless_round() {
let body = build_body( let body = build_body(
@@ -840,6 +931,9 @@ mod tests {
assert_unique_tool_use_ids_per_message(&body); assert_unique_tool_use_ids_per_message(&body);
} }
/// Legacy sessions recorded before round tracking carry no
/// `round_starts`; without a text marker or an ID collision their
/// textless rounds still serialize merged (the pre-round-tracking shape).
#[test] #[test]
fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() { fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() {
let body = build_body( let body = build_body(
+66
View File
@@ -216,6 +216,16 @@ pub struct MessageContentToolCalls {
pub tool_results: Vec<ToolResult>, pub tool_results: Vec<ToolResult>,
pub text: String, pub text: String,
pub sequence: bool, pub sequence: bool,
/// Indices into `tool_results` where a new tool-call round begins (the
/// first round implicitly starts at index 0). Recorded by `merge()` on
/// every merge, so serializers can emit one assistant/user pair per
/// round even when the model produced no narration text for a round.
/// Keeping each round in its own message pair makes the request history
/// append-only across tool-loop iterations, which prefix-based prompt
/// caches require to produce growing cache reads instead of rewriting
/// the whole accumulated history on every request.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub round_starts: Vec<usize>,
} }
impl MessageContentToolCalls { impl MessageContentToolCalls {
@@ -224,6 +234,7 @@ impl MessageContentToolCalls {
tool_results, tool_results,
text, text,
sequence: false, sequence: false,
round_starts: vec![],
} }
} }
@@ -233,6 +244,9 @@ impl MessageContentToolCalls {
{ {
first.text = Some(text); first.text = Some(text);
} }
if !self.tool_results.is_empty() && !tool_results.is_empty() {
self.round_starts.push(self.tool_results.len());
}
self.tool_results.extend(tool_results); self.tool_results.extend(tool_results);
self.sequence = true; self.sequence = true;
} }
@@ -358,4 +372,56 @@ mod tests {
let restored: Vec<ThinkingBlock> = serde_yaml::from_str(&yaml).unwrap(); let restored: Vec<ThinkingBlock> = serde_yaml::from_str(&yaml).unwrap();
assert_eq!(serde_json::to_value(&restored).unwrap(), json); assert_eq!(serde_json::to_value(&restored).unwrap(), json);
} }
/// Session files recorded before round tracking carry no `round_starts`
/// key — they must deserialize to an empty default, and an empty
/// `round_starts` must be skipped on serialize so new files keep the old
/// shape until a merge actually records a round boundary.
#[test]
fn tool_calls_round_starts_defaults_and_skips_when_empty() {
let json = json!({
"tool_results": [],
"text": "t",
"sequence": true,
});
let parsed: MessageContentToolCalls = serde_json::from_value(json).unwrap();
assert!(parsed.round_starts.is_empty());
let serialized = serde_json::to_value(&parsed).unwrap();
assert!(
serialized.get("round_starts").is_none(),
"serialized: {serialized}"
);
}
#[test]
fn tool_calls_round_starts_round_trip_through_yaml() {
let mut tool_calls = MessageContentToolCalls::new(vec![], "t".to_string());
tool_calls.round_starts = vec![1, 3];
let yaml = serde_yaml::to_string(&tool_calls).unwrap();
let restored: MessageContentToolCalls = serde_yaml::from_str(&yaml).unwrap();
assert_eq!(restored.round_starts, vec![1, 3]);
}
/// Each `merge()` call is one tool-loop round: the boundary index must be
/// recorded even when the model produced no narration text for the round.
#[test]
fn merge_records_round_start_per_call() {
let tool_result = |id: &str| {
crate::function::ToolResult::new(
crate::function::ToolCall::new("t".into(), json!({}), Some(id.to_string())),
json!("ok"),
)
};
let mut tool_calls = MessageContentToolCalls::new(vec![tool_result("a")], String::new());
tool_calls.merge(vec![tool_result("b"), tool_result("c")], String::new());
tool_calls.merge(vec![tool_result("d")], String::new());
assert!(tool_calls.sequence);
assert_eq!(tool_calls.round_starts, vec![1, 3]);
}
} }
+4 -1
View File
@@ -481,6 +481,7 @@ pub fn openai_build_chat_completions_body(data: ChatCompletionsData, model: &Mod
tool_results, tool_results,
text, text,
sequence, sequence,
..
}) => { }) => {
// Empty tool_results (reachable via deserialized sessions) must not emit empty tool_calls. // Empty tool_results (reachable via deserialized sessions) must not emit empty tool_calls.
if tool_results.is_empty() { if tool_results.is_empty() {
@@ -681,7 +682,7 @@ pub fn openai_build_responses_body(data: ChatCompletionsData, model: &Model) ->
MessageContent::ToolCalls(MessageContentToolCalls { MessageContent::ToolCalls(MessageContentToolCalls {
tool_results, tool_results,
text, text,
sequence: _, ..
}) => tool_results }) => tool_results
.into_iter() .into_iter()
.enumerate() .enumerate()
@@ -993,6 +994,7 @@ mod tests {
tool_results: vec![], tool_results: vec![],
text: "leftover text".to_string(), text: "leftover text".to_string(),
sequence, sequence,
round_starts: vec![],
}), }),
), ),
], ],
@@ -1057,6 +1059,7 @@ mod tests {
tool_results, tool_results,
text: "first round".to_string(), text: "first round".to_string(),
sequence: false, sequence: false,
round_starts: vec![],
}), }),
), ),
], ],
+1
View File
@@ -600,6 +600,7 @@ mod tests {
tool_results: vec![], tool_results: vec![],
text: "leftover text".to_string(), text: "leftover text".to_string(),
sequence: false, sequence: false,
round_starts: vec![],
}), }),
), ),
], ],
+24
View File
@@ -1023,6 +1023,7 @@ mod tests {
assert!(!tool_calls.sequence); assert!(!tool_calls.sequence);
assert_eq!(tool_calls.tool_results.len(), 1); assert_eq!(tool_calls.tool_results.len(), 1);
assert!(tool_calls.tool_results[0].text.is_none()); assert!(tool_calls.tool_results[0].text.is_none());
assert!(tool_calls.round_starts.is_empty());
} }
#[test] #[test]
@@ -1044,6 +1045,29 @@ mod tests {
tool_calls.tool_results[1].text, tool_calls.tool_results[1].text,
Some("second text".to_string()) Some("second text".to_string())
); );
assert_eq!(tool_calls.round_starts, vec![1]);
}
/// Regression: rounds where the model emitted only tool calls (no
/// narration text) carry no text marker, so `round_starts` is the only
/// record of the round boundary — it must be set on every merge.
#[test]
fn merge_tool_results_records_round_starts_for_textless_rounds() {
let ctx = create_test_ctx();
let input = Input::from_str(&ctx, "test", None)
.unwrap()
.merge_tool_results(String::new(), vec![tool_result("id-1", "ok")])
.merge_tool_results(
String::new(),
vec![tool_result("id-2", "ok2"), tool_result("id-3", "ok3")],
)
.merge_tool_results(String::new(), vec![tool_result("id-4", "ok4")]);
let tool_calls = input.tool_calls().as_ref().unwrap();
assert!(tool_calls.sequence);
assert_eq!(tool_calls.tool_results.len(), 4);
assert!(tool_calls.tool_results.iter().all(|t| t.text.is_none()));
assert_eq!(tool_calls.round_starts, vec![1, 3]);
} }
#[test] #[test]