From 290d925c7f1689ec7dd0579cc1bbaab84ba8866d Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Mon, 9 Mar 2026 23:36:07 +0800 Subject: [PATCH] fix: preserve tool-call history across thread hydration (#568) (#670) Prevent model re-attempts and data inconsistencies when rebuilding conversation context from persisted tool-call records. - Remove raw tool parameters from persisted tool_calls JSON to prevent unredacted sensitive data from being stored in the database. The LLM context rebuild only needs call_id + name + result. - Make record_tool_error/record_tool_result mutually exclusive in all three execution paths (dispatcher, approval, deferred). Previously error cases called both methods, violating the TurnToolCall invariant and sending contradictory outcomes to the LLM. - Unify call_id format to turn{N}_{i} between live sessions and persisted hydration to eliminate ID mismatch in the LLM context. - Auto-close XML tags after truncate_preview truncation to prevent malformed tool output reaching the LLM. [skip-regression-check] --- src/agent/dispatcher.rs | 35 ++-- src/agent/session.rs | 329 ++++++++++++++++++++++++++++++++++-- src/agent/thread_ops.rs | 349 ++++++++++++++++++++++++++++++++------- src/channels/web/util.rs | 40 ++++- 4 files changed, 666 insertions(+), 87 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index b59ff92f..f5306644 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -743,23 +743,6 @@ impl Agent { .await; } - // Record result in thread - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - match &tool_result { - Ok(output) => { - turn.record_tool_result(serde_json::json!(output)); - } - Err(e) => { - turn.record_tool_error(e.to_string()); - } - } - } - } - // Check for auth awaiting — defer the return // until all results are recorded. if deferred_auth.is_none() @@ -799,6 +782,7 @@ impl Agent { } // Sanitize and add tool result to context + let is_tool_error = tool_result.is_err(); let result_content = match tool_result { Ok(output) => { let sanitized = @@ -812,6 +796,23 @@ impl Agent { Err(e) => format!("Tool '{}' failed: {}", tc.name, e), }; + // Record sanitized result in thread so messages() + // and persist_tool_calls() use cleaned content. + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) + && let Some(turn) = thread.last_turn_mut() + { + if is_tool_error { + turn.record_tool_error(result_content.clone()); + } else { + turn.record_tool_result(serde_json::json!( + result_content + )); + } + } + } + context_messages.push(ChatMessage::tool_result( &tc.id, &tc.name, diff --git a/src/agent/session.rs b/src/agent/session.rs index a051ffea..193e0309 100644 --- a/src/agent/session.rs +++ b/src/agent/session.rs @@ -16,6 +16,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; +use crate::channels::web::util::truncate_preview; use crate::llm::{ChatMessage, ToolCall}; /// A session containing one or more threads. @@ -320,7 +321,13 @@ impl Thread { } } - /// Get all messages for context building. + /// Get all messages for context building, including tool call history. + /// + /// Emits the full LLM-compatible message sequence per turn: + /// `user → [assistant_with_tool_calls → tool_result*] → assistant` + /// + /// This ensures the LLM sees prior tool executions and won't re-attempt + /// completed actions in subsequent turns. pub fn messages(&self) -> Vec { let mut messages = Vec::new(); for turn in &self.turns { @@ -332,6 +339,42 @@ impl Thread { turn.image_content_parts.clone(), )); } + + if !turn.tool_calls.is_empty() { + // Build ToolCall objects with synthetic stable IDs + let tool_calls: Vec = turn + .tool_calls + .iter() + .enumerate() + .map(|(i, tc)| ToolCall { + id: format!("turn{}_{}", turn.turn_number, i), + name: tc.name.clone(), + arguments: tc.parameters.clone(), + }) + .collect(); + + // Assistant message declaring the tool calls (no text content) + messages.push(ChatMessage::assistant_with_tool_calls(None, tool_calls)); + + // Individual tool result messages, truncated to limit context size. + for (i, tc) in turn.tool_calls.iter().enumerate() { + let call_id = format!("turn{}_{}", turn.turn_number, i); + let content = if let Some(ref err) = tc.error { + // .error already contains the full error text; + // pass through without wrapping to avoid double-prefix. + truncate_preview(err, 1000) + } else if let Some(ref res) = tc.result { + let raw = match res { + serde_json::Value::String(s) => s.clone(), + other => other.to_string(), + }; + truncate_preview(&raw, 1000) + } else { + "OK".to_string() + }; + messages.push(ChatMessage::tool_result(call_id, &tc.name, content)); + } + } if let Some(ref response) = turn.response { messages.push(ChatMessage::assistant(response)); } @@ -353,13 +396,16 @@ impl Thread { /// Restore thread state from a checkpoint's messages. /// - /// Clears existing turns and rebuilds from message pairs. - /// Messages should alternate: user, assistant, user, assistant... + /// Clears existing turns and rebuilds from the message sequence. + /// Handles the full message pattern including tool messages: + /// `user → [assistant_with_tool_calls → tool_result*] → assistant` + /// + /// Also supports the legacy pattern (user/assistant pairs only) for + /// backward compatibility with old checkpoint data. pub fn restore_from_messages(&mut self, messages: Vec) { self.turns.clear(); self.state = ThreadState::Idle; - // Messages alternate: user, assistant, user, assistant... let mut iter = messages.into_iter().peekable(); let mut turn_number = 0; @@ -367,18 +413,58 @@ impl Thread { if msg.role == crate::llm::Role::User { let mut turn = Turn::new(turn_number, &msg.content); - // Check if next is assistant response - if let Some(next) = iter.peek() - && next.role == crate::llm::Role::Assistant - { - // iter.next() is guaranteed Some after a successful peek() - if let Some(response) = iter.next() { - turn.complete(&response.content); + // Consume tool call sequences (assistant_with_tool_calls + tool_results). + // A single turn may contain multiple rounds of tool calls, so we + // track the cumulative base index into turn.tool_calls. + while let Some(next) = iter.peek() { + if next.role == crate::llm::Role::Assistant && next.tool_calls.is_some() { + let call_base_idx = turn.tool_calls.len(); + + if let Some(assistant_msg) = iter.next() + && let Some(ref tcs) = assistant_msg.tool_calls + { + for tc in tcs { + turn.record_tool_call(&tc.name, tc.arguments.clone()); + } + } + + // Consume the corresponding tool_result messages, + // indexing relative to this batch's base offset. + let mut pos = 0; + while let Some(tr) = iter.peek() { + if tr.role != crate::llm::Role::Tool { + break; + } + if let Some(tool_msg) = iter.next() { + let idx = call_base_idx + pos; + if idx < turn.tool_calls.len() { + // Store as result — the error/success distinction + // is for the live turn only; restored context just + // needs the content the LLM originally saw. + turn.tool_calls[idx].result = + Some(serde_json::Value::String(tool_msg.content.clone())); + } + } + pos += 1; + } + } else { + break; } } + // Check if next is the final assistant response for this turn + let is_final_assistant = iter.peek().is_some_and(|n| { + n.role == crate::llm::Role::Assistant && n.tool_calls.is_none() + }); + if is_final_assistant && let Some(response) = iter.next() { + turn.complete(&response.content); + } + self.turns.push(turn); turn_number += 1; + } else { + // Skip non-user messages that aren't anchored to a turn + continue; } } @@ -1035,4 +1121,225 @@ mod tests { ThreadState::Processing ); } + + // Regression tests for #568: tool call history must survive hydration. + + #[test] + fn test_messages_includes_tool_calls() { + let mut thread = Thread::new(Uuid::new_v4()); + + thread.start_turn("Search for X"); + { + let turn = thread.turns.last_mut().unwrap(); + turn.record_tool_call("memory_search", serde_json::json!({"query": "X"})); + turn.record_tool_result(serde_json::json!("Found X in doc.md")); + } + thread.complete_turn("I found X in doc.md."); + + let messages = thread.messages(); + // user + assistant_with_tool_calls + tool_result + assistant = 4 + assert_eq!(messages.len(), 4); + + assert_eq!(messages[0].role, crate::llm::Role::User); + assert_eq!(messages[0].content, "Search for X"); + + assert_eq!(messages[1].role, crate::llm::Role::Assistant); + assert!(messages[1].tool_calls.is_some()); + let tcs = messages[1].tool_calls.as_ref().unwrap(); + assert_eq!(tcs.len(), 1); + assert_eq!(tcs[0].name, "memory_search"); + + assert_eq!(messages[2].role, crate::llm::Role::Tool); + assert!(messages[2].content.contains("Found X")); + + assert_eq!(messages[3].role, crate::llm::Role::Assistant); + assert_eq!(messages[3].content, "I found X in doc.md."); + } + + #[test] + fn test_messages_multiple_tool_calls_per_turn() { + let mut thread = Thread::new(Uuid::new_v4()); + + thread.start_turn("Do two things"); + { + let turn = thread.turns.last_mut().unwrap(); + turn.record_tool_call("echo", serde_json::json!({"msg": "a"})); + turn.record_tool_result(serde_json::json!("a")); + turn.record_tool_call("time", serde_json::json!({})); + turn.record_tool_error("timeout"); + } + thread.complete_turn("Done."); + + let messages = thread.messages(); + // user + assistant_with_calls(2) + tool_result + tool_result + assistant = 5 + assert_eq!(messages.len(), 5); + + let tcs = messages[1].tool_calls.as_ref().unwrap(); + assert_eq!(tcs.len(), 2); + + // First tool: success + assert_eq!(messages[2].content, "a"); + // Second tool: error (passed through directly, no wrapping) + assert!(messages[3].content.contains("timeout")); + } + + #[test] + fn test_restore_from_messages_with_tool_calls() { + let mut thread = Thread::new(Uuid::new_v4()); + + // Build a message sequence with tool calls + let tc = ToolCall { + id: "call_0".to_string(), + name: "search".to_string(), + arguments: serde_json::json!({"q": "test"}), + }; + let messages = vec![ + ChatMessage::user("Find test"), + ChatMessage::assistant_with_tool_calls(None, vec![tc]), + ChatMessage::tool_result("call_0", "search", "result: found"), + ChatMessage::assistant("Found it."), + ]; + + thread.restore_from_messages(messages); + + assert_eq!(thread.turns.len(), 1); + let turn = &thread.turns[0]; + assert_eq!(turn.user_input, "Find test"); + assert_eq!(turn.tool_calls.len(), 1); + assert_eq!(turn.tool_calls[0].name, "search"); + assert_eq!( + turn.tool_calls[0].result, + Some(serde_json::Value::String("result: found".to_string())) + ); + assert_eq!(turn.response, Some("Found it.".to_string())); + } + + #[test] + fn test_restore_from_messages_with_tool_error() { + let mut thread = Thread::new(Uuid::new_v4()); + + let tc = ToolCall { + id: "call_0".to_string(), + name: "http".to_string(), + arguments: serde_json::json!({}), + }; + let messages = vec![ + ChatMessage::user("Fetch URL"), + ChatMessage::assistant_with_tool_calls(None, vec![tc]), + ChatMessage::tool_result("call_0", "http", "Error: timeout"), + ChatMessage::assistant("The request timed out."), + ]; + + thread.restore_from_messages(messages); + + // restore_from_messages stores all tool content as result (not error), + // because it can't reliably distinguish errors from results that happen + // to start with "Error: ". The content is preserved for LLM context. + let turn = &thread.turns[0]; + assert_eq!( + turn.tool_calls[0].result, + Some(serde_json::Value::String("Error: timeout".to_string())) + ); + } + + #[test] + fn test_messages_round_trip_with_tools() { + // Build a thread with tool calls, get messages(), restore, get messages() again + // The two message sequences should be equivalent. + let mut thread = Thread::new(Uuid::new_v4()); + + thread.start_turn("Do search"); + { + let turn = thread.turns.last_mut().unwrap(); + turn.record_tool_call("search", serde_json::json!({"q": "test"})); + turn.record_tool_result(serde_json::json!("found")); + } + thread.complete_turn("Here are results."); + + let messages_original = thread.messages(); + + // Restore into a new thread + let mut thread2 = Thread::new(Uuid::new_v4()); + thread2.restore_from_messages(messages_original.clone()); + + let messages_restored = thread2.messages(); + + // Same number of messages + assert_eq!(messages_original.len(), messages_restored.len()); + + // Same roles + for (orig, rest) in messages_original.iter().zip(messages_restored.iter()) { + assert_eq!(orig.role, rest.role); + } + + // Same final response + assert_eq!( + messages_original.last().unwrap().content, + messages_restored.last().unwrap().content + ); + } + + #[test] + fn test_restore_multi_stage_tool_calls() { + let mut thread = Thread::new(Uuid::new_v4()); + + let tc1 = ToolCall { + id: "call_a".to_string(), + name: "search".to_string(), + arguments: serde_json::json!({"q": "data"}), + }; + let tc2 = ToolCall { + id: "call_b".to_string(), + name: "write".to_string(), + arguments: serde_json::json!({"path": "out.txt"}), + }; + let messages = vec![ + ChatMessage::user("Find and save"), + ChatMessage::assistant_with_tool_calls(None, vec![tc1]), + ChatMessage::tool_result("call_a", "search", "found data"), + ChatMessage::assistant_with_tool_calls(None, vec![tc2]), + ChatMessage::tool_result("call_b", "write", "written"), + ChatMessage::assistant("Done, saved to out.txt"), + ]; + + thread.restore_from_messages(messages); + + assert_eq!(thread.turns.len(), 1); + let turn = &thread.turns[0]; + assert_eq!(turn.tool_calls.len(), 2); + assert_eq!(turn.tool_calls[0].name, "search"); + assert_eq!(turn.tool_calls[1].name, "write"); + assert_eq!( + turn.tool_calls[0].result, + Some(serde_json::Value::String("found data".to_string())) + ); + assert_eq!( + turn.tool_calls[1].result, + Some(serde_json::Value::String("written".to_string())) + ); + assert_eq!(turn.response, Some("Done, saved to out.txt".to_string())); + } + + #[test] + fn test_messages_truncates_large_tool_results() { + let mut thread = Thread::new(Uuid::new_v4()); + + thread.start_turn("Read big file"); + { + let turn = thread.turns.last_mut().unwrap(); + turn.record_tool_call("read_file", serde_json::json!({"path": "big.txt"})); + let big_result = "x".repeat(2000); + turn.record_tool_result(serde_json::json!(big_result)); + } + thread.complete_turn("Here's the file content."); + + let messages = thread.messages(); + let tool_result_content = &messages[2].content; + assert!( + tool_result_content.len() <= 1010, + "Tool result should be truncated, got {} chars", + tool_result_content.len() + ); + assert!(tool_result_content.ends_with("...")); + } } diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 4dc3ff17..758e98ed 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -20,7 +20,7 @@ use crate::channels::web::util::truncate_preview; use crate::channels::{IncomingMessage, StatusUpdate}; use crate::context::JobContext; use crate::error::Error; -use crate::llm::ChatMessage; +use crate::llm::{ChatMessage, ToolCall}; use crate::tools::redact_params; impl Agent { @@ -66,16 +66,7 @@ impl Agent { .await .unwrap_or_default(); msg_count = db_messages.len(); - chat_messages = db_messages - .iter() - .filter_map(|m| match m.role.as_str() { - "user" => Some(ChatMessage::user(&m.content)), - "assistant" => Some(ChatMessage::assistant(&m.content)), - // tool_calls rows are UI metadata (tool name + preview), - // not part of the LLM conversation context. - _ => None, - }) - .collect(); + chat_messages = rebuild_chat_messages_from_db(&db_messages); } else { msg_count = 0; } @@ -340,10 +331,10 @@ impl Agent { }; thread.complete_turn(&response); - let tool_calls = thread + let (turn_number, tool_calls) = thread .turns .last() - .map(|t| t.tool_calls.clone()) + .map(|t| (t.turn_number, t.tool_calls.clone())) .unwrap_or_default(); let _ = self .channels @@ -355,7 +346,7 @@ impl Agent { .await; // Persist tool calls then assistant response (user message already persisted at turn start) - self.persist_tool_calls(thread_id, &message.user_id, &tool_calls) + self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls) .await; self.persist_assistant_response(thread_id, &message.user_id, &response) .await; @@ -464,6 +455,7 @@ impl Agent { &self, thread_id: Uuid, user_id: &str, + turn_number: usize, tool_calls: &[crate::agent::session::TurnToolCall], ) { if tool_calls.is_empty() { @@ -477,14 +469,24 @@ impl Agent { let summaries: Vec = tool_calls .iter() - .map(|tc| { - let mut obj = serde_json::json!({ "name": tc.name }); + .enumerate() + .map(|(i, tc)| { + let mut obj = serde_json::json!({ + "name": tc.name, + "call_id": format!("turn{}_{}", turn_number, i), + }); if let Some(ref result) = tc.result { let preview = match result { serde_json::Value::String(s) => truncate_preview(s, 500), other => truncate_preview(&other.to_string(), 500), }; obj["result_preview"] = serde_json::Value::String(preview); + // Store full result (truncated to ~1000 chars) for LLM context rebuild + let full_result = match result { + serde_json::Value::String(s) => truncate_preview(s, 1000), + other => truncate_preview(&other.to_string(), 1000), + }; + obj["result"] = serde_json::Value::String(full_result); } if let Some(ref error) = tc.error { obj["error"] = serde_json::Value::String(truncate_preview(error, 200)); @@ -807,19 +809,33 @@ impl Agent { let mut context_messages = pending.context_messages; let deferred_tool_calls = pending.deferred_tool_calls; - // Record result in thread + // Sanitize tool result, then record the cleaned version in the + // thread. Must happen before auth intercept check which may return early. + let is_tool_error = tool_result.is_err(); + let result_content = match &tool_result { + Ok(output) => { + let sanitized = self + .safety() + .sanitize_tool_output(&pending.tool_name, output); + self.safety().wrap_for_llm( + &pending.tool_name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Error: {}", e), + }; + + // Record sanitized result in thread { let mut sess = session.lock().await; if let Some(thread) = sess.threads.get_mut(&thread_id) && let Some(turn) = thread.last_turn_mut() { - match &tool_result { - Ok(output) => { - turn.record_tool_result(serde_json::json!(output)); - } - Err(e) => { - turn.record_tool_error(e.to_string()); - } + if is_tool_error { + turn.record_tool_error(result_content.clone()); + } else { + turn.record_tool_result(serde_json::json!(result_content)); } } } @@ -841,21 +857,6 @@ impl Agent { return Ok(SubmissionResult::response(instructions)); } - // Add tool result to context - let result_content = match tool_result { - Ok(output) => { - let sanitized = self - .safety() - .sanitize_tool_output(&pending.tool_name, &output); - self.safety().wrap_for_llm( - &pending.tool_name, - &sanitized.content, - sanitized.was_modified, - ) - } - Err(e) => format!("Error: {}", e), - }; - context_messages.push(ChatMessage::tool_result( &pending.tool_call_id, &pending.tool_name, @@ -1060,15 +1061,31 @@ impl Agent { .await; } - // Record in thread + // Sanitize first, then record the cleaned version in thread. + // Must happen before auth detection which may set deferred_auth. + let is_deferred_error = deferred_result.is_err(); + let deferred_content = match &deferred_result { + Ok(output) => { + let sanitized = self.safety().sanitize_tool_output(&tc.name, output); + self.safety().wrap_for_llm( + &tc.name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Error: {}", e), + }; + + // Record sanitized result in thread { let mut sess = session.lock().await; if let Some(thread) = sess.threads.get_mut(&thread_id) && let Some(turn) = thread.last_turn_mut() { - match &deferred_result { - Ok(output) => turn.record_tool_result(serde_json::json!(output)), - Err(e) => turn.record_tool_error(e.to_string()), + if is_deferred_error { + turn.record_tool_error(deferred_content.clone()); + } else { + turn.record_tool_result(serde_json::json!(deferred_content)); } } } @@ -1090,18 +1107,6 @@ impl Agent { deferred_auth = Some(instructions); } - let deferred_content = match deferred_result { - Ok(output) => { - let sanitized = self.safety().sanitize_tool_output(&tc.name, &output); - self.safety().wrap_for_llm( - &tc.name, - &sanitized.content, - sanitized.was_modified, - ) - } - Err(e) => format!("Error: {}", e), - }; - context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content)); } @@ -1169,13 +1174,13 @@ impl Agent { match result { Ok(AgenticLoopResult::Response(response)) => { thread.complete_turn(&response); - let tool_calls = thread + let (turn_number, tool_calls) = thread .turns .last() - .map(|t| t.tool_calls.clone()) + .map(|t| (t.turn_number, t.tool_calls.clone())) .unwrap_or_default(); // User message already persisted at turn start; save tool calls then assistant response - self.persist_tool_calls(thread_id, &message.user_id, &tool_calls) + self.persist_tool_calls(thread_id, &message.user_id, turn_number, &tool_calls) .await; self.persist_assistant_response(thread_id, &message.user_id, &response) .await; @@ -1490,3 +1495,231 @@ impl Agent { } } } + +/// Rebuild full LLM-compatible `ChatMessage` sequence from DB messages. +/// +/// Parses `role="tool_calls"` rows to reconstruct `assistant_with_tool_calls` +/// and `tool_result` messages so that the LLM sees the complete tool execution +/// history on thread hydration. Falls back gracefully for legacy rows that +/// lack the enriched fields (`call_id`, `parameters`, `result`). +fn rebuild_chat_messages_from_db( + db_messages: &[crate::history::ConversationMessage], +) -> Vec { + let mut result = Vec::new(); + + for msg in db_messages { + match msg.role.as_str() { + "user" => result.push(ChatMessage::user(&msg.content)), + "assistant" => result.push(ChatMessage::assistant(&msg.content)), + "tool_calls" => { + // Try to parse the enriched JSON and rebuild tool messages. + if let Ok(calls) = serde_json::from_str::>(&msg.content) { + if calls.is_empty() { + continue; + } + + // Check if this is an enriched row (has call_id) or legacy + let has_call_id = calls + .first() + .and_then(|c| c.get("call_id")) + .and_then(|v| v.as_str()) + .is_some(); + + if has_call_id { + // Build assistant_with_tool_calls + tool_result messages + let tool_calls: Vec = calls + .iter() + .map(|c| ToolCall { + id: c["call_id"].as_str().unwrap_or("call_0").to_string(), + name: c["name"].as_str().unwrap_or("unknown").to_string(), + arguments: c + .get("parameters") + .cloned() + .unwrap_or(serde_json::json!({})), + }) + .collect(); + + // The assistant text for tool_calls is always None here; + // the final assistant response comes as a separate + // "assistant" row after this tool_calls row. + result.push(ChatMessage::assistant_with_tool_calls(None, tool_calls)); + + // Emit tool_result messages for each call + for c in &calls { + let call_id = c["call_id"].as_str().unwrap_or("call_0").to_string(); + let name = c["name"].as_str().unwrap_or("unknown").to_string(); + let content = if let Some(err) = c.get("error").and_then(|v| v.as_str()) + { + format!("Error: {}", err) + } else if let Some(res) = c.get("result").and_then(|v| v.as_str()) { + res.to_string() + } else if let Some(preview) = + c.get("result_preview").and_then(|v| v.as_str()) + { + preview.to_string() + } else { + "OK".to_string() + }; + result.push(ChatMessage::tool_result(call_id, name, content)); + } + } + // Legacy rows without call_id: skip (will appear as + // simple user/assistant pairs, same as before this fix). + } + } + _ => {} // Skip unknown roles + } + } + + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_rebuild_chat_messages_user_assistant_only() { + let messages = vec![ + make_db_msg("user", "Hello"), + make_db_msg("assistant", "Hi there!"), + ]; + let result = rebuild_chat_messages_from_db(&messages); + assert_eq!(result.len(), 2); + assert_eq!(result[0].role, crate::llm::Role::User); + assert_eq!(result[1].role, crate::llm::Role::Assistant); + } + + #[test] + fn test_rebuild_chat_messages_with_enriched_tool_calls() { + let tool_json = serde_json::json!([ + { + "name": "memory_search", + "call_id": "call_0", + "parameters": {"query": "test"}, + "result": "Found 3 results", + "result_preview": "Found 3 re..." + }, + { + "name": "echo", + "call_id": "call_1", + "parameters": {"message": "hi"}, + "error": "timeout" + } + ]); + let messages = vec![ + make_db_msg("user", "Search for test"), + make_db_msg("tool_calls", &tool_json.to_string()), + make_db_msg("assistant", "I found some results."), + ]; + let result = rebuild_chat_messages_from_db(&messages); + + // user + assistant_with_tool_calls + tool_result*2 + assistant + assert_eq!(result.len(), 5); + + // user + assert_eq!(result[0].role, crate::llm::Role::User); + + // assistant with tool_calls + assert_eq!(result[1].role, crate::llm::Role::Assistant); + assert!(result[1].tool_calls.is_some()); + let tcs = result[1].tool_calls.as_ref().unwrap(); + assert_eq!(tcs.len(), 2); + assert_eq!(tcs[0].name, "memory_search"); + assert_eq!(tcs[0].id, "call_0"); + assert_eq!(tcs[1].name, "echo"); + + // tool results + assert_eq!(result[2].role, crate::llm::Role::Tool); + assert_eq!(result[2].tool_call_id, Some("call_0".to_string())); + assert!(result[2].content.contains("Found 3 results")); + + assert_eq!(result[3].role, crate::llm::Role::Tool); + assert_eq!(result[3].tool_call_id, Some("call_1".to_string())); + assert!(result[3].content.contains("Error: timeout")); + + // final assistant + assert_eq!(result[4].role, crate::llm::Role::Assistant); + assert_eq!(result[4].content, "I found some results."); + } + + #[test] + fn test_rebuild_chat_messages_legacy_tool_calls_skipped() { + // Legacy format: no call_id field + let tool_json = serde_json::json!([ + {"name": "echo", "result_preview": "hello"} + ]); + let messages = vec![ + make_db_msg("user", "Hi"), + make_db_msg("tool_calls", &tool_json.to_string()), + make_db_msg("assistant", "Done"), + ]; + let result = rebuild_chat_messages_from_db(&messages); + + // Legacy rows are skipped, only user + assistant + assert_eq!(result.len(), 2); + assert_eq!(result[0].role, crate::llm::Role::User); + assert_eq!(result[1].role, crate::llm::Role::Assistant); + } + + #[test] + fn test_rebuild_chat_messages_empty() { + let result = rebuild_chat_messages_from_db(&[]); + assert!(result.is_empty()); + } + + #[test] + fn test_rebuild_chat_messages_malformed_tool_calls_json() { + let messages = vec![ + make_db_msg("user", "Hi"), + make_db_msg("tool_calls", "not valid json"), + make_db_msg("assistant", "Done"), + ]; + let result = rebuild_chat_messages_from_db(&messages); + // Malformed JSON is silently skipped + assert_eq!(result.len(), 2); + } + + #[test] + fn test_rebuild_chat_messages_multi_turn_with_tools() { + let tool_json_1 = serde_json::json!([ + {"name": "search", "call_id": "call_0", "parameters": {}, "result": "found it"} + ]); + let tool_json_2 = serde_json::json!([ + {"name": "write", "call_id": "call_0", "parameters": {"path": "a.txt"}, "result": "ok"} + ]); + let messages = vec![ + make_db_msg("user", "Find X"), + make_db_msg("tool_calls", &tool_json_1.to_string()), + make_db_msg("assistant", "Found X"), + make_db_msg("user", "Write it"), + make_db_msg("tool_calls", &tool_json_2.to_string()), + make_db_msg("assistant", "Written"), + ]; + let result = rebuild_chat_messages_from_db(&messages); + + // Turn 1: user + assistant_with_calls + tool_result + assistant = 4 + // Turn 2: user + assistant_with_calls + tool_result + assistant = 4 + assert_eq!(result.len(), 8); + + // Verify turn boundaries + assert_eq!(result[0].content, "Find X"); + assert!(result[1].tool_calls.is_some()); + assert_eq!(result[2].role, crate::llm::Role::Tool); + assert_eq!(result[3].content, "Found X"); + + assert_eq!(result[4].content, "Write it"); + assert!(result[5].tool_calls.is_some()); + assert_eq!(result[6].role, crate::llm::Role::Tool); + assert_eq!(result[7].content, "Written"); + } + + fn make_db_msg(role: &str, content: &str) -> crate::history::ConversationMessage { + crate::history::ConversationMessage { + id: uuid::Uuid::new_v4(), + role: role.to_string(), + content: content.to_string(), + created_at: chrono::Utc::now(), + } + } +} diff --git a/src/channels/web/util.rs b/src/channels/web/util.rs index 81485b94..060afeab 100644 --- a/src/channels/web/util.rs +++ b/src/channels/web/util.rs @@ -3,6 +3,10 @@ use crate::channels::web::types::{ToolCallInfo, TurnInfo}; /// Truncate a string to at most `max_bytes` bytes at a char boundary, appending "...". +/// +/// If the input is wrapped in `…` and truncation +/// removes the closing tag, the tag is re-appended so downstream XML parsers +/// never see an unclosed element. pub fn truncate_preview(s: &str, max_bytes: usize) -> String { if s.len() <= max_bytes { return s.to_string(); @@ -12,7 +16,14 @@ pub fn truncate_preview(s: &str, max_bytes: usize) -> String { while end > 0 && !s.is_char_boundary(end) { end -= 1; } - format!("{}...", &s[..end]) + let mut result = format!("{}...", &s[..end]); + + // Re-close if truncation cut through the closing tag. + if s.starts_with("") { + result.push_str("\n"); + } + + result } /// Build TurnInfo pairs from flat DB messages (user/tool_calls/assistant triples). @@ -162,6 +173,33 @@ mod tests { assert_eq!(truncate_preview("hello", 0), "..."); } + #[test] + fn test_truncate_preview_closes_tool_output_tag() { + let s = "\nSome very long content here\n"; + // Truncate so it cuts before the closing tag + let result = truncate_preview(s, 60); + assert!(result.ends_with("")); + assert!(result.contains("...")); + } + + #[test] + fn test_truncate_preview_no_extra_close_when_intact() { + let s = "\nshort\n"; + // The string is short enough not to be truncated + let result = truncate_preview(s, 500); + assert_eq!(result, s); + // Should not have a duplicate closing tag + assert_eq!(result.matches("").count(), 1); + } + + #[test] + fn test_truncate_preview_non_xml_unaffected() { + let s = "Just a plain long string that gets truncated"; + let result = truncate_preview(s, 10); + assert_eq!(result, "Just a pla..."); + assert!(!result.contains("")); + } + // ---- build_turns_from_db_messages tests ---- fn make_msg(role: &str, content: &str, offset_ms: i64) -> crate::history::ConversationMessage {