diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index ae9f300..022588f 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,21 @@ 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, + /// The divergence tree node. Pass as `node_id` to the branches endpoint + /// for a list scoped to this exact fork. + pub node_id: u64, } #[derive(Debug, Serialize)] @@ -1414,31 +1434,50 @@ 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, + node_id: fi.node_id, + }) }) - .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 +1491,118 @@ pub async fn chat_history_handler( } } +/// Query for GET /insights/chat/branches. +#[derive(Debug, Deserialize)] +pub struct ChatBranchesQuery { + pub path: String, + #[serde(default)] + pub library: Option, + /// When set (from `ForkInfo.node_id`), scope the list to the sibling + /// branches diverging at this tree node instead of every leaf. + #[serde(default)] + pub node_id: Option, + /// The leaf the client is currently viewing. Scoped options in the + /// same subtree anchor to it so the client can identify its own branch. + /// Defaults to the active leaf when absent. + #[serde(default)] + pub viewing_branch_id: Option, +} + +/// GET /insights/chat/branches — return the branch list for a photo's +/// conversation tree. Without `node_id`, each entry is a leaf node; +/// with `node_id`, entries are the position-ranked sibling branches at +/// that divergence point. Use an entry's `id` with `branch_id` on the +/// history endpoint (or the switch-branch endpoint) to load that 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, + query.node_id, + query.viewing_branch_id, + ) { + 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 if msg.contains("not found in tree") { + return HttpResponse::BadRequest().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..75d050c 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,125 @@ 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 metadata and the active leaf ID for a photo's + /// conversation tree. + /// + /// With `node_id = None` this lists every leaf in the tree. With + /// `node_id = Some(fork)` (from `ForkInfo.node_id`) it lists only the + /// sibling branches diverging at that node, position-ranked; options in + /// the same subtree as `viewing_leaf` anchor to that leaf so the client + /// can recognise "the branch I'm on". + pub fn get_branches( + &self, + library_id: i32, + file_path: &str, + node_id: Option, + viewing_leaf: Option, + ) -> 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))? + }; + + let list = match node_id { + Some(fork_node) => { + if !store.nodes.iter().any(|n| n.id == fork_node) { + bail!("node_id {} not found in tree", fork_node); + } + store.branch_options_at(fork_node, viewing_leaf.unwrap_or(store.active_leaf_id)) + } + None => store.leaves_with_info(), + }; + Ok((list, 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 +929,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 +1045,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 +2289,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 +2315,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 +2476,126 @@ 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, + /// The divergence node itself (the parent whose children fork). Pass as + /// `node_id` to the branches endpoint to list the sibling branches at + /// this exact point rather than every leaf in the tree. + pub node_id: u64, +} + +/// 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 non-rendered + // nodes (e.g. tool-dispatch) so the fork is attributed to the next + // rendered message. Cleared after the first rendered message consumes it, + // so the chip appears only at the actual divergence point. + 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. + // (fork_at_node only fires for nodes with a parent, so the pair + // pattern below can't drop a real fork.) + if let (Some(node_id), Some((position, total))) = + (node.parent_id, store.fork_at_node(node.id)) + { + last_fork = Some(ForkInfo { + position, + total, + node_id, + }); + } + 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.take()); + } + "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.take()); + } + _ => continue, + } + } + + (rendered, assistant_turns_seen, node_ids, fork_info) +} + /// View returned to clients for chat-UI rendering. #[derive(Debug)] pub struct HistoryView { @@ -2356,6 +2603,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 +2837,94 @@ 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)); + assert_eq!(f.node_id, u, "divergence node is the forked user turn"); + } + + #[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)); + assert_eq!(f.node_id, u, "divergence node is the forked user turn"); + } + + #[test] + fn render_tree_path_nested_forks_only_emits_at_divergence() { + // Regression: with .take() instead of .clone(), fork_info should only + // appear at the actual divergence point, not propagate downstream. + // Structure: + // u (user "q1") + // ├── a1 (assistant "answer 1") + // └── a2 (assistant "answer 2") + // ├── q2 (user "q2") + // └── r2 (assistant "reply 2") <- active leaf + // + // Expected fork_info: [None, Some, None, None] + // (only a2 has fork info since it's the divergence from a1) + 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")); + let q2 = store.append_node(Some(a2), ChatMessage::user("q2")); + let r2 = store.append_node(Some(q2), assistant_text("reply 2")); + store.active_leaf_id = r2; + + let path = store.path_to_leaf(r2).expect("path to active leaf"); + let (rendered, _turns, _node_ids, fork_info) = render_tree_path(&store, &path); + + assert_eq!(rendered.len(), 4, "user + asst + user + asst"); + assert!(fork_info[0].is_none(), "initial user has no fork"); + assert!(fork_info[1].is_some(), "a2 is the fork point"); + assert!(fork_info[2].is_none(), "q2 should NOT carry fork forward"); + assert!(fork_info[3].is_none(), "r2 should NOT carry fork forward"); + } + #[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..ca4ee58 100644 --- a/src/ai/llm_client.rs +++ b/src/ai/llm_client.rs @@ -171,6 +171,292 @@ 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 within the subtree rooted at `node_id` (the node + /// itself when it has no children). Empty when the node doesn't exist. + pub fn descendant_leaves(&self, node_id: u64) -> Vec { + if !self.nodes.iter().any(|n| n.id == node_id) { + return Vec::new(); + } + let mut result = Vec::new(); + let mut stack = vec![node_id]; + while let Some(id) = stack.pop() { + let children = self.children_of(id); + if children.is_empty() { + result.push(id); + } else { + stack.extend(children.iter().map(|c| c.id)); + } + } + result + } + + /// Branch options at a divergence point: one entry per direct child of + /// `fork_node_id`, ordered by id (chronological — matches the stable + /// position rank shown in "X/Y" fork chips). + /// + /// Each entry's `id` is the leaf to switch to for that branch: the + /// `preferred_leaf` when it lives in that child's subtree (so the + /// caller's current view maps onto its own option and can be filtered + /// out client-side), otherwise the most recent (max-id) descendant leaf. + /// + /// Snippets preview the first visible message that *differs* between + /// the branches, not blindly the first child: Rewind & Regenerate + /// resends the identical question, so the first child is often shared + /// across siblings and would make every option read the same. + pub fn branch_options_at(&self, fork_node_id: u64, preferred_leaf: u64) -> Vec { + struct Candidate { + rep: u64, + message_count: usize, + /// User/assistant contents after the divergence, tool + /// scaffolding (tool results, empty dispatch turns) skipped. + visible: Vec, + } + let candidates: Vec = self + .children_of(fork_node_id) + .iter() + .map(|child| { + let leaves = self.descendant_leaves(child.id); + let rep = if leaves.contains(&preferred_leaf) { + preferred_leaf + } else { + leaves.iter().copied().max().unwrap_or(child.id) + }; + let path = self.path_to_leaf(rep).unwrap_or_default(); + let message_count = path.len(); + let visible: Vec = path + .iter() + .skip_while(|n| n.id != child.id) + .filter(|n| { + (n.message.role == "user" || n.message.role == "assistant") + && !n.message.content.trim().is_empty() + }) + .map(|n| n.message.content.clone()) + .collect(); + Candidate { + rep, + message_count, + visible, + } + }) + .collect(); + + // Walk forward past positions where every branch has the same + // content; stop at the first divergence (or when any branch runs + // out — a length difference is itself distinguishing). + let mut snippet_idx = 0usize; + loop { + let contents: Vec> = candidates + .iter() + .map(|c| c.visible.get(snippet_idx)) + .collect(); + let all_present_and_equal = + contents.iter().all(|c| c.is_some()) && contents.windows(2).all(|w| w[0] == w[1]); + if !all_present_and_equal { + break; + } + snippet_idx += 1; + } + + candidates + .into_iter() + .enumerate() + .map(|(i, c)| { + // A branch exhausted before the divergence index falls back + // to its own last message so it still shows *something*. + let snippet = c + .visible + .get(snippet_idx) + .or_else(|| c.visible.last()) + .map(|s| s.chars().take(100).collect::()) + .unwrap_or_default(); + BranchLeafInfo { + id: c.rep, + snippet, + message_count: c.message_count, + position: Some(i + 1), + } + }) + .collect() + } + + /// 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, + position: None, + } + }) + .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, + /// 1-based sibling rank at the divergence point. Only present for + /// scoped queries (`branch_options_at`); tree-wide leaf listings have + /// no single divergence to rank against. + #[serde(skip_serializing_if = "Option::is_none")] + pub position: Option, +} + /// Strip a leading `…` reasoning block from model output. /// /// Thinking models sometimes emit chain-of-thought inside think tags before @@ -221,4 +507,284 @@ 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_descendant_leaves_walks_subtree() { + 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 d2 = store.append_node(Some(c1), mk_msg("assistant", "d2")); + // c1's subtree has two leaves; c2 is its own leaf; root sees all. + let mut c1_leaves = store.descendant_leaves(c1); + c1_leaves.sort_unstable(); + assert_eq!(c1_leaves, vec![d1, d2]); + assert_eq!(store.descendant_leaves(c2), vec![c2]); + let mut all = store.descendant_leaves(root); + all.sort_unstable(); + assert_eq!(all, vec![c2, d1, d2]); + assert!(store.descendant_leaves(999).is_empty()); + } + + #[test] + fn branch_options_ranks_children_and_prefers_viewing_leaf() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + // First branch: two leaves (older + newer). + let c1 = store.append_node(Some(root), mk_msg("assistant", "reply one")); + let old_leaf = store.append_node(Some(c1), mk_msg("user", "old follow-up")); + let new_leaf = store.append_node(Some(c1), mk_msg("user", "new follow-up")); + // Second branch: single leaf. + let c2 = store.append_node(Some(root), mk_msg("assistant", "reply two")); + + // Viewing the *older* leaf of branch 1: its option must anchor to + // that exact leaf, not the max-id one, so the client's + // "hide the branch I'm on" filter matches. + let opts = store.branch_options_at(root, old_leaf); + assert_eq!(opts.len(), 2); + assert_eq!(opts[0].id, old_leaf); + assert_eq!(opts[0].position, Some(1)); + assert_eq!(opts[0].snippet, "reply one"); + assert_eq!(opts[1].id, c2); + assert_eq!(opts[1].position, Some(2)); + + // Viewing branch 2: branch 1 falls back to its most recent leaf. + let opts = store.branch_options_at(root, c2); + assert_eq!(opts[0].id, new_leaf); + assert_eq!(opts[1].id, c2); + } + + #[test] + fn branch_options_snippet_skips_shared_resent_question() { + // Rewind & Regenerate resends the identical question, so both + // branches' first child matches — the preview must advance to the + // first *differing* message (the assistant replies) or every + // option would read the same. + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("assistant", "prior reply")); + let q1 = store.append_node(Some(root), mk_msg("user", "same question")); + let _a1 = store.append_node(Some(q1), mk_msg("assistant", "first answer")); + let q2 = store.append_node(Some(root), mk_msg("user", "same question")); + let a2 = store.append_node(Some(q2), mk_msg("assistant", "second answer")); + + let opts = store.branch_options_at(root, a2); + assert_eq!(opts.len(), 2); + assert_eq!(opts[0].snippet, "first answer"); + assert_eq!(opts[1].snippet, "second answer"); + } + + #[test] + fn branch_options_snippet_falls_back_when_branch_exhausted() { + // One branch is just the shared question (no reply yet); the other + // continues past it. The shorter branch falls back to its own last + // message rather than showing an empty preview. + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("assistant", "prior reply")); + let _q1 = store.append_node(Some(root), mk_msg("user", "same question")); + let q2 = store.append_node(Some(root), mk_msg("user", "same question")); + let a2 = store.append_node(Some(q2), mk_msg("assistant", "second answer")); + + let opts = store.branch_options_at(root, a2); + assert_eq!(opts[0].snippet, "same question"); + assert_eq!(opts[1].snippet, "second answer"); + } + + #[test] + fn branch_options_snippet_skips_tool_scaffolding() { + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "root")); + // Branch whose diverging child is an empty tool-dispatch turn. + let dispatch = store.append_node( + Some(root), + ChatMessage { + role: "assistant".to_string(), + content: String::new(), + tool_calls: None, + images: None, + }, + ); + let tool = store.append_node(Some(dispatch), mk_msg("tool", "raw tool output")); + let reply = store.append_node(Some(tool), mk_msg("assistant", "final answer")); + let _other = store.append_node(Some(root), mk_msg("assistant", "direct reply")); + + let opts = store.branch_options_at(root, reply); + assert_eq!(opts[0].snippet, "final answer"); + assert_eq!(opts[1].snippet, "direct reply"); + } + + #[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); + } + + #[test] + fn branch_options_handles_nested_forks() { + // Regression: verify sub-fork display when there are multiple forks. + // Structure: + // root (user "q1") + // ├── a1 (assistant "answer 1") + // └── a2 (assistant "answer 2") + // ├── q2 (user "q2") + // │ ├── r1 (assistant "reply 1") + // │ └── r2 (assistant "reply 2") <- active leaf + // + // branch_options_at(root) should return [a1, a2] (2 branches) + // branch_options_at(q2) should return [r1, r2] (2 sub-branches) + let mut store = ChatHistoryStore::empty(); + let root = store.append_node(None, mk_msg("user", "q1")); + let a1 = store.append_node(Some(root), mk_msg("assistant", "answer 1")); + let a2 = store.append_node(Some(root), mk_msg("assistant", "answer 2")); + let q2 = store.append_node(Some(a2), mk_msg("user", "q2")); + let r1 = store.append_node(Some(q2), mk_msg("assistant", "reply 1")); + let r2 = store.append_node(Some(q2), mk_msg("assistant", "reply 2")); + store.active_leaf_id = r2; + + // Root fork: should show [a1, a2's subtree] + let opts = store.branch_options_at(root, r2); + assert_eq!(opts.len(), 2); + assert_eq!(opts[0].id, a1); + assert_eq!(opts[0].position, Some(1)); + assert_eq!(opts[0].snippet, "answer 1"); + // a2's subtree includes r2, so the rep should be r2 (preferred leaf) + assert_eq!(opts[1].id, r2); + assert_eq!(opts[1].position, Some(2)); + assert_eq!(opts[1].snippet, "answer 2"); + + // Sub-fork at q2: should show [r1, r2] + let opts = store.branch_options_at(q2, r2); + assert_eq!(opts.len(), 2); + assert_eq!(opts[0].id, r1); + assert_eq!(opts[0].position, Some(1)); + assert_eq!(opts[0].snippet, "reply 1"); + assert_eq!(opts[1].id, r2); + assert_eq!(opts[1].position, Some(2)); + assert_eq!(opts[1].snippet, "reply 2"); + } } 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) diff --git a/src/thumbnails.rs b/src/thumbnails.rs index a7334b0..75a456f 100644 --- a/src/thumbnails.rs +++ b/src/thumbnails.rs @@ -98,7 +98,7 @@ pub fn generate_image_thumbnail(src: &Path, thumb_path: &Path) -> std::io::Resul } if file_types::needs_ffmpeg_thumbnail(src) { - return generate_image_thumbnail_ffmpeg(src, thumb_path); + return generate_image_thumbnail_ffmpeg(src, thumb_path, orientation); } let img = image::open(src).map_err(|e| { @@ -140,7 +140,7 @@ pub fn generate_large_preview(src: &Path, dest: &Path) -> std::io::Result<()> { } if file_types::needs_ffmpeg_thumbnail(src) { - return generate_large_preview_ffmpeg(src, dest); + return generate_large_preview_ffmpeg(src, dest, orientation); } let img = image::open(src).map_err(|e| { @@ -175,14 +175,34 @@ fn encode_large_jpeg(img: image::DynamicImage, dest: &Path) -> std::io::Result<( /// ffmpeg path for HEIC/HEIF (image crate can't decode these). Mirrors /// [`crate::video::actors::generate_image_thumbnail_ffmpeg`] but scales /// to the large-preview cap instead of 200. -fn generate_large_preview_ffmpeg(src: &Path, dest: &Path) -> std::io::Result<()> { - // scale=W:-1 with force_original_aspect_ratio=decrease + the min(iw,W) - // trick caps the long edge regardless of orientation, mirroring what - // image::thumbnail does for the non-ffmpeg branch. - let vf = format!( +fn generate_large_preview_ffmpeg( + src: &Path, + dest: &Path, + orientation: i32, +) -> std::io::Result<()> { + // Rotation + scale + colorspace. HEIC sources use Display P3; without + // colorspace=bt709 the mjpeg encoder treats P3 values as sRGB, producing + // warm/oversaturated output. The min(iw,cap) trick caps the long edge + // regardless of orientation, mirroring image::thumbnail. + let rotation = match orientation { + 2 => "hflip", + 3 => "transpose=2", + 4 => "vflip", + 5 => "transpose=0,hflip", + 6 => "transpose=0", + 7 => "transpose=1,hflip", + 8 => "transpose=1", + _ => "", + }; + let scale_expr = format!( "scale='if(gt(iw,ih),min(iw,{cap}),-1)':'if(gt(iw,ih),-1,min(ih,{cap}))'", cap = LARGE_PREVIEW_MAX_DIM ); + let vf = if rotation.is_empty() { + format!("{},colorspace=bt709", scale_expr) + } else { + format!("{},{},colorspace=bt709", rotation, scale_expr) + }; let output = Command::new("ffmpeg") .arg("-y") .arg("-i") @@ -232,7 +252,7 @@ pub fn generate_xlarge_preview(src: &Path, dest: &Path) -> std::io::Result<()> { } if file_types::needs_ffmpeg_thumbnail(src) { - return generate_xlarge_preview_ffmpeg(src, dest); + return generate_xlarge_preview_ffmpeg(src, dest, orientation); } let img = image::open(src).map_err(|e| { @@ -260,11 +280,30 @@ fn encode_xlarge_jpeg(img: image::DynamicImage, dest: &Path) -> std::io::Result< Ok(()) } -fn generate_xlarge_preview_ffmpeg(src: &Path, dest: &Path) -> std::io::Result<()> { - let vf = format!( +fn generate_xlarge_preview_ffmpeg( + src: &Path, + dest: &Path, + orientation: i32, +) -> std::io::Result<()> { + let rotation = match orientation { + 2 => "hflip", + 3 => "transpose=2", + 4 => "vflip", + 5 => "transpose=0,hflip", + 6 => "transpose=0", + 7 => "transpose=1,hflip", + 8 => "transpose=1", + _ => "", + }; + let scale_expr = format!( "scale='if(gt(iw,ih),min(iw,{cap}),-1)':'if(gt(iw,ih),-1,min(ih,{cap}))'", cap = XLARGE_PREVIEW_MAX_DIM ); + let vf = if rotation.is_empty() { + format!("{},colorspace=bt709", scale_expr) + } else { + format!("{},{},colorspace=bt709", rotation, scale_expr) + }; let output = Command::new("ffmpeg") .arg("-y") .arg("-i") diff --git a/src/video/actors.rs b/src/video/actors.rs index 22ec1ac..ba44329 100644 --- a/src/video/actors.rs +++ b/src/video/actors.rs @@ -90,10 +90,39 @@ pub fn generate_video_thumbnail(path: &Path, destination: &Path) -> std::io::Res Ok(()) } -/// Use ffmpeg to extract a 200px-wide thumbnail from formats the `image` crate -/// can't decode (RAW: NEF/ARW, HEIC/HEIF). Writes JPEG bytes to `destination` -/// regardless of its extension. -pub fn generate_image_thumbnail_ffmpeg(path: &Path, destination: &Path) -> std::io::Result<()> { +/// Build the ffmpeg filter chain for image thumbnails: rotation (from EXIF +/// orientation) + scale + color-space conversion. HEIC sources use Display P3 +/// primaries; without `colorspace=bt709` the mjpeg encoder treats P3 values +/// as sRGB, producing warm/oversaturated output. +fn build_image_thumb_filter(orientation: i32, scale_w: u32) -> String { + let rotation = match orientation { + 2 => "hflip", + 3 => "transpose=2", + 4 => "vflip", + 5 => "transpose=0,hflip", + 6 => "transpose=0", + 7 => "transpose=1,hflip", + 8 => "transpose=1", + _ => "", // orientation 1 or unknown — no rotation needed + }; + if rotation.is_empty() { + format!("scale={}:{{-1}},colorspace=bt709", scale_w) + } else { + format!("{},scale={}:{{-1}},colorspace=bt709", rotation, scale_w) + } +} + +/// Use ffmpeg to extract a thumbnail from formats the `image` crate can't +/// decode (HEIC/HEIF, RAW: NEF/ARW). `orientation` is the EXIF Orientation +/// tag value (1..=8) — baked into the pixels so the saved JPEG is +/// canonically oriented. Writes JPEG bytes to `destination` regardless of +/// its extension. +pub fn generate_image_thumbnail_ffmpeg( + path: &Path, + destination: &Path, + orientation: i32, +) -> std::io::Result<()> { + let vf = build_image_thumb_filter(orientation, 200); let output = Command::new("ffmpeg") .arg("-y") .arg("-i") @@ -101,7 +130,7 @@ pub fn generate_image_thumbnail_ffmpeg(path: &Path, destination: &Path) -> std:: .arg("-vframes") .arg("1") .arg("-vf") - .arg("scale=200:-1") + .arg(&vf) .arg("-f") .arg("image2") .arg("-c:v")