From 7ed641941978520e3dc1df718bc2e513931dd966 Mon Sep 17 00:00:00 2001 From: Ben Gao Date: Thu, 18 Jun 2026 15:24:14 +0800 Subject: [PATCH 1/2] feat(tui): preserve thinking/tool blocks when seeding thread from session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the text-only `seed_thread_from_messages` with a block-type-aware implementation that preserves ContentBlock variants (Thinking, ToolUse, ToolResult) as distinct TurnItem entries, so `loadHistory` / `messages_from_thread_detail` can reconstruct the full conversation including process information. - Add `session_id` field to `ThreadRecord` for session→thread linking - Add `SeedItem` / `TurnSeed` helpers for block-type turn construction - Add `set_thread_session_id` to associate a session with a thread - Rewrite `messages_from_thread_detail` to emit typed ContentBlocks - Link session→thread on both `resume_session` and `create_session_from_thread` - Fall back to turn-based reconstruction when session file is missing --- crates/tui/src/runtime_api.rs | 157 ++++++++-- crates/tui/src/runtime_threads.rs | 493 +++++++++++++++++++++++++----- 2 files changed, 552 insertions(+), 98 deletions(-) diff --git a/crates/tui/src/runtime_api.rs b/crates/tui/src/runtime_api.rs index ec02d8eeea..93adde9909 100644 --- a/crates/tui/src/runtime_api.rs +++ b/crates/tui/src/runtime_api.rs @@ -940,6 +940,16 @@ async fn resume_session_thread( .await .map_err(|e| ApiError::internal(format!("Failed to seed thread history: {e}")))?; + // Link the session to the new thread so that `ensure_engine_loaded` + // can restore the full message history from the session file. + if let Err(e) = state + .runtime_threads + .set_thread_session_id(&thread.id, &id) + .await + { + tracing::warn!("Failed to link session {id} to thread {}: {e}", thread.id); + } + let summary = format!( "Resumed session '{}' ({} messages) into thread {}", session.metadata.title, msg_count, thread.id @@ -1014,6 +1024,19 @@ async fn create_session_from_thread( .save_session(&session) .map_err(|e| ApiError::internal(format!("Failed to save session: {e}")))?; + // Link the session to the thread so that `ensure_engine_loaded` can + // restore the full message history from the session file. + if let Err(e) = state + .runtime_threads + .set_thread_session_id(&detail.thread.id, &session_id) + .await + { + tracing::warn!( + "Failed to link session {session_id} to thread {}: {e}", + detail.thread.id + ); + } + Ok(( StatusCode::CREATED, Json(CreateSessionResponse { @@ -1048,29 +1071,115 @@ fn messages_from_thread_detail(detail: &ThreadDetail) -> Vec { let mut messages = Vec::new(); for turn in &detail.turns { + // Collect content blocks for the current assistant message. + // Multiple items (AgentMessage, AgentReasoning, ToolCall) may + // belong to the same assistant message, so we batch them. + let mut assistant_blocks: Vec = Vec::new(); + let flush_assistant = |blocks: &mut Vec, msgs: &mut Vec| { + if !blocks.is_empty() { + msgs.push(Message { + role: "assistant".to_string(), + content: std::mem::take(blocks), + }); + } + }; + for item_id in &turn.item_ids { let Some(item) = items_by_id.get(item_id.as_str()) else { continue; }; - let role = match item.kind { - TurnItemKind::UserMessage => "user", - TurnItemKind::AgentMessage => "assistant", - _ => continue, - }; - let Some(text) = item.detail.as_deref().map(str::trim) else { - continue; - }; - if text.is_empty() { - continue; + match item.kind { + TurnItemKind::UserMessage => { + // Flush any pending assistant blocks before starting a + // new user message. + flush_assistant(&mut assistant_blocks, &mut messages); + + let text = item.detail.as_deref().map(str::trim).unwrap_or(""); + if !text.is_empty() { + messages.push(Message { + role: "user".to_string(), + content: vec![ContentBlock::Text { + text: text.to_string(), + cache_control: None, + }], + }); + } + } + TurnItemKind::AgentMessage => { + let text = item.detail.as_deref().map(str::trim).unwrap_or(""); + if !text.is_empty() { + assistant_blocks.push(ContentBlock::Text { + text: text.to_string(), + cache_control: None, + }); + } + } + TurnItemKind::AgentReasoning => { + let thinking = item.detail.as_deref().map(str::trim).unwrap_or(""); + if !thinking.is_empty() { + assistant_blocks.push(ContentBlock::Thinking { + thinking: thinking.to_string(), + signature: None, + }); + } + } + TurnItemKind::ToolCall => { + // Check metadata to distinguish tool_use from tool_result. + let meta = item.metadata.as_ref(); + let is_tool_result = meta.and_then(|m| m.get("tool_result_for")).is_some(); + if is_tool_result { + // tool_result blocks go in a user message. + // Flush any pending assistant blocks first. + flush_assistant(&mut assistant_blocks, &mut messages); + + let tool_use_id = meta + .and_then(|m| m.get("tool_result_for")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let content = item.detail.as_deref().unwrap_or("").to_string(); + let is_error = meta + .and_then(|m| m.get("is_error")) + .and_then(|v| v.as_bool()) + .unwrap_or(false); + messages.push(Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id, + content, + is_error: if is_error { Some(true) } else { None }, + content_blocks: None, + }], + }); + } else { + // tool_use block — part of assistant message. + let tool_use_id = meta + .and_then(|m| m.get("tool_use_id")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let tool_name = meta + .and_then(|m| m.get("tool_name")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let input_str = item.detail.as_deref().unwrap_or("{}"); + let input: serde_json::Value = + serde_json::from_str(input_str).unwrap_or(serde_json::Value::Null); + assistant_blocks.push(ContentBlock::ToolUse { + id: tool_use_id, + name: tool_name, + input, + caller: None, + }); + } + } + // Skip other item kinds (file_change, command_execution, etc.) + _ => {} } - messages.push(Message { - role: role.to_string(), - content: vec![ContentBlock::Text { - text: text.to_string(), - cache_control: None, - }], - }); } + // Flush any remaining assistant blocks. + flush_assistant(&mut assistant_blocks, &mut messages); } messages @@ -1193,8 +1302,20 @@ async fn save_current_session( .save_session(&session) .map_err(|e| ApiError::internal(format!("Failed to save session: {e}")))?; + // Link the session to the thread so that `ensure_engine_loaded` can + // restore the full message history (including thinking/tool blocks) + // from the session file instead of reconstructing from turns. + let session_id = session.metadata.id.clone(); + if let Err(e) = state + .runtime_threads + .set_thread_session_id(&thread_id, &session_id) + .await + { + tracing::warn!("Failed to link session {session_id} to thread {thread_id}: {e}"); + } + Ok(Json(SaveSessionResponse { - session_id: session.metadata.id.clone(), + session_id, session: session_to_detail(session), })) } diff --git a/crates/tui/src/runtime_threads.rs b/crates/tui/src/runtime_threads.rs index 0555e9be8a..81d2df35ca 100644 --- a/crates/tui/src/runtime_threads.rs +++ b/crates/tui/src/runtime_threads.rs @@ -151,6 +151,11 @@ pub struct ThreadRecord { /// additive metadata — older readers ignore it without misinterpretation. #[serde(default, skip_serializing_if = "Option::is_none")] pub title: Option, + /// The session ID associated with this thread. When set, `ensure_engine_loaded` + /// loads the full message history (including thinking/tool blocks) from the + /// session file instead of reconstructing from turns (which loses process info). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_id: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -794,6 +799,31 @@ pub struct RuntimeThreadManager { pending_dynamic_tools: Arc>>>, } +/// Helper types for `seed_thread_from_messages` — intermediate representation +/// of a turn being built from session messages before persisting as items. +/// +/// A single content block extracted from an assistant message. +enum SeedItem { + Text(String), + Thinking(String), + ToolUse { + id: String, + name: String, + input: serde_json::Value, + }, + ToolResult { + tool_use_id: String, + content: String, + is_error: bool, + }, +} + +/// A turn being assembled from session messages. +struct TurnSeed { + user_text: String, + items: Vec, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RuntimeApprovalDecision { ApproveTool, @@ -1078,6 +1108,7 @@ impl RuntimeThreadManager { system_prompt: req.system_prompt, task_id: req.task_id, title: None, + session_id: None, }; self.store.save_thread(&thread)?; self.emit_event( @@ -1324,6 +1355,28 @@ impl RuntimeThreadManager { Ok(thread) } + /// Link a session to a thread so that `ensure_engine_loaded` can restore + /// the full message history (including thinking/tool blocks) from the + /// session file instead of reconstructing from turns. + pub async fn set_thread_session_id(&self, thread_id: &str, session_id: &str) -> Result<()> { + let mut thread = self.get_thread(thread_id).await?; + if thread.session_id.as_deref() == Some(session_id) { + return Ok(()); + } + thread.session_id = Some(session_id.to_string()); + thread.updated_at = Utc::now(); + self.store.save_thread(&thread)?; + self.emit_event( + thread_id, + None, + None, + "thread.updated", + json!({ "thread": thread, "changes": { "session_id": session_id } }), + ) + .await?; + Ok(()) + } + async fn ensure_thread_has_no_active_turn(&self, thread_id: &str) -> Result<()> { let active = self.active.lock().await; if active @@ -1559,6 +1612,11 @@ impl RuntimeThreadManager { /// Seed a thread with messages from a saved session so subsequent turns /// continue with the prior conversation context. + /// + /// Unlike the old text-only implementation, this preserves all content + /// block types (thinking, tool_use, tool_result, etc.) as separate turn + /// items so that `loadHistory` in the GUI can reconstruct the full + /// conversation including process information. pub async fn seed_thread_from_messages( &self, thread_id: &str, @@ -1567,44 +1625,110 @@ impl RuntimeThreadManager { let mut thread = self.get_thread(thread_id).await?; let now = Utc::now(); - let mut user_buf: Vec = Vec::new(); - let mut pending_pairs: Vec<(String, Option)> = Vec::new(); + // Group messages into turns. A turn starts with a user message and + // includes all subsequent assistant messages (which may contain + // thinking, tool_use, tool_result blocks) until the next user message. + let mut turns: Vec = Vec::new(); + let mut current_turn: Option = None; for msg in messages { - let text = msg - .content - .iter() - .filter_map(|block| match block { - ContentBlock::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect::>() - .join("\n"); - if text.trim().is_empty() { - continue; - } - if msg.role == "user" { - user_buf.push(text); - } else if msg.role == "assistant" { - let user_text = if user_buf.is_empty() { - String::new() - } else { - std::mem::take(&mut user_buf).join("\n") - }; - pending_pairs.push((user_text, Some(text))); + match msg.role.as_str() { + "user" => { + // Flush any pending turn before starting a new one. + if let Some(t) = current_turn.take() { + turns.push(t); + } + let mut turn = TurnSeed { + user_text: String::new(), + items: Vec::new(), + }; + // Extract text from user message content blocks. + // Tool result blocks in user messages are part of the + // tool loop and should be stored as tool_call items. + for block in &msg.content { + match block { + ContentBlock::Text { text, .. } if !text.trim().is_empty() => { + if !turn.user_text.is_empty() { + turn.user_text.push('\n'); + } + turn.user_text.push_str(text); + } + ContentBlock::ToolResult { + tool_use_id, + content, + is_error, + .. + } => { + turn.items.push(SeedItem::ToolResult { + tool_use_id: tool_use_id.clone(), + content: content.clone(), + is_error: is_error.unwrap_or(false), + }); + } + // Other block types in user messages are rare; + // skip them gracefully. + _ => {} + } + } + current_turn = Some(turn); + } + "assistant" => { + // If no current turn exists (e.g. session starts with + // an assistant message), create a placeholder turn. + let turn = current_turn.get_or_insert_with(|| TurnSeed { + user_text: String::new(), + items: Vec::new(), + }); + for block in &msg.content { + match block { + ContentBlock::Text { text, .. } if !text.trim().is_empty() => { + turn.items.push(SeedItem::Text(text.clone())); + } + ContentBlock::Thinking { thinking, .. } + if !thinking.trim().is_empty() => + { + turn.items.push(SeedItem::Thinking(thinking.clone())); + } + ContentBlock::ToolUse { + id, name, input, .. + } => { + turn.items.push(SeedItem::ToolUse { + id: id.clone(), + name: name.clone(), + input: input.clone(), + }); + } + ContentBlock::ServerToolUse { + id, name, input, .. + } => { + turn.items.push(SeedItem::ToolUse { + id: id.clone(), + name: name.clone(), + input: input.clone(), + }); + } + // Skip other block types (image_url, etc.) + _ => {} + } + } + } + // System messages and other roles are ignored for turn seeding. + _ => {} } } - if !user_buf.is_empty() { - let user_text = std::mem::take(&mut user_buf).join("\n"); - pending_pairs.push((user_text, None)); + // Flush the last turn. + if let Some(t) = current_turn.take() { + turns.push(t); } - for (user_text, assistant_text) in pending_pairs { + for turn_seed in turns { let turn_id = format!("turn_{}", &Uuid::new_v4().to_string()[..8]); - let summary = crate::utils::truncate_with_ellipsis(&user_text, SUMMARY_LIMIT, "..."); + let summary = + crate::utils::truncate_with_ellipsis(&turn_seed.user_text, SUMMARY_LIMIT, "..."); let mut item_ids = Vec::new(); - if !user_text.is_empty() { + // Save user message item. + if !turn_seed.user_text.is_empty() { let item_id = format!("item_{}", &Uuid::new_v4().to_string()[..8]); self.store.save_item(&TurnItemRecord { schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, @@ -1613,7 +1737,7 @@ impl RuntimeThreadManager { kind: TurnItemKind::UserMessage, status: TurnItemLifecycleStatus::Completed, summary: summary.clone(), - detail: Some(user_text), + detail: Some(turn_seed.user_text.clone()), metadata: None, artifact_refs: Vec::new(), started_at: Some(now), @@ -1622,47 +1746,148 @@ impl RuntimeThreadManager { item_ids.push(item_id); } - if let Some(assistant_text) = assistant_text { - let asst_summary = if assistant_text.len() > SUMMARY_LIMIT { - crate::utils::truncate_with_ellipsis(&assistant_text, SUMMARY_LIMIT, "...") - } else { - assistant_text.clone() - }; + // Save assistant content items in order. + for seed_item in &turn_seed.items { let item_id = format!("item_{}", &Uuid::new_v4().to_string()[..8]); - self.store.save_item(&TurnItemRecord { + match seed_item { + SeedItem::Text(text) => { + let asst_summary = if text.len() > SUMMARY_LIMIT { + crate::utils::truncate_with_ellipsis(text, SUMMARY_LIMIT, "...") + } else { + text.clone() + }; + self.store.save_item(&TurnItemRecord { + schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, + id: item_id.clone(), + turn_id: turn_id.clone(), + kind: TurnItemKind::AgentMessage, + status: TurnItemLifecycleStatus::Completed, + summary: asst_summary, + detail: Some(text.clone()), + metadata: None, + artifact_refs: Vec::new(), + started_at: Some(now), + ended_at: Some(now), + })?; + } + SeedItem::Thinking(thinking) => { + let thinking_summary = if thinking.len() > SUMMARY_LIMIT { + crate::utils::truncate_with_ellipsis(thinking, SUMMARY_LIMIT, "...") + } else { + thinking.clone() + }; + self.store.save_item(&TurnItemRecord { + schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, + id: item_id.clone(), + turn_id: turn_id.clone(), + kind: TurnItemKind::AgentReasoning, + status: TurnItemLifecycleStatus::Completed, + summary: thinking_summary, + detail: Some(thinking.clone()), + metadata: None, + artifact_refs: Vec::new(), + started_at: Some(now), + ended_at: Some(now), + })?; + } + SeedItem::ToolUse { + id: tool_id, + name, + input, + } => { + let input_str = + serde_json::to_string(input).unwrap_or_else(|_| input.to_string()); + let tool_summary = format!("{name}({})", { + let s = &input_str; + if s.len() > 80 { + crate::utils::truncate_with_ellipsis(s, 80, "...") + } else { + s.clone() + } + }); + self.store.save_item(&TurnItemRecord { + schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, + id: item_id.clone(), + turn_id: turn_id.clone(), + kind: TurnItemKind::ToolCall, + status: TurnItemLifecycleStatus::Completed, + summary: tool_summary, + detail: Some(input_str), + metadata: Some(serde_json::Value::Object( + serde_json::json!({ + "tool_use_id": tool_id, + "tool_name": name, + }) + .as_object() + .unwrap() + .clone(), + )), + artifact_refs: Vec::new(), + started_at: Some(now), + ended_at: Some(now), + })?; + } + SeedItem::ToolResult { + tool_use_id, + content, + is_error, + } => { + let result_summary = if content.len() > SUMMARY_LIMIT { + crate::utils::truncate_with_ellipsis(content, SUMMARY_LIMIT, "...") + } else { + content.clone() + }; + self.store.save_item(&TurnItemRecord { + schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, + id: item_id.clone(), + turn_id: turn_id.clone(), + kind: TurnItemKind::ToolCall, + status: if *is_error { + TurnItemLifecycleStatus::Failed + } else { + TurnItemLifecycleStatus::Completed + }, + summary: result_summary, + detail: Some(content.clone()), + metadata: Some(serde_json::Value::Object( + serde_json::json!({ + "tool_result_for": tool_use_id, + "is_error": is_error, + }) + .as_object() + .unwrap() + .clone(), + )), + artifact_refs: Vec::new(), + started_at: Some(now), + ended_at: Some(now), + })?; + } + } + item_ids.push(item_id); + } + + // Only create a turn if there's content. + if !item_ids.is_empty() { + self.store.save_turn(&TurnRecord { schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, - id: item_id.clone(), - turn_id: turn_id.clone(), - kind: TurnItemKind::AgentMessage, - status: TurnItemLifecycleStatus::Completed, - summary: asst_summary, - detail: Some(assistant_text), - metadata: None, - artifact_refs: Vec::new(), + id: turn_id.clone(), + thread_id: thread_id.to_string(), + status: RuntimeTurnStatus::Completed, + input_summary: summary, + created_at: now, started_at: Some(now), ended_at: Some(now), + duration_ms: Some(0), + usage: None, + error: None, + item_ids, + steer_count: 0, })?; - item_ids.push(item_id); - } - self.store.save_turn(&TurnRecord { - schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, - id: turn_id.clone(), - thread_id: thread_id.to_string(), - status: RuntimeTurnStatus::Completed, - input_summary: summary, - created_at: now, - started_at: Some(now), - ended_at: Some(now), - duration_ms: Some(0), - usage: None, - error: None, - item_ids, - steer_count: 0, - })?; - - thread.latest_turn_id = Some(turn_id); - thread.updated_at = now; + thread.latest_turn_id = Some(turn_id); + thread.updated_at = now; + } } self.store.save_thread(&thread)?; @@ -2226,8 +2451,46 @@ impl RuntimeThreadManager { let engine = spawn_engine(engine_cfg, &self.config); - let turns = self.store.list_turns_for_thread(&thread.id)?; - let session_messages = self.reconstruct_messages_from_turns(&turns)?; + // When the thread has an associated session, load the full message history + // (including thinking/tool blocks) from the session file. This preserves + // process information that `reconstruct_messages_from_turns` would lose. + let session_messages = if let Some(ref sid) = thread.session_id { + match crate::session_manager::default_sessions_dir() { + Ok(sessions_dir) => { + match crate::session_manager::SessionManager::new(sessions_dir) { + Ok(manager) => match manager.load_session(sid) { + Ok(session) => session.messages, + Err(e) => { + tracing::warn!( + "Failed to load session {} for thread {}: {e}; falling back to turn reconstruction", + sid, + thread.id + ); + let turns = self.store.list_turns_for_thread(&thread.id)?; + self.reconstruct_messages_from_turns(&turns)? + } + }, + Err(e) => { + tracing::warn!( + "Failed to open sessions dir: {e}; falling back to turn reconstruction" + ); + let turns = self.store.list_turns_for_thread(&thread.id)?; + self.reconstruct_messages_from_turns(&turns)? + } + } + } + Err(e) => { + tracing::warn!( + "Failed to resolve sessions dir: {e}; falling back to turn reconstruction" + ); + let turns = self.store.list_turns_for_thread(&thread.id)?; + self.reconstruct_messages_from_turns(&turns)? + } + } + } else { + let turns = self.store.list_turns_for_thread(&thread.id)?; + self.reconstruct_messages_from_turns(&turns)? + }; let sys_prompt = thread .system_prompt .as_ref() @@ -2235,7 +2498,7 @@ impl RuntimeThreadManager { if !session_messages.is_empty() || sys_prompt.is_some() { engine .send(Op::SyncSession { - session_id: None, + session_id: thread.session_id.clone(), messages: session_messages, system_prompt: sys_prompt, system_prompt_override: thread.system_prompt.is_some(), @@ -2274,31 +2537,99 @@ impl RuntimeThreadManager { let mut messages = Vec::new(); for turn in turns { let items = self.store.list_items_for_turn(&turn.id)?; + // Collect content blocks for the current assistant message. + let mut assistant_blocks: Vec = Vec::new(); + let flush_assistant = |blocks: &mut Vec, msgs: &mut Vec| { + if !blocks.is_empty() { + msgs.push(Message { + role: "assistant".to_string(), + content: std::mem::take(blocks), + }); + } + }; for item in items { match item.kind { TurnItemKind::UserMessage => { + flush_assistant(&mut assistant_blocks, &mut messages); let text = item.detail.unwrap_or(item.summary); - messages.push(Message { - role: "user".to_string(), - content: vec![ContentBlock::Text { - text, - cache_control: None, - }], - }); + if !text.trim().is_empty() { + messages.push(Message { + role: "user".to_string(), + content: vec![ContentBlock::Text { + text, + cache_control: None, + }], + }); + } } TurnItemKind::AgentMessage => { let text = item.detail.unwrap_or(item.summary); - messages.push(Message { - role: "assistant".to_string(), - content: vec![ContentBlock::Text { + if !text.trim().is_empty() { + assistant_blocks.push(ContentBlock::Text { text, cache_control: None, - }], - }); + }); + } + } + TurnItemKind::AgentReasoning => { + let thinking = item.detail.unwrap_or(item.summary); + if !thinking.trim().is_empty() { + assistant_blocks.push(ContentBlock::Thinking { + thinking, + signature: None, + }); + } + } + TurnItemKind::ToolCall => { + let meta = item.metadata.as_ref(); + let is_tool_result = meta.and_then(|m| m.get("tool_result_for")).is_some(); + if is_tool_result { + flush_assistant(&mut assistant_blocks, &mut messages); + let tool_use_id = meta + .and_then(|m| m.get("tool_result_for")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let content = item.detail.unwrap_or_default(); + let is_error = meta + .and_then(|m| m.get("is_error")) + .and_then(|v| v.as_bool()) + .unwrap_or(false); + messages.push(Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id, + content, + is_error: if is_error { Some(true) } else { None }, + content_blocks: None, + }], + }); + } else { + let tool_use_id = meta + .and_then(|m| m.get("tool_use_id")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let tool_name = meta + .and_then(|m| m.get("tool_name")) + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let input_str = item.detail.unwrap_or_default(); + let input: serde_json::Value = + serde_json::from_str(&input_str).unwrap_or(serde_json::Value::Null); + assistant_blocks.push(ContentBlock::ToolUse { + id: tool_use_id, + name: tool_name, + input, + caller: None, + }); + } } _ => {} } } + flush_assistant(&mut assistant_blocks, &mut messages); } Ok(messages) } @@ -3503,6 +3834,7 @@ mod tests { system_prompt: None, task_id: None, title: None, + session_id: None, } } @@ -5559,6 +5891,7 @@ mod tests { system_prompt: None, task_id: None, title: None, + session_id: None, }; manager.store.save_thread(&thread)?; From 5ce2ebcb38147918475a6301186bf600310c851a Mon Sep 17 00:00:00 2001 From: Ben Gao Date: Fri, 19 Jun 2026 15:10:49 +0800 Subject: [PATCH 2/2] Preserve session process history on thread resume Keep runtime thread resume aligned with the configured sessions directory instead of always reading the default location. Also batch consecutive user tool_result blocks during thread seeding and turn reconstruction, and add regression tests covering parallel tool results plus custom sessions-dir resume. --- crates/tui/src/runtime_threads.rs | 432 +++++++++++++++++++++++++----- 1 file changed, 366 insertions(+), 66 deletions(-) diff --git a/crates/tui/src/runtime_threads.rs b/crates/tui/src/runtime_threads.rs index 81d2df35ca..ecbe73ad56 100644 --- a/crates/tui/src/runtime_threads.rs +++ b/crates/tui/src/runtime_threads.rs @@ -584,6 +584,7 @@ impl RuntimeThreadStore { pub struct RuntimeThreadManagerConfig { pub data_dir: PathBuf, pub task_data_dir: PathBuf, + pub sessions_dir: Option, pub max_active_threads: usize, } @@ -602,6 +603,7 @@ impl RuntimeThreadManagerConfig { Self { data_dir, task_data_dir, + sessions_dir: crate::session_manager::default_sessions_dir().ok(), max_active_threads: MAX_ACTIVE_THREADS_DEFAULT, } } @@ -824,6 +826,14 @@ struct TurnSeed { items: Vec, } +impl TurnSeed { + fn has_assistant_content(&self) -> bool { + self.items + .iter() + .any(|item| !matches!(item, SeedItem::ToolResult { .. })) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RuntimeApprovalDecision { ApproveTool, @@ -1628,20 +1638,28 @@ impl RuntimeThreadManager { // Group messages into turns. A turn starts with a user message and // includes all subsequent assistant messages (which may contain // thinking, tool_use, tool_result blocks) until the next user message. + // + // Consecutive user messages (e.g. parallel tool call results) are + // merged into the same turn to avoid violating LLM provider contracts + // that forbid back-to-back user messages. let mut turns: Vec = Vec::new(); let mut current_turn: Option = None; for msg in messages { match msg.role.as_str() { "user" => { - // Flush any pending turn before starting a new one. - if let Some(t) = current_turn.take() { - turns.push(t); + // Only flush when the current turn has assistant content; + // otherwise merge consecutive user messages into one turn. + if current_turn + .as_ref() + .is_some_and(TurnSeed::has_assistant_content) + { + turns.push(current_turn.take().unwrap()); } - let mut turn = TurnSeed { + let turn = current_turn.get_or_insert_with(|| TurnSeed { user_text: String::new(), items: Vec::new(), - }; + }); // Extract text from user message content blocks. // Tool result blocks in user messages are part of the // tool loop and should be stored as tool_call items. @@ -1670,7 +1688,6 @@ impl RuntimeThreadManager { _ => {} } } - current_turn = Some(turn); } "assistant" => { // If no current turn exists (e.g. session starts with @@ -1721,14 +1738,23 @@ impl RuntimeThreadManager { turns.push(t); } - for turn_seed in turns { + for (offset, turn_seed) in turns.into_iter().enumerate() { + let created_at = now + chrono::Duration::seconds(offset as i64); let turn_id = format!("turn_{}", &Uuid::new_v4().to_string()[..8]); let summary = crate::utils::truncate_with_ellipsis(&turn_seed.user_text, SUMMARY_LIMIT, "..."); let mut item_ids = Vec::new(); + let mut item_offset_ms = 0_i64; + + let mut next_item_time = || { + let ts = created_at + chrono::Duration::milliseconds(item_offset_ms); + item_offset_ms += 1; + ts + }; // Save user message item. if !turn_seed.user_text.is_empty() { + let item_time = next_item_time(); let item_id = format!("item_{}", &Uuid::new_v4().to_string()[..8]); self.store.save_item(&TurnItemRecord { schema_version: CURRENT_RUNTIME_SCHEMA_VERSION, @@ -1740,14 +1766,15 @@ impl RuntimeThreadManager { detail: Some(turn_seed.user_text.clone()), metadata: None, artifact_refs: Vec::new(), - started_at: Some(now), - ended_at: Some(now), + started_at: Some(item_time), + ended_at: Some(item_time), })?; item_ids.push(item_id); } // Save assistant content items in order. for seed_item in &turn_seed.items { + let item_time = next_item_time(); let item_id = format!("item_{}", &Uuid::new_v4().to_string()[..8]); match seed_item { SeedItem::Text(text) => { @@ -1766,8 +1793,8 @@ impl RuntimeThreadManager { detail: Some(text.clone()), metadata: None, artifact_refs: Vec::new(), - started_at: Some(now), - ended_at: Some(now), + started_at: Some(item_time), + ended_at: Some(item_time), })?; } SeedItem::Thinking(thinking) => { @@ -1786,8 +1813,8 @@ impl RuntimeThreadManager { detail: Some(thinking.clone()), metadata: None, artifact_refs: Vec::new(), - started_at: Some(now), - ended_at: Some(now), + started_at: Some(item_time), + ended_at: Some(item_time), })?; } SeedItem::ToolUse { @@ -1823,8 +1850,8 @@ impl RuntimeThreadManager { .clone(), )), artifact_refs: Vec::new(), - started_at: Some(now), - ended_at: Some(now), + started_at: Some(item_time), + ended_at: Some(item_time), })?; } SeedItem::ToolResult { @@ -1859,8 +1886,8 @@ impl RuntimeThreadManager { .clone(), )), artifact_refs: Vec::new(), - started_at: Some(now), - ended_at: Some(now), + started_at: Some(item_time), + ended_at: Some(item_time), })?; } } @@ -1875,9 +1902,9 @@ impl RuntimeThreadManager { thread_id: thread_id.to_string(), status: RuntimeTurnStatus::Completed, input_summary: summary, - created_at: now, - started_at: Some(now), - ended_at: Some(now), + created_at, + started_at: Some(created_at), + ended_at: Some(created_at), duration_ms: Some(0), usage: None, error: None, @@ -2455,38 +2482,31 @@ impl RuntimeThreadManager { // (including thinking/tool blocks) from the session file. This preserves // process information that `reconstruct_messages_from_turns` would lose. let session_messages = if let Some(ref sid) = thread.session_id { - match crate::session_manager::default_sessions_dir() { - Ok(sessions_dir) => { - match crate::session_manager::SessionManager::new(sessions_dir) { - Ok(manager) => match manager.load_session(sid) { - Ok(session) => session.messages, - Err(e) => { - tracing::warn!( - "Failed to load session {} for thread {}: {e}; falling back to turn reconstruction", - sid, - thread.id - ); - let turns = self.store.list_turns_for_thread(&thread.id)?; - self.reconstruct_messages_from_turns(&turns)? - } - }, - Err(e) => { - tracing::warn!( - "Failed to open sessions dir: {e}; falling back to turn reconstruction" - ); - let turns = self.store.list_turns_for_thread(&thread.id)?; - self.reconstruct_messages_from_turns(&turns)? - } - } - } - Err(e) => { - tracing::warn!( - "Failed to resolve sessions dir: {e}; falling back to turn reconstruction" - ); - let turns = self.store.list_turns_for_thread(&thread.id)?; - self.reconstruct_messages_from_turns(&turns)? - } - } + let fallback = || -> Result> { + let turns = self.store.list_turns_for_thread(&thread.id)?; + self.reconstruct_messages_from_turns(&turns) + }; + + self.manager_cfg + .sessions_dir + .clone() + .ok_or_else(|| anyhow!("No sessions dir configured")) + .and_then(|dir| { + crate::session_manager::SessionManager::new(dir).map_err(anyhow::Error::from) + }) + .and_then(|manager| manager.load_session(sid).map_err(anyhow::Error::from)) + .map(|session| session.messages) + .map_or_else( + |e| { + tracing::warn!( + "Failed to load session {sid} for thread {}: {e}; \ + falling back to turn reconstruction", + thread.id, + ); + fallback() + }, + Ok, + )? } else { let turns = self.store.list_turns_for_thread(&thread.id)?; self.reconstruct_messages_from_turns(&turns)? @@ -2537,8 +2557,17 @@ impl RuntimeThreadManager { let mut messages = Vec::new(); for turn in turns { let items = self.store.list_items_for_turn(&turn.id)?; + let mut user_blocks: Vec = Vec::new(); // Collect content blocks for the current assistant message. let mut assistant_blocks: Vec = Vec::new(); + let flush_user = |blocks: &mut Vec, msgs: &mut Vec| { + if !blocks.is_empty() { + msgs.push(Message { + role: "user".to_string(), + content: std::mem::take(blocks), + }); + } + }; let flush_assistant = |blocks: &mut Vec, msgs: &mut Vec| { if !blocks.is_empty() { msgs.push(Message { @@ -2553,16 +2582,14 @@ impl RuntimeThreadManager { flush_assistant(&mut assistant_blocks, &mut messages); let text = item.detail.unwrap_or(item.summary); if !text.trim().is_empty() { - messages.push(Message { - role: "user".to_string(), - content: vec![ContentBlock::Text { - text, - cache_control: None, - }], + user_blocks.push(ContentBlock::Text { + text, + cache_control: None, }); } } TurnItemKind::AgentMessage => { + flush_user(&mut user_blocks, &mut messages); let text = item.detail.unwrap_or(item.summary); if !text.trim().is_empty() { assistant_blocks.push(ContentBlock::Text { @@ -2572,6 +2599,7 @@ impl RuntimeThreadManager { } } TurnItemKind::AgentReasoning => { + flush_user(&mut user_blocks, &mut messages); let thinking = item.detail.unwrap_or(item.summary); if !thinking.trim().is_empty() { assistant_blocks.push(ContentBlock::Thinking { @@ -2595,16 +2623,14 @@ impl RuntimeThreadManager { .and_then(|m| m.get("is_error")) .and_then(|v| v.as_bool()) .unwrap_or(false); - messages.push(Message { - role: "user".to_string(), - content: vec![ContentBlock::ToolResult { - tool_use_id, - content, - is_error: if is_error { Some(true) } else { None }, - content_blocks: None, - }], + user_blocks.push(ContentBlock::ToolResult { + tool_use_id, + content, + is_error: if is_error { Some(true) } else { None }, + content_blocks: None, }); } else { + flush_user(&mut user_blocks, &mut messages); let tool_use_id = meta .and_then(|m| m.get("tool_use_id")) .and_then(|v| v.as_str()) @@ -2630,6 +2656,7 @@ impl RuntimeThreadManager { } } flush_assistant(&mut assistant_blocks, &mut messages); + flush_user(&mut user_blocks, &mut messages); } Ok(messages) } @@ -3790,6 +3817,8 @@ mod tests { use super::*; use crate::core::engine::{MockApprovalEvent, mock_engine_handle}; use crate::core::events::{Event as EngineEvent, TurnOutcomeStatus}; + use crate::session_manager::{SessionManager, create_saved_session_with_id_and_mode}; + use std::fs; use std::time::{Duration, Instant}; use tokio::sync::oneshot; use tokio::time::sleep; @@ -3803,6 +3832,19 @@ mod tests { RuntimeThreadManagerConfig { task_data_dir: data_dir.clone(), data_dir, + sessions_dir: None, + max_active_threads: 4, + } + } + + fn test_manager_config_with_sessions_dir( + data_dir: PathBuf, + sessions_dir: PathBuf, + ) -> RuntimeThreadManagerConfig { + RuntimeThreadManagerConfig { + task_data_dir: data_dir.clone(), + data_dir, + sessions_dir: Some(sessions_dir), max_active_threads: 4, } } @@ -3815,6 +3857,17 @@ mod tests { ) } + fn test_manager_with_sessions_dir( + data_dir: PathBuf, + sessions_dir: PathBuf, + ) -> Result { + RuntimeThreadManager::open( + Config::default(), + PathBuf::from("."), + test_manager_config_with_sessions_dir(data_dir, sessions_dir), + ) + } + fn sample_thread(thread_id: &str) -> ThreadRecord { let now = Utc::now(); ThreadRecord { @@ -4908,6 +4961,253 @@ mod tests { Ok(()) } + #[tokio::test] + async fn seed_thread_from_messages_keeps_parallel_tool_results_in_one_turn() -> Result<()> { + let manager = test_manager(test_runtime_dir())?; + let thread = manager + .create_thread(CreateThreadRequest { + model: None, + workspace: None, + mode: None, + allow_shell: None, + trust_mode: None, + auto_approve: None, + archived: false, + system_prompt: None, + task_id: None, + ..Default::default() + }) + .await?; + + let messages = vec![ + Message { + role: "user".to_string(), + content: vec![ContentBlock::Text { + text: "Calculate both files".to_string(), + cache_control: None, + }], + }, + Message { + role: "assistant".to_string(), + content: vec![ + ContentBlock::ToolUse { + id: "tool_a".to_string(), + name: "read_file".to_string(), + input: serde_json::json!({"path":"a.txt"}), + caller: None, + }, + ContentBlock::ToolUse { + id: "tool_b".to_string(), + name: "read_file".to_string(), + input: serde_json::json!({"path":"b.txt"}), + caller: None, + }, + ], + }, + Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id: "tool_a".to_string(), + content: "A".to_string(), + is_error: None, + content_blocks: None, + }], + }, + Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id: "tool_b".to_string(), + content: "B".to_string(), + is_error: None, + content_blocks: None, + }], + }, + Message { + role: "assistant".to_string(), + content: vec![ContentBlock::Text { + text: "Done".to_string(), + cache_control: None, + }], + }, + ]; + + manager + .seed_thread_from_messages(&thread.id, &messages) + .await?; + + let detail = manager.get_thread_detail(&thread.id).await?; + assert_eq!(detail.turns.len(), 2); + let merged_tool_result_turn_found = + detail.turns.iter().try_fold(false, |found, turn| { + let items = manager.store.list_items_for_turn(&turn.id)?; + let tool_result_count = items + .iter() + .filter(|item| { + item.metadata + .as_ref() + .and_then(|meta| meta.get("tool_result_for")) + .is_some() + }) + .count(); + let has_done = items.iter().any(|item| { + item.kind == TurnItemKind::AgentMessage + && item.detail.as_deref() == Some("Done") + }); + Ok::(found || (tool_result_count == 2 && has_done)) + })?; + assert!(merged_tool_result_turn_found); + Ok(()) + } + + #[tokio::test] + async fn reconstruct_messages_from_turns_batches_parallel_tool_results() -> Result<()> { + let manager = test_manager(test_runtime_dir())?; + let thread = manager + .create_thread(CreateThreadRequest { + model: None, + workspace: None, + mode: None, + allow_shell: None, + trust_mode: None, + auto_approve: None, + archived: false, + system_prompt: None, + task_id: None, + ..Default::default() + }) + .await?; + + let messages = vec![ + Message { + role: "user".to_string(), + content: vec![ContentBlock::Text { + text: "Calculate both files".to_string(), + cache_control: None, + }], + }, + Message { + role: "assistant".to_string(), + content: vec![ + ContentBlock::ToolUse { + id: "tool_a".to_string(), + name: "read_file".to_string(), + input: serde_json::json!({"path":"a.txt"}), + caller: None, + }, + ContentBlock::ToolUse { + id: "tool_b".to_string(), + name: "read_file".to_string(), + input: serde_json::json!({"path":"b.txt"}), + caller: None, + }, + ], + }, + Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id: "tool_a".to_string(), + content: "A".to_string(), + is_error: None, + content_blocks: None, + }], + }, + Message { + role: "user".to_string(), + content: vec![ContentBlock::ToolResult { + tool_use_id: "tool_b".to_string(), + content: "B".to_string(), + is_error: None, + content_blocks: None, + }], + }, + Message { + role: "assistant".to_string(), + content: vec![ContentBlock::Text { + text: "Done".to_string(), + cache_control: None, + }], + }, + ]; + + manager + .seed_thread_from_messages(&thread.id, &messages) + .await?; + + let turns = manager.store.list_turns_for_thread(&thread.id)?; + let reconstructed = manager.reconstruct_messages_from_turns(&turns)?; + let roles: Vec<&str> = reconstructed.iter().map(|msg| msg.role.as_str()).collect(); + assert_eq!(roles, vec!["user", "assistant", "user", "assistant"]); + assert_eq!(reconstructed[2].content.len(), 2); + assert!( + reconstructed[2] + .content + .iter() + .all(|block| matches!(block, ContentBlock::ToolResult { .. })) + ); + Ok(()) + } + + #[tokio::test] + async fn ensure_engine_loaded_uses_configured_sessions_dir() -> Result<()> { + let data_dir = test_runtime_dir(); + let sessions_dir = data_dir.join("sessions-custom"); + fs::create_dir_all(&sessions_dir)?; + let manager = test_manager_with_sessions_dir(data_dir, sessions_dir.clone())?; + let thread = manager + .create_thread(CreateThreadRequest { + model: None, + workspace: None, + mode: None, + allow_shell: None, + trust_mode: None, + auto_approve: None, + archived: false, + system_prompt: None, + task_id: None, + ..Default::default() + }) + .await?; + + let messages = vec![ + Message { + role: "user".to_string(), + content: vec![ContentBlock::Text { + text: "restore me".to_string(), + cache_control: None, + }], + }, + Message { + role: "assistant".to_string(), + content: vec![ContentBlock::Thinking { + thinking: "hidden chain".to_string(), + signature: None, + }], + }, + ]; + let session = create_saved_session_with_id_and_mode( + "sess_custom".to_string(), + &messages, + DEFAULT_TEXT_MODEL, + Path::new("."), + 42, + None, + Some("agent"), + ); + SessionManager::new(sessions_dir)?.save_session(&session)?; + manager + .set_thread_session_id(&thread.id, &session.metadata.id) + .await?; + + let engine = manager.get_engine(&thread.id).await?; + let snapshot = engine.get_session_snapshot().await?; + assert_eq!(snapshot.messages.len(), 2); + assert!(matches!( + &snapshot.messages[1].content[0], + ContentBlock::Thinking { thinking, .. } if thinking == "hidden chain" + )); + Ok(()) + } + #[tokio::test] async fn interrupt_turn_marks_interrupted_after_cleanup() -> Result<()> { let manager = test_manager(test_runtime_dir())?;