diff --git a/src/client/bedrock.rs b/src/client/bedrock.rs index a8da844..1a1e63d 100644 --- a/src/client/bedrock.rs +++ b/src/client/bedrock.rs @@ -423,6 +423,7 @@ fn build_chat_completions_body( tool_results, text, sequence, + round_starts, }) => { if !sequence { let mut assistant_parts = vec![]; @@ -472,11 +473,15 @@ fn build_chat_completions_body( } } else { // 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 - // rounds where the model emitted only tool calls (no narration) - // carry no text marker. As a backstop, also split whenever a - // toolUseId would repeat within the current assistant message. - // Converse rejects duplicate toolUseIds in a single message. + // Round boundaries are recorded in `round_starts` at merge time. + // Splitting per round also keeps the serialized history append-only + // across tool-loop iterations, so prompt caching keeps hitting on + // earlier rounds instead of rewriting them every request. + // 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 assistant_parts: Vec = vec![]; let mut user_parts: Vec = vec![]; @@ -487,7 +492,11 @@ fn build_chat_completions_body( .id .as_deref() .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( json!({ "role": "assistant", "content": assistant_parts }), ); @@ -896,6 +905,7 @@ mod tests { tool_results, text: text.to_string(), sequence, + round_starts: vec![], }), ), ], @@ -912,6 +922,29 @@ mod tests { build_body_with(tool_results, "", true) } + fn build_body_rounds(tool_results: Vec, round_starts: Vec) -> 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 { let body = build_body(vec![tool_result(output)]); body["messages"][2]["content"][0]["toolResult"]["content"][0].clone() @@ -1067,6 +1100,9 @@ mod tests { 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] fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() { let body = build_body(vec![ @@ -1081,6 +1117,64 @@ mod tests { 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] fn non_sequence_emits_assistant_user_pair() { let body = build_body_with( diff --git a/src/client/claude.rs b/src/client/claude.rs index db31271..78f12a1 100644 --- a/src/client/claude.rs +++ b/src/client/claude.rs @@ -401,6 +401,7 @@ pub fn claude_build_chat_completions_body( tool_results, text, sequence, + round_starts, }) => { if !sequence { let mut assistant_parts = vec![]; @@ -456,14 +457,18 @@ pub fn claude_build_chat_completions_body( } } else { // 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 - // rounds where the model emitted only tool calls (no narration) - // carry no text marker. As a backstop, also split whenever a - // tool_use ID would repeat within the current assistant message — - // the API rejects duplicate tool_use IDs in a single message. + // Round boundaries are recorded in `round_starts` at merge time. + // Splitting per round also keeps the serialized history append-only + // across tool-loop iterations, so prefix-based prompt caching keeps + // hitting on earlier rounds instead of rewriting them every request. + // 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 assistant_parts: Vec = vec![]; - let mut user_parts: Vec = vec![]; + let mut assistant_parts: Vec = vec![]; + let mut user_parts: Vec = vec![]; let mut chunk_ids: HashSet<&str> = HashSet::new(); for (index, tool_result) in tool_results.iter().enumerate() { let id_collision = tool_result @@ -471,7 +476,11 @@ pub fn claude_build_chat_completions_body( .id .as_deref() .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( json!({ "role": "assistant", "content": assistant_parts }), ); @@ -735,6 +744,7 @@ mod tests { tool_results, text: text.to_string(), sequence, + round_starts: vec![], }), ), ], @@ -752,6 +762,30 @@ mod tests { build_body_with(tool_results, "", true, prompt_cache) } + fn build_body_rounds(tool_results: Vec, round_starts: Vec) -> 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>) -> ChatCompletionsData { ChatCompletionsData { messages: vec![ @@ -823,6 +857,63 @@ mod tests { 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] fn sequence_splits_on_reused_id_in_textless_round() { let body = build_body( @@ -840,6 +931,9 @@ mod tests { 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] fn sequence_keeps_textless_rounds_merged_when_ids_are_unique() { let body = build_body( diff --git a/src/client/message.rs b/src/client/message.rs index 46e6b7d..5e095c4 100644 --- a/src/client/message.rs +++ b/src/client/message.rs @@ -216,6 +216,16 @@ pub struct MessageContentToolCalls { pub tool_results: Vec, pub text: String, 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, } impl MessageContentToolCalls { @@ -224,6 +234,7 @@ impl MessageContentToolCalls { tool_results, text, sequence: false, + round_starts: vec![], } } @@ -233,6 +244,9 @@ impl MessageContentToolCalls { { 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.sequence = true; } @@ -358,4 +372,56 @@ mod tests { let restored: Vec = serde_yaml::from_str(&yaml).unwrap(); 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]); + } } diff --git a/src/client/openai.rs b/src/client/openai.rs index ffc0933..bf78a4a 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -481,6 +481,7 @@ pub fn openai_build_chat_completions_body(data: ChatCompletionsData, model: &Mod tool_results, text, sequence, + .. }) => { // Empty tool_results (reachable via deserialized sessions) must not emit empty tool_calls. if tool_results.is_empty() { @@ -681,7 +682,7 @@ pub fn openai_build_responses_body(data: ChatCompletionsData, model: &Model) -> MessageContent::ToolCalls(MessageContentToolCalls { tool_results, text, - sequence: _, + .. }) => tool_results .into_iter() .enumerate() @@ -993,6 +994,7 @@ mod tests { tool_results: vec![], text: "leftover text".to_string(), sequence, + round_starts: vec![], }), ), ], @@ -1057,6 +1059,7 @@ mod tests { tool_results, text: "first round".to_string(), sequence: false, + round_starts: vec![], }), ), ], diff --git a/src/client/vertexai.rs b/src/client/vertexai.rs index d9f3fed..39eb39d 100644 --- a/src/client/vertexai.rs +++ b/src/client/vertexai.rs @@ -600,6 +600,7 @@ mod tests { tool_results: vec![], text: "leftover text".to_string(), sequence: false, + round_starts: vec![], }), ), ], diff --git a/src/config/input.rs b/src/config/input.rs index 59dd634..eb1fedc 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -1023,6 +1023,7 @@ mod tests { assert!(!tool_calls.sequence); assert_eq!(tool_calls.tool_results.len(), 1); assert!(tool_calls.tool_results[0].text.is_none()); + assert!(tool_calls.round_starts.is_empty()); } #[test] @@ -1044,6 +1045,29 @@ mod tests { tool_calls.tool_results[1].text, 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]