fix: improved caching by simply keeping groupings of tool calls by terminating a round on each tool merge into the bucket, rather than waiting for LLM narration to end a round.

2026-09-17 16:32:21 -06:00
1 parent dc6c176b4f
commit 3ee78a46df
@@ -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<Value> = vec![];
                        let mut user_parts: Vec<Value> = 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<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 {
        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(
@@ -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<serde_json::Value> = vec![];
-                        let mut user_parts: Vec<serde_json::Value> = vec![];
+                        let mut assistant_parts: Vec<Value> = vec![];
+                        let mut user_parts: Vec<Value> = 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<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 {
        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(
@@ -216,6 +216,16 @@ pub struct MessageContentToolCalls {
    pub tool_results: Vec<ToolResult>,
    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<usize>,
 }

 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<ThinkingBlock> = 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]);
+    }
 }
@@ -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![],
                    }),
                ),
            ],
@@ -600,6 +600,7 @@ mod tests {
                        tool_results: vec![],
                        text: "leftover text".to_string(),
                        sequence: false,
+                        round_starts: vec![],
                    }),
                ),
            ],