From 71d3d2d350fb997861ec59ceed9eabf9a91f0d82 Mon Sep 17 00:00:00 2001
From: Steven Enamakel <enamakel@tinyhumans.ai>
Date: Fri, 25 Sep 2026 17:50:20 +0530
Subject: [PATCH 1/2] fix: keep tool result and active request in one turn

---
 crates/tinyinference-llm/src/prompt_tools/mod.rs | 16 +++++++++++-----
 .../tinyinference-llm/src/prompt_tools/test.rs   |  7 +++++--
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/crates/tinyinference-llm/src/prompt_tools/mod.rs b/crates/tinyinference-llm/src/prompt_tools/mod.rs
index 358af54a..2b95d803 100644
--- a/crates/tinyinference-llm/src/prompt_tools/mod.rs
+++ b/crates/tinyinference-llm/src/prompt_tools/mod.rs
@@ -232,8 +232,10 @@ pub fn ensure_resolvable_user_turn(messages: &[Message]) -> Vec<Message> {
 /// Keep the active request at the end of a prompt-guided tool continuation.
 /// Some chat templates treat the synthetic tool-result user turn as transport
 /// data and resolve an earlier user message as the query. Repeating the latest
-/// real text request after a terminal result gives those templates the correct
-/// query without changing the durable transcript.
+/// real text request alongside a terminal result gives those templates the
+/// correct query without changing the durable transcript. Both belong in one
+/// user turn: a separate trailing request makes a model restart discovery and
+/// ignore the result it just received.
 #[must_use]
 pub fn anchor_user_request_after_tool_result(messages: &[Message]) -> Vec<Message> {
     if !messages.last().is_some_and(|message| {
@@ -258,9 +260,13 @@ pub fn anchor_user_request_after_tool_result(messages: &[Message]) -> Vec<Messag
         return messages.to_vec();
     }
     let mut out = messages.to_vec();
-    out.push(Message::user(format!(
-        "Continue the latest user request using the tool result above. Latest user request:\n{request}"
-    )));
+    let result = out.last().expect("terminal tool result exists").text();
+    if let Some(Message::User(last)) = out.last_mut() {
+        last.content = vec![ContentBlock::Text(format!(
+            "Continue the active user request using the completed tool result below. \
+             Do not repeat the completed tool call.\nActive user request:\n{request}\n\n{result}"
+        ))];
+    }
     out
 }
 
diff --git a/crates/tinyinference-llm/src/prompt_tools/test.rs b/crates/tinyinference-llm/src/prompt_tools/test.rs
index 1b9609a3..b670305f 100644
--- a/crates/tinyinference-llm/src/prompt_tools/test.rs
+++ b/crates/tinyinference-llm/src/prompt_tools/test.rs
@@ -177,9 +177,11 @@ fn tool_continuation_anchors_the_latest_request_after_search_results() {
         Message::tool("search-1", "GMAIL_FETCH_EMAILS schema"),
     ]);
     let anchored = anchor_user_request_after_tool_result(&messages);
-    assert_eq!(anchored.len(), messages.len() + 1);
+    assert_eq!(anchored.len(), messages.len());
     assert!(
-        anchored[anchored.len() - 2]
+        anchored
+            .last()
+            .unwrap()
             .text()
             .contains("GMAIL_FETCH_EMAILS")
     );
@@ -191,6 +193,7 @@ fn tool_continuation_anchors_the_latest_request_after_search_results() {
             .contains("fetch my latest email")
     );
     assert!(!anchored.last().unwrap().text().contains("Hey! What's up?"));
+    assert!(anchored.last().unwrap().text().contains("Do not repeat"));
     assert_eq!(anchor_user_request_after_tool_result(&anchored), anchored);
 }
 

From 6aff95c549357d082f27bb32e492994df379c5d6 Mon Sep 17 00:00:00 2001
From: Steven Enamakel <enamakel@tinyhumans.ai>
Date: Fri, 25 Sep 2026 17:56:25 +0530
Subject: [PATCH 2/2] fix: preserve tool result blocks in continuation

---
 .../tinyinference-llm/src/prompt_tools/mod.rs | 16 +++++++++-----
 .../src/prompt_tools/test.rs                  | 21 +++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/crates/tinyinference-llm/src/prompt_tools/mod.rs b/crates/tinyinference-llm/src/prompt_tools/mod.rs
index 2b95d803..08ca0542 100644
--- a/crates/tinyinference-llm/src/prompt_tools/mod.rs
+++ b/crates/tinyinference-llm/src/prompt_tools/mod.rs
@@ -260,12 +260,18 @@ pub fn anchor_user_request_after_tool_result(messages: &[Message]) -> Vec<Messag
         return messages.to_vec();
     }
     let mut out = messages.to_vec();
-    let result = out.last().expect("terminal tool result exists").text();
+    // `coalesce_tool_results` supplies a user-role result here. Keep every
+    // original block intact: direct callers can also pass JSON or media blocks
+    // alongside the result text.
     if let Some(Message::User(last)) = out.last_mut() {
-        last.content = vec![ContentBlock::Text(format!(
-            "Continue the active user request using the completed tool result below. \
-             Do not repeat the completed tool call.\nActive user request:\n{request}\n\n{result}"
-        ))];
+        last.content.insert(
+            0,
+            ContentBlock::Text(format!(
+                "Continue the active user request using the completed tool result below. \
+             Do not repeat the completed tool call unless new information requires it.\n\
+             Active user request:\n{request}\n\n"
+            )),
+        );
     }
     out
 }
diff --git a/crates/tinyinference-llm/src/prompt_tools/test.rs b/crates/tinyinference-llm/src/prompt_tools/test.rs
index b670305f..ad9ebc57 100644
--- a/crates/tinyinference-llm/src/prompt_tools/test.rs
+++ b/crates/tinyinference-llm/src/prompt_tools/test.rs
@@ -207,6 +207,27 @@ fn tool_continuation_without_a_real_user_request_stays_unchanged() {
     assert_eq!(anchor_user_request_after_tool_result(&messages), messages);
 }
 
+#[test]
+fn tool_continuation_preserves_non_text_result_blocks() {
+    let mut result = Message::user("[Tool results]\nlookup completed");
+    let Message::User(user) = &mut result else {
+        unreachable!()
+    };
+    let json_block = ContentBlock::Json(serde_json::json!({"count": 1}));
+    user.content.push(json_block.clone());
+    let anchored =
+        anchor_user_request_after_tool_result(&[Message::user("find the count"), result]);
+    let Message::User(last) = anchored.last().unwrap() else {
+        panic!("tool result remains a user turn");
+    };
+    assert_eq!(last.content.last(), Some(&json_block));
+    assert!(
+        last.content[0]
+            .as_text()
+            .is_some_and(|text| text.contains("find the count"))
+    );
+}
+
 #[test]
 fn user_turn_normalization_ignores_blank_and_accepts_non_text_turns() {
     let out = ensure_resolvable_user_turn(&[Message::system("system"), Message::user("   ")]);