diff --git a/src/ai/handlers.rs b/src/ai/handlers.rs index 6cd4b48..022588f 100644 --- a/src/ai/handlers.rs +++ b/src/ai/handlers.rs @@ -1335,6 +1335,9 @@ pub struct ChatHistoryHttpResponse { 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)] @@ -1443,6 +1446,7 @@ pub async fn chat_history_handler( f.map(|fi| ChatForkInfo { position: fi.position, total: fi.total, + node_id: fi.node_id, }) }) .collect(); @@ -1487,14 +1491,32 @@ 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. +/// 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, + query: web::Query, app_state: web::Data, ) -> impl Responder { let library = libraries::resolve_library_param_state(&app_state, query.library.as_deref()) @@ -1502,21 +1524,27 @@ pub async fn chat_branches_handler( .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 })); - } + 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, diff --git a/src/ai/insight_chat.rs b/src/ai/insight_chat.rs index af9acf9..58cce7a 100644 --- a/src/ai/insight_chat.rs +++ b/src/ai/insight_chat.rs @@ -717,12 +717,20 @@ impl InsightChatService { Ok(()) } - /// Return branch leaf metadata (id, snippet, message_count) and the - /// active leaf ID for a photo's conversation tree. + /// 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(); @@ -745,7 +753,16 @@ impl InsightChatService { .map_err(|e| anyhow!("failed to deserialize chat history: {}", e))? }; - Ok((store.leaves_with_info(), store.active_leaf_id)) + 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 @@ -2465,6 +2482,10 @@ pub(crate) fn restore_system_prompt_override( 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, @@ -2498,8 +2519,16 @@ fn render_tree_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 }); + // (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() { @@ -2839,6 +2868,7 @@ mod tests { .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] @@ -2859,6 +2889,7 @@ mod tests { .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] diff --git a/src/ai/llm_client.rs b/src/ai/llm_client.rs index 8309cde..bdaa463 100644 --- a/src/ai/llm_client.rs +++ b/src/ai/llm_client.rs @@ -288,6 +288,67 @@ impl ChatHistoryStore { 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. + pub fn branch_options_at(&self, fork_node_id: u64, preferred_leaf: u64) -> Vec { + self.children_of(fork_node_id) + .iter() + .enumerate() + .map(|(i, 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(); + // Snippet: first visible message *after* the divergence point — + // skip tool scaffolding (tool results, empty dispatch turns). + let snippet = path + .iter() + .skip_while(|n| n.id != child.id) + .find(|n| { + (n.message.role == "user" || n.message.role == "assistant") + && !n.message.content.trim().is_empty() + }) + .map(|n| n.message.content.chars().take(100).collect::()) + .unwrap_or_default(); + BranchLeafInfo { + id: rep, + snippet, + 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 = @@ -328,6 +389,7 @@ impl ChatHistoryStore { id: leaf_id, snippet, message_count, + position: None, } }) .collect() @@ -341,6 +403,11 @@ pub struct BranchLeafInfo { #[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. @@ -515,6 +582,76 @@ mod tests { 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_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();