diff --git a/crates/code_assistant/src/acp/agent.rs b/crates/code_assistant/src/acp/agent.rs index d9d05b2d..a97d584f 100644 --- a/crates/code_assistant/src/acp/agent.rs +++ b/crates/code_assistant/src/acp/agent.rs @@ -646,10 +646,12 @@ impl acp::Agent for ACPAgentImpl { &arguments.session_id.0, Some(session_model_config.clone()), )?; + manager .start_agent_for_message( &arguments.session_id.0, content_blocks, + None, // ACP doesn't support branching yet llm_client, project_manager, command_executor, diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index 5b183798..2a9092e3 100644 --- a/crates/code_assistant/src/acp/ui.rs +++ b/crates/code_assistant/src/acp/ui.rs @@ -556,6 +556,7 @@ impl UserInterface for ACPUserUI { UiEvent::DisplayUserInput { content, attachments, + node_id: _, // ACP doesn't use node_id for branching } => { // Send user message content self.send_session_update(acp::SessionUpdate::UserMessageChunk( @@ -713,7 +714,12 @@ impl UserInterface for ACPUserUI { | UiEvent::UpdateCurrentModel { .. } | UiEvent::UpdateSandboxPolicy { .. } | UiEvent::CancelSubAgent { .. } - | UiEvent::HiddenToolCompleted => { + | UiEvent::HiddenToolCompleted + | UiEvent::StartMessageEdit { .. } + | UiEvent::SwitchBranch { .. } + | UiEvent::MessageEditReady { .. } + | UiEvent::BranchSwitched { .. } + | UiEvent::UpdateBranchInfo { .. } => { // These are UI management events, not relevant for ACP } UiEvent::DisplayError { message } => { @@ -1026,6 +1032,7 @@ mod tests { let send_future = ui.send_event(UiEvent::DisplayUserInput { content: "Hello".into(), attachments: vec![], + node_id: None, }); let receive_future = async { let (notification, ack) = rx.recv().await.expect("session update"); diff --git a/crates/code_assistant/src/agent/persistence.rs b/crates/code_assistant/src/agent/persistence.rs index b96eeedd..07f033e2 100644 --- a/crates/code_assistant/src/agent/persistence.rs +++ b/crates/code_assistant/src/agent/persistence.rs @@ -137,7 +137,10 @@ impl AgentStatePersistence for FileStatePersistence { let SessionState { session_id, name, - messages, + message_nodes, + active_path, + next_node_id, + messages: _, tool_executions, plan, config, @@ -152,7 +155,15 @@ impl AgentStatePersistence for FileStatePersistence { // Create a ChatSession with the current state let mut session = ChatSession::new_empty(session_id, name, config, model_config); - session.messages = messages; + + // Store tree structure + session.message_nodes = message_nodes; + session.active_path = active_path; + session.next_node_id = next_node_id; + + // Clear legacy messages (tree is authoritative) + session.messages.clear(); + session.tool_executions = serialized_executions; session.plan = plan; session.next_request_id = next_request_id.unwrap_or(0); diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index 38ad7e26..0c241790 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -58,8 +58,24 @@ pub struct Agent { permission_handler: Option>, sub_agent_runner: Option>, - // Store all messages exchanged + + // ======================================================================== + // Branching: Tree-based message storage + // ======================================================================== + /// All message nodes in the session (tree structure) + message_nodes: + std::collections::BTreeMap, + /// The currently active path through the tree + active_path: crate::persistence::ConversationPath, + /// Counter for generating unique node IDs + next_node_id: crate::persistence::NodeId, + + // ======================================================================== + // Legacy: Linearized message history (derived from active_path) + // ======================================================================== + /// Store all messages exchanged (kept in sync with active_path) message_history: Vec, + // Store the history of tool executions tool_executions: Vec, // Cached system prompts keyed by model hint @@ -129,6 +145,11 @@ impl Agent { state_persistence, permission_handler, sub_agent_runner, + // Branching tree structure + message_nodes: std::collections::BTreeMap::new(), + active_path: Vec::new(), + next_node_id: 1, + // Linearized message history message_history: Vec::new(), tool_executions: Vec::new(), cached_system_prompts: HashMap::new(), @@ -295,14 +316,18 @@ impl Agent { /// Save the current state (message history and tool executions) fn save_state(&mut self) -> Result<()> { trace!( - "saving {} messages to persistence", - self.message_history.len() + "saving {} messages to persistence (tree nodes: {})", + self.message_history.len(), + self.message_nodes.len() ); if let Some(session_id) = &self.session_id { let session_state = crate::session::SessionState { session_id: session_id.clone(), name: self.session_name.clone(), + message_nodes: self.message_nodes.clone(), + active_path: self.active_path.clone(), + next_node_id: self.next_node_id, messages: self.message_history.clone(), tool_executions: self.tool_executions.clone(), plan: self.plan.clone(), @@ -329,13 +354,53 @@ impl Agent { Ok(()) } - /// Adds a message to the history and saves the state + /// Adds a message to the history and saves the state. + /// This adds the message to both the tree structure and the linearized history. pub fn append_message(&mut self, message: Message) -> Result<()> { + // Add to tree structure + let parent_id = self.active_path.last().copied(); + let node_id = self.next_node_id; + self.next_node_id += 1; + + let node = crate::persistence::MessageNode { + id: node_id, + message: message.clone(), + parent_id, + created_at: std::time::SystemTime::now(), + plan_snapshot: None, + }; + + self.message_nodes.insert(node_id, node); + self.active_path.push(node_id); + + // Also add to linearized history self.message_history.push(message); + self.save_state()?; Ok(()) } + /// Save the current plan state as a snapshot on the last assistant message in the tree. + /// This is called after update_plan tool execution to enable correct plan reconstruction + /// when switching branches. + fn save_plan_snapshot_to_last_assistant_message(&mut self) { + // Find the last assistant message in the active path + for &node_id in self.active_path.iter().rev() { + if let Some(node) = self.message_nodes.get(&node_id) { + if node.message.role == llm::MessageRole::Assistant { + // Found it - set the snapshot + if let Some(node_mut) = self.message_nodes.get_mut(&node_id) { + node_mut.plan_snapshot = Some(self.plan.clone()); + trace!("Saved plan snapshot to assistant message node {}", node_id); + } + return; + } + } + } + // No assistant message found - this shouldn't happen in normal flow + trace!("No assistant message found to save plan snapshot"); + } + /// Run a single iteration of the agent loop without waiting for user input /// This is used in the new on-demand agent architecture pub async fn run_single_iteration(&mut self) -> Result<()> { @@ -350,6 +415,7 @@ impl Agent { .send_event(UiEvent::DisplayUserInput { content: pending_message, attachments: Vec::new(), + node_id: None, // Pending messages don't have node_id yet }) .await?; } @@ -438,17 +504,25 @@ impl Agent { } } - /// Load state from session state (for backward compatibility) + /// Load state from session state pub async fn load_from_session_state( &mut self, session_state: crate::session::SessionState, ) -> Result<()> { // Restore all state components self.session_id = Some(session_state.session_id); + + // Load tree structure + self.message_nodes = session_state.message_nodes; + self.active_path = session_state.active_path; + self.next_node_id = session_state.next_node_id; + + // Use the linearized messages from session state (derived from active_path) self.message_history = session_state.messages; debug!( - "loaded {} messages from session", - self.message_history.len() + "loaded {} messages from session (tree nodes: {})", + self.message_history.len(), + self.message_nodes.len() ); self.tool_executions = session_state.tool_executions; self.plan = session_state.plan.clone(); @@ -944,6 +1018,7 @@ impl Agent { .send_event(UiEvent::DisplayUserInput { content: task.clone(), attachments: Vec::new(), + node_id: None, // Initial task message }) .await?; @@ -1743,6 +1818,12 @@ impl Agent { // Store the execution record self.tool_executions.push(tool_execution); + // If this was an update_plan tool, save plan snapshot to the last assistant message + // This enables correct plan reconstruction when switching branches + if tool_request.name == "update_plan" && success { + self.save_plan_snapshot_to_last_assistant_message(); + } + // Update message history if input was modified if input_modified { if !is_hidden { diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index f04b167b..eafc6756 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -1977,16 +1977,12 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { ]) .with_request_id(1); - let session_state = SessionState { - session_id: "native-session".to_string(), - name: "Native Session".to_string(), - messages: vec![user_message, assistant_message], - tool_executions: Vec::new(), - plan: PlanState::default(), - config: session_config.clone(), - next_request_id: Some(2), - model_config: None, - }; + let session_state = SessionState::from_messages( + "native-session", + "Native Session", + vec![user_message, assistant_message], + session_config.clone(), + ); agent.load_from_session_state(session_state).await?; @@ -2033,20 +2029,16 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() .with_request_id(1); let followup_user_message = Message::new_user("Also check the contributing guide."); - let session_state = SessionState { - session_id: "native-session".to_string(), - name: "Native Session".to_string(), - messages: vec![ + let session_state = SessionState::from_messages( + "native-session", + "Native Session", + vec![ user_message.clone(), assistant_message, followup_user_message.clone(), ], - tool_executions: Vec::new(), - plan: PlanState::default(), - config: session_config.clone(), - next_request_id: Some(3), - model_config: None, - }; + session_config.clone(), + ); agent.load_from_session_state(session_state).await?; @@ -2099,16 +2091,12 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { Message::new_assistant_content(vec![ContentBlock::new_text(assistant_text)]) .with_request_id(1); - let session_state = SessionState { - session_id: "xml-session".to_string(), - name: "XML Session".to_string(), - messages: vec![user_message, assistant_message], - tool_executions: Vec::new(), - plan: PlanState::default(), - config: session_config.clone(), - next_request_id: Some(2), - model_config: None, - }; + let session_state = SessionState::from_messages( + "xml-session", + "XML Session", + vec![user_message, assistant_message], + session_config.clone(), + ); agent.load_from_session_state(session_state).await?; @@ -2143,16 +2131,12 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() let user_message = Message::new_user("Summarize the repo."); let assistant_message = Message::new_assistant("Here is a summary of the repository."); - let session_state = SessionState { - session_id: "text-session".to_string(), - name: "Regular Session".to_string(), - messages: vec![user_message.clone(), assistant_message.clone()], - tool_executions: Vec::new(), - plan: PlanState::default(), - config: session_config.clone(), - next_request_id: Some(3), - model_config: None, - }; + let session_state = SessionState::from_messages( + "text-session", + "Regular Session", + vec![user_message.clone(), assistant_message.clone()], + session_config.clone(), + ); agent.load_from_session_state(session_state).await?; diff --git a/crates/code_assistant/src/app/gpui.rs b/crates/code_assistant/src/app/gpui.rs index 389b433b..e2542735 100644 --- a/crates/code_assistant/src/app/gpui.rs +++ b/crates/code_assistant/src/app/gpui.rs @@ -104,6 +104,7 @@ pub fn run(config: AgentRunConfig) -> Result<()> { .start_agent_for_message( &session_id, vec![llm::ContentBlock::new_text(initial_task)], + None, // Initial task is not a branch llm_client, project_manager, command_executor, diff --git a/crates/code_assistant/src/persistence.rs b/crates/code_assistant/src/persistence.rs index b807e75d..04e124b3 100644 --- a/crates/code_assistant/src/persistence.rs +++ b/crates/code_assistant/src/persistence.rs @@ -1,7 +1,7 @@ use anyhow::Result; use llm::Message; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -11,6 +11,50 @@ use crate::session::SessionConfig; use crate::tools::ToolRequest; use crate::types::{PlanState, ToolSyntax}; +// ============================================================================ +// Session Branching Types +// ============================================================================ + +/// Unique identifier for a message node within a session +pub type NodeId = u64; + +/// A path through the conversation tree (list of node IDs from root to leaf) +pub type ConversationPath = Vec; + +/// A single message node in the conversation tree +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct MessageNode { + /// Unique ID within this session + pub id: NodeId, + + /// The actual message content + pub message: Message, + + /// Parent node ID (None for root/first message) + pub parent_id: Option, + + /// Creation timestamp (for ordering siblings) + pub created_at: SystemTime, + + /// Plan state snapshot (only set if plan changed in this message's response) + /// Used for efficient plan reconstruction when switching branches + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plan_snapshot: Option, +} + +/// Information about a branch point in the conversation (for UI) +#[derive(Debug, Clone, PartialEq)] +pub struct BranchInfo { + /// Node ID where the branch occurs (the node that has multiple children) + pub parent_node_id: Option, + + /// All sibling node IDs at this branch point (different continuations) + pub sibling_ids: Vec, + + /// Index of the currently active sibling (0-based) + pub active_index: usize, +} + /// Model configuration for a session #[derive(Debug, Serialize, Deserialize, Clone)] pub struct SessionModelConfig { @@ -35,11 +79,34 @@ pub struct ChatSession { pub created_at: SystemTime, /// Last updated timestamp pub updated_at: SystemTime, - /// Message history + + // ======================================================================== + // Branching: Tree-based message storage + // ======================================================================== + /// All message nodes in the session (tree structure) + /// Key: NodeId, Value: MessageNode + #[serde(default)] + pub message_nodes: BTreeMap, + + /// The currently active path through the tree + /// This determines which messages are shown and sent to LLM + #[serde(default)] + pub active_path: ConversationPath, + + /// Counter for generating unique node IDs + #[serde(default = "default_next_node_id")] + pub next_node_id: NodeId, + + // ======================================================================== + // Legacy: Linear message list (for migration from old sessions) + // ======================================================================== + /// Legacy linear message history - migrated to message_nodes on first load + #[serde(default, skip_serializing_if = "Vec::is_empty")] pub messages: Vec, + /// Serialized tool execution results pub tool_executions: Vec, - /// Current session plan + /// Current session plan (for the active path) #[serde(default)] pub plan: PlanState, /// Persistent session configuration @@ -79,6 +146,10 @@ pub struct ChatSession { _legacy_working_memory: serde_json::Value, } +fn default_next_node_id() -> NodeId { + 1 +} + impl ChatSession { /// Merge any legacy top-level fields into the nested SessionConfig. pub fn ensure_config(&mut self) -> Result<()> { @@ -96,6 +167,10 @@ impl ChatSession { if let Some(use_diff_blocks) = self.legacy_use_diff_blocks.take() { self.config.use_diff_blocks = use_diff_blocks; } + + // Migrate linear messages to tree structure if needed + self.migrate_to_tree_structure(); + Ok(()) } @@ -111,6 +186,9 @@ impl ChatSession { name, created_at: SystemTime::now(), updated_at: SystemTime::now(), + message_nodes: BTreeMap::new(), + active_path: Vec::new(), + next_node_id: 1, messages: Vec::new(), tool_executions: Vec::new(), plan: PlanState::default(), @@ -124,6 +202,264 @@ impl ChatSession { _legacy_working_memory: serde_json::Value::Null, } } + + // ======================================================================== + // Migration + // ======================================================================== + + /// Migrate legacy linear messages to tree structure. + /// Called automatically by ensure_config() on session load. + fn migrate_to_tree_structure(&mut self) { + if self.message_nodes.is_empty() && !self.messages.is_empty() { + debug!( + "Migrating session {} from linear to tree structure ({} messages)", + self.id, + self.messages.len() + ); + + let mut parent_id: Option = None; + + for message in self.messages.drain(..) { + let node_id = self.next_node_id; + self.next_node_id += 1; + + let node = MessageNode { + id: node_id, + message, + parent_id, + created_at: SystemTime::now(), + plan_snapshot: None, + }; + + self.message_nodes.insert(node_id, node); + self.active_path.push(node_id); + parent_id = Some(node_id); + } + + debug!( + "Migration complete: {} nodes, active_path length: {}", + self.message_nodes.len(), + self.active_path.len() + ); + } + } + + // ======================================================================== + // Tree Navigation & Query + // ======================================================================== + + /// Get the linearized message history for the active path. + /// This is what gets sent to the LLM. + pub fn get_active_messages(&self) -> Vec<&Message> { + self.active_path + .iter() + .filter_map(|id| self.message_nodes.get(id)) + .map(|node| &node.message) + .collect() + } + + /// Get owned copies of messages for the active path. + pub fn get_active_messages_cloned(&self) -> Vec { + self.active_path + .iter() + .filter_map(|id| self.message_nodes.get(id)) + .map(|node| node.message.clone()) + .collect() + } + + /// Get all direct children of a node. + pub fn get_children(&self, parent_id: Option) -> Vec<&MessageNode> { + self.message_nodes + .values() + .filter(|node| node.parent_id == parent_id) + .collect() + } + + /// Get children sorted by creation time (oldest first). + pub fn get_children_sorted(&self, parent_id: Option) -> Vec<&MessageNode> { + let mut children = self.get_children(parent_id); + children.sort_by_key(|n| n.created_at); + children + } + + /// Get branch info for a specific node (if it's part of a branch). + /// Returns None if the node has no siblings (no branching at this point). + pub fn get_branch_info(&self, node_id: NodeId) -> Option { + let node = self.message_nodes.get(&node_id)?; + let siblings: Vec = self + .get_children_sorted(node.parent_id) + .into_iter() + .map(|n| n.id) + .collect(); + + if siblings.len() <= 1 { + return None; // No branching here + } + + let active_index = siblings.iter().position(|&id| id == node_id)?; + + Some(BranchInfo { + parent_node_id: node.parent_id, + sibling_ids: siblings, + active_index, + }) + } + + /// Find the plan state for the active path by walking backwards + /// to find the most recent plan_snapshot. + pub fn get_plan_for_active_path(&self) -> PlanState { + for &node_id in self.active_path.iter().rev() { + if let Some(node) = self.message_nodes.get(&node_id) { + if let Some(plan) = &node.plan_snapshot { + return plan.clone(); + } + } + } + PlanState::default() + } + + // ======================================================================== + // Tree Modification + // ======================================================================== + + /// Add a new message as a child of the last node in the active path. + /// Updates active_path to include the new node. + /// Returns the new node ID. + pub fn add_message(&mut self, message: Message) -> NodeId { + self.add_message_with_parent(message, self.active_path.last().copied()) + } + + /// Add a new message as a child of a specific parent node. + /// Updates active_path to follow this new branch. + /// Returns the new node ID. + pub fn add_message_with_parent( + &mut self, + message: Message, + parent_id: Option, + ) -> NodeId { + let node_id = self.next_node_id; + self.next_node_id += 1; + + let node = MessageNode { + id: node_id, + message, + parent_id, + created_at: SystemTime::now(), + plan_snapshot: None, + }; + + self.message_nodes.insert(node_id, node); + + // Update active_path: build path to parent and add new node + if let Some(parent) = parent_id { + if let Some(parent_pos) = self.active_path.iter().position(|&id| id == parent) { + // Parent is in current active path - just truncate + self.active_path.truncate(parent_pos + 1); + } else { + // Parent is NOT in current active path (we're branching from a different branch) + // Rebuild the path from root to parent + self.active_path = self.build_path_to_node(parent); + } + } else { + // No parent means this is a root node + self.active_path.clear(); + } + self.active_path.push(node_id); + + self.updated_at = SystemTime::now(); + node_id + } + + /// Switch to a different branch by making a different sibling node active. + /// Updates active_path to follow the new branch to its deepest descendant. + pub fn switch_branch(&mut self, new_node_id: NodeId) -> Result<()> { + let node = self + .message_nodes + .get(&new_node_id) + .ok_or_else(|| anyhow::anyhow!("Node not found: {}", new_node_id))?; + + // Find where in active_path the parent is + if let Some(parent_id) = node.parent_id { + if let Some(parent_pos) = self.active_path.iter().position(|&id| id == parent_id) { + // Truncate path after parent + self.active_path.truncate(parent_pos + 1); + } else { + // Parent not in active path - this shouldn't happen in normal use + // but we handle it by rebuilding the path from root + self.active_path = self.build_path_to_node(parent_id); + } + } else { + // Switching to a root node + self.active_path.clear(); + } + + // Extend path from new node to deepest descendant + self.extend_active_path_from(new_node_id); + + // Update the plan to match the new active path + self.plan = self.get_plan_for_active_path(); + + Ok(()) + } + + /// Build the path from root to a specific node. + fn build_path_to_node(&self, target_id: NodeId) -> ConversationPath { + let mut path = Vec::new(); + let mut current_id = Some(target_id); + + // Walk up to root, collecting node IDs + while let Some(id) = current_id { + path.push(id); + current_id = self.message_nodes.get(&id).and_then(|n| n.parent_id); + } + + // Reverse to get root-to-target order + path.reverse(); + path + } + + /// Extend active_path from a given node, following the most recent child at each step. + fn extend_active_path_from(&mut self, start_node_id: NodeId) { + self.active_path.push(start_node_id); + + let mut current_id = start_node_id; + loop { + // Collect child IDs to avoid borrowing issues + let mut child_ids: Vec<(NodeId, SystemTime)> = self + .message_nodes + .values() + .filter(|node| node.parent_id == Some(current_id)) + .map(|node| (node.id, node.created_at)) + .collect(); + + if child_ids.is_empty() { + break; + } + + // Sort by creation time and take the most recent (last) + child_ids.sort_by_key(|(_, created_at)| *created_at); + let next_id = child_ids.last().unwrap().0; + + self.active_path.push(next_id); + current_id = next_id; + } + } + + /// Get the total number of messages (nodes) in the session. + pub fn message_count(&self) -> usize { + self.message_nodes.len() + } + + /// Check if the session has any branches. + #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 + pub fn has_branches(&self) -> bool { + // A session has branches if any node has more than one child + let mut child_counts: HashMap, usize> = HashMap::new(); + for node in self.message_nodes.values() { + *child_counts.entry(node.parent_id).or_insert(0) += 1; + } + child_counts.values().any(|&count| count > 1) + } } impl SessionModelConfig { @@ -250,7 +586,7 @@ impl FileSessionPersistence { name: session.name.clone(), created_at: session.created_at, updated_at: session.updated_at, - message_count: session.messages.len(), + message_count: session.message_count(), total_usage, last_usage, tokens_limit, @@ -383,7 +719,7 @@ impl FileSessionPersistence { for session_id in candidate_ids { // Safety check: load the session and verify it's actually empty match self.load_chat_session(&session_id) { - Ok(Some(session)) if session.messages.is_empty() => { + Ok(Some(session)) if session.message_count() == 0 => { if let Err(e) = self.delete_chat_session(&session_id) { warn!("Failed to delete empty session {}: {}", session_id, e); } else { @@ -458,7 +794,7 @@ impl FileSessionPersistence { name: session.name.clone(), created_at: session.created_at, updated_at: session.updated_at, - message_count: session.messages.len(), + message_count: session.message_count(), total_usage, last_usage, tokens_limit, @@ -513,14 +849,22 @@ pub fn generate_session_id() -> String { format!("chat_{timestamp:x}_{random_part:x}") } -/// Calculate usage information from session messages +/// Calculate usage information from session messages. +/// Uses tree structure (message_nodes) if available, falls back to legacy messages. fn calculate_session_usage(session: &ChatSession) -> (llm::Usage, llm::Usage, Option) { let mut total_usage = llm::Usage::zero(); let mut last_usage = llm::Usage::zero(); let tokens_limit = None; + // Get messages from tree structure (active path) or legacy list + let messages: Vec<&Message> = if !session.message_nodes.is_empty() { + session.get_active_messages() + } else { + session.messages.iter().collect() + }; + // Calculate total usage and find most recent assistant message usage - for message in &session.messages { + for message in messages { if let Some(usage) = &message.usage { // Add to total usage total_usage.input_tokens += usage.input_tokens; @@ -862,4 +1206,283 @@ mod tests { assert_eq!(deleted_count, 0); } + + // ======================================================================== + // Session Branching Tests + // ======================================================================== + + #[test] + fn test_add_message_creates_tree_structure() { + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + // Add first message + let node1 = session.add_message(Message::new_user("Hello")); + assert_eq!(node1, 1); + assert_eq!(session.active_path, vec![1]); + assert_eq!(session.message_nodes.len(), 1); + + // Add second message + let node2 = session.add_message(Message::new_assistant("Hi there!")); + assert_eq!(node2, 2); + assert_eq!(session.active_path, vec![1, 2]); + assert_eq!(session.message_nodes.len(), 2); + + // Verify parent relationships + assert_eq!(session.message_nodes.get(&1).unwrap().parent_id, None); + assert_eq!(session.message_nodes.get(&2).unwrap().parent_id, Some(1)); + } + + #[test] + fn test_add_message_with_parent_creates_branch() { + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + // Create initial conversation + let _node1 = session.add_message(Message::new_user("Hello")); + let _node2 = session.add_message(Message::new_assistant("Hi!")); + let _node3 = session.add_message(Message::new_user("How are you?")); + + // Now create a branch from node 2 (after the first assistant response) + let branch_node = session.add_message_with_parent( + Message::new_user("What's the weather like?"), + Some(2), // Branch from node 2 + ); + + assert_eq!(branch_node, 4); + // Active path should now follow the new branch + assert_eq!(session.active_path, vec![1, 2, 4]); + + // Both node 3 and node 4 should have parent_id = 2 + assert_eq!(session.message_nodes.get(&3).unwrap().parent_id, Some(2)); + assert_eq!(session.message_nodes.get(&4).unwrap().parent_id, Some(2)); + + // Session should detect branches + assert!(session.has_branches()); + } + + #[test] + fn test_switch_branch() { + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + // Create initial conversation + session.add_message(Message::new_user("Hello")); // node 1 + session.add_message(Message::new_assistant("Hi!")); // node 2 + session.add_message(Message::new_user("Original followup")); // node 3 + + // Create a branch from node 2 + session.add_message_with_parent(Message::new_user("Alternative followup"), Some(2)); // node 4 + + // Add continuation on the branch + session.add_message(Message::new_assistant("Alternative response")); // node 5 + + // Active path should be: 1 -> 2 -> 4 -> 5 + assert_eq!(session.active_path, vec![1, 2, 4, 5]); + + // Switch back to node 3 (the original branch) + session.switch_branch(3).expect("switch branch"); + + // Active path should now be: 1 -> 2 -> 3 + assert_eq!(session.active_path, vec![1, 2, 3]); + + // Verify linearized messages + let messages = session.get_active_messages(); + assert_eq!(messages.len(), 3); + } + + #[test] + fn test_get_branch_info() { + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + // Create initial conversation + session.add_message(Message::new_user("Hello")); // node 1 + session.add_message(Message::new_assistant("Hi!")); // node 2 + session.add_message(Message::new_user("Followup A")); // node 3 + + // No branch info for node 3 yet (only child of node 2) + assert!(session.get_branch_info(3).is_none()); + + // Create branches from node 2 + session.add_message_with_parent(Message::new_user("Followup B"), Some(2)); // node 4 + session.add_message_with_parent(Message::new_user("Followup C"), Some(2)); // node 5 + + // Now node 3, 4, 5 are siblings + let info_3 = session.get_branch_info(3).expect("should have branch info"); + assert_eq!(info_3.parent_node_id, Some(2)); + assert_eq!(info_3.sibling_ids.len(), 3); + + let info_5 = session.get_branch_info(5).expect("should have branch info"); + assert_eq!(info_5.parent_node_id, Some(2)); + assert_eq!(info_5.active_index, 2); // node 5 is the third sibling + } + + #[test] + fn test_migration_from_linear_messages() { + // Create a session with legacy linear messages + let mut session = ChatSession { + id: "test".to_string(), + name: "Test".to_string(), + created_at: SystemTime::now(), + updated_at: SystemTime::now(), + message_nodes: BTreeMap::new(), + active_path: Vec::new(), + next_node_id: 1, + messages: vec![ + Message::new_user("Hello"), + Message::new_assistant("Hi!"), + Message::new_user("How are you?"), + ], + tool_executions: Vec::new(), + plan: PlanState::default(), + config: SessionConfig::default(), + next_request_id: 1, + model_config: None, + legacy_init_path: None, + legacy_initial_project: None, + legacy_tool_syntax: None, + legacy_use_diff_blocks: None, + _legacy_working_memory: serde_json::Value::Null, + }; + + // Run migration + session.ensure_config().expect("migration should succeed"); + + // Verify migration + assert_eq!(session.message_nodes.len(), 3); + assert_eq!(session.active_path, vec![1, 2, 3]); + assert!(session.messages.is_empty()); // Legacy messages should be cleared + + // Verify tree structure + assert_eq!(session.message_nodes.get(&1).unwrap().parent_id, None); + assert_eq!(session.message_nodes.get(&2).unwrap().parent_id, Some(1)); + assert_eq!(session.message_nodes.get(&3).unwrap().parent_id, Some(2)); + + // Verify messages are accessible + let messages = session.get_active_messages(); + assert_eq!(messages.len(), 3); + } + + #[test] + fn test_get_active_messages_cloned() { + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + session.add_message(Message::new_user("Hello")); + session.add_message(Message::new_assistant("Hi!")); + + let messages = session.get_active_messages_cloned(); + assert_eq!(messages.len(), 2); + + // Verify content + match &messages[0].content { + llm::MessageContent::Text(text) => assert_eq!(text, "Hello"), + _ => panic!("Expected text content"), + } + } + + #[test] + fn test_nested_branching() { + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + // Level 1: Initial message + session.add_message(Message::new_user("Start")); // 1 + + // Level 2: Two branches from node 1 + session.add_message(Message::new_assistant("Response A")); // 2 + session.add_message_with_parent(Message::new_assistant("Response B"), Some(1)); // 3 + + // Level 3: Two branches from node 2 + session.switch_branch(2).unwrap(); + session.add_message(Message::new_user("Follow A1")); // 4 + session.add_message_with_parent(Message::new_user("Follow A2"), Some(2)); // 5 + + // Verify structure + assert_eq!(session.message_nodes.len(), 5); + + // Check node 1 has two children (2 and 3) + let children_of_1 = session.get_children(Some(1)); + assert_eq!(children_of_1.len(), 2); + + // Check node 2 has two children (4 and 5) + let children_of_2 = session.get_children(Some(2)); + assert_eq!(children_of_2.len(), 2); + + // Navigate to different paths and verify + session.switch_branch(3).unwrap(); + assert_eq!(session.active_path, vec![1, 3]); + + session.switch_branch(5).unwrap(); + assert_eq!(session.active_path, vec![1, 2, 5]); + } + + #[test] + fn test_branch_from_different_branch() { + // This tests the scenario where we create a branch while on a different branch + // i.e., the parent_id is NOT in the current active_path + let mut session = ChatSession::new_empty( + "test".to_string(), + "Test".to_string(), + SessionConfig::default(), + None, + ); + + // Create initial conversation on main branch + session.add_message(Message::new_user("User 1")); // node 1 + session.add_message(Message::new_assistant("Asst 1")); // node 2 + session.add_message(Message::new_user("User 2")); // node 3 + session.add_message(Message::new_assistant("Asst 2")); // node 4 + + // active_path: [1, 2, 3, 4] + assert_eq!(session.active_path, vec![1, 2, 3, 4]); + + // Create branch 2 from node 2 (alternative User 2) + session.add_message_with_parent(Message::new_user("User 2 alt"), Some(2)); // node 5 + session.add_message(Message::new_assistant("Asst 2 alt")); // node 6 + + // active_path: [1, 2, 5, 6] (we're now on branch 2) + assert_eq!(session.active_path, vec![1, 2, 5, 6]); + + // Now while on branch 2, create a new branch from node 4 (which is on branch 1) + // This should properly switch to branch 1's path and then add the new node + let new_node = + session.add_message_with_parent(Message::new_user("User 3 on branch 1"), Some(4)); // node 7 + + assert_eq!(new_node, 7); + // active_path should be: [1, 2, 3, 4, 7] - NOT [1, 2, 5, 6, 7] + assert_eq!(session.active_path, vec![1, 2, 3, 4, 7]); + + // Verify parent relationship + assert_eq!(session.message_nodes.get(&7).unwrap().parent_id, Some(4)); + + // Verify we can still switch back to branch 2 + session.switch_branch(6).unwrap(); + assert_eq!(session.active_path, vec![1, 2, 5, 6]); + } } diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 175aa54e..3a51b3fe 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -6,14 +6,14 @@ use tokio::task::JoinHandle; // Agent instances are created on-demand, no need to import use crate::agent::SubAgentCancellationRegistry; -use crate::persistence::{ChatMetadata, ChatSession}; +use crate::persistence::{ChatMetadata, ChatSession, NodeId}; use crate::ui::gpui::elements::MessageRole; use crate::ui::streaming::create_stream_processor; use crate::ui::ui_events::{MessageData, UiEvent}; use crate::ui::{DisplayFragment, UIError, UserInterface}; use async_trait::async_trait; use sandbox::SandboxContext; -use tracing::{debug, error, trace}; +use tracing::{debug, error}; /// Represents the current activity state of a session #[derive(Debug, Clone, PartialEq, Default)] @@ -123,15 +123,28 @@ impl SessionInstance { } } - /// Add a message to the session - pub fn add_message(&mut self, message: Message) { - self.session.messages.push(message); - self.session.updated_at = std::time::SystemTime::now(); - } + /// Add a message with optional branching support. + /// If `branch_parent_id` is Some, creates a new branch from that parent. + /// If `branch_parent_id` is None, appends to the end of the active path. + pub fn add_message_with_branch( + &mut self, + message: Message, + branch_parent_id: Option, + ) -> Result { + let node_id = if let Some(parent_id) = branch_parent_id { + // Branching: create new message as child of specified parent + debug!( + "Creating new branch from parent {} in session {}", + parent_id, self.session.id + ); + self.session + .add_message_with_parent(message, Some(parent_id)) + } else { + // Normal append: add to end of active path + self.session.add_message(message) + }; - /// Get all messages in the session - pub fn messages(&self) -> &[Message] { - &self.session.messages + Ok(node_id) } /// Get the current context size (input tokens + cache reads from most recent assistant message) @@ -229,6 +242,8 @@ impl SessionInstance { let incomplete_message = MessageData { role: crate::ui::gpui::elements::MessageRole::Assistant, fragments: buffered_fragments, + node_id: None, // Streaming message doesn't have a node yet + branch_info: None, // No branch info for incomplete message }; messages_data.push(incomplete_message); } @@ -277,6 +292,17 @@ impl SessionInstance { pub fn convert_messages_to_ui_data( &self, tool_syntax: crate::types::ToolSyntax, + ) -> Result, anyhow::Error> { + self.convert_messages_to_ui_data_until(tool_syntax, None) + } + + /// Convert session messages to UI MessageData format, stopping at a specific node + /// If `until_node_id` is Some, includes all messages up to and including that node. + /// If `until_node_id` is None, includes all messages (same as convert_messages_to_ui_data). + pub fn convert_messages_to_ui_data_until( + &self, + tool_syntax: crate::types::ToolSyntax, + until_node_id: Option, ) -> Result, anyhow::Error> { // Create dummy UI for stream processor struct DummyUI; @@ -298,13 +324,8 @@ impl SessionInstance { fn should_streaming_continue(&self) -> bool { true } - fn notify_rate_limit(&self, _seconds_remaining: u64) { - // No-op for dummy UI - } - fn clear_rate_limit(&self) { - // No-op for dummy UI - } - + fn notify_rate_limit(&self, _seconds_remaining: u64) {} + fn clear_rate_limit(&self) {} fn as_any(&self) -> &dyn std::any::Any { self } @@ -315,12 +336,31 @@ impl SessionInstance { let mut messages_data = Vec::new(); - trace!( - "preparing {} messages for event", - self.session.messages.len() - ); + // Build message iterator from tree or legacy messages + let message_iter: Vec<(Option, &llm::Message)> = + if !self.session.message_nodes.is_empty() { + // Use active path from tree, but stop at until_node_id + let mut iter = Vec::new(); + for &node_id in &self.session.active_path { + if let Some(node) = self.session.message_nodes.get(&node_id) { + iter.push((Some(node_id), &node.message)); + // Stop after adding the until_node_id + if until_node_id == Some(node_id) { + break; + } + } + } + iter + } else { + // Fall back to legacy linear messages (no until_node_id support) + self.session + .messages + .iter() + .map(|msg| (None, msg)) + .collect() + }; - for message in &self.session.messages { + for (node_id, message) in message_iter { if message.is_compaction_summary { let summary = match &message.content { llm::MessageContent::Text(text) => text.trim().to_string(), @@ -336,34 +376,29 @@ impl SessionInstance { .trim() .to_string(), }; + messages_data.push(MessageData { role: MessageRole::User, fragments: vec![crate::ui::DisplayFragment::CompactionDivider { summary }], + node_id, + branch_info: node_id.and_then(|id| self.session.get_branch_info(id)), }); continue; } - // Filter out tool-result user messages (they have tool IDs in structured content) + + // Filter out tool-result user messages if message.role == llm::MessageRole::User { match &message.content { - llm::MessageContent::Text(text) if text.trim().is_empty() => { - // Skip empty user messages (legacy tool results in XML mode) - continue; - } + llm::MessageContent::Text(text) if text.trim().is_empty() => continue, llm::MessageContent::Structured(blocks) => { - // Check if this is a tool-result message by looking for ToolResult blocks let has_tool_results = blocks .iter() .any(|block| matches!(block, llm::ContentBlock::ToolResult { .. })); - if has_tool_results { - // Skip tool-result user messages (they shouldn't be shown in UI) continue; } - // Otherwise, this is a real structured user message, process it - } - _ => { - // This is a real user message, process it } + _ => {} } } @@ -373,22 +408,24 @@ impl SessionInstance { llm::MessageRole::User => MessageRole::User, llm::MessageRole::Assistant => MessageRole::Assistant, }; - messages_data.push(MessageData { role, fragments }); + messages_data.push(MessageData { + role, + fragments, + node_id, + branch_info: node_id.and_then(|id| self.session.get_branch_info(id)), + }); } Err(e) => { error!("Failed to extract fragments from message: {}", e); - // Continue with other messages even if one fails } } } - trace!("prepared {} message data for event", messages_data.len()); - Ok(messages_data) } /// Convert tool executions to UI tool result data - fn convert_tool_executions_to_ui_data( + pub fn convert_tool_executions_to_ui_data( &self, ) -> Result, anyhow::Error> { use crate::tools::core::ResourcesTracker; diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index b90b423c..6f6d3761 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -222,13 +222,103 @@ impl SessionManager { Ok(ui_events) } - /// Start an agent for a session with a user message + /// Add a user message to a session and return the new node_id. + /// This is used to add the message before displaying it in the UI, + /// ensuring the node_id is available for the edit button. + pub fn add_user_message( + &mut self, + session_id: &str, + content_blocks: Vec, + branch_parent_id: Option, + ) -> Result { + let session_instance = self + .active_sessions + .get_mut(session_id) + .ok_or_else(|| anyhow::anyhow!("Session not found: {session_id}"))?; + + // Make sure the session instance is not stale + session_instance.reload_from_persistence(&self.persistence)?; + + // Add structured user message to session, optionally creating a branch + let node_id = session_instance + .add_message_with_branch(Message::new_user_content(content_blocks), branch_parent_id)?; + + // Save the session state with the new message + self.persistence + .save_chat_session(&session_instance.session)?; + + Ok(node_id) + } + + /// Get branch info for all siblings of a given node. + /// Used after creating a new branch to update the UI for all sibling nodes. + pub fn get_sibling_branch_infos( + &self, + session_id: &str, + node_id: crate::persistence::NodeId, + ) -> Vec<(crate::persistence::NodeId, crate::persistence::BranchInfo)> { + let Some(session_instance) = self.active_sessions.get(session_id) else { + return Vec::new(); + }; + + // Get the branch info for the new node (which contains all siblings) + let Some(branch_info) = session_instance.session.get_branch_info(node_id) else { + return Vec::new(); + }; + + // Return branch info for each sibling (they all have the same siblings but different active_index) + branch_info + .sibling_ids + .iter() + .enumerate() + .map(|(idx, &sibling_id)| { + let sibling_branch_info = crate::persistence::BranchInfo { + parent_node_id: branch_info.parent_node_id, + sibling_ids: branch_info.sibling_ids.clone(), + active_index: idx, + }; + (sibling_id, sibling_branch_info) + }) + .collect() + } + + /// Start an agent for a session (message must already be added via add_user_message) /// This is the key method - agents run on-demand for specific messages + /// + /// Convenience method that adds a user message and starts the agent in one call. + /// This is for callers that don't need the node_id (e.g., ACP, initial task). #[allow(clippy::too_many_arguments)] pub async fn start_agent_for_message( &mut self, session_id: &str, content_blocks: Vec, + branch_parent_id: Option, + llm_provider: Box, + project_manager: Box, + command_executor: Box, + ui: Arc, + permission_handler: Option>, + ) -> Result<()> { + // Add the message first + self.add_user_message(session_id, content_blocks, branch_parent_id)?; + + // Then start the agent + self.start_agent_for_session( + session_id, + llm_provider, + project_manager, + command_executor, + ui, + permission_handler, + ) + .await + } + + /// Start an agent for a session (message must already be added via add_user_message) + #[allow(clippy::too_many_arguments)] + pub async fn start_agent_for_session( + &mut self, + session_id: &str, llm_provider: Box, project_manager: Box, command_executor: Box, @@ -252,8 +342,7 @@ impl SessionManager { // Make sure the session instance is not stale session_instance.reload_from_persistence(&self.persistence)?; - // Add structured user message to session - session_instance.add_message(Message::new_user_content(content_blocks)); + // Note: User message should already be added via add_user_message() // Clone all needed data to avoid borrowing conflicts let name = session_instance.session.name.clone(); @@ -265,7 +354,10 @@ impl SessionManager { let session_state = crate::session::SessionState { session_id: session_id.to_string(), name, - messages: session_instance.messages().to_vec(), + message_nodes: session_instance.session.message_nodes.clone(), + active_path: session_instance.session.active_path.clone(), + next_node_id: session_instance.session.next_node_id, + messages: session_instance.session.get_active_messages_cloned(), tool_executions: session_instance .session .tool_executions @@ -568,6 +660,17 @@ impl SessionManager { self.persistence.get_chat_session_metadata(session_id) } + /// Save the current state of an active session to persistence + pub fn save_session(&mut self, session_id: &str) -> Result<()> { + let session_instance = self + .active_sessions + .get(session_id) + .ok_or_else(|| anyhow::anyhow!("Session not found: {}", session_id))?; + + self.persistence + .save_chat_session(&session_instance.session) + } + /// Save agent state to a specific session pub fn save_session_state(&mut self, state: SessionState) -> Result<()> { let mut session = self @@ -577,7 +680,14 @@ impl SessionManager { // Update session with current state session.name = state.name; - session.messages = state.messages; + + // Update tree structure + session.message_nodes = state.message_nodes; + session.active_path = state.active_path; + session.next_node_id = state.next_node_id; + + // Clear legacy messages (tree is now authoritative) + session.messages.clear(); session.tool_executions = state .tool_executions diff --git a/crates/code_assistant/src/session/mod.rs b/crates/code_assistant/src/session/mod.rs index c2521b7a..58e660ea 100644 --- a/crates/code_assistant/src/session/mod.rs +++ b/crates/code_assistant/src/session/mod.rs @@ -1,9 +1,10 @@ use crate::agent::ToolExecution; -use crate::persistence::SessionModelConfig; +use crate::persistence::{ConversationPath, MessageNode, NodeId, SessionModelConfig}; use crate::types::{PlanState, ToolSyntax}; use llm::Message; use sandbox::SandboxPolicy; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::path::PathBuf; // New session management architecture @@ -44,15 +45,93 @@ impl Default for SessionConfig { } } -/// State data needed to restore an agent session +/// State data needed to restore an agent session. +/// +/// This struct supports both the new tree-based branching structure and +/// a legacy linear message list for backward compatibility. #[derive(Debug, Clone)] pub struct SessionState { pub session_id: String, pub name: String, + + // ======================================================================== + // Branching: Tree-based message storage + // ======================================================================== + /// All message nodes in the session (tree structure) + pub message_nodes: BTreeMap, + + /// The currently active path through the tree + pub active_path: ConversationPath, + + /// Counter for generating unique node IDs + pub next_node_id: NodeId, + + // ======================================================================== + // Legacy: For backward compatibility during transition + // ======================================================================== + /// Linearized message history (derived from active_path for convenience) + /// This is kept in sync with the tree and used by the agent loop pub messages: Vec, + pub tool_executions: Vec, pub plan: PlanState, pub config: SessionConfig, pub next_request_id: Option, pub model_config: Option, } + +#[cfg(test)] +impl SessionState { + /// Create a SessionState from a linear list of messages. + /// This is primarily for tests and backward compatibility. + /// The messages are converted to a tree structure with a single linear path. + pub fn from_messages( + session_id: impl Into, + name: impl Into, + messages: Vec, + config: SessionConfig, + ) -> Self { + let mut message_nodes = BTreeMap::new(); + let mut active_path = Vec::new(); + let mut next_node_id: NodeId = 1; + let mut parent_id: Option = None; + + // Infer next_request_id from messages + let max_request_id = messages + .iter() + .filter_map(|m| m.request_id) + .max() + .unwrap_or(0); + + for message in &messages { + let node_id = next_node_id; + next_node_id += 1; + + let node = crate::persistence::MessageNode { + id: node_id, + message: message.clone(), + parent_id, + created_at: std::time::SystemTime::now(), + plan_snapshot: None, + }; + + message_nodes.insert(node_id, node); + active_path.push(node_id); + parent_id = Some(node_id); + } + + Self { + session_id: session_id.into(), + name: name.into(), + message_nodes, + active_path, + next_node_id, + messages, + tool_executions: Vec::new(), + plan: PlanState::default(), + config, + next_request_id: Some(max_request_id + 1), + model_config: None, + } + } +} diff --git a/crates/code_assistant/src/ui/backend.rs b/crates/code_assistant/src/ui/backend.rs index a3511c5b..043facc4 100644 --- a/crates/code_assistant/src/ui/backend.rs +++ b/crates/code_assistant/src/ui/backend.rs @@ -33,6 +33,8 @@ pub enum BackendEvent { session_id: String, message: String, attachments: Vec, + /// If set, creates a new branch from this parent node instead of appending to active path + branch_parent_id: Option, }, QueueUserMessage { session_id: String, @@ -58,6 +60,16 @@ pub enum BackendEvent { session_id: String, tool_id: String, }, + + // Session branching + StartMessageEdit { + session_id: String, + node_id: crate::persistence::NodeId, + }, + SwitchBranch { + session_id: String, + new_node_id: crate::persistence::NodeId, + }, } // Response from backend to UI @@ -94,10 +106,28 @@ pub enum BackendResponse { session_id: String, policy: SandboxPolicy, }, + SubAgentCancelled { session_id: String, tool_id: String, }, + + // Session branching responses + MessageEditReady { + session_id: String, + content: String, + attachments: Vec, + branch_parent_id: Option, + /// Messages up to (but not including) the message being edited + messages: Vec, + tool_results: Vec, + }, + BranchSwitched { + session_id: String, + messages: Vec, + tool_results: Vec, + plan: crate::types::PlanState, + }, } #[derive(Debug, Clone)] @@ -138,12 +168,14 @@ pub async fn handle_backend_events( session_id, message, attachments, + branch_parent_id, } => { handle_send_user_message( &multi_session_manager, &session_id, &message, &attachments, + branch_parent_id, runtime_options.as_ref(), &ui, ) @@ -180,6 +212,18 @@ pub async fn handle_backend_events( session_id, tool_id, } => Some(handle_cancel_sub_agent(&multi_session_manager, &session_id, &tool_id).await), + + BackendEvent::StartMessageEdit { + session_id, + node_id, + } => { + Some(handle_start_message_edit(&multi_session_manager, &session_id, node_id).await) + } + + BackendEvent::SwitchBranch { + session_id, + new_node_id, + } => Some(handle_switch_branch(&multi_session_manager, &session_id, new_node_id).await), }; // Send response back to UI only if there is one @@ -304,31 +348,69 @@ async fn handle_send_user_message( session_id: &str, message: &str, attachments: &[DraftAttachment], + branch_parent_id: Option, runtime_options: &BackendRuntimeOptions, ui: &Arc, ) -> Option { debug!( - "User message for session {}: {} (with {} attachments)", + "User message for session {}: {} (with {} attachments, branch_parent: {:?})", session_id, message, - attachments.len() + attachments.len(), + branch_parent_id ); // Convert DraftAttachments to ContentBlocks let content_blocks = content_blocks_from(message, attachments); - // Display the user message with attachments in the UI + // First, add the user message to the session and get the new node_id + let (new_node_id, branch_info_updates) = { + let mut manager = multi_session_manager.lock().await; + match manager.add_user_message(session_id, content_blocks.clone(), branch_parent_id) { + Ok(node_id) => { + // If we created a branch, get branch info updates for all siblings + let updates = if branch_parent_id.is_some() { + manager.get_sibling_branch_infos(session_id, node_id) + } else { + Vec::new() + }; + (Some(node_id), updates) + } + Err(e) => { + error!("Failed to add user message to session: {}", e); + return Some(BackendResponse::Error { + message: format!("Failed to add user message: {e}"), + }); + } + } + }; + + // Now display the user message with the correct node_id if let Err(e) = ui .send_event(crate::ui::UiEvent::DisplayUserInput { content: message.to_string(), attachments: attachments.to_vec(), + node_id: new_node_id, }) .await { error!("Failed to display user message with attachments: {}", e); } - // Start the agent with structured content + // Send branch info updates for all siblings (so they show the branch switcher) + for (sibling_node_id, branch_info) in branch_info_updates { + if let Err(e) = ui + .send_event(crate::ui::UiEvent::UpdateBranchInfo { + node_id: sibling_node_id, + branch_info, + }) + .await + { + error!("Failed to send branch info update: {}", e); + } + } + + // Start the agent (message already added) let result = { let project_manager = Box::new(DefaultProjectManager::new()); let command_executor = Box::new(DefaultCommandExecutor); @@ -368,10 +450,10 @@ async fn handle_send_user_message( ); Err(e) } else { + // Message already added via add_user_message above manager - .start_agent_for_message( + .start_agent_for_session( session_id, - content_blocks, client, project_manager, command_executor, @@ -599,6 +681,7 @@ async fn handle_cancel_sub_agent( tool_id: tool_id.to_string(), } } + Err(e) => { error!( "Failed to cancel sub-agent {} for session {}: {}", @@ -610,3 +693,177 @@ async fn handle_cancel_sub_agent( } } } + +// ============================================================================ +// Session Branching Handlers +// ============================================================================ + +async fn handle_start_message_edit( + multi_session_manager: &Arc>, + session_id: &str, + node_id: crate::persistence::NodeId, +) -> BackendResponse { + debug!( + "Starting message edit for session {} node {}", + session_id, node_id + ); + + let result = { + let manager = multi_session_manager.lock().await; + if let Some(session_instance) = manager.get_session(session_id) { + // Get the message node + if let Some(node) = session_instance.session.message_nodes.get(&node_id) { + // Extract content from message + let content = match &node.message.content { + llm::MessageContent::Text(text) => text.clone(), + llm::MessageContent::Structured(blocks) => { + // Extract text content from structured blocks + blocks + .iter() + .filter_map(|block| match block { + llm::ContentBlock::Text { text, .. } => Some(text.clone()), + _ => None, + }) + .collect::>() + .join("\n") + } + }; + + // Extract attachments (images) from message + let attachments = match &node.message.content { + llm::MessageContent::Structured(blocks) => blocks + .iter() + .filter_map(|block| match block { + llm::ContentBlock::Image { + media_type, data, .. + } => Some(DraftAttachment::Image { + content: data.clone(), + mime_type: media_type.clone(), + }), + _ => None, + }) + .collect(), + _ => Vec::new(), + }; + + // The branch parent is the parent of the node being edited + let branch_parent_id = node.parent_id; + + // Generate truncated messages (up to but not including the message being edited) + let messages = session_instance + .convert_messages_to_ui_data_until( + session_instance.session.config.tool_syntax, + branch_parent_id, + ) + .unwrap_or_default(); + + let tool_results = session_instance + .convert_tool_executions_to_ui_data() + .unwrap_or_default(); + + Ok(( + content, + attachments, + branch_parent_id, + messages, + tool_results, + )) + } else { + Err(anyhow::anyhow!("Message node {} not found", node_id)) + } + } else { + Err(anyhow::anyhow!("Session {} not found", session_id)) + } + }; + + match result { + Ok((content, attachments, branch_parent_id, messages, tool_results)) => { + BackendResponse::MessageEditReady { + session_id: session_id.to_string(), + content, + attachments, + branch_parent_id, + messages, + tool_results, + } + } + Err(e) => { + error!("Failed to start message edit: {}", e); + BackendResponse::Error { + message: format!("Failed to start message edit: {e}"), + } + } + } +} + +async fn handle_switch_branch( + multi_session_manager: &Arc>, + session_id: &str, + new_node_id: crate::persistence::NodeId, +) -> BackendResponse { + debug!( + "Switching branch for session {} to node {}", + session_id, new_node_id + ); + + let mut manager = multi_session_manager.lock().await; + + let Some(session_instance) = manager.get_session_mut(session_id) else { + return BackendResponse::Error { + message: format!("Session {} not found", session_id), + }; + }; + + // Perform the branch switch + if let Err(e) = session_instance.session.switch_branch(new_node_id) { + error!("Failed to switch branch: {}", e); + return BackendResponse::Error { + message: format!("Failed to switch branch: {e}"), + }; + } + + // Persist the updated active_path + if let Err(e) = manager.save_session(session_id) { + error!("Failed to save session after branch switch: {}", e); + // Continue anyway - the switch worked in memory + } + + // Re-get session reference after save (borrow checker) + let Some(session_instance) = manager.get_session(session_id) else { + return BackendResponse::Error { + message: format!("Session {} not found after save", session_id), + }; + }; + + // Generate new messages for UI + let messages_data = match session_instance + .convert_messages_to_ui_data(session_instance.session.config.tool_syntax) + { + Ok(data) => data, + Err(e) => { + error!("Failed to convert messages: {}", e); + return BackendResponse::Error { + message: format!("Failed to convert messages: {e}"), + }; + } + }; + + let tool_results = match session_instance.convert_tool_executions_to_ui_data() { + Ok(results) => results, + Err(e) => { + error!("Failed to convert tool results: {}", e); + return BackendResponse::Error { + message: format!("Failed to convert tool results: {e}"), + }; + } + }; + + let plan = session_instance.session.plan.clone(); + + BackendResponse::BranchSwitched { + session_id: session_id.to_string(), + messages: messages_data, + tool_results, + plan, + } +} diff --git a/crates/code_assistant/src/ui/gpui/branch_switcher.rs b/crates/code_assistant/src/ui/gpui/branch_switcher.rs new file mode 100644 index 00000000..1f3af163 --- /dev/null +++ b/crates/code_assistant/src/ui/gpui/branch_switcher.rs @@ -0,0 +1,166 @@ +use crate::persistence::BranchInfo; +use gpui::{div, prelude::*, px, App, CursorStyle, MouseButton, SharedString, Window}; +use gpui_component::{ActiveTheme, Icon}; + +/// A stateless branch navigation component styled as a bubble on the message border. +/// Displayed at the bottom-right corner of user messages where branches exist. +#[derive(Clone, IntoElement)] +pub struct BranchSwitcherElement { + branch_info: BranchInfo, + session_id: String, +} + +impl BranchSwitcherElement { + pub fn new(branch_info: BranchInfo, session_id: String) -> Self { + Self { + branch_info, + session_id, + } + } +} + +impl RenderOnce for BranchSwitcherElement { + fn render(self, _window: &mut Window, cx: &mut App) -> impl IntoElement { + let has_prev = self.branch_info.active_index > 0; + let has_next = self.branch_info.active_index + 1 < self.branch_info.sibling_ids.len(); + let current = self.branch_info.active_index + 1; + let total = self.branch_info.sibling_ids.len(); + + let muted_color = cx.theme().muted_foreground; + let active_color = cx.theme().foreground; + + // Get sibling IDs for navigation + let prev_node_id = if has_prev { + self.branch_info + .sibling_ids + .get(self.branch_info.active_index - 1) + .copied() + } else { + None + }; + + let next_node_id = if has_next { + self.branch_info + .sibling_ids + .get(self.branch_info.active_index + 1) + .copied() + } else { + None + }; + + let session_id = self.session_id.clone(); + let session_id_for_next = session_id.clone(); + + // Outer container for positioning + div().flex().w_full().justify_end().child( + // The bubble itself - positioned to overlap the border + div() + .flex() + .flex_row() + .items_center() + .gap(px(2.)) + .px_2() + .py(px(2.)) + .mr_2() + .mt(px(-12.)) // Pull up to sit on the border + .bg(cx.theme().background) + .border_1() + .border_color(cx.theme().border) + .rounded_full() + .shadow_xs() + .text_xs() + .children(vec![ + // Previous button + { + let base = div() + .id("branch-prev") + .flex() + .items_center() + .justify_center() + .size(px(18.)) + .rounded_full() + .cursor(if has_prev { + CursorStyle::PointingHand + } else { + CursorStyle::default() + }) + .child( + Icon::default() + .path(SharedString::from("icons/chevron_left.svg")) + .text_color(if has_prev { active_color } else { muted_color }) + .size(px(14.)), + ); + + if has_prev { + base.hover(|s| s.bg(cx.theme().accent.opacity(0.15))) + .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { + if let Some(node_id) = prev_node_id { + if let Some(sender) = + cx.try_global::() + { + let _ = sender.0.try_send( + crate::ui::UiEvent::SwitchBranch { + session_id: session_id.clone(), + new_node_id: node_id, + }, + ); + } + } + }) + .into_any_element() + } else { + base.into_any_element() + } + }, + // Current position display + div() + .px_1() + .text_color(muted_color) + .child(format!("{}/{}", current, total)) + .into_any_element(), + // Next button + { + let base = div() + .id("branch-next") + .flex() + .items_center() + .justify_center() + .size(px(18.)) + .rounded_full() + .cursor(if has_next { + CursorStyle::PointingHand + } else { + CursorStyle::default() + }) + .child( + Icon::default() + .path(SharedString::from("icons/chevron_right.svg")) + .text_color(if has_next { active_color } else { muted_color }) + .size(px(14.)), + ); + + if has_next { + base.hover(|s| s.bg(cx.theme().accent.opacity(0.25))) + .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { + if let Some(node_id) = next_node_id { + if let Some(sender) = + cx.try_global::() + { + let _ = sender.0.try_send( + crate::ui::UiEvent::SwitchBranch { + session_id: session_id_for_next.clone(), + new_node_id: node_id, + }, + ); + } + } + }) + .into_any_element() + } else { + base.into_any_element() + } + }, + ]), + ) + } +} diff --git a/crates/code_assistant/src/ui/gpui/elements.rs b/crates/code_assistant/src/ui/gpui/elements.rs index ee5057ae..5ed04143 100644 --- a/crates/code_assistant/src/ui/gpui/elements.rs +++ b/crates/code_assistant/src/ui/gpui/elements.rs @@ -1,3 +1,4 @@ +use crate::persistence::{BranchInfo, NodeId}; use crate::ui::gpui::file_icons; use crate::ui::gpui::image; use crate::ui::gpui::parameter_renderers::ParameterRendererRegistry; @@ -82,6 +83,11 @@ pub struct MessageContainer { last_block_type_for_hidden_tool: Arc>>, /// Flag indicating a hidden tool completed and we may need a paragraph break needs_paragraph_break_after_hidden_tool: Arc>, + + /// Node ID for this message (for branching support) + node_id: Arc>>, + /// Branch info if this message is part of a branch point + branch_info: Arc>>, } /// Tracks the last block type for paragraph breaks after hidden tools @@ -100,6 +106,8 @@ impl MessageContainer { current_project: Arc::new(Mutex::new(String::new())), last_block_type_for_hidden_tool: Arc::new(Mutex::new(None)), needs_paragraph_break_after_hidden_tool: Arc::new(Mutex::new(false)), + node_id: Arc::new(Mutex::new(None)), + branch_info: Arc::new(Mutex::new(None)), } } @@ -113,6 +121,26 @@ impl MessageContainer { *self.current_project.lock().unwrap() = project; } + /// Set the node ID for this message (for branching support) + pub fn set_node_id(&self, node_id: Option) { + *self.node_id.lock().unwrap() = node_id; + } + + /// Get the node ID for this message + pub fn node_id(&self) -> Option { + *self.node_id.lock().unwrap() + } + + /// Set the branch info for this message + pub fn set_branch_info(&self, branch_info: Option) { + *self.branch_info.lock().unwrap() = branch_info; + } + + /// Get the branch info for this message + pub fn branch_info(&self) -> Option { + self.branch_info.lock().unwrap().clone() + } + /// Mark that a hidden tool completed - paragraph break may be needed before next text pub fn mark_hidden_tool_completed(&self, _cx: &mut Context) { *self.needs_paragraph_break_after_hidden_tool.lock().unwrap() = true; diff --git a/crates/code_assistant/src/ui/gpui/input_area.rs b/crates/code_assistant/src/ui/gpui/input_area.rs index 5717554b..e742273e 100644 --- a/crates/code_assistant/src/ui/gpui/input_area.rs +++ b/crates/code_assistant/src/ui/gpui/input_area.rs @@ -2,7 +2,7 @@ use super::attachment::{AttachmentEvent, AttachmentView}; use super::file_icons; use super::model_selector::{ModelSelector, ModelSelectorEvent}; use super::sandbox_selector::{SandboxSelector, SandboxSelectorEvent}; -use crate::persistence::DraftAttachment; +use crate::persistence::{DraftAttachment, NodeId}; use base64::Engine; use gpui::{ div, prelude::*, px, ClipboardEntry, Context, CursorStyle, Entity, EventEmitter, FocusHandle, @@ -19,6 +19,8 @@ pub enum InputAreaEvent { MessageSubmitted { content: String, attachments: Vec, + /// If set, this message creates a new branch from this parent node + branch_parent_id: Option, }, /// Content changed (for draft saving) ContentChanged { @@ -51,6 +53,11 @@ pub struct InputArea { // Agent state for button rendering agent_is_running: bool, cancel_enabled: bool, + + // Branch editing state + /// When editing a message, this is the parent node ID where the new branch will be created + branch_parent_id: Option, + // Subscriptions _input_subscription: Subscription, _model_selector_subscription: Subscription, @@ -92,6 +99,9 @@ impl InputArea { agent_is_running: false, cancel_enabled: false, + + branch_parent_id: None, + _input_subscription: input_subscription, _model_selector_subscription: model_selector_subscription, _sandbox_selector_subscription: sandbox_selector_subscription, @@ -116,6 +126,25 @@ impl InputArea { self.rebuild_attachment_views(cx); } + /// Set content for editing a message (creates a branch) + pub fn set_content_for_edit( + &mut self, + text: String, + attachments: Vec, + branch_parent_id: Option, + window: &mut Window, + cx: &mut Context, + ) { + self.set_content(text, attachments, window, cx); + self.branch_parent_id = branch_parent_id; + cx.notify(); + } + + /// Clear the edit mode state + fn clear_edit_mode(&mut self) { + self.branch_parent_id = None; + } + /// Sync the dropdown with the current model selection pub fn set_current_model( &mut self, @@ -167,6 +196,9 @@ impl InputArea { // Clear attachments self.attachments.clear(); self.attachment_views.clear(); + + // Clear edit mode + self.clear_edit_mode(); } /// Get current content (text and attachments) @@ -270,6 +302,7 @@ impl InputArea { cx.emit(InputAreaEvent::FocusRequested); } InputEvent::Blur => {} + InputEvent::PressEnter { secondary } => { // Only send message on plain ENTER (not with modifiers) if !secondary { @@ -278,6 +311,9 @@ impl InputArea { // Remove trailing newline if present (from ENTER key press) let cleaned_text = current_text.trim_end_matches('\n').to_string(); + // Capture branch_parent_id before clearing + let branch_parent_id = self.branch_parent_id; + // FIRST: Clear draft before doing anything else cx.emit(InputAreaEvent::ClearDraftRequested); @@ -285,6 +321,7 @@ impl InputArea { cx.emit(InputAreaEvent::MessageSubmitted { content: cleaned_text, attachments: self.attachments.clone(), + branch_parent_id, }); // Clear the input and attachments @@ -335,6 +372,9 @@ impl InputArea { let content = self.text_input.read(cx).value().to_string(); if !content.trim().is_empty() || !self.attachments.is_empty() { + // Capture branch_parent_id before clearing + let branch_parent_id = self.branch_parent_id; + // FIRST: Clear draft before doing anything else cx.emit(InputAreaEvent::ClearDraftRequested); @@ -342,6 +382,7 @@ impl InputArea { cx.emit(InputAreaEvent::MessageSubmitted { content: content.clone(), attachments: self.attachments.clone(), + branch_parent_id, }); // Clear the input and attachments diff --git a/crates/code_assistant/src/ui/gpui/messages.rs b/crates/code_assistant/src/ui/gpui/messages.rs index 941a76b0..b6b7e3b9 100644 --- a/crates/code_assistant/src/ui/gpui/messages.rs +++ b/crates/code_assistant/src/ui/gpui/messages.rs @@ -1,6 +1,10 @@ +use super::branch_switcher::BranchSwitcherElement; use super::elements::MessageContainer; -use gpui::{div, prelude::*, px, rgb, App, Context, Entity, FocusHandle, Focusable, Window}; -use gpui_component::{v_flex, ActiveTheme}; +use gpui::{ + div, prelude::*, px, rgb, App, Context, CursorStyle, Entity, FocusHandle, Focusable, + MouseButton, SharedString, Window, +}; +use gpui_component::{v_flex, ActiveTheme, Icon}; use std::sync::{Arc, Mutex}; /// MessagesView - Component responsible for displaying the message history @@ -8,6 +12,7 @@ pub struct MessagesView { message_queue: Arc>>>, current_pending_message: Arc>>, current_project: Arc>, + current_session_id: Arc>>, focus_handle: FocusHandle, } @@ -20,10 +25,21 @@ impl MessagesView { message_queue, current_pending_message: Arc::new(Mutex::new(None)), current_project: Arc::new(Mutex::new(String::new())), + current_session_id: Arc::new(Mutex::new(None)), focus_handle: cx.focus_handle(), } } + /// Update the current session ID + pub fn set_current_session_id(&self, session_id: Option) { + *self.current_session_id.lock().unwrap() = session_id; + } + + /// Get the current session ID + fn get_current_session_id(&self) -> Option { + self.current_session_id.lock().unwrap().clone() + } + /// Group consecutive image blocks into horizontal galleries for user messages fn group_user_message_elements( elements: Vec>, @@ -107,6 +123,7 @@ impl Render for MessagesView { // Get current project for parameter filtering let current_project = self.get_current_project(); + let current_session_id = self.get_current_session_id(); // Collect all message elements first let message_elements: Vec<_> = messages @@ -114,10 +131,15 @@ impl Render for MessagesView { .map(|msg| { // Update the message container with current project msg.read(cx).set_current_project(current_project.clone()); + + let is_user_message = msg.read(cx).is_user_message(); + let node_id = msg.read(cx).node_id(); + let branch_info = msg.read(cx).branch_info(); + // Create message container with appropriate styling based on role let mut message_container = div().p_3(); - if msg.read(cx).is_user_message() { + if is_user_message { message_container = message_container .m_3() .bg(cx.theme().muted) // Use opaque muted color (darker than card background) @@ -127,30 +149,77 @@ impl Render for MessagesView { .shadow_xs(); } - // Create message container with user badge if needed - let message_container = if msg.read(cx).is_user_message() { - message_container.child( - div() - .flex() - .flex_row() - .items_center() - .gap_2() - .children(vec![ - super::file_icons::render_icon_container( - &super::file_icons::get() - .get_type_icon(super::file_icons::TOOL_USER_INPUT), - 16.0, - user_accent, // Use themed user accent color - "πŸ‘€", - ) - .into_any_element(), - div() - .font_weight(gpui::FontWeight(600.0)) - .text_color(user_accent) // Use themed user accent color - .child("You") + // Create message container with user badge and edit button if needed + let message_container = if is_user_message { + // Build header row with user badge and edit button + let session_id_for_edit = current_session_id.clone(); + let node_id_for_edit = node_id; + + let header_row = div() + .flex() + .flex_row() + .items_center() + .justify_between() + .w_full() + .child( + // Left side: User badge + div() + .flex() + .flex_row() + .items_center() + .gap_2() + .children(vec![ + super::file_icons::render_icon_container( + &super::file_icons::get() + .get_type_icon(super::file_icons::TOOL_USER_INPUT), + 16.0, + user_accent, + "πŸ‘€", + ) .into_any_element(), - ]), - ) + div() + .font_weight(gpui::FontWeight(600.0)) + .text_color(user_accent) + .child("You") + .into_any_element(), + ]), + ) + .when(node_id.is_some(), |el| { + // Right side: Edit button (only shown when node_id is present) + el.child( + div() + .id("edit-message-btn") + .p_1() + .rounded_sm() + .cursor(CursorStyle::PointingHand) + .hover(|s| s.bg(cx.theme().accent.opacity(0.25))) + .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { + if let (Some(session_id), Some(node_id)) = + (session_id_for_edit.clone(), node_id_for_edit) + { + // Send StartMessageEdit event + if let Some(sender) = + cx.try_global::() + { + let _ = sender.0.try_send( + crate::ui::UiEvent::StartMessageEdit { + session_id, + node_id, + }, + ); + } + } + }) + .child( + Icon::default() + .path(SharedString::from("icons/pencil.svg")) + .text_color(cx.theme().muted_foreground) + .size_4(), + ), + ) + }); + + message_container.child(header_row) } else { message_container }; @@ -158,7 +227,7 @@ impl Render for MessagesView { // Render all block elements with special handling for user messages let elements = msg.read(cx).elements(); - if msg.read(cx).is_user_message() { + let message_container = if is_user_message { // For user messages, group consecutive image blocks into horizontal galleries let container_children = Self::group_user_message_elements(elements, cx); message_container.children(container_children) @@ -169,7 +238,23 @@ impl Render for MessagesView { .map(|element| element.into_any_element()) .collect(); message_container.children(container_children) + }; + + // Add branch switcher if branch_info is present (only for user messages) + if is_user_message { + if let (Some(branch_info), Some(session_id)) = + (branch_info, current_session_id.clone()) + { + // Only show if there are multiple siblings (actual branches) + if branch_info.sibling_ids.len() > 1 { + return message_container + .child(BranchSwitcherElement::new(branch_info, session_id)) + .into_any_element(); + } + } } + + message_container.into_any_element() }) .collect(); diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index 605f15a6..f72c6ca3 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -1,6 +1,7 @@ pub mod assets; pub mod attachment; pub mod auto_scroll; +pub mod branch_switcher; pub mod chat_sidebar; pub mod content_renderer; pub mod diff_renderer; @@ -101,6 +102,17 @@ pub struct Gpui { current_model: Arc>>, // Current sandbox selection current_sandbox_policy: Arc>>, + + // Pending message edit state (for branching) + pending_edit: Arc>>, +} + +/// State for a pending message edit (for branching) +#[derive(Clone, Debug)] +pub struct PendingEdit { + pub content: String, + pub attachments: Vec, + pub branch_parent_id: Option, } fn init(cx: &mut App) { @@ -303,6 +315,9 @@ impl Gpui { current_model: Arc::new(Mutex::new(None)), // Current sandbox selection current_sandbox_policy: Arc::new(Mutex::new(None)), + + // Pending message edit state + pending_edit: Arc::new(Mutex::new(None)), } } @@ -471,11 +486,15 @@ impl Gpui { UiEvent::DisplayUserInput { content, attachments, + node_id, } => { let mut queue = self.message_queue.lock().unwrap(); let result = cx.new(|cx| { let new_message = MessageContainer::with_role(MessageRole::User, cx); + // Set node_id for edit button support + new_message.set_node_id(node_id); + // Add text content if not empty if !content.is_empty() { new_message.add_text_block(&content, cx); @@ -618,9 +637,11 @@ impl Gpui { warn!("Using initial project: '{}'", current_project); - // Update MessagesView with current project + // Update MessagesView with current project and session ID + let session_id_for_messages = session_id.clone(); self.update_messages_view(cx, |messages_view, _cx| { messages_view.set_current_project(current_project.clone()); + messages_view.set_current_session_id(Some(session_id_for_messages)); }); } @@ -648,6 +669,8 @@ impl Gpui { let mut queue = self.message_queue.lock().unwrap(); // Check if we can reuse the last container (same role) + // Note: For user messages, we always create a new container to preserve + // node_id and branch_info for each message let needs_new_container = if let Some(last_container) = queue.last() { let last_role = cx .update_entity(last_container, |container, _cx| { @@ -658,6 +681,7 @@ impl Gpui { } }) .expect("Failed to get container role"); + // User messages always get their own container (for branching) last_role == MessageRole::User || last_role != message_data.role } else { true @@ -666,12 +690,18 @@ impl Gpui { if needs_new_container { // Create new container for this role let container = cx - .new(|cx| MessageContainer::with_role(message_data.role, cx)) + .new(|cx| { + MessageContainer::with_role(message_data.role.clone(), cx) + }) .expect("Failed to create message container"); - // Set current project on the new container + // Set current project, node_id, and branch_info on the new container + let node_id = message_data.node_id; + let branch_info = message_data.branch_info.clone(); self.update_container(&container, cx, |container, _cx| { container.set_current_project(current_project.clone()); + container.set_node_id(node_id); + container.set_branch_info(branch_info); }); queue.push(container.clone()); @@ -804,16 +834,19 @@ impl Gpui { queue.clear(); cx.refresh().expect("Failed to refresh windows"); } + UiEvent::SendUserMessage { message, session_id, attachments, + branch_parent_id, } => { debug!( - "UI: SendUserMessage event for session {}: {} (with {} attachments)", + "UI: SendUserMessage event for session {}: {} (with {} attachments, branch_parent: {:?})", session_id, message, - attachments.len() + attachments.len(), + branch_parent_id ); // Clear any existing error when user sends a new message *self.current_error.lock().unwrap() = None; @@ -823,6 +856,7 @@ impl Gpui { session_id, message, attachments, + branch_parent_id, }); } else { warn!("UI: No backend event sender available"); @@ -1020,6 +1054,7 @@ impl Gpui { path.display() ); } + UiEvent::CancelSubAgent { tool_id } => { debug!("UI: CancelSubAgent event for tool_id: {}", tool_id); // Forward to backend with current session ID @@ -1034,6 +1069,234 @@ impl Gpui { warn!("UI: CancelSubAgent requested but no active session"); } } + + // === Session Branching Events === + UiEvent::StartMessageEdit { + session_id, + node_id, + } => { + debug!( + "UI: StartMessageEdit event for session {} node {}", + session_id, node_id + ); + // Forward to backend to get message content + if let Some(sender) = self.backend_event_sender.lock().unwrap().as_ref() { + let _ = sender.try_send(BackendEvent::StartMessageEdit { + session_id, + node_id, + }); + } + } + UiEvent::SwitchBranch { + session_id, + new_node_id, + } => { + debug!( + "UI: SwitchBranch event for session {} to node {}", + session_id, new_node_id + ); + // Forward to backend to perform branch switch + if let Some(sender) = self.backend_event_sender.lock().unwrap().as_ref() { + let _ = sender.try_send(BackendEvent::SwitchBranch { + session_id, + new_node_id, + }); + } + } + + UiEvent::MessageEditReady { + content, + attachments, + branch_parent_id, + messages, + tool_results, + } => { + debug!( + "UI: MessageEditReady event - content len: {}, attachments: {}, parent: {:?}, {} messages", + content.len(), + attachments.len(), + branch_parent_id, + messages.len() + ); + + // Get current session_id without holding lock during SetMessages processing + let session_id = self.current_session_id.lock().unwrap().clone(); + + // Get current project for new containers + let current_project = if let Some(ref session_id) = session_id { + let sessions = self.chat_sessions.lock().unwrap(); + sessions + .iter() + .find(|s| s.id == *session_id) + .map(|s| s.initial_project.clone()) + .unwrap_or_default() + } else { + String::new() + }; + + // Clear existing messages and rebuild with truncated set + // (Inline version of SetMessages logic to avoid recursive call) + { + let mut queue = self.message_queue.lock().unwrap(); + queue.clear(); + } + + // Update MessagesView with current project and session ID + if let Some(ref session_id) = session_id { + let session_id_for_messages = session_id.clone(); + self.update_messages_view(cx, |messages_view, _cx| { + messages_view.set_current_project(current_project.clone()); + messages_view.set_current_session_id(Some(session_id_for_messages)); + }); + } + + // Process message data + for message_data in messages { + let current_container = { + let mut queue = self.message_queue.lock().unwrap(); + + let needs_new_container = if let Some(last_container) = queue.last() { + let last_role = cx + .update_entity(last_container, |container, _cx| { + if container.is_user_message() { + MessageRole::User + } else { + MessageRole::Assistant + } + }) + .expect("Failed to get container role"); + last_role == MessageRole::User || last_role != message_data.role + } else { + true + }; + + if needs_new_container { + let container = cx + .new(|cx| { + MessageContainer::with_role(message_data.role.clone(), cx) + }) + .expect("Failed to create message container"); + + let node_id = message_data.node_id; + let branch_info = message_data.branch_info.clone(); + self.update_container(&container, cx, |container, _cx| { + container.set_current_project(current_project.clone()); + container.set_node_id(node_id); + container.set_branch_info(branch_info); + }); + + queue.push(container.clone()); + container + } else { + let container = queue.last().unwrap().clone(); + self.update_container(&container, cx, |container, _cx| { + container.set_current_project(current_project.clone()); + }); + container + } + }; + + self.process_fragments_for_container( + ¤t_container, + message_data.fragments, + cx, + ); + } + + // Apply tool results + for tool_result in tool_results { + self.update_all_messages(cx, |message_container, cx| { + message_container.update_tool_status( + &tool_result.tool_id, + tool_result.status, + tool_result.message.clone(), + tool_result.output.clone(), + cx, + ); + }); + } + + // Ensure we end with an Assistant container for the edit response + { + let mut queue = self.message_queue.lock().unwrap(); + let needs_assistant_container = if let Some(last) = queue.last() { + cx.update_entity(last, |message, _cx| message.is_user_message()) + .expect("Failed to check container role") + } else { + true + }; + + if needs_assistant_container { + let assistant_container = cx + .new(|cx| MessageContainer::with_role(MessageRole::Assistant, cx)) + .expect("Failed to create assistant container"); + queue.push(assistant_container); + } + } + + // Store pending edit state for RootView to pick up on refresh + self.set_pending_edit(PendingEdit { + content, + attachments, + branch_parent_id, + }); + + // Refresh UI to trigger RootView to process the pending edit + cx.refresh().expect("Failed to refresh windows"); + } + UiEvent::BranchSwitched { + session_id, + messages, + tool_results, + plan, + } => { + debug!( + "UI: BranchSwitched event for session {} with {} messages", + session_id, + messages.len() + ); + // TODO Phase 4: Update messages display with new branch content + // For now, we can reuse the SetMessages logic + self.process_ui_event_async( + UiEvent::SetMessages { + messages, + session_id: Some(session_id), + tool_results, + }, + cx, + ); + self.process_ui_event_async(UiEvent::UpdatePlan { plan }, cx); + } + + UiEvent::UpdateBranchInfo { + node_id, + branch_info, + } => { + debug!( + "UI: UpdateBranchInfo for node {} with {} siblings", + node_id, + branch_info.sibling_ids.len() + ); + + // Find the message container with this node_id and update its branch_info + let queue = self.message_queue.lock().unwrap(); + for container in queue.iter() { + let container_node_id = cx + .update_entity(container, |c, _cx| c.node_id()) + .ok() + .flatten(); + + if container_node_id == Some(node_id) { + let branch_info_clone = branch_info.clone(); + self.update_container(container, cx, |c, _cx| { + c.set_branch_info(Some(branch_info_clone)); + }); + break; + } + } + + cx.refresh().expect("Failed to refresh windows"); + } } } @@ -1176,6 +1439,16 @@ impl Gpui { self.current_sandbox_policy.lock().unwrap().clone() } + /// Get and clear pending edit (used by RootView to pick up edit state) + pub fn take_pending_edit(&self) -> Option { + self.pending_edit.lock().unwrap().take() + } + + /// Set pending edit state + pub fn set_pending_edit(&self, edit: PendingEdit) { + *self.pending_edit.lock().unwrap() = Some(edit); + } + // Extended draft management methods with attachments pub fn save_draft_for_session( &self, @@ -1286,7 +1559,7 @@ impl Gpui { } // Handle backend responses - fn handle_backend_response(&self, response: BackendResponse, _cx: &mut AsyncApp) { + fn handle_backend_response(&self, response: BackendResponse, cx: &mut AsyncApp) { match response { BackendResponse::SessionCreated { session_id } => { debug!("Received BackendResponse::SessionCreated"); @@ -1378,6 +1651,7 @@ impl Gpui { ); } } + BackendResponse::SubAgentCancelled { session_id, tool_id, @@ -1389,6 +1663,58 @@ impl Gpui { // The sub-agent will update its own UI state via the normal tool output mechanism // No additional UI update needed here } + + // Session branching responses + BackendResponse::MessageEditReady { + session_id, + content, + attachments, + branch_parent_id, + messages, + tool_results, + } => { + debug!( + "Received BackendResponse::MessageEditReady for session {} with {} chars, {} attachments, {} messages", + session_id, + content.len(), + attachments.len(), + messages.len() + ); + + // Forward to UI as event + self.process_ui_event_async( + UiEvent::MessageEditReady { + content: content.clone(), + attachments: attachments.clone(), + branch_parent_id, + messages: messages.clone(), + tool_results: tool_results.clone(), + }, + cx, + ); + } + BackendResponse::BranchSwitched { + session_id, + messages, + tool_results, + plan, + } => { + debug!( + "Received BackendResponse::BranchSwitched for session {} with {} messages", + session_id, + messages.len() + ); + // Forward to UI as event + self.process_ui_event_async( + UiEvent::BranchSwitched { + session_id: session_id.clone(), + messages: messages.clone(), + tool_results: tool_results.clone(), + plan: plan.clone(), + }, + cx, + ); + } } } } diff --git a/crates/code_assistant/src/ui/gpui/root.rs b/crates/code_assistant/src/ui/gpui/root.rs index 0166283f..b4d8ae13 100644 --- a/crates/code_assistant/src/ui/gpui/root.rs +++ b/crates/code_assistant/src/ui/gpui/root.rs @@ -170,9 +170,16 @@ impl RootView { InputAreaEvent::MessageSubmitted { content, attachments, + branch_parent_id, } => { if let Some(session_id) = self.current_session_id.clone() { - self.send_message(&session_id, content.clone(), attachments.clone(), cx); + self.send_message( + &session_id, + content.clone(), + attachments.clone(), + *branch_parent_id, + cx, + ); } } InputAreaEvent::ContentChanged { @@ -278,6 +285,7 @@ impl RootView { session_id: &str, content: String, attachments: Vec, + branch_parent_id: Option, cx: &mut Context, ) { if content.trim().is_empty() && attachments.is_empty() { @@ -299,6 +307,17 @@ impl RootView { false }; + // Log branch editing info + if branch_parent_id.is_some() { + tracing::info!( + "RootView: Sending edited message (branch from {:?}) to session {}: {} (with {} attachments)", + branch_parent_id, + session_id, + content, + attachments.len() + ); + } + if agent_is_running { // Queue the message for the running agent tracing::info!( @@ -315,20 +334,37 @@ impl RootView { } else { // Send message normally (agent is idle) tracing::info!( - "RootView: Sending user message to session {}: {} (with {} attachments)", + "RootView: Sending user message to session {}: {} (with {} attachments, branch_parent: {:?})", session_id, content, - attachments.len() + attachments.len(), + branch_parent_id ); let _ = sender.0.try_send(UiEvent::SendUserMessage { message: content.clone(), session_id: session_id.to_string(), attachments: attachments.clone(), + branch_parent_id, }); } } } + /// Handle message edit ready event - load content into input area + pub fn handle_message_edit_ready( + &mut self, + content: String, + attachments: Vec, + branch_parent_id: Option, + window: &mut gpui::Window, + cx: &mut Context, + ) { + self.input_area.update(cx, |input_area, cx| { + input_area.set_content_for_edit(content, attachments, branch_parent_id, window, cx); + }); + cx.notify(); + } + fn save_draft_for_session( &self, session_id: &str, @@ -656,6 +692,19 @@ impl Render for RootView { } } + // Check for pending edit (message editing for branching) + if let Some(gpui) = cx.try_global::() { + if let Some(pending_edit) = gpui.take_pending_edit() { + self.handle_message_edit_ready( + pending_edit.content, + pending_edit.attachments, + pending_edit.branch_parent_id, + window, + cx, + ); + } + } + // Ensure InputArea stays in sync with the current model let selected_model = self.input_area.read(cx).current_model(); if selected_model != current_model { diff --git a/crates/code_assistant/src/ui/terminal/app.rs b/crates/code_assistant/src/ui/terminal/app.rs index 20482c38..e8439e8d 100644 --- a/crates/code_assistant/src/ui/terminal/app.rs +++ b/crates/code_assistant/src/ui/terminal/app.rs @@ -137,6 +137,7 @@ async fn event_loop( session_id, message, attachments: Vec::new(), + branch_parent_id: None, // Terminal UI doesn't support branching yet } } _ => BackendEvent::QueueUserMessage { @@ -442,6 +443,7 @@ impl TerminalTuiApp { policy ))); } + BackendResponse::SubAgentCancelled { session_id: _, tool_id: _, @@ -449,6 +451,10 @@ impl TerminalTuiApp { // Sub-agent cancellation handled; the sub-agent will // update its tool output via the normal mechanism } + BackendResponse::MessageEditReady { .. } + | BackendResponse::BranchSwitched { .. } => { + // Session branching not supported in terminal UI + } } } }); @@ -508,6 +514,7 @@ impl TerminalTuiApp { session_id: session_id.clone(), message: task.clone(), attachments: Vec::new(), + branch_parent_id: None, }); } diff --git a/crates/code_assistant/src/ui/terminal/ui.rs b/crates/code_assistant/src/ui/terminal/ui.rs index be3ef547..68cb27f9 100644 --- a/crates/code_assistant/src/ui/terminal/ui.rs +++ b/crates/code_assistant/src/ui/terminal/ui.rs @@ -171,9 +171,11 @@ impl UserInterface for TerminalTuiUI { renderer_guard.clear_all_messages(); } } + UiEvent::DisplayUserInput { content, attachments, + node_id: _, // Terminal UI doesn't support branching } => { debug!("Displaying user input: {}", content); diff --git a/crates/code_assistant/src/ui/ui_events.rs b/crates/code_assistant/src/ui/ui_events.rs index 6415e1bd..4530608a 100644 --- a/crates/code_assistant/src/ui/ui_events.rs +++ b/crates/code_assistant/src/ui/ui_events.rs @@ -1,4 +1,4 @@ -use crate::persistence::{ChatMetadata, DraftAttachment}; +use crate::persistence::{BranchInfo, ChatMetadata, DraftAttachment, NodeId}; use crate::session::instance::SessionActivityState; use crate::types::PlanState; use crate::ui::gpui::elements::MessageRole; @@ -11,6 +11,10 @@ use std::path::PathBuf; pub struct MessageData { pub role: MessageRole, pub fragments: Vec, + /// Optional node ID for branching support + pub node_id: Option, + /// Optional branch info if this message is part of a branch + pub branch_info: Option, } /// Tool execution result data for UI updates @@ -29,6 +33,8 @@ pub enum UiEvent { DisplayUserInput { content: String, attachments: Vec, + /// Node ID for this message (for edit button support) + node_id: Option, }, /// Display a system-generated compaction divider message DisplayCompactionSummary { summary: String }, @@ -88,6 +94,8 @@ pub enum UiEvent { message: String, session_id: String, attachments: Vec, + /// If set, creates a new branch from this parent node instead of appending to active path + branch_parent_id: Option, }, /// Update metadata for a single session without refreshing the entire list UpdateSessionMetadata { metadata: ChatMetadata }, @@ -121,9 +129,61 @@ pub enum UiEvent { UpdateCurrentModel { model_name: String }, /// Update the current sandbox selection in the UI UpdateSandboxPolicy { policy: SandboxPolicy }, + /// Cancel a running sub-agent by its tool id CancelSubAgent { tool_id: String }, + // === Session Branching Events === + /// Request to start editing a message (creates a branch point) + /// UI should load the message content into the input area + StartMessageEdit { + session_id: String, + /// The node ID of the message being edited + node_id: NodeId, + }, + + /// Switch to a different branch at a branch point + SwitchBranch { + session_id: String, + /// The node ID to switch to (a sibling of the current node at a branch point) + new_node_id: NodeId, + }, + + /// Response: Message content loaded for editing + /// Sent in response to StartMessageEdit + MessageEditReady { + /// The text content of the message + content: String, + /// Any attachments from the original message + attachments: Vec, + /// The parent node ID where the new branch will be created + branch_parent_id: Option, + /// Messages up to (but not including) the message being edited + messages: Vec, + /// Tool results for the truncated path + tool_results: Vec, + }, + + /// Response: Branch switch completed, new messages to display + BranchSwitched { + session_id: String, + /// Full message list for the new active path + messages: Vec, + /// Tool results for the new path + tool_results: Vec, + /// Updated plan for the new path + plan: PlanState, + }, + + /// Update the branch info for a specific message node + /// Used when a new branch is created to update the UI for the parent message + UpdateBranchInfo { + /// The node ID whose branch info should be updated + node_id: NodeId, + /// The updated branch info (siblings at this branch point) + branch_info: BranchInfo, + }, + // === Resource Events (for tool operations) === /// A file was loaded/read by a tool ResourceLoaded { project: String, path: PathBuf }, diff --git a/docs/session-branching.md b/docs/session-branching.md new file mode 100644 index 00000000..f6f5dbfd --- /dev/null +++ b/docs/session-branching.md @@ -0,0 +1,715 @@ +# Session Branching Feature + +## Overview + +Session Branching ermΓΆglicht das Erstellen von alternativen GesprΓ€chsverlΓ€ufen durch Editieren bereits abgeschickter User-Nachrichten. Statt den ursprΓΌnglichen Verlauf zu lΓΆschen, entstehen "Abzweigungen" (Branches), zwischen denen der User umschalten kann. + +### Ziele + +1. **Exploration**: "Was wΓ€re wenn ich etwas anders formuliert hΓ€tte?" +2. **Preservation**: Bestehende VerlΓ€ufe bleiben erhalten +3. **Eleganz**: Einfache, mΓ€chtige Datenstruktur + +## Architektur-Design + +### Kernkonzept: Baum statt Liste + +Die aktuelle Architektur speichert Messages als lineare Liste (`Vec`). Das neue Design verwendet einen **gerichteten Baum** (DAG - Directed Acyclic Graph): + +``` + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Root (Start) β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ User: "Hi" β”‚ node_id: 1 + β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Assistant... β”‚ node_id: 2 + β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ β”‚ β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ User: "Fix A" β”‚ β”‚ User: "Fix B" β”‚ β”‚ User: "Fix C" β”‚ + β”‚ (original) β”‚ β”‚ (branch 1) β”‚ β”‚ (branch 2) β”‚ + β”‚ node_id: 3 β”‚ β”‚ node_id: 5 β”‚ β”‚ node_id: 7 β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ β”‚ β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β” + β”‚ Assistant... β”‚ β”‚ Assistant... β”‚ β”‚ Assistant... β”‚ + β”‚ node_id: 4 β”‚ β”‚ node_id: 6 β”‚ β”‚ node_id: 8 β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +### Datenstrukturen + +#### MessageNode + +Ersetzt `Message` als Speichereinheit innerhalb einer Session: + +```rust +// In crates/code_assistant/src/persistence.rs + +/// Unique identifier for a message node within a session +pub type NodeId = u64; + +/// A single message node in the conversation tree +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct MessageNode { + /// Unique ID within this session + pub id: NodeId, + + /// The actual message content + pub message: Message, + + /// Parent node ID (None for root/first message) + pub parent_id: Option, + + /// Creation timestamp (for ordering siblings) + pub created_at: SystemTime, + + /// Plan state snapshot (only set if plan changed in this message's response) + /// Used for efficient plan reconstruction when switching branches + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plan_snapshot: Option, +} + +/// A path through the conversation tree (list of node IDs from root to leaf) +pub type ConversationPath = Vec; +``` + +#### Erweiterte ChatSession + +```rust +// In crates/code_assistant/src/persistence.rs + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ChatSession { + // ... existing fields ... + + /// All message nodes in the session (tree structure) + /// Key: NodeId, Value: MessageNode + /// Using BTreeMap for ordered iteration and efficient lookup + #[serde(default)] + pub message_nodes: BTreeMap, + + /// The currently active path through the tree + /// This determines which messages are shown and sent to LLM + #[serde(default)] + pub active_path: ConversationPath, + + /// Counter for generating unique node IDs + #[serde(default)] + pub next_node_id: NodeId, + + /// Legacy: Old linear message list (for migration) + /// Will be migrated to message_nodes on first load + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub messages: Vec, +} +``` + +#### Branch-Metadaten fΓΌr UI + +```rust +// In crates/code_assistant/src/ui/ui_events.rs + +/// Information about a branch point in the conversation +#[derive(Debug, Clone)] +pub struct BranchInfo { + /// Node ID where the branch occurs (parent node) + pub branch_point_id: NodeId, + + /// All sibling node IDs (different continuations) + pub sibling_ids: Vec, + + /// Index of the currently active sibling (0-based) + pub active_index: usize, + + /// Total number of branches at this point + pub total_branches: usize, +} +``` + +### Migration bestehender Sessions + +```rust +impl ChatSession { + /// Migrate legacy linear messages to tree structure + pub fn migrate_to_tree_structure(&mut self) -> Result<()> { + if self.message_nodes.is_empty() && !self.messages.is_empty() { + // Convert linear messages to tree + let mut parent_id: Option = None; + + for message in self.messages.drain(..) { + let node_id = self.next_node_id; + self.next_node_id += 1; + + let node = MessageNode { + id: node_id, + message, + parent_id, + created_at: SystemTime::now(), + }; + + self.message_nodes.insert(node_id, node); + self.active_path.push(node_id); + parent_id = Some(node_id); + } + } + Ok(()) + } +} +``` + +## UI-Γ„nderungen + +### 1. Edit-Button auf User-Nachrichten + +In `crates/code_assistant/src/ui/gpui/messages.rs`: + +```rust +// FΓΌr User-Nachrichten einen Edit-IconButton hinzufΓΌgen +if msg.read(cx).is_user_message() { + message_container = message_container.child( + div() + .flex() + .flex_row() + .items_center() + .justify_between() // Verteilt Inhalt zwischen links und rechts + .gap_2() + .children(vec![ + // Linke Seite: User Badge (wie bisher) + div().flex().flex_row().items_center().gap_2() + .child(/* user icon */) + .child(/* "You" label */), + + // Rechte Seite: Edit Button (nur bei Hover sichtbar) + IconButton::new("edit") + .icon(IconName::Pencil) + .tooltip("Edit this message and create a new branch") + .on_click(cx.listener(move |this, _, window, cx| { + this.start_message_edit(node_id, window, cx); + })) + ]) + ) +} +``` + +### 2. Branch-Switcher unter User-Nachrichten + +Wenn ein Punkt mehrere Branches hat, wird ein kleiner Umschalter angezeigt: + +```rust +/// Branch switcher component +pub struct BranchSwitcher { + branch_info: BranchInfo, + on_switch: Callback, // Called with new index +} + +impl Render for BranchSwitcher { + fn render(&mut self, window: &mut Window, cx: &mut Context) -> impl IntoElement { + div() + .flex() + .flex_row() + .items_center() + .gap_1() + .text_xs() + .text_color(cx.theme().muted_foreground) + .children([ + // Left arrow button (disabled if at index 0) + IconButton::new("prev") + .icon(IconName::ChevronLeft) + .disabled(self.branch_info.active_index == 0) + .on_click(/* switch to prev */), + + // "2/3" indicator + div().child(format!( + "{}/{}", + self.branch_info.active_index + 1, + self.branch_info.total_branches + )), + + // Right arrow button + IconButton::new("next") + .icon(IconName::ChevronRight) + .disabled(self.branch_info.active_index >= self.branch_info.total_branches - 1) + .on_click(/* switch to next */), + ]) + } +} +``` + +**Darstellung:** +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ πŸ‘€ You ✏️ β”‚ +β”‚ β”‚ +β”‚ "Please fix the bug in file.rs" β”‚ +β”‚ β”‚ +β”‚ β—€ 2/3 β–Ά β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +### 3. Edit-Workflow + +1. User klickt Edit-Button auf einer User-Nachricht +2. Die Message-Inhalte und Attachments werden in den Input-Bereich geladen +3. UI markiert den "Branch-Punkt" (die Node, nach der die neue Nachricht eingefΓΌgt wird) +4. User editiert und sendet +5. System erstellt neue MessageNode mit gleichem `parent_id` wie das Original +6. `active_path` wird aktualisiert auf den neuen Branch + +```rust +// In input_area.rs oder messages.rs +pub fn start_message_edit(&mut self, node_id: NodeId, cx: &mut Context) { + // 1. Find the node + let node = self.get_node(node_id); + + // 2. Load content into input area + let content = node.message.extract_text_content(); + let attachments = node.message.extract_attachments(); + + // 3. Set editing state (parent of the node being edited) + self.editing_branch_point = node.parent_id; + + // 4. Populate input + self.input_area.set_content(content); + self.input_area.set_attachments(attachments); + + // 5. Focus input + cx.focus(&self.input_area); +} +``` + +## Neue UI-Events + +```rust +// In crates/code_assistant/src/ui/ui_events.rs + +pub enum UiEvent { + // ... existing events ... + + /// Start editing a message (load content to input, set branch point) + StartMessageEdit { + node_id: NodeId, + session_id: String, + }, + + /// Switch to a different branch at a branch point + SwitchBranch { + branch_point_id: NodeId, + new_active_sibling_id: NodeId, + session_id: String, + }, + + /// Notify UI about branch info for a node (for rendering branch switcher) + UpdateBranchInfo { + node_id: NodeId, + branch_info: BranchInfo, + }, + + /// Set messages with branch information + SetMessagesWithBranches { + messages: Vec, + branch_infos: Vec, + session_id: Option, + tool_results: Vec, + }, +} +``` + +## Session-Management Γ„nderungen + +### ChatSession Methoden + +```rust +impl ChatSession { + /// Get the linearized message history for the active path + pub fn get_active_messages(&self) -> Vec<&Message> { + self.active_path + .iter() + .filter_map(|id| self.message_nodes.get(id)) + .map(|node| &node.message) + .collect() + } + + /// Get all children of a node (for branch detection) + pub fn get_children(&self, parent_id: Option) -> Vec<&MessageNode> { + self.message_nodes + .values() + .filter(|node| node.parent_id == parent_id) + .collect() + } + + /// Check if a node has multiple children (is a branch point) + pub fn is_branch_point(&self, node_id: NodeId) -> bool { + self.get_children(Some(node_id)).len() > 1 + } + + /// Get branch info for a specific node + pub fn get_branch_info(&self, node_id: NodeId) -> Option { + let node = self.message_nodes.get(&node_id)?; + let siblings: Vec<_> = self.get_children(node.parent_id) + .into_iter() + .map(|n| n.id) + .collect(); + + if siblings.len() <= 1 { + return None; // No branching here + } + + let active_index = siblings.iter().position(|&id| id == node_id)?; + + Some(BranchInfo { + branch_point_id: node.parent_id.unwrap_or(0), + sibling_ids: siblings, + active_index, + total_branches: siblings.len(), + }) + } + + /// Add a new message as a child of the given parent + /// Returns the new node ID + pub fn add_message(&mut self, message: Message, parent_id: Option) -> NodeId { + let node_id = self.next_node_id; + self.next_node_id += 1; + + let node = MessageNode { + id: node_id, + message, + parent_id, + created_at: SystemTime::now(), + }; + + self.message_nodes.insert(node_id, node); + node_id + } + + /// Switch to a different branch at a branch point + /// Updates active_path to follow the new branch + pub fn switch_branch(&mut self, new_node_id: NodeId) -> Result<()> { + let node = self.message_nodes.get(&new_node_id) + .ok_or_else(|| anyhow::anyhow!("Node not found: {}", new_node_id))?; + + // Find where in active_path the parent is + if let Some(parent_id) = node.parent_id { + if let Some(parent_pos) = self.active_path.iter().position(|&id| id == parent_id) { + // Truncate path after parent + self.active_path.truncate(parent_pos + 1); + + // Build path from new node to deepest descendant on current active path + self.extend_active_path_from(new_node_id); + } + } else { + // Switching root node + self.active_path.clear(); + self.extend_active_path_from(new_node_id); + } + + Ok(()) + } + + /// Extend active_path from a given node, following the "most recent" child + fn extend_active_path_from(&mut self, start_node_id: NodeId) { + self.active_path.push(start_node_id); + + let mut current_id = start_node_id; + loop { + // Find children of current node + let mut children: Vec<_> = self.get_children(Some(current_id)) + .into_iter() + .collect(); + + if children.is_empty() { + break; + } + + // Sort by created_at, take most recent (or first in existing path) + children.sort_by_key(|n| n.created_at); + + // Prefer child that was in the original active_path, otherwise most recent + let next_child = children.last().unwrap(); + self.active_path.push(next_child.id); + current_id = next_child.id; + } + } +} +``` + +### SessionState Γ„nderungen + +```rust +// In crates/code_assistant/src/session/mod.rs + +pub struct SessionState { + pub session_id: String, + pub name: String, + + // Replace Vec with tree data + pub message_nodes: BTreeMap, + pub active_path: ConversationPath, + pub next_node_id: NodeId, + + pub tool_executions: Vec, + pub plan: PlanState, + pub config: SessionConfig, + pub next_request_id: Option, + pub model_config: Option, +} + +impl SessionState { + /// Get linearized messages for the active path (for LLM requests) + pub fn get_active_messages(&self) -> Vec { + self.active_path + .iter() + .filter_map(|id| self.message_nodes.get(id)) + .map(|node| node.message.clone()) + .collect() + } +} +``` + +## Agent-Γ„nderungen + +### Message-Handling im Agent + +```rust +// In crates/code_assistant/src/agent/runner.rs + +impl Agent { + // Die message_history bleibt als "flache" Ansicht fΓΌr die aktuelle Iteration + // Sie wird aus dem aktiven Pfad der Session rekonstruiert + + /// Load session state including branch information + pub async fn load_from_session_state(&mut self, state: SessionState) -> Result<()> { + // ... existing code ... + + // Load the linearized active path as message history + self.message_history = state.get_active_messages(); + + // Store tree structure for state saving + self.message_nodes = state.message_nodes; + self.active_path = state.active_path; + self.next_node_id = state.next_node_id; + + // ... rest of loading ... + } + + /// Save state with tree structure + fn save_state(&mut self) -> Result<()> { + let session_state = SessionState { + session_id: self.session_id.clone().unwrap_or_default(), + name: self.session_name.clone(), + message_nodes: self.message_nodes.clone(), + active_path: self.active_path.clone(), + next_node_id: self.next_node_id, + // ... other fields ... + }; + + self.state_persistence.save_agent_state(session_state)?; + Ok(()) + } + + /// Append a message to the active path + pub fn append_message(&mut self, message: Message) -> Result<()> { + // Get parent (last node in active path) + let parent_id = self.active_path.last().copied(); + + // Create new node + let node_id = self.next_node_id; + self.next_node_id += 1; + + let node = MessageNode { + id: node_id, + message: message.clone(), + parent_id, + created_at: SystemTime::now(), + }; + + // Add to tree + self.message_nodes.insert(node_id, node); + + // Extend active path + self.active_path.push(node_id); + + // Keep linearized history in sync + self.message_history.push(message); + + self.save_state()?; + Ok(()) + } +} +``` + + +## Implementierungsplan + +### Phase 1: Datenstruktur (Backend) βœ… COMPLETE + +1. **Neue Typen definieren** (`persistence.rs`) βœ… + - `NodeId` type alias + - `MessageNode` struct + - `ConversationPath` type alias + - `BranchInfo` struct + +2. **ChatSession erweitern** (`persistence.rs`) βœ… + - Neue Felder: `message_nodes`, `active_path`, `next_node_id` + - Migration-Logik fΓΌr bestehende Sessions + - Methoden: `get_active_messages()`, `add_message()`, `switch_branch()`, etc. + +3. **SessionState anpassen** (`session/mod.rs`) βœ… + - Gleiche Struktur wie ChatSession + - Hilfsmethoden fΓΌr Linearisierung + +4. **Tests schreiben** βœ… + - Migration von linear zu tree + - Branch-Erstellung + - Branch-Switching + - Pfad-Berechnung + +### Phase 2: Agent-Anpassungen βœ… COMPLETE + +1. **Agent-State erweitern** (`agent/runner.rs`) βœ… + - Neue Felder fΓΌr Tree-Struktur + - `append_message()` anpassen + +2. **State-Laden/Speichern** (`agent/persistence.rs`) βœ… + - Tree-Struktur in SessionState + +3. **LLM-Request-Building** βœ… + - Nur aktiven Pfad an LLM senden + +### Phase 3: UI-Backend-Kommunikation βœ… COMPLETE + +1. **Neue UiEvents** (`ui/ui_events.rs`) βœ… + - `StartMessageEdit` + - `SwitchBranch` + - `MessageEditReady` + - `BranchSwitched` + +2. **SessionInstance-Methoden** (`session/instance.rs`) βœ… + - `convert_messages_to_ui_data()` with branch info + - `convert_tool_executions_to_ui_data()` + +3. **Backend handlers** (`ui/backend.rs`) βœ… + - `handle_start_message_edit()` + - `handle_switch_branch()` + +### Phase 4: UI (GPUI) βœ… COMPLETE + +1. **Edit-Button** (`gpui/messages.rs`) βœ… + - Shown on user messages + - Click-Handler sends `StartMessageEdit` event + +2. **BranchSwitcher-Komponente** (`gpui/branch_switcher.rs` - neu) βœ… + - Compact-Darstellung: "β—€ 2/3 β–Ά" + - Click-Handler fΓΌr Navigation + - Both entity-based (`BranchSwitcher`) and element-based (`BranchSwitcherElement`) + +3. **MessagesView erweitern** (`gpui/messages.rs`) βœ… + - Branch-Info pro Message + - BranchSwitcher rendern wo nΓΆtig + +4. **InputArea-Anpassungen** (`gpui/input_area.rs`) βœ… + - `branch_parent_id` tracking for edit mode + - `set_content_for_edit()` method + - MessageSubmitted event includes `branch_parent_id` + +5. **Root/App-Integration** (`gpui/root.rs`, `gpui/mod.rs`) βœ… + - Event-Handling fΓΌr Branch-Events + - `MessageEditReady` loads content into InputArea via `PendingEdit` state + - `BranchSwitched` updates messages display + + +### Phase 5: Testing & Polish πŸ”„ PENDING + +1. **Integrationstests** + - Branch-Erstellung durch Edit + - Branch-Wechsel + - Agent-Loop mit Branches + +2. **UI-Polish** + - Animationen fΓΌr Branch-Wechsel + - Visual Feedback beim Editieren + - Tooltips + +3. **Edge Cases** + - Leere Sessions + - Single-Message Sessions + - Tiefe Verschachtelung + +4. **Additional Work Needed** + - ~~Backend support for creating new branches from edited messages (when `branch_parent_id` is set)~~ βœ… DONE + - Terminal UI support for branching (currently ignores branch events) + +## Datei-Γ„nderungen Übersicht + +| Datei | Art | Beschreibung | +|-------|-----|--------------| +| `persistence.rs` | Modify | NodeId, MessageNode, ChatSession-Erweiterung | +| `session/mod.rs` | Modify | SessionState-Anpassung | +| `session/instance.rs` | Modify | Branch-Info-Generierung | +| `session/manager.rs` | Modify | Branch-Switch-Methoden | +| `agent/runner.rs` | Modify | Tree-basierte Message-History | +| `ui/ui_events.rs` | Modify | Neue Branch-Events | +| `ui/gpui/messages.rs` | Modify | Edit-Button, BranchSwitcher-Integration | +| `ui/gpui/branch_switcher.rs` | New | BranchSwitcher-Komponente | +| `ui/gpui/mod.rs` | Modify | Event-Handling | +| `ui/gpui/input_area.rs` | Modify | Edit-Mode | + +## Design-Entscheidungen + +1. **Tool-Executions bei Branches**: Tool-Executions bleiben wie bisher global in der Session gespeichert (`Vec`). Sie werden per `tool_id` auf den konkreten Tool-Aufruf in einer Message gemappt. Da die `tool_id` eindeutig ist, funktioniert das auch mit Branches. + +2. **Plan-State bei Branches**: Der Plan ist pfad-spezifisch und wird bei Branch-Wechseln aus dem neuen aktiven Pfad rekonstruiert. Dazu geht man die Message-History rΓΌckwΓ€rts durch bis zur ersten `MessageNode` mit gesetztem `plan_snapshot` Feld. Das `plan_snapshot` wird beim Speichern einer Assistant-Message gesetzt, wenn diese einen `update_plan` Tool-Aufruf enthielt. Der `plan`-Field in `ChatSession` speichert immer den Plan des aktuell aktiven Pfads. + +3. **Compaction bei Branches**: Compaction-Messages sind normale Messages im Baum (nur mit `is_compaction_summary: true` Flag). Wenn Branch A eine Compaction hatte und Branch B nicht, wird bei Wechsel zu B die volle History an das LLM geschickt. Im Agent gilt weiterhin: Messages ab der letzten Compaction-Message im aktiven Pfad werden ans API gesendet. + +4. **Maximale Tiefe/Breite**: Keine Limits. + +## Beispiel-Szenarien + +### Szenario 1: Einfaches Branching + +``` +1. User: "Fix the bug" +2. Assistant: [fixes bug incorrectly] +3. User klickt Edit auf Nachricht 1 +4. User: "Fix the bug in utils.rs, line 42" +5. Assistant: [fixes correct bug] +``` + +Baum: +``` + [1: "Fix the bug"] + β”‚ + [2: Assistant] + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β” +[3: "Fix bug"] [4: "Fix bug in utils.rs"] + β”‚ β”‚ +[*: Assistant] [5: Assistant] +``` + +### Szenario 2: Verschachteltes Branching + +User probiert mehrere AnsΓ€tze in verschiedenen Zweigen: + +``` + [User: "Build feature X"] + β”‚ + [Assistant: Plan A] + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + [User: "Use approach A"] [User: "Use approach B"] + β”‚ β”‚ + [Assistant...] [Assistant...] + β”‚ β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β” ... +[User: A1] [User: A2] +``` + +Der User kann jederzeit zwischen allen Pfaden wechseln.