From 1512fc5bdc3f67b6a4c794486264f2c4cb436f27 Mon Sep 17 00:00:00 2001 From: Cameron Cordes Date: Fri, 7 Aug 2026 19:19:20 -0400 Subject: [PATCH] Add conversation branching to insight chat (rewind preserves forks) Store chat history as a tree in training_messages instead of a flat array. Rewind now sets active_leaf_id to the target node rather than truncating, so discarded paths survive as alternate branches. Fork indicators ("X/Y") mark divergence points and let the client load or switch to alternate paths. - llm_client: StoredChatNode/ChatHistoryStore with tree traversal helpers (path_to_leaf, children_of, fork_at_node, leaves*, etc.) and from_flat_array for backward-compatible reads of the old flat format. - insight_chat: load_history/chat_turn/rewind_history operate on the tree; new switch_branch and get_branches. render_tree_path walks every node (including non-rendered tool-dispatch nodes) when computing fork info, so a fork whose diverging child is a tool call still surfaces on the following rendered message. - handlers: GET /insights/chat/branches, POST /insights/chat/switch-branch, and a branch_id param on the history endpoint. Backward-compatible: flat arrays are converted to a tree lazily on read. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/ai/handlers.rs | 169 ++++++++++++-- src/ai/insight_chat.rs | 508 ++++++++++++++++++++++++++++++++--------- src/ai/llm_client.rs | 303 ++++++++++++++++++++++++ src/ai/mod.rs | 12 +- src/main.rs | 2 + 5 files changed, 851 insertions(+), 143 deletions(-) diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index ae9f300..6cd4b48 100644 --- a/src/ai/handlers.rs +++ b/src/ai/handlers.rs @@ -1310,6 +1310,11 @@ pub struct ChatHistoryQuery { pub path: String, #[serde(default)] pub library: Option, + /// When set, load the branch anchored at this leaf ID instead of the + /// active branch. Allows the client to preview alternate conversation + /// paths without making them active. + #[serde(default)] + pub branch_id: Option, } #[derive(Debug, Serialize)] @@ -1318,6 +1323,18 @@ pub struct ChatHistoryHttpResponse { pub turn_count: usize, pub model_version: String, pub backend: String, + pub active_leaf_id: u64, + /// The branch leaf ID being viewed. Equals `active_leaf_id` when viewing + /// the active branch, or the `branch_id` query param for alternates. + pub viewing_branch_id: u64, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub fork_info: Vec>, +} + +#[derive(Debug, Clone, Serialize)] +pub struct ChatForkInfo { + pub position: usize, + pub total: usize, } #[derive(Debug, Serialize)] @@ -1414,31 +1431,49 @@ pub async fn chat_history_handler( .flatten() .unwrap_or_else(|| app_state.primary_library()); - match app_state.insight_chat.load_history(library.id, &query.path) { - Ok(view) => HttpResponse::Ok().json(ChatHistoryHttpResponse { - messages: view - .messages + match app_state + .insight_chat + .load_history(library.id, &query.path, query.branch_id) + { + Ok(view) => { + let fork_info: Vec> = view + .fork_info .into_iter() - .map(|m| RenderedHistoryMessage { - role: m.role, - content: m.content, - is_initial: m.is_initial, - tools: m - .tools - .into_iter() - .map(|t| HistoryToolInvocation { - name: t.name, - arguments: t.arguments, - result: t.result, - result_truncated: t.result_truncated, - }) - .collect(), + .map(|f| { + f.map(|fi| ChatForkInfo { + position: fi.position, + total: fi.total, + }) }) - .collect(), - turn_count: view.turn_count, - model_version: view.model_version, - backend: view.backend, - }), + .collect(); + HttpResponse::Ok().json(ChatHistoryHttpResponse { + messages: view + .messages + .into_iter() + .map(|m| RenderedHistoryMessage { + role: m.role, + content: m.content, + is_initial: m.is_initial, + tools: m + .tools + .into_iter() + .map(|t| HistoryToolInvocation { + name: t.name, + arguments: t.arguments, + result: t.result, + result_truncated: t.result_truncated, + }) + .collect(), + }) + .collect(), + turn_count: view.turn_count, + model_version: view.model_version, + backend: view.backend, + active_leaf_id: view.active_leaf_id, + viewing_branch_id: view.viewing_branch_id, + fork_info, + }) + } Err(e) => { let msg = format!("{}", e); if msg.contains("no insight found") { @@ -1452,6 +1487,94 @@ pub async fn chat_history_handler( } } +/// GET /insights/chat/branches — return the list of branch leaf IDs for a +/// photo's conversation tree. Each branch is a leaf node; the active branch +/// is the one matching `active_leaf_id`. Use with `branch_id` on the +/// history endpoint to load an alternate conversation path. +#[get("/insights/chat/branches")] +pub async fn chat_branches_handler( + _claims: Claims, + query: web::Query, + app_state: web::Data, +) -> impl Responder { + let library = libraries::resolve_library_param_state(&app_state, query.library.as_deref()) + .ok() + .flatten() + .unwrap_or_else(|| app_state.primary_library()); + + let (branches, active_leaf_id) = + match app_state.insight_chat.get_branches(library.id, &query.path) { + Ok(result) => result, + Err(e) => { + let msg = format!("{}", e); + if msg.contains("no insight found") { + return HttpResponse::NotFound().json(serde_json::json!({ "error": msg })); + } else if msg.contains("no chat history") { + return HttpResponse::Conflict().json(serde_json::json!({ "error": msg })); + } else { + return HttpResponse::InternalServerError() + .json(serde_json::json!({ "error": msg })); + } + } + }; + + HttpResponse::Ok().json(serde_json::json!({ + "branches": branches, + "active_leaf_id": active_leaf_id, + })) +} + +/// POST /insights/chat/switch-branch — switch the active branch to a +/// different conversation path. The new branch becomes the active +/// conversation; the previous active branch becomes a regular fork. +#[post("/insights/chat/switch-branch")] +pub async fn chat_switch_branch_handler( + _claims: Claims, + request: web::Json, + app_state: web::Data, +) -> impl Responder { + let library = + match libraries::resolve_library_param_state(&app_state, request.library.as_deref()) { + Ok(Some(lib)) => lib, + Ok(None) => app_state.primary_library(), + Err(e) => { + return HttpResponse::BadRequest().json(serde_json::json!({ + "error": format!("invalid library: {}", e) + })); + } + }; + + match app_state + .insight_chat + .switch_branch(library.id, &request.file_path, request.branch_id) + .await + { + Ok(()) => HttpResponse::Ok().json(serde_json::json!({ "success": true })), + Err(e) => { + let msg = format!("{}", e); + log::error!("Switch branch failed: {}", msg); + if msg.contains("no insight found") { + HttpResponse::NotFound().json(serde_json::json!({ "error": msg })) + } else if msg.contains("no chat history") { + HttpResponse::Conflict().json(serde_json::json!({ "error": msg })) + } else if msg.contains("not found") || msg.contains("not a leaf") { + HttpResponse::BadRequest().json(serde_json::json!({ "error": msg })) + } else { + HttpResponse::InternalServerError().json(serde_json::json!({ "error": msg })) + } + } + } +} + +#[derive(Debug, Deserialize)] +pub struct ChatSwitchBranchRequest { + pub file_path: String, + #[serde(default)] + pub library: Option, + /// The leaf node ID to switch to. Must be a valid leaf in the tree. + pub branch_id: u64, +} + /// POST /insights/chat/stream — streaming variant of /insights/chat. /// Returns `text/event-stream` with one event per chat stream event. #[post("/insights/chat/stream")] diff --git a/src/ai/insight_chat.rs b/src/ai/insight_chat.rs index 1a67540..af9acf9 100644 --- a/src/ai/insight_chat.rs +++ b/src/ai/insight_chat.rs @@ -8,7 +8,9 @@ use tokio::sync::Mutex as TokioMutex; use crate::ai::backend::{BackendKind, ResolvedBackend, SamplingOverrides}; use crate::ai::insight_generator::InsightGenerator; -use crate::ai::llm_client::{ChatMessage, LlmStreamEvent, Tool}; +use crate::ai::llm_client::{ + BranchLeafInfo, ChatHistoryStore, ChatMessage, LlmStreamEvent, StoredChatNode, Tool, +}; use crate::ai::turn_registry::TurnEntry; use crate::ai::turn_registry::TurnRegistry; use crate::database::InsightDao; @@ -137,19 +139,20 @@ impl InsightChatService { &self.insight_dao } - /// Load the rendered transcript for chat-UI display. Filters internal - /// scaffolding (system message, tool turns, tool-dispatch-only assistant - /// messages) and drops base64 images from user turns to keep payloads - /// small. The first remaining user message is flagged `is_initial`. + /// Load the rendered transcript for chat-UI display. Deserializes the + /// `training_messages` tree, traverses to `active_leaf_id`, and renders + /// the path from root to that leaf. Backward-compatible: if the stored + /// value is a flat array (old format), converts to a tree automatically. /// /// `library_id` scopes the lookup to one library — without it, a /// regenerate on lib1 can be shadowed on the next refresh by an /// untouched `is_current=true` row in lib2 for the same rel_path. - /// Falls back to the cross-library `get_insight` only when the - /// scoped lookup misses, preserving the cross-library "show this - /// photo's primary insight" merge for the case where the active - /// library has no insight but another library does. - pub fn load_history(&self, library_id: i32, file_path: &str) -> Result { + pub fn load_history( + &self, + library_id: i32, + file_path: &str, + branch_id: Option, + ) -> Result { let normalized = normalize_path(file_path); let cx = opentelemetry::Context::new(); let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); @@ -168,98 +171,32 @@ impl InsightChatService { .training_messages .as_ref() .ok_or_else(|| anyhow!("insight has no chat history (pre-agentic insight)"))?; - let messages: Vec = serde_json::from_str(raw) - .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?; - let mut rendered = Vec::new(); - let mut user_turns_seen = 0usize; - let mut assistant_turns_seen = 0usize; + // Backward-compatible deserialization: flat array (old) vs tree (new). + let store: ChatHistoryStore = if let Ok(arr) = serde_json::from_str::>(raw) + { + ChatHistoryStore::from_flat_array(arr) + } else { + serde_json::from_str(raw) + .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? + }; - // Accumulate tool invocations seen since the last user turn. An - // invocation is: one assistant tool_call message (which may hold - // multiple calls) + the N following tool-role messages (one per call, - // in order). They attach to the next assistant-with-content, which - // is the "final" reply for the current turn. - // - // Wire shape from the model: - // assistant { tool_calls: [A, B], content: "" } - // tool { content: "result of A" } - // tool { content: "result of B" } - // assistant { content: "here's the answer" } ← rendered as final - let mut pending_tools: Vec = Vec::new(); - // Queue of (name, arguments) awaiting a tool_result to pair with. - let mut pending_calls: std::collections::VecDeque<(String, serde_json::Value)> = - std::collections::VecDeque::new(); + // Use branch_id if provided, otherwise use active_leaf_id. + let target_leaf = branch_id.unwrap_or(store.active_leaf_id); + let path = store + .path_to_leaf(target_leaf) + .ok_or_else(|| anyhow!("branch_id {} not found in tree", target_leaf))?; - for msg in &messages { - match msg.role.as_str() { - "system" => continue, - "tool" => { - if let Some((name, arguments)) = pending_calls.pop_front() { - let (result, result_truncated) = truncate_tool_result(&msg.content); - pending_tools.push(ToolInvocation { - name, - arguments, - result, - result_truncated, - }); - } - // If there's no pending call, the tool message is an - // orphan (shouldn't happen in practice) — skip silently. - } - "assistant" => { - let has_tool_calls = msg - .tool_calls - .as_ref() - .map(|c| !c.is_empty()) - .unwrap_or(false); - if has_tool_calls && msg.content.trim().is_empty() { - // Tool-dispatch turn: enqueue calls, wait for tool - // results on subsequent messages. - if let Some(ref tcs) = msg.tool_calls { - for tc in tcs { - pending_calls.push_back(( - tc.function.name.clone(), - tc.function.arguments.clone(), - )); - } - } - continue; - } - // Final assistant reply for this turn — drain accumulated - // tools into it. - assistant_turns_seen += 1; - let tools = std::mem::take(&mut pending_tools); - pending_calls.clear(); // any leftover unpaired calls are dropped - rendered.push(RenderedMessage { - role: "assistant".to_string(), - content: msg.content.clone(), - is_initial: false, - tools, - }); - } - "user" => { - let is_initial = user_turns_seen == 0; - user_turns_seen += 1; - // New user turn resets any in-flight tool state. - pending_tools.clear(); - pending_calls.clear(); - rendered.push(RenderedMessage { - role: "user".to_string(), - content: msg.content.clone(), - is_initial, - tools: Vec::new(), - }); - } - _ => continue, - } - } + let (rendered, turn_count, _node_ids, fork_info) = render_tree_path(&store, &path); Ok(HistoryView { messages: rendered, - turn_count: assistant_turns_seen, + turn_count, model_version: insight.model_version, backend: insight.backend, + active_leaf_id: store.active_leaf_id, + viewing_branch_id: target_leaf, + fork_info, }) } @@ -315,8 +252,21 @@ impl InsightChatService { anyhow!("insight has no chat history; regenerate this insight in agentic mode") })? .clone(); - let mut messages: Vec = serde_json::from_str(&raw_history) - .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?; + + // Backward-compatible: flat array (old) vs tree (new). + let mut store: ChatHistoryStore = + if let Ok(arr) = serde_json::from_str::>(&raw_history) { + ChatHistoryStore::from_flat_array(arr) + } else { + serde_json::from_str(&raw_history) + .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? + }; + + // Build messages from the path to the active leaf. + let path = store + .path_to_leaf(store.active_leaf_id) + .ok_or_else(|| anyhow!("active_leaf_id {} not found in tree", store.active_leaf_id))?; + let mut messages: Vec = path.iter().map(|n| n.message.clone()).collect(); // 3. Resolve effective backend. Reject the unsupported switch. let stored_backend = insight.backend.clone(); @@ -509,8 +459,22 @@ impl InsightChatService { // 9. Persist. Append mode rewrites the JSON blob in place; amend // mode regenerates the title and inserts a new insight row, // relying on store_insight to flip prior rows' is_current=false. - let json = serde_json::to_string(&messages) - .map_err(|e| anyhow!("failed to serialize chat history: {}", e))?; + // Append new messages (this turn's user + assistant exchanges) as + // tree nodes chained from the previous active_leaf_id. + let path_len = path.len(); + let new_messages = messages[path_len..].to_vec(); + let mut parent_id = Some(store.active_leaf_id); + for msg in &new_messages { + let new_id = store.append_node(parent_id, msg.clone()); + parent_id = Some(new_id); + } + // Update active_leaf_id to the last node appended this turn. + if let Some(last_id) = parent_id { + store.active_leaf_id = last_id; + } + + let json = serde_json::to_string(&store) + .map_err(|e| anyhow!("failed to serialize chat tree: {}", e))?; let mut amended_insight_id: Option = None; if req.amend { @@ -597,11 +561,14 @@ impl InsightChatService { }) } - /// Truncate the stored conversation so the rendered message at - /// `discard_from_rendered_index` (and everything after it — including - /// the tool-call scaffolding that produced a discarded assistant reply) - /// is removed. The initial user turn cannot be discarded; attempting to - /// do so returns an error. + /// Rewind the conversation to the rendered message at + /// `discard_from_rendered_index` by setting `active_leaf_id` to the + /// node just before the discarded message. Unlike the legacy flat-array + /// approach, this preserves all fork branches — the discarded path + /// remains accessible as an alternate branch. + /// + /// The initial user turn cannot be discarded; attempting to do so + /// returns an error. /// /// Holds the per-file chat mutex so it serialises with `chat_turn`. pub async fn rewind_history( @@ -636,15 +603,30 @@ impl InsightChatService { .training_messages .as_ref() .ok_or_else(|| anyhow!("insight has no chat history"))?; - let messages: Vec = serde_json::from_str(raw_history) - .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?; - let cut_at = find_raw_cut(&messages, discard_from_rendered_index) + let mut store: ChatHistoryStore = + if let Ok(arr) = serde_json::from_str::>(raw_history) { + ChatHistoryStore::from_flat_array(arr) + } else { + serde_json::from_str(raw_history) + .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? + }; + + let path = store + .path_to_leaf(store.active_leaf_id) + .ok_or_else(|| anyhow!("active_leaf_id not found in tree"))?; + + let (_rendered, _turn_count, node_ids, _fork_info) = render_tree_path(&store, &path); + + // The last kept rendered message is at index `discard_from_rendered_index - 1`. + let last_kept_idx = discard_from_rendered_index - 1; + let new_active_leaf_id = *node_ids + .get(last_kept_idx) .ok_or_else(|| anyhow!("discard_from_rendered_index out of range"))?; - let truncated = &messages[..cut_at]; - let json = serde_json::to_string(truncated) - .map_err(|e| anyhow!("failed to serialize truncated history: {}", e))?; + store.active_leaf_id = new_active_leaf_id; + let json = serde_json::to_string(&store) + .map_err(|e| anyhow!("failed to serialize tree: {}", e))?; let dao = self.insight_dao.clone(); let rows = retry_with_backoff("update_training_messages", 3, || { @@ -652,7 +634,7 @@ impl InsightChatService { let mut d = dao.lock().expect("Unable to lock InsightDao"); d.update_training_messages(&cx, library_id, &normalized, &json) }) - .map_err(|e| anyhow!("failed to persist truncated history: {:?}", e))?; + .map_err(|e| anyhow!("failed to persist rewound history: {:?}", e))?; if rows == 0 { log::warn!( "update_training_messages (rewind) updated 0 rows for {} (lib {}), \ @@ -664,6 +646,108 @@ impl InsightChatService { Ok(()) } + /// Switch the active branch to the given leaf ID. The new branch becomes + /// the active conversation path; the previous active branch becomes a + /// regular fork. If the leaf ID doesn't exist, returns an error. + pub async fn switch_branch( + &self, + library_id: i32, + file_path: &str, + leaf_id: u64, + ) -> Result<()> { + let normalized = normalize_path(file_path); + + let lock_key = (library_id, normalized.clone()); + let entry_lock = { + let mut locks = self.chat_locks.lock().await; + locks + .entry(lock_key.clone()) + .or_insert_with(|| Arc::new(TokioMutex::new(()))) + .clone() + }; + let _guard = entry_lock.lock().await; + + let insight = { + let cx = opentelemetry::Context::new(); + let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); + dao.get_current_insight_for_library(&cx, library_id, &normalized) + .map_err(|e| anyhow!("failed to load insight: {:?}", e))? + .ok_or_else(|| anyhow!("no insight found for path"))? + }; + let raw_history = insight + .training_messages + .as_ref() + .ok_or_else(|| anyhow!("insight has no chat history"))?; + + let mut store: ChatHistoryStore = + if let Ok(arr) = serde_json::from_str::>(raw_history) { + ChatHistoryStore::from_flat_array(arr) + } else { + serde_json::from_str(raw_history) + .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? + }; + + // Validate the leaf_id exists and is a leaf. + if !store.nodes.iter().any(|n| n.id == leaf_id) { + bail!("branch_id {} not found in tree", leaf_id); + } + if !store.children_of(leaf_id).is_empty() { + bail!("branch_id {} is not a leaf node", leaf_id); + } + + store.active_leaf_id = leaf_id; + let json = serde_json::to_string(&store) + .map_err(|e| anyhow!("failed to serialize tree: {}", e))?; + + let dao = self.insight_dao.clone(); + let rows = retry_with_backoff("update_training_messages", 3, || { + let cx = opentelemetry::Context::new(); + let mut d = dao.lock().expect("Unable to lock InsightDao"); + d.update_training_messages(&cx, library_id, &normalized, &json) + }) + .map_err(|e| anyhow!("failed to persist branch switch: {:?}", e))?; + if rows == 0 { + log::warn!( + "update_training_messages (switch_branch) updated 0 rows for {} (lib {}), \ + concurrent regenerate likely flipped is_current", + normalized, + library_id + ); + } + Ok(()) + } + + /// Return branch leaf metadata (id, snippet, message_count) and the + /// active leaf ID for a photo's conversation tree. + pub fn get_branches( + &self, + library_id: i32, + file_path: &str, + ) -> Result<(Vec, u64)> { + let normalized = normalize_path(file_path); + let cx = opentelemetry::Context::new(); + let mut dao = self.insight_dao.lock().expect("Unable to lock InsightDao"); + let insight = dao + .get_current_insight_for_library(&cx, library_id, &normalized) + .map_err(|e| anyhow!("failed to load insight: {:?}", e))? + .ok_or_else(|| anyhow!("no insight found for path"))?; + + let raw = insight + .training_messages + .as_ref() + .ok_or_else(|| anyhow!("insight has no chat history"))?; + + let store: ChatHistoryStore = if let Ok(arr) = serde_json::from_str::>(raw) + { + ChatHistoryStore::from_flat_array(arr) + } else { + serde_json::from_str(raw) + .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? + }; + + Ok((store.leaves_with_info(), store.active_leaf_id)) + } + /// Streaming variant of `chat_turn`. Emits user-facing events as the /// conversation progresses: iteration starts, tool dispatch + result, /// text deltas from the final assistant reply, and a terminal `Done` @@ -828,8 +912,21 @@ impl InsightChatService { let raw_history = insight.training_messages.as_ref().ok_or_else(|| { anyhow!("insight has no chat history; regenerate this insight in agentic mode") })?; - let mut messages: Vec = serde_json::from_str(raw_history) - .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))?; + + // Backward-compatible: flat array (old) vs tree (new). + let mut store: ChatHistoryStore = + if let Ok(arr) = serde_json::from_str::>(raw_history) { + ChatHistoryStore::from_flat_array(arr) + } else { + serde_json::from_str(raw_history) + .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? + }; + + let path = store + .path_to_leaf(store.active_leaf_id) + .ok_or_else(|| anyhow!("active_leaf_id {} not found in tree", store.active_leaf_id))?; + let path_len = path.len(); + let mut messages: Vec = path.iter().map(|n| n.message.clone()).collect(); let stored_backend = insight.backend.clone(); let effective_backend = req @@ -931,8 +1028,19 @@ impl InsightChatService { restore_system_prompt_override(&mut messages, override_stash); } - let json = serde_json::to_string(&messages) - .map_err(|e| anyhow!("failed to serialize chat history: {}", e))?; + // Append new messages as tree nodes. + let new_messages = messages[path_len..].to_vec(); + let mut parent_id = Some(store.active_leaf_id); + for msg in &new_messages { + let new_id = store.append_node(parent_id, msg.clone()); + parent_id = Some(new_id); + } + if let Some(last_id) = parent_id { + store.active_leaf_id = last_id; + } + + let json = serde_json::to_string(&store) + .map_err(|e| anyhow!("failed to serialize chat tree: {}", e))?; let mut amended_insight_id: Option = None; if req.amend { @@ -2164,6 +2272,7 @@ pub enum ChatStreamEvent { /// Is this raw message visible in the rendered transcript? Must match /// `load_history`'s filter exactly — `find_raw_cut` depends on it to map /// rendered indices back to raw positions. +#[allow(dead_code)] fn is_rendered(m: &ChatMessage) -> bool { match m.role.as_str() { "user" => true, @@ -2189,6 +2298,7 @@ fn is_rendered(m: &ChatMessage) -> bool { /// regenerating after a failed turn — its optimistic user bubble lives at /// the index just past the server's persisted history. Strictly past the end /// (`discard > rendered_count`) returns `None`. +#[allow(dead_code)] pub(crate) fn find_raw_cut( messages: &[ChatMessage], discard_from_rendered_index: usize, @@ -2349,6 +2459,112 @@ pub(crate) fn restore_system_prompt_override( } } +/// Fork indicator for a rendered message: position and total branches at +/// the nearest divergence point upstream in the conversation tree. +#[derive(Debug, Clone, serde::Serialize)] +pub struct ForkInfo { + pub position: usize, + pub total: usize, +} + +/// Render a path of tree nodes into the UI-friendly `RenderedMessage` format, +/// matching the same filtering logic as the legacy flat-array `load_history`. +/// Returns `(rendered_messages, turn_count, node_ids, fork_info)` where +/// `node_ids[i]` is the tree node ID that produced `rendered_messages[i]` and +/// `fork_info[i]` carries the nearest divergence point at or upstream of it. +/// +/// Fork detection walks *every* node in the path — including the non-rendered +/// tool-dispatch/tool/system nodes — so a fork whose diverging child is a +/// tool-dispatch assistant (empty content + tool_calls) is still attributed to +/// the next rendered message. Detecting forks only on rendered nodes would miss +/// regenerations where the model replied with a tool call. +fn render_tree_path( + store: &ChatHistoryStore, + path: &[&StoredChatNode], +) -> (Vec, usize, Vec, Vec>) { + let mut rendered = Vec::new(); + let mut user_turns_seen = 0usize; + let mut assistant_turns_seen = 0usize; + let mut pending_tools: Vec = Vec::new(); + let mut pending_calls: std::collections::VecDeque<(String, serde_json::Value)> = + std::collections::VecDeque::new(); + let mut node_ids: Vec = Vec::new(); + let mut fork_info: Vec> = Vec::new(); + // Nearest divergence seen so far, carried forward across nodes (rendered or + // not) so every message at/after a fork shows the indicator. + let mut last_fork: Option = None; + + for node in path { + // Update the carried fork BEFORE rendering: a non-rendered fork child + // (e.g. a tool-dispatch assistant) must colour the rendered message + // that follows it. + if let Some((position, total)) = store.fork_at_node(node.id) { + last_fork = Some(ForkInfo { position, total }); + } + let msg = &node.message; + match msg.role.as_str() { + "system" => continue, + "tool" => { + if let Some((name, arguments)) = pending_calls.pop_front() { + let (result, result_truncated) = truncate_tool_result(&msg.content); + pending_tools.push(ToolInvocation { + name, + arguments, + result, + result_truncated, + }); + } + } + "assistant" => { + let has_tool_calls = msg + .tool_calls + .as_ref() + .map(|c| !c.is_empty()) + .unwrap_or(false); + if has_tool_calls && msg.content.trim().is_empty() { + if let Some(ref tcs) = msg.tool_calls { + for tc in tcs { + pending_calls.push_back(( + tc.function.name.clone(), + tc.function.arguments.clone(), + )); + } + } + continue; + } + assistant_turns_seen += 1; + let tools = std::mem::take(&mut pending_tools); + pending_calls.clear(); + rendered.push(RenderedMessage { + role: "assistant".to_string(), + content: msg.content.clone(), + is_initial: false, + tools, + }); + node_ids.push(node.id); + fork_info.push(last_fork.clone()); + } + "user" => { + let is_initial = user_turns_seen == 0; + user_turns_seen += 1; + pending_tools.clear(); + pending_calls.clear(); + rendered.push(RenderedMessage { + role: "user".to_string(), + content: msg.content.clone(), + is_initial, + tools: Vec::new(), + }); + node_ids.push(node.id); + fork_info.push(last_fork.clone()); + } + _ => continue, + } + } + + (rendered, assistant_turns_seen, node_ids, fork_info) +} + /// View returned to clients for chat-UI rendering. #[derive(Debug)] pub struct HistoryView { @@ -2356,6 +2572,15 @@ pub struct HistoryView { pub turn_count: usize, pub model_version: String, pub backend: String, + /// ID of the leaf node the current view is anchored to. + pub active_leaf_id: u64, + /// ID of the branch leaf being viewed. Equals `active_leaf_id` when + /// viewing the active branch, or the `branch_id` parameter when + /// viewing an alternate branch. + pub viewing_branch_id: u64, + /// Fork info for each rendered message. `None` means no divergence at + /// or before that point in the tree. + pub fork_info: Vec>, } #[derive(Debug)] @@ -2581,6 +2806,61 @@ mod tests { assert!(!dropped); } + #[test] + fn render_tree_path_detects_fork_at_tool_dispatch_child() { + // Regression: a regeneration where the model replies with a tool call + // forks at the user turn, but the diverging children are non-rendered + // tool-dispatch assistant nodes. The indicator must still surface on + // the rendered final reply that follows — earlier code only inspected + // rendered nodes and silently dropped the fork, orphaning the branch. + let mut store = ChatHistoryStore::empty(); + let u = store.append_node(None, ChatMessage::user("q1")); + // Original branch: tool-dispatch → tool result → final reply. + let d1 = store.append_node(Some(u), assistant_with_tool_call("lookup")); + let t1 = store.append_node(Some(d1), ChatMessage::tool_result("data1")); + let _a1 = store.append_node(Some(t1), assistant_text("answer 1")); + // Regenerated branch: also a tool reply, forking at the user turn. + let d2 = store.append_node(Some(u), assistant_with_tool_call("lookup")); + let t2 = store.append_node(Some(d2), ChatMessage::tool_result("data2")); + let a2 = store.append_node(Some(t2), assistant_text("answer 2")); + store.active_leaf_id = a2; + + let path = store.path_to_leaf(a2).expect("path to active leaf"); + let (rendered, _turns, node_ids, fork_info) = render_tree_path(&store, &path); + + assert_eq!(rendered.len(), 2, "user turn + final reply are rendered"); + assert_eq!(rendered[1].content, "answer 2"); + assert_eq!(node_ids, vec![u, a2]); + assert!( + fork_info[0].is_none(), + "user turn sits before the divergence" + ); + let f = fork_info[1] + .as_ref() + .expect("fork indicator carried onto the rendered reply"); + assert_eq!((f.position, f.total), (2, 2)); + } + + #[test] + fn render_tree_path_detects_fork_at_rendered_child() { + // Non-tool case: two direct assistant replies to the same user turn. + // The diverging child is itself rendered, so the indicator lands on it. + let mut store = ChatHistoryStore::empty(); + let u = store.append_node(None, ChatMessage::user("q1")); + let _a1 = store.append_node(Some(u), assistant_text("answer 1")); + let a2 = store.append_node(Some(u), assistant_text("answer 2")); + store.active_leaf_id = a2; + + let path = store.path_to_leaf(a2).expect("path to active leaf"); + let (rendered, _turns, _node_ids, fork_info) = render_tree_path(&store, &path); + + assert_eq!(rendered.len(), 2); + let f = fork_info[1] + .as_ref() + .expect("fork indicator on the divergent rendered reply"); + assert_eq!((f.position, f.total), (2, 2)); + } + #[test] fn rewind_strips_assistant_and_tool_scaffolding() { // Rendered: [user1, asst1, user2, asst2] → cut at rendered index 3 diff --git a/src/ai/llm_client.rs b/src/ai/llm_client.rs index a50a6d8..8309cde 100644 --- a/src/ai/llm_client.rs +++ b/src/ai/llm_client.rs @@ -171,6 +171,178 @@ pub struct ModelCapabilities { pub has_tool_calling: bool, } +/// A single node in the conversation tree. Wraps ChatMessage with structural +/// metadata so the tree can be persisted as JSON in `training_messages`. +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct StoredChatNode { + pub id: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_id: Option, + pub message: ChatMessage, +} + +/// Top-level container for tree-based chat history. Stored as JSON in the +/// `training_messages` column. Backward-compatible reader detects whether +/// the stored value is a flat array (old format) or this object (new format). +#[derive(Serialize, Deserialize, Clone, Debug)] +pub struct ChatHistoryStore { + pub nodes: Vec, + pub active_leaf_id: u64, +} + +impl ChatHistoryStore { + /// Create an empty store. Test-only: production stores are always built + /// from `from_flat_array` or deserialized from `training_messages`. + #[cfg(test)] + pub fn empty() -> Self { + Self { + nodes: Vec::new(), + active_leaf_id: 0, + } + } + + /// Convert a flat `Vec` (old format) into a tree. + /// Each message becomes a node chained by parent_id. + pub fn from_flat_array(messages: Vec) -> Self { + let mut nodes = Vec::with_capacity(messages.len()); + let mut prev_id: Option = None; + for (i, msg) in messages.into_iter().enumerate() { + let id = (i + 1) as u64; + nodes.push(StoredChatNode { + id, + parent_id: prev_id, + message: msg, + }); + prev_id = Some(id); + } + let active_leaf_id = prev_id.unwrap_or(0); + Self { + nodes, + active_leaf_id, + } + } + + /// Allocate the next node ID (max existing + 1, or 1 if empty). + pub fn next_id(&self) -> u64 { + self.nodes.iter().map(|n| n.id).max().unwrap_or(0) + 1 + } + + /// Append a new node as a child of the given parent. Returns the new node's id. + pub fn append_node(&mut self, parent_id: Option, message: ChatMessage) -> u64 { + let id = self.next_id(); + self.nodes.push(StoredChatNode { + id, + parent_id, + message, + }); + id + } + + /// Traverse from root to the given leaf_id, returning the path in order. + /// Returns None if the leaf_id doesn't exist. + pub fn path_to_leaf(&self, leaf_id: u64) -> Option> { + // Build a map from id to node for fast lookup. + let node_map: std::collections::HashMap = self + .nodes + .iter() + .enumerate() + .map(|(i, n)| (n.id, i)) + .collect(); + let &_leaf_idx = node_map.get(&leaf_id)?; + + let mut path_indices = Vec::new(); + let mut current_id = leaf_id; + while let Some(&idx) = node_map.get(¤t_id) { + path_indices.push(idx); + match self.nodes[idx].parent_id { + Some(pid) => current_id = pid, + None => break, + } + } + path_indices.reverse(); + Some(path_indices.into_iter().map(|i| &self.nodes[i]).collect()) + } + + /// All direct children of a given node, ordered by id (chronological). + pub fn children_of(&self, node_id: u64) -> Vec<&StoredChatNode> { + self.nodes + .iter() + .filter(|n| n.parent_id == Some(node_id)) + .collect() + } + + /// For a given node, return (position, total) where: + /// - total = number of children of the parent of this node (i.e., siblings + self) + /// - position = 1-based rank of this node among siblings, ordered by id + /// + /// Returns None if the node has no parent (root) or the parent has 0 or 1 child. + pub fn fork_at_node(&self, node_id: u64) -> Option<(usize, usize)> { + let node = self.nodes.iter().find(|n| n.id == node_id)?; + let parent_id = node.parent_id?; + let children = self.children_of(parent_id); + let total = children.len(); + if total <= 1 { + return None; + } + let position = children.iter().position(|c| c.id == node_id)? + 1; + Some((position, total)) + } + + /// All leaf node IDs (nodes with no children). + pub fn leaves(&self) -> Vec { + let child_ids: std::collections::HashSet = + self.nodes.iter().filter_map(|n| n.parent_id).collect(); + self.nodes + .iter() + .filter(|n| !child_ids.contains(&n.id)) + .map(|n| n.id) + .collect() + } + + /// The parent of the given node, if any. Test-only helper. + #[cfg(test)] + pub fn parent_of(&self, node_id: u64) -> Option<&StoredChatNode> { + let node = self.nodes.iter().find(|n| n.id == node_id)?; + let parent_id = node.parent_id?; + self.nodes.iter().find(|n| n.id == parent_id) + } + + /// Return all leaf nodes with metadata: id, snippet (first 100 chars of + /// the first user-visible message after the divergence point), and + /// message_count (number of nodes in the path from root to leaf). + pub fn leaves_with_info(&self) -> Vec { + let leaf_ids = self.leaves(); + leaf_ids + .into_iter() + .map(|leaf_id| { + let path = self.path_to_leaf(leaf_id).unwrap_or_default(); + let message_count = path.len(); + // Snippet: first 100 chars of the first user/assistant message + // (skip system messages). + let snippet = path + .iter() + .find(|n| n.message.role != "system") + .map(|n| n.message.content.chars().take(100).collect::()) + .unwrap_or_default(); + BranchLeafInfo { + id: leaf_id, + snippet, + message_count, + } + }) + .collect() + } +} + +/// Metadata about a branch leaf, returned by the branches endpoint. +#[derive(Serialize, Clone, Debug)] +pub struct BranchLeafInfo { + pub id: u64, + #[serde(skip_serializing_if = "String::is_empty")] + pub snippet: String, + pub message_count: usize, +} + /// Strip a leading `…` reasoning block from model output. /// /// Thinking models sometimes emit chain-of-thought inside think tags before @@ -221,4 +393,135 @@ mod tests { let raw = "thinking forever"; assert_eq!(strip_think_blocks(raw), raw); } + + // ── ChatHistoryStore tree traversal tests ────────────────────────── + + fn mk_msg(role: &str, content: &str) -> ChatMessage { + match role { + "system" => ChatMessage::system(content), + "user" => ChatMessage::user(content), + _ => ChatMessage { + role: role.to_string(), + content: content.to_string(), + tool_calls: None, + images: None, + }, + } + } + + #[test] + fn store_from_flat_array_chains_parent_ids() { + let msgs = vec![ + mk_msg("system", "sys"), + mk_msg("user", "hello"), + mk_msg("assistant", "hi there"), + ]; + let store = ChatHistoryStore::from_flat_array(msgs); + assert_eq!(store.nodes.len(), 3); + assert_eq!(store.nodes[0].id, 1); + assert_eq!(store.nodes[0].parent_id, None); + assert_eq!(store.nodes[1].parent_id, Some(1)); + assert_eq!(store.nodes[2].parent_id, Some(2)); + assert_eq!(store.active_leaf_id, 3); + } + + #[test] + fn store_path_to_leaf_returns_ordered_path() { + let mut store = ChatHistoryStore::empty(); + let n1 = store.append_node(None, mk_msg("user", "u1")); + let n2 = store.append_node(Some(n1), mk_msg("assistant", "a1")); + let n3 = store.append_node(Some(n2), mk_msg("user", "u2")); + let path = store.path_to_leaf(n3).unwrap(); + assert_eq!(path.len(), 3); + assert_eq!(path[0].id, n1); + assert_eq!(path[1].id, n2); + assert_eq!(path[2].id, n3); + } + + #[test] + fn store_path_to_leaf_returns_none_for_missing_id() { + let store = ChatHistoryStore::empty(); + assert!(store.path_to_leaf(999).is_none()); + } + + #[test] + fn store_children_of_returns_direct_children() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + let c1 = store.append_node(Some(root), mk_msg("assistant", "c1")); + let c2 = store.append_node(Some(root), mk_msg("user", "c2")); + let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1")); + let children = store.children_of(root); + assert_eq!(children.len(), 2); + assert_eq!(children[0].id, c1); + assert_eq!(children[1].id, c2); + } + + #[test] + fn store_fork_at_node_detects_fork() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + let c1 = store.append_node(Some(root), mk_msg("assistant", "c1")); + let c2 = store.append_node(Some(root), mk_msg("user", "c2")); + let c3 = store.append_node(Some(root), mk_msg("user", "c3")); + // c1 is position 1 of 3 siblings + assert_eq!(store.fork_at_node(c1), Some((1, 3))); + assert_eq!(store.fork_at_node(c2), Some((2, 3))); + assert_eq!(store.fork_at_node(c3), Some((3, 3))); + // root has no parent, so no fork + assert!(store.fork_at_node(root).is_none()); + } + + #[test] + fn store_fork_at_node_returns_none_for_single_child() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + let _only = store.append_node(Some(root), mk_msg("assistant", "only")); + assert!(store.fork_at_node(2).is_none()); + } + + #[test] + fn store_leaves_returns_leaf_ids() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + let c1 = store.append_node(Some(root), mk_msg("assistant", "c1")); + let c2 = store.append_node(Some(root), mk_msg("user", "c2")); + let _d1 = store.append_node(Some(c1), mk_msg("assistant", "d1")); + let leaves = store.leaves(); + assert_eq!(leaves.len(), 2); + assert!(leaves.contains(&c2)); + assert!(leaves.contains(&_d1)); + assert!(!leaves.contains(&root)); + assert!(!leaves.contains(&c1)); + } + + #[test] + fn store_parent_of_returns_parent() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + let child = store.append_node(Some(root), mk_msg("assistant", "child")); + let parent = store.parent_of(child).unwrap(); + assert_eq!(parent.id, root); + assert!(store.parent_of(root).is_none()); + } + + #[test] + fn store_next_id_increments() { + let mut store = ChatHistoryStore::empty(); + assert_eq!(store.next_id(), 1); + store.append_node(None, mk_msg("user", "a")); + assert_eq!(store.next_id(), 2); + store.append_node(Some(1), mk_msg("assistant", "b")); + assert_eq!(store.next_id(), 3); + } + + #[test] + fn store_append_node_sets_parent() { + let mut store = ChatHistoryStore::empty(); + let parent = store.append_node(None, mk_msg("user", "parent")); + let child = store.append_node(Some(parent), mk_msg("assistant", "child")); + assert_eq!(store.nodes.len(), 2); + assert_eq!(store.nodes[1].parent_id, Some(parent)); + assert_eq!(store.nodes[1].id, child); + } } diff --git a/src/ai/mod.rs b/src/ai/mod.rs index 7d0802e..3f17064 100644 --- a/src/ai/mod.rs +++ b/src/ai/mod.rs @@ -25,12 +25,12 @@ pub use daily_summary_job::{ generate_daily_summaries, strip_summary_boilerplate, }; pub use handlers::{ - cancel_generation_handler, cancel_turn_handler, chat_history_handler, chat_rewind_handler, - chat_stream_handler, chat_turn_handler, delete_insight_handler, export_training_data_handler, - generate_agentic_insight_handler, generate_insight_handler, generation_status_handler, - get_all_insights_handler, get_available_models_handler, get_insight_handler, - get_insight_history_handler, get_openrouter_models_handler, rate_insight_handler, - turn_async_handler, turn_replay_handler, + cancel_generation_handler, cancel_turn_handler, chat_branches_handler, chat_history_handler, + chat_rewind_handler, chat_stream_handler, chat_switch_branch_handler, chat_turn_handler, + delete_insight_handler, export_training_data_handler, generate_agentic_insight_handler, + generate_insight_handler, generation_status_handler, get_all_insights_handler, + get_available_models_handler, get_insight_handler, get_insight_history_handler, + get_openrouter_models_handler, rate_insight_handler, turn_async_handler, turn_replay_handler, }; pub use insight_generator::InsightGenerator; pub use llamacpp::LlamaCppClient; diff --git a/src/main.rs b/src/main.rs index 7faa959..9d27f2e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -377,6 +377,8 @@ fn main() -> std::io::Result<()> { .service(ai::chat_stream_handler) .service(ai::chat_history_handler) .service(ai::chat_rewind_handler) + .service(ai::chat_branches_handler) + .service(ai::chat_switch_branch_handler) .service(ai::turn_async_handler) .service(ai::turn_replay_handler) .service(ai::cancel_turn_handler)