Implement session branching

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,

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");

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);

crates/code_assistant/src/agent/runner.rs

@@ -58,8 +58,24 @@ pub struct Agent {
 
     permission_handler: Option<Arc<dyn PermissionMediator>>,
     sub_agent_runner: Option<Arc<dyn crate::agent::SubAgentRunner>>,
-    // Store all messages exchanged
+
+    // ========================================================================
+    // Branching: Tree-based message storage
+    // ========================================================================
+    /// All message nodes in the session (tree structure)
+    message_nodes:
+        std::collections::BTreeMap<crate::persistence::NodeId, crate::persistence::MessageNode>,
+    /// 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<Message>,
+
     // Store the history of tool executions
     tool_executions: Vec<crate::agent::types::ToolExecution>,
     // 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 {

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?;
 

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,