From 3b5e537e2693810c8e66c26db6fad59cbeb99c29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 11 Jan 2026 10:46:19 +0100 Subject: [PATCH 01/14] Plan for session branching sorry.. German --- docs/session-branching.md | 705 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 705 insertions(+) create mode 100644 docs/session-branching.md diff --git a/docs/session-branching.md b/docs/session-branching.md new file mode 100644 index 00000000..4c7a86e9 --- /dev/null +++ b/docs/session-branching.md @@ -0,0 +1,705 @@ +# 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) + +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 + +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 + +1. **Neue UiEvents** (`ui/ui_events.rs`) + - `StartMessageEdit` + - `SwitchBranch` + - `UpdateBranchInfo` + +2. **SessionInstance-Methoden** (`session/instance.rs`) + - `get_branch_info_for_path()` + - `generate_session_connect_events()` erweitern + +3. **SessionManager-Methoden** (`session/manager.rs`) + - `switch_branch()` + - `start_message_edit()` + +### Phase 4: UI (GPUI) + +1. **Edit-Button** (`gpui/messages.rs`) + - Auf User-Nachrichten, bei Hover sichtbar + - Click-Handler + +2. **BranchSwitcher-Komponente** (`gpui/branch_switcher.rs` - neu) + - Compact-Darstellung: "◀ 2/3 ▶" + - Click-Handler für Navigation + +3. **MessagesView erweitern** (`gpui/messages.rs`) + - Branch-Info pro Message + - BranchSwitcher rendern wo nötig + +4. **InputArea-Anpassungen** (`gpui/input_area.rs`) + - "Edit mode" state + - Branch-Point tracking + +5. **Root/App-Integration** (`gpui/root.rs`, `gpui/mod.rs`) + - Event-Handling für Branch-Events + - State-Management + +### Phase 5: Testing & Polish + +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 + +## 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. From 9a6ecd39e3871ceea5b632f9b8eaec09e40809d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 11 Jan 2026 12:52:50 +0100 Subject: [PATCH 02/14] phase 1-3 implemented --- crates/code_assistant/src/acp/ui.rs | 6 +- .../code_assistant/src/agent/persistence.rs | 15 +- crates/code_assistant/src/agent/runner.rs | 66 +- crates/code_assistant/src/agent/tests.rs | 80 +-- crates/code_assistant/src/persistence.rs | 640 +++++++++++++++++- crates/code_assistant/src/session/instance.rs | 51 +- crates/code_assistant/src/session/manager.rs | 14 +- crates/code_assistant/src/session/mod.rs | 99 ++- crates/code_assistant/src/ui/backend.rs | 177 +++++ crates/code_assistant/src/ui/gpui/mod.rs | 128 +++- crates/code_assistant/src/ui/terminal/app.rs | 5 + crates/code_assistant/src/ui/ui_events.rs | 51 +- 12 files changed, 1258 insertions(+), 74 deletions(-) diff --git a/crates/code_assistant/src/acp/ui.rs b/crates/code_assistant/src/acp/ui.rs index 5b183798..6c2e95b1 100644 --- a/crates/code_assistant/src/acp/ui.rs +++ b/crates/code_assistant/src/acp/ui.rs @@ -713,7 +713,11 @@ impl UserInterface for ACPUserUI { | UiEvent::UpdateCurrentModel { .. } | UiEvent::UpdateSandboxPolicy { .. } | UiEvent::CancelSubAgent { .. } - | UiEvent::HiddenToolCompleted => { + | UiEvent::HiddenToolCompleted + | UiEvent::StartMessageEdit { .. } + | UiEvent::SwitchBranch { .. } + | UiEvent::MessageEditReady { .. } + | UiEvent::BranchSwitched { .. } => { // These are UI management events, not relevant for ACP } UiEvent::DisplayError { message } => { 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..9200a5ad 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,9 +354,28 @@ 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(()) } @@ -438,17 +482,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(); diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index f04b167b..2b8250ec 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -1977,16 +1977,16 @@ 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".to_string(), + "Native Session".to_string(), + vec![user_message, assistant_message], + Vec::new(), + PlanState::default(), + session_config.clone(), + Some(2), + None, + ); agent.load_from_session_state(session_state).await?; @@ -2033,20 +2033,20 @@ 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".to_string(), + "Native Session".to_string(), + 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, - }; + Vec::new(), + PlanState::default(), + session_config.clone(), + Some(3), + None, + ); agent.load_from_session_state(session_state).await?; @@ -2099,16 +2099,16 @@ 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".to_string(), + "XML Session".to_string(), + vec![user_message, assistant_message], + Vec::new(), + PlanState::default(), + session_config.clone(), + Some(2), + None, + ); agent.load_from_session_state(session_state).await?; @@ -2143,16 +2143,16 @@ 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".to_string(), + "Regular Session".to_string(), + vec![user_message.clone(), assistant_message.clone()], + Vec::new(), + PlanState::default(), + session_config.clone(), + Some(3), + None, + ); agent.load_from_session_state(session_state).await?; diff --git a/crates/code_assistant/src/persistence.rs b/crates/code_assistant/src/persistence.rs index b807e75d..9038edff 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,59 @@ 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)] +#[allow(dead_code)] // Will be used by UI in Phase 4 +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, +} + +impl BranchInfo { + /// Total number of branches at this point + #[allow(dead_code)] // Will be used by UI in Phase 4 + pub fn total_branches(&self) -> usize { + self.sibling_ids.len() + } +} + /// Model configuration for a session #[derive(Debug, Serialize, Deserialize, Clone)] pub struct SessionModelConfig { @@ -35,11 +88,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 +155,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 +176,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 +195,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 +211,299 @@ 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. + #[allow(dead_code)] // Will be used by UI in Phase 4 + 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). + #[allow(dead_code)] // Will be used by UI in Phase 4 + 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 + } + + /// Check if a node has multiple children (is a branch point). + #[allow(dead_code)] // Will be used by UI in Phase 4 + 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 (if it's part of a branch). + /// Returns None if the node has no siblings (no branching at this point). + #[allow(dead_code)] // Will be used by UI in Phase 4 + 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, + }) + } + + /// Get branch infos for all branch points in the active path. + #[allow(dead_code)] // Will be used by UI in Phase 4 + pub fn get_all_branch_infos(&self) -> Vec<(NodeId, BranchInfo)> { + self.active_path + .iter() + .filter_map(|&node_id| self.get_branch_info(node_id).map(|info| (node_id, info))) + .collect() + } + + /// Find the plan state for the active path by walking backwards + /// to find the most recent plan_snapshot. + #[allow(dead_code)] // Will be used by UI in Phase 4 + 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. + #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 + 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. + #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 + 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: truncate 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) { + self.active_path.truncate(parent_pos + 1); + } + } 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 + } + + /// Add a message and store a plan snapshot with it. + #[allow(dead_code)] // Will be used by UI in Phase 4 + pub fn add_message_with_plan(&mut self, message: Message, plan: PlanState) -> NodeId { + let node_id = self.add_message(message); + if let Some(node) = self.message_nodes.get_mut(&node_id) { + node.plan_snapshot = Some(plan); + } + node_id + } + + /// Update the plan snapshot for an existing node. + #[allow(dead_code)] // Will be used by UI in Phase 4 + pub fn set_plan_snapshot(&mut self, node_id: NodeId, plan: PlanState) { + if let Some(node) = self.message_nodes.get_mut(&node_id) { + node.plan_snapshot = Some(plan); + } + } + + /// Switch to a different branch by making a different sibling node active. + /// Updates active_path to follow the new branch to its deepest descendant. + #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 + 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 +630,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 +763,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 +838,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 +893,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 +1250,240 @@ 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); + assert_eq!(info_3.total_branches(), 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]); + } } diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 175aa54e..c5886155 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -129,9 +129,10 @@ impl SessionInstance { self.session.updated_at = std::time::SystemTime::now(); } - /// Get all messages in the session - pub fn messages(&self) -> &[Message] { - &self.session.messages + /// Get all messages in the session (linearized from active path) + #[allow(dead_code)] // Kept for backward compatibility + pub fn messages(&self) -> Vec { + self.session.get_active_messages_cloned() } /// Get the current context size (input tokens + cache reads from most recent assistant message) @@ -229,6 +230,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); } @@ -315,12 +318,32 @@ impl SessionInstance { let mut messages_data = Vec::new(); - trace!( - "preparing {} messages for event", - self.session.messages.len() - ); + // Use tree structure if available, otherwise fall back to legacy messages + let message_iter: Vec<(Option, &llm::Message)> = + if !self.session.message_nodes.is_empty() { + // Use active path from tree + self.session + .active_path + .iter() + .filter_map(|&node_id| { + self.session + .message_nodes + .get(&node_id) + .map(|node| (Some(node_id), &node.message)) + }) + .collect() + } else { + // Fall back to legacy linear messages + self.session + .messages + .iter() + .map(|msg| (None, msg)) + .collect() + }; - for message in &self.session.messages { + trace!("preparing {} messages for event", message_iter.len()); + + 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,9 +359,12 @@ 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; } @@ -373,7 +399,12 @@ 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); @@ -388,7 +419,7 @@ impl SessionInstance { } /// 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..0e444fda 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -265,7 +265,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 @@ -577,7 +580,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..a4be2c57 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,109 @@ 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, } + +impl SessionState { + /// Get the linearized message history for the active path. + /// This returns cloned messages from the tree structure. + #[allow(dead_code)] // Will be used more extensively in Phase 4 + 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() + } + + /// Ensure the messages vec is in sync with the tree. + /// Call this after modifying the tree structure. + #[allow(dead_code)] // Will be used more extensively in Phase 4 + pub fn sync_messages_from_tree(&mut self) { + self.messages = self.get_active_messages(); + } + + /// 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. + #[allow(dead_code)] // Used by tests + #[allow(clippy::too_many_arguments)] + pub fn from_messages( + session_id: String, + name: String, + messages: Vec, + tool_executions: Vec, + plan: PlanState, + config: SessionConfig, + next_request_id: Option, + model_config: Option, + ) -> 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; + + 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, + name, + message_nodes, + active_path, + next_node_id, + messages, + tool_executions, + plan, + config, + next_request_id, + model_config, + } + } +} diff --git a/crates/code_assistant/src/ui/backend.rs b/crates/code_assistant/src/ui/backend.rs index a3511c5b..ff779875 100644 --- a/crates/code_assistant/src/ui/backend.rs +++ b/crates/code_assistant/src/ui/backend.rs @@ -58,6 +58,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 +104,25 @@ 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, + }, + BranchSwitched { + session_id: String, + messages: Vec, + tool_results: Vec, + plan: crate::types::PlanState, + }, } #[derive(Debug, Clone)] @@ -180,6 +205,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 @@ -599,6 +636,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 +648,142 @@ 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; + + Ok((content, attachments, branch_parent_id)) + } 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)) => BackendResponse::MessageEditReady { + session_id: session_id.to_string(), + content, + attachments, + branch_parent_id, + }, + 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}"), + }; + } + + // 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/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index 605f15a6..37232c3d 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -1020,6 +1020,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 +1035,82 @@ 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, + } => { + debug!( + "UI: MessageEditReady event - content len: {}, attachments: {}, parent: {:?}", + content.len(), + attachments.len(), + branch_parent_id + ); + // TODO Phase 4: Load content into input area + // self.update_input_area(cx, |input, cx| { + // input.set_content(content); + // input.set_attachments(attachments); + // input.set_branch_parent_id(branch_parent_id); + // cx.notify(); + // }); + } + 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); + } } } @@ -1286,7 +1363,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 +1455,7 @@ impl Gpui { ); } } + BackendResponse::SubAgentCancelled { session_id, tool_id, @@ -1389,6 +1467,54 @@ 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, + } => { + debug!( + "Received BackendResponse::MessageEditReady for session {} with {} chars, {} attachments", + session_id, + content.len(), + attachments.len() + ); + + // Forward to UI as event + self.process_ui_event_async( + UiEvent::MessageEditReady { + content: content.clone(), + attachments: attachments.clone(), + + branch_parent_id, + }, + 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/terminal/app.rs b/crates/code_assistant/src/ui/terminal/app.rs index 20482c38..bd132f7b 100644 --- a/crates/code_assistant/src/ui/terminal/app.rs +++ b/crates/code_assistant/src/ui/terminal/app.rs @@ -442,6 +442,7 @@ impl TerminalTuiApp { policy ))); } + BackendResponse::SubAgentCancelled { session_id: _, tool_id: _, @@ -449,6 +450,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 + } } } }); diff --git a/crates/code_assistant/src/ui/ui_events.rs b/crates/code_assistant/src/ui/ui_events.rs index 6415e1bd..f725c1e2 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,12 @@ use std::path::PathBuf; pub struct MessageData { pub role: MessageRole, pub fragments: Vec, + /// Optional node ID for branching support + #[allow(dead_code)] // Will be used in Phase 4 + pub node_id: Option, + /// Optional branch info if this message is part of a branch + #[allow(dead_code)] // Will be used in Phase 4 + pub branch_info: Option, } /// Tool execution result data for UI updates @@ -121,9 +127,52 @@ 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 + #[allow(dead_code)] // Will be used in Phase 4 + StartMessageEdit { + session_id: String, + /// The node ID of the message being edited + node_id: NodeId, + }, + + /// Switch to a different branch at a branch point + #[allow(dead_code)] // Will be used in Phase 4 + 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 + #[allow(dead_code)] // Will be used in Phase 4 + 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, + }, + + /// Response: Branch switch completed, new messages to display + #[allow(dead_code)] // Will be used in Phase 4 + 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, + }, + // === Resource Events (for tool operations) === /// A file was loaded/read by a tool ResourceLoaded { project: String, path: PathBuf }, From b0160d9297c7b614b5e6e467565e182c8014715c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Sun, 11 Jan 2026 16:00:31 +0100 Subject: [PATCH 03/14] Phase 4 completed (GPUI) --- .../src/ui/gpui/branch_switcher.rs | 303 ++++++++++++++++++ crates/code_assistant/src/ui/gpui/elements.rs | 28 ++ .../code_assistant/src/ui/gpui/input_area.rs | 55 +++- crates/code_assistant/src/ui/gpui/messages.rs | 137 ++++++-- crates/code_assistant/src/ui/gpui/mod.rs | 58 +++- crates/code_assistant/src/ui/gpui/root.rs | 52 ++- docs/session-branching.md | 69 ++-- 7 files changed, 634 insertions(+), 68 deletions(-) create mode 100644 crates/code_assistant/src/ui/gpui/branch_switcher.rs 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..8c7ecf9e --- /dev/null +++ b/crates/code_assistant/src/ui/gpui/branch_switcher.rs @@ -0,0 +1,303 @@ +use crate::persistence::{BranchInfo, NodeId}; +use gpui::{ + div, prelude::*, px, App, Context, CursorStyle, EventEmitter, FocusHandle, Focusable, + MouseButton, MouseUpEvent, Render, Window, +}; +use gpui_component::ActiveTheme; + +/// Events emitted by the BranchSwitcher component +#[derive(Clone, Debug)] +pub enum BranchSwitcherEvent { + /// User clicked to switch to a different branch + #[allow(dead_code)] // Will be used when entity-based BranchSwitcher is used + SwitchToBranch { node_id: NodeId }, +} + +/// A compact branch navigation component that shows "◀ 2/3 ▶" +/// Displayed below user messages where branches exist +#[allow(dead_code)] // Entity-based component - will be used when needed +pub struct BranchSwitcher { + branch_info: BranchInfo, + session_id: String, + focus_handle: FocusHandle, +} + +impl BranchSwitcher { + #[allow(dead_code)] // Will be used when entity-based BranchSwitcher is used + pub fn new(branch_info: BranchInfo, session_id: String, cx: &mut Context) -> Self { + Self { + branch_info, + session_id, + focus_handle: cx.focus_handle(), + } + } + + /// Update the branch info (e.g., after a branch switch) + #[allow(dead_code)] // Will be used when entity-based BranchSwitcher is used + pub fn set_branch_info(&mut self, branch_info: BranchInfo) { + self.branch_info = branch_info; + } + + fn on_prev_click(&mut self, _: &MouseUpEvent, _window: &mut Window, cx: &mut Context) { + if let Some(prev_node_id) = self.get_previous_sibling() { + cx.emit(BranchSwitcherEvent::SwitchToBranch { + node_id: prev_node_id, + }); + } + } + + fn on_next_click(&mut self, _: &MouseUpEvent, _window: &mut Window, cx: &mut Context) { + if let Some(next_node_id) = self.get_next_sibling() { + cx.emit(BranchSwitcherEvent::SwitchToBranch { + node_id: next_node_id, + }); + } + } + + fn get_previous_sibling(&self) -> Option { + if self.branch_info.active_index > 0 { + self.branch_info + .sibling_ids + .get(self.branch_info.active_index - 1) + .copied() + } else { + None + } + } + + fn get_next_sibling(&self) -> Option { + if self.branch_info.active_index + 1 < self.branch_info.sibling_ids.len() { + self.branch_info + .sibling_ids + .get(self.branch_info.active_index + 1) + .copied() + } else { + None + } + } + + fn has_previous(&self) -> bool { + self.branch_info.active_index > 0 + } + + fn has_next(&self) -> bool { + self.branch_info.active_index + 1 < self.branch_info.sibling_ids.len() + } + + /// Get the session ID this switcher belongs to + #[allow(dead_code)] + pub fn session_id(&self) -> &str { + &self.session_id + } +} + +impl EventEmitter for BranchSwitcher {} + +impl Focusable for BranchSwitcher { + fn focus_handle(&self, _: &App) -> FocusHandle { + self.focus_handle.clone() + } +} + +impl Render for BranchSwitcher { + fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { + let has_prev = self.has_previous(); + let has_next = self.has_next(); + let current = self.branch_info.active_index + 1; + let total = self.branch_info.sibling_ids.len(); + + let text_color = cx.theme().muted_foreground; + let active_color = cx.theme().foreground; + let hover_bg = cx.theme().muted; + + div() + .flex() + .flex_row() + .items_center() + .gap_1() + .text_xs() + .children(vec![ + // Previous button + div() + .id("branch-prev") + .px_1() + .py(px(2.)) + .rounded_sm() + .cursor(if has_prev { + CursorStyle::PointingHand + } else { + CursorStyle::OperationNotAllowed + }) + .text_color(if has_prev { active_color } else { text_color }) + .when(has_prev, |el| { + el.hover(|s| s.bg(hover_bg)) + .on_mouse_up(MouseButton::Left, cx.listener(Self::on_prev_click)) + }) + .child("◀") + .into_any_element(), + // Current position display + div() + .px_1() + .text_color(text_color) + .child(format!("{}/{}", current, total)) + .into_any_element(), + // Next button + div() + .id("branch-next") + .px_1() + .py(px(2.)) + .rounded_sm() + .cursor(if has_next { + CursorStyle::PointingHand + } else { + CursorStyle::OperationNotAllowed + }) + .text_color(if has_next { active_color } else { text_color }) + .when(has_next, |el| { + el.hover(|s| s.bg(hover_bg)) + .on_mouse_up(MouseButton::Left, cx.listener(Self::on_next_click)) + }) + .child("▶") + .into_any_element(), + ]) + } +} + +/// A stateless version of BranchSwitcher that can be rendered directly in elements +/// without needing an Entity. Used when rendering message lists. +#[derive(Clone, IntoElement)] +pub struct BranchSwitcherElement { + branch_info: BranchInfo, + session_id: String, + #[allow(dead_code)] // Reserved for future use + node_id: NodeId, +} + +impl BranchSwitcherElement { + pub fn new(branch_info: BranchInfo, session_id: String, node_id: NodeId) -> Self { + Self { + branch_info, + session_id, + node_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 text_color = cx.theme().muted_foreground; + let active_color = cx.theme().foreground; + let hover_bg = cx.theme().muted; + + // 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(); + + div() + .flex() + .flex_row() + .items_center() + .gap_1() + .text_xs() + .mt_1() + .children(vec![ + // Previous button + { + let base = div() + .id("branch-prev") + .px_1() + .py(px(2.)) + .rounded_sm() + .cursor(if has_prev { + CursorStyle::PointingHand + } else { + CursorStyle::OperationNotAllowed + }) + .text_color(if has_prev { active_color } else { text_color }) + .child("◀"); + + if has_prev { + base.hover(|s| s.bg(hover_bg)) + .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { + if let Some(node_id) = prev_node_id { + // Send event through UiEventSender global + 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(text_color) + .child(format!("{}/{}", current, total)) + .into_any_element(), + // Next button + { + let base = div() + .id("branch-next") + .px_1() + .py(px(2.)) + .rounded_sm() + .cursor(if has_next { + CursorStyle::PointingHand + } else { + CursorStyle::OperationNotAllowed + }) + .text_color(if has_next { active_color } else { text_color }) + .child("▶"); + + if has_next { + base.hover(|s| s.bg(hover_bg)) + .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { + if let Some(node_id) = next_node_id { + // Send event through UiEventSender global + 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..293cdf14 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,37 @@ 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(); + } + + /// Get the current branch parent ID (for editing) + #[allow(dead_code)] + pub fn branch_parent_id(&self) -> Option { + self.branch_parent_id + } + + /// Check if we're currently in edit mode (have a branch parent) + #[allow(dead_code)] // Will be used for edit mode visual indicator + pub fn is_editing(&self) -> bool { + self.branch_parent_id.is_some() + } + + /// 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 +208,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 +314,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 +323,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 +333,7 @@ impl InputArea { cx.emit(InputAreaEvent::MessageSubmitted { content: cleaned_text, attachments: self.attachments.clone(), + branch_parent_id, }); // Clear the input and attachments @@ -335,6 +384,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 +394,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..9d923096 100644 --- a/crates/code_assistant/src/ui/gpui/messages.rs +++ b/crates/code_assistant/src/ui/gpui/messages.rs @@ -1,5 +1,9 @@ +use super::branch_switcher::BranchSwitcherElement; use super::elements::MessageContainer; -use gpui::{div, prelude::*, px, rgb, App, Context, Entity, FocusHandle, Focusable, Window}; +use gpui::{ + div, prelude::*, px, rgb, App, Context, CursorStyle, Entity, FocusHandle, Focusable, + MouseButton, Window, +}; use gpui_component::{v_flex, ActiveTheme}; use std::sync::{Arc, Mutex}; @@ -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") + .px_2() + .py_1() + .rounded_sm() + .cursor(CursorStyle::PointingHand) + .text_xs() + .text_color(cx.theme().muted_foreground) + .hover(|s| { + s.bg(cx.theme().muted).text_color(cx.theme().foreground) + }) + .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("Edit"), + ) + }); + + 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(node_id), Some(session_id)) = + (branch_info, node_id, 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, node_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 37232c3d..efef9ef7 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)), } } @@ -618,9 +633,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 +665,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 +677,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 +686,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()); @@ -1069,6 +1095,7 @@ impl Gpui { }); } } + UiEvent::MessageEditReady { content, attachments, @@ -1080,13 +1107,14 @@ impl Gpui { attachments.len(), branch_parent_id ); - // TODO Phase 4: Load content into input area - // self.update_input_area(cx, |input, cx| { - // input.set_content(content); - // input.set_attachments(attachments); - // input.set_branch_parent_id(branch_parent_id); - // cx.notify(); - // }); + // 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, @@ -1253,6 +1281,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, diff --git a/crates/code_assistant/src/ui/gpui/root.rs b/crates/code_assistant/src/ui/gpui/root.rs index 0166283f..7f0869a3 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!( @@ -314,6 +333,9 @@ impl RootView { }); } else { // Send message normally (agent is idle) + // TODO: When branch_parent_id is set, we need to create a new branch + // For now, we send as a regular message - full branching requires + // additional backend support to create the branch tracing::info!( "RootView: Sending user message to session {}: {} (with {} attachments)", session_id, @@ -329,6 +351,21 @@ impl RootView { } } + /// 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 +693,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/docs/session-branching.md b/docs/session-branching.md index 4c7a86e9..8fb51234 100644 --- a/docs/session-branching.md +++ b/docs/session-branching.md @@ -544,81 +544,86 @@ impl Agent { } ``` + ## Implementierungsplan -### Phase 1: Datenstruktur (Backend) +### Phase 1: Datenstruktur (Backend) ✅ COMPLETE -1. **Neue Typen definieren** (`persistence.rs`) +1. **Neue Typen definieren** (`persistence.rs`) ✅ - `NodeId` type alias - `MessageNode` struct - `ConversationPath` type alias - `BranchInfo` struct -2. **ChatSession erweitern** (`persistence.rs`) +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`) +3. **SessionState anpassen** (`session/mod.rs`) ✅ - Gleiche Struktur wie ChatSession - Hilfsmethoden für Linearisierung -4. **Tests schreiben** +4. **Tests schreiben** ✅ - Migration von linear zu tree - Branch-Erstellung - Branch-Switching - Pfad-Berechnung -### Phase 2: Agent-Anpassungen +### Phase 2: Agent-Anpassungen ✅ COMPLETE -1. **Agent-State erweitern** (`agent/runner.rs`) +1. **Agent-State erweitern** (`agent/runner.rs`) ✅ - Neue Felder für Tree-Struktur - `append_message()` anpassen -2. **State-Laden/Speichern** (`agent/persistence.rs`) +2. **State-Laden/Speichern** (`agent/persistence.rs`) ✅ - Tree-Struktur in SessionState -3. **LLM-Request-Building** +3. **LLM-Request-Building** ✅ - Nur aktiven Pfad an LLM senden -### Phase 3: UI-Backend-Kommunikation +### Phase 3: UI-Backend-Kommunikation ✅ COMPLETE -1. **Neue UiEvents** (`ui/ui_events.rs`) +1. **Neue UiEvents** (`ui/ui_events.rs`) ✅ - `StartMessageEdit` - `SwitchBranch` - - `UpdateBranchInfo` + - `MessageEditReady` + - `BranchSwitched` -2. **SessionInstance-Methoden** (`session/instance.rs`) - - `get_branch_info_for_path()` - - `generate_session_connect_events()` erweitern +2. **SessionInstance-Methoden** (`session/instance.rs`) ✅ + - `convert_messages_to_ui_data()` with branch info + - `convert_tool_executions_to_ui_data()` -3. **SessionManager-Methoden** (`session/manager.rs`) - - `switch_branch()` - - `start_message_edit()` +3. **Backend handlers** (`ui/backend.rs`) ✅ + - `handle_start_message_edit()` + - `handle_switch_branch()` -### Phase 4: UI (GPUI) +### Phase 4: UI (GPUI) ✅ COMPLETE -1. **Edit-Button** (`gpui/messages.rs`) - - Auf User-Nachrichten, bei Hover sichtbar - - Click-Handler +1. **Edit-Button** (`gpui/messages.rs`) ✅ + - Shown on user messages + - Click-Handler sends `StartMessageEdit` event -2. **BranchSwitcher-Komponente** (`gpui/branch_switcher.rs` - neu) +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`) +3. **MessagesView erweitern** (`gpui/messages.rs`) ✅ - Branch-Info pro Message - BranchSwitcher rendern wo nötig -4. **InputArea-Anpassungen** (`gpui/input_area.rs`) - - "Edit mode" state - - Branch-Point tracking +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`) +5. **Root/App-Integration** (`gpui/root.rs`, `gpui/mod.rs`) ✅ - Event-Handling für Branch-Events - - State-Management + - `MessageEditReady` loads content into InputArea via `PendingEdit` state + - `BranchSwitched` updates messages display -### Phase 5: Testing & Polish +### Phase 5: Testing & Polish 🔄 PENDING 1. **Integrationstests** - Branch-Erstellung durch Edit @@ -635,6 +640,10 @@ impl Agent { - Single-Message Sessions - Tiefe Verschachtelung +4. **Additional Work Needed** + - Backend support for creating new branches from edited messages (when `branch_parent_id` is set) + - Terminal UI support for branching (currently ignores branch events) + ## Datei-Änderungen Übersicht | Datei | Art | Beschreibung | From db429b51f97bb3ed975376eb59c5f762bbae5b07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 10:19:33 +0100 Subject: [PATCH 04/14] Various fixes after implementing all phases --- crates/code_assistant/src/acp/agent.rs | 2 + crates/code_assistant/src/acp/ui.rs | 5 +- crates/code_assistant/src/agent/runner.rs | 2 + crates/code_assistant/src/app/gpui.rs | 1 + crates/code_assistant/src/persistence.rs | 51 ++++- crates/code_assistant/src/session/instance.rs | 161 +++++++++++++++- crates/code_assistant/src/session/manager.rs | 106 ++++++++++- crates/code_assistant/src/ui/backend.rs | 106 +++++++++-- crates/code_assistant/src/ui/gpui/mod.rs | 176 +++++++++++++++++- crates/code_assistant/src/ui/gpui/root.rs | 9 +- crates/code_assistant/src/ui/terminal/app.rs | 2 + crates/code_assistant/src/ui/terminal/ui.rs | 2 + crates/code_assistant/src/ui/ui_events.rs | 17 ++ docs/session-branching.md | 3 +- 14 files changed, 608 insertions(+), 35 deletions(-) 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 6c2e95b1..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( @@ -717,7 +718,8 @@ impl UserInterface for ACPUserUI { | UiEvent::StartMessageEdit { .. } | UiEvent::SwitchBranch { .. } | UiEvent::MessageEditReady { .. } - | UiEvent::BranchSwitched { .. } => { + | UiEvent::BranchSwitched { .. } + | UiEvent::UpdateBranchInfo { .. } => { // These are UI management events, not relevant for ACP } UiEvent::DisplayError { message } => { @@ -1030,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/runner.rs b/crates/code_assistant/src/agent/runner.rs index 9200a5ad..f3bb1a3c 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -394,6 +394,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?; } @@ -996,6 +997,7 @@ impl Agent { .send_event(UiEvent::DisplayUserInput { content: task.clone(), attachments: Vec::new(), + node_id: None, // Initial task message }) .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 9038edff..40532f08 100644 --- a/crates/code_assistant/src/persistence.rs +++ b/crates/code_assistant/src/persistence.rs @@ -380,10 +380,15 @@ impl ChatSession { self.message_nodes.insert(node_id, node); - // Update active_path: truncate to parent and add new 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 @@ -1486,4 +1491,48 @@ mod tests { 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 c5886155..63935d5f 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -6,7 +6,7 @@ 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}; @@ -123,10 +123,35 @@ impl SessionInstance { } } - /// Add a message to the session + /// Add a message to the session (appends to active path) + #[allow(dead_code)] // Kept for backward compatibility / internal use pub fn add_message(&mut self, message: Message) { - self.session.messages.push(message); - self.session.updated_at = std::time::SystemTime::now(); + // Use the tree-based method which handles active_path correctly + self.session.add_message(message); + } + + /// 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) + }; + + Ok(node_id) } /// Get all messages in the session (linearized from active path) @@ -418,6 +443,134 @@ impl SessionInstance { Ok(messages_data) } + /// 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; + #[async_trait::async_trait] + impl crate::ui::UserInterface for DummyUI { + async fn send_event( + &self, + _event: crate::ui::UiEvent, + ) -> Result<(), crate::ui::UIError> { + Ok(()) + } + + fn display_fragment( + &self, + _fragment: &crate::ui::DisplayFragment, + ) -> Result<(), crate::ui::UIError> { + Ok(()) + } + fn should_streaming_continue(&self) -> bool { + true + } + fn notify_rate_limit(&self, _seconds_remaining: u64) {} + fn clear_rate_limit(&self) {} + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + let dummy_ui: std::sync::Arc = std::sync::Arc::new(DummyUI); + let mut processor = create_stream_processor(tool_syntax, dummy_ui, 0); + + let mut messages_data = Vec::new(); + + // 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 (node_id, message) in message_iter { + if message.is_compaction_summary { + let summary = match &message.content { + llm::MessageContent::Text(text) => text.trim().to_string(), + llm::MessageContent::Structured(blocks) => blocks + .iter() + .filter_map(|block| match block { + llm::ContentBlock::Text { text, .. } => Some(text.as_str()), + llm::ContentBlock::Thinking { thinking, .. } => Some(thinking.as_str()), + _ => None, + }) + .collect::>() + .join("\n") + .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 + if message.role == llm::MessageRole::User { + match &message.content { + llm::MessageContent::Text(text) if text.trim().is_empty() => continue, + llm::MessageContent::Structured(blocks) => { + let has_tool_results = blocks + .iter() + .any(|block| matches!(block, llm::ContentBlock::ToolResult { .. })); + if has_tool_results { + continue; + } + } + _ => {} + } + } + + match processor.extract_fragments_from_message(message) { + Ok(fragments) => { + let role = match message.role { + llm::MessageRole::User => MessageRole::User, + llm::MessageRole::Assistant => MessageRole::Assistant, + }; + 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); + } + } + } + + Ok(messages_data) + } + /// Convert tool executions to UI tool result data pub fn convert_tool_executions_to_ui_data( &self, diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index 0e444fda..3d0edc52 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(); @@ -571,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 diff --git a/crates/code_assistant/src/ui/backend.rs b/crates/code_assistant/src/ui/backend.rs index ff779875..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, @@ -116,6 +118,9 @@ pub enum BackendResponse { 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, @@ -163,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, ) @@ -341,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); @@ -405,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, @@ -704,7 +749,25 @@ async fn handle_start_message_edit( // The branch parent is the parent of the node being edited let branch_parent_id = node.parent_id; - Ok((content, attachments, branch_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)) } @@ -714,12 +777,16 @@ async fn handle_start_message_edit( }; match result { - Ok((content, attachments, branch_parent_id)) => BackendResponse::MessageEditReady { - session_id: session_id.to_string(), - content, - attachments, - branch_parent_id, - }, + 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 { @@ -755,6 +822,19 @@ async fn handle_switch_branch( }; } + // 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) diff --git a/crates/code_assistant/src/ui/gpui/mod.rs b/crates/code_assistant/src/ui/gpui/mod.rs index efef9ef7..f72c6ca3 100644 --- a/crates/code_assistant/src/ui/gpui/mod.rs +++ b/crates/code_assistant/src/ui/gpui/mod.rs @@ -486,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); @@ -830,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; @@ -849,6 +856,7 @@ impl Gpui { session_id, message, attachments, + branch_parent_id, }); } else { warn!("UI: No backend event sender available"); @@ -1100,19 +1108,139 @@ impl Gpui { content, attachments, branch_parent_id, + messages, + tool_results, } => { debug!( - "UI: MessageEditReady event - content len: {}, attachments: {}, parent: {:?}", + "UI: MessageEditReady event - content len: {}, attachments: {}, parent: {:?}, {} messages", content.len(), attachments.len(), - branch_parent_id + 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"); } @@ -1139,6 +1267,36 @@ impl Gpui { ); 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"); + } } } @@ -1512,12 +1670,15 @@ impl Gpui { content, attachments, branch_parent_id, + messages, + tool_results, } => { debug!( - "Received BackendResponse::MessageEditReady for session {} with {} chars, {} attachments", + "Received BackendResponse::MessageEditReady for session {} with {} chars, {} attachments, {} messages", session_id, content.len(), - attachments.len() + attachments.len(), + messages.len() ); // Forward to UI as event @@ -1525,8 +1686,9 @@ impl Gpui { UiEvent::MessageEditReady { content: content.clone(), attachments: attachments.clone(), - branch_parent_id, + messages: messages.clone(), + tool_results: tool_results.clone(), }, cx, ); diff --git a/crates/code_assistant/src/ui/gpui/root.rs b/crates/code_assistant/src/ui/gpui/root.rs index 7f0869a3..b4d8ae13 100644 --- a/crates/code_assistant/src/ui/gpui/root.rs +++ b/crates/code_assistant/src/ui/gpui/root.rs @@ -333,19 +333,18 @@ impl RootView { }); } else { // Send message normally (agent is idle) - // TODO: When branch_parent_id is set, we need to create a new branch - // For now, we send as a regular message - full branching requires - // additional backend support to create the branch 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, }); } } diff --git a/crates/code_assistant/src/ui/terminal/app.rs b/crates/code_assistant/src/ui/terminal/app.rs index bd132f7b..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 { @@ -513,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 f725c1e2..d582d95f 100644 --- a/crates/code_assistant/src/ui/ui_events.rs +++ b/crates/code_assistant/src/ui/ui_events.rs @@ -35,6 +35,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 }, @@ -94,6 +96,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 }, @@ -159,6 +163,10 @@ pub enum UiEvent { 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 @@ -173,6 +181,15 @@ pub enum UiEvent { 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 index 8fb51234..f6f5dbfd 100644 --- a/docs/session-branching.md +++ b/docs/session-branching.md @@ -623,6 +623,7 @@ impl Agent { - `MessageEditReady` loads content into InputArea via `PendingEdit` state - `BranchSwitched` updates messages display + ### Phase 5: Testing & Polish 🔄 PENDING 1. **Integrationstests** @@ -641,7 +642,7 @@ impl Agent { - Tiefe Verschachtelung 4. **Additional Work Needed** - - Backend support for creating new branches from edited messages (when `branch_parent_id` is set) + - ~~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 From b8a13c3a36a6a917aa382eb0e0ed69063ec239af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 11:05:43 +0100 Subject: [PATCH 05/14] Fix empty line in doc comment --- crates/code_assistant/src/session/manager.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/code_assistant/src/session/manager.rs b/crates/code_assistant/src/session/manager.rs index 3d0edc52..6f6d3761 100644 --- a/crates/code_assistant/src/session/manager.rs +++ b/crates/code_assistant/src/session/manager.rs @@ -284,7 +284,7 @@ impl SessionManager { /// 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)] From 1dc6c0e7aeb2079df4c700cabb0e0ea55cce55ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 11:25:54 +0100 Subject: [PATCH 06/14] Move more defaults into new constructor --- crates/code_assistant/src/agent/tests.rs | 32 ++++++------------------ crates/code_assistant/src/session/mod.rs | 28 +++++++++++---------- 2 files changed, 23 insertions(+), 37 deletions(-) diff --git a/crates/code_assistant/src/agent/tests.rs b/crates/code_assistant/src/agent/tests.rs index 2b8250ec..eafc6756 100644 --- a/crates/code_assistant/src/agent/tests.rs +++ b/crates/code_assistant/src/agent/tests.rs @@ -1978,14 +1978,10 @@ async fn test_load_normalizes_native_dangling_tool_request() -> Result<()> { .with_request_id(1); let session_state = SessionState::from_messages( - "native-session".to_string(), - "Native Session".to_string(), + "native-session", + "Native Session", vec![user_message, assistant_message], - Vec::new(), - PlanState::default(), session_config.clone(), - Some(2), - None, ); agent.load_from_session_state(session_state).await?; @@ -2034,18 +2030,14 @@ async fn test_load_normalizes_native_dangling_tool_request_with_followup_user() let followup_user_message = Message::new_user("Also check the contributing guide."); let session_state = SessionState::from_messages( - "native-session".to_string(), - "Native Session".to_string(), + "native-session", + "Native Session", vec![ user_message.clone(), assistant_message, followup_user_message.clone(), ], - Vec::new(), - PlanState::default(), session_config.clone(), - Some(3), - None, ); agent.load_from_session_state(session_state).await?; @@ -2100,14 +2092,10 @@ async fn test_load_normalizes_xml_dangling_tool_request() -> Result<()> { .with_request_id(1); let session_state = SessionState::from_messages( - "xml-session".to_string(), - "XML Session".to_string(), + "xml-session", + "XML Session", vec![user_message, assistant_message], - Vec::new(), - PlanState::default(), session_config.clone(), - Some(2), - None, ); agent.load_from_session_state(session_state).await?; @@ -2144,14 +2132,10 @@ async fn test_load_keeps_assistant_messages_without_tool_requests() -> Result<() let assistant_message = Message::new_assistant("Here is a summary of the repository."); let session_state = SessionState::from_messages( - "text-session".to_string(), - "Regular Session".to_string(), + "text-session", + "Regular Session", vec![user_message.clone(), assistant_message.clone()], - Vec::new(), - PlanState::default(), session_config.clone(), - Some(3), - None, ); agent.load_from_session_state(session_state).await?; diff --git a/crates/code_assistant/src/session/mod.rs b/crates/code_assistant/src/session/mod.rs index a4be2c57..2da5775a 100644 --- a/crates/code_assistant/src/session/mod.rs +++ b/crates/code_assistant/src/session/mod.rs @@ -103,22 +103,24 @@ impl SessionState { /// This is primarily for tests and backward compatibility. /// The messages are converted to a tree structure with a single linear path. #[allow(dead_code)] // Used by tests - #[allow(clippy::too_many_arguments)] pub fn from_messages( - session_id: String, - name: String, + session_id: impl Into, + name: impl Into, messages: Vec, - tool_executions: Vec, - plan: PlanState, config: SessionConfig, - next_request_id: Option, - model_config: Option, ) -> 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; @@ -137,17 +139,17 @@ impl SessionState { } Self { - session_id, - name, + session_id: session_id.into(), + name: name.into(), message_nodes, active_path, next_node_id, messages, - tool_executions, - plan, + tool_executions: Vec::new(), + plan: PlanState::default(), config, - next_request_id, - model_config, + next_request_id: Some(max_request_id + 1), + model_config: None, } } } From b5fdef8636af5d602d4bc8ef141cbacdf726603c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 11:34:24 +0100 Subject: [PATCH 07/14] DRY --- crates/code_assistant/src/session/instance.rs | 138 +----------------- 1 file changed, 2 insertions(+), 136 deletions(-) diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 63935d5f..207d20e9 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -13,7 +13,7 @@ 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)] @@ -306,141 +306,7 @@ impl SessionInstance { &self, tool_syntax: crate::types::ToolSyntax, ) -> Result, anyhow::Error> { - // Create dummy UI for stream processor - struct DummyUI; - #[async_trait::async_trait] - impl crate::ui::UserInterface for DummyUI { - async fn send_event( - &self, - _event: crate::ui::UiEvent, - ) -> Result<(), crate::ui::UIError> { - Ok(()) - } - - fn display_fragment( - &self, - _fragment: &crate::ui::DisplayFragment, - ) -> Result<(), crate::ui::UIError> { - Ok(()) - } - 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 as_any(&self) -> &dyn std::any::Any { - self - } - } - - let dummy_ui: std::sync::Arc = std::sync::Arc::new(DummyUI); - let mut processor = create_stream_processor(tool_syntax, dummy_ui, 0); - - let mut messages_data = Vec::new(); - - // Use tree structure if available, otherwise fall back to legacy messages - let message_iter: Vec<(Option, &llm::Message)> = - if !self.session.message_nodes.is_empty() { - // Use active path from tree - self.session - .active_path - .iter() - .filter_map(|&node_id| { - self.session - .message_nodes - .get(&node_id) - .map(|node| (Some(node_id), &node.message)) - }) - .collect() - } else { - // Fall back to legacy linear messages - self.session - .messages - .iter() - .map(|msg| (None, msg)) - .collect() - }; - - trace!("preparing {} messages for event", message_iter.len()); - - 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(), - llm::MessageContent::Structured(blocks) => blocks - .iter() - .filter_map(|block| match block { - llm::ContentBlock::Text { text, .. } => Some(text.as_str()), - llm::ContentBlock::Thinking { thinking, .. } => Some(thinking.as_str()), - _ => None, - }) - .collect::>() - .join("\n") - .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) - 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::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 - } - } - } - - match processor.extract_fragments_from_message(message) { - Ok(fragments) => { - let role = match message.role { - llm::MessageRole::User => MessageRole::User, - llm::MessageRole::Assistant => MessageRole::Assistant, - }; - 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) + self.convert_messages_to_ui_data_until(tool_syntax, None) } /// Convert session messages to UI MessageData format, stopping at a specific node From b36431d16a104da7ca47c5f55bfeabb8c69fdb47 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 11:36:42 +0100 Subject: [PATCH 08/14] Remove dead code --- crates/code_assistant/src/session/instance.rs | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/crates/code_assistant/src/session/instance.rs b/crates/code_assistant/src/session/instance.rs index 207d20e9..3a51b3fe 100644 --- a/crates/code_assistant/src/session/instance.rs +++ b/crates/code_assistant/src/session/instance.rs @@ -123,13 +123,6 @@ impl SessionInstance { } } - /// Add a message to the session (appends to active path) - #[allow(dead_code)] // Kept for backward compatibility / internal use - pub fn add_message(&mut self, message: Message) { - // Use the tree-based method which handles active_path correctly - self.session.add_message(message); - } - /// 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. @@ -154,12 +147,6 @@ impl SessionInstance { Ok(node_id) } - /// Get all messages in the session (linearized from active path) - #[allow(dead_code)] // Kept for backward compatibility - pub fn messages(&self) -> Vec { - self.session.get_active_messages_cloned() - } - /// Get the current context size (input tokens + cache reads from most recent assistant message) /// This represents the total tokens being processed in the current LLM request #[allow(dead_code)] From 1606e8d654e969cb2d26e058e0e74d113d44bbb8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 11:44:17 +0100 Subject: [PATCH 09/14] Cleanup --- crates/code_assistant/src/session/mod.rs | 20 +------------------- 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/crates/code_assistant/src/session/mod.rs b/crates/code_assistant/src/session/mod.rs index 2da5775a..58e660ea 100644 --- a/crates/code_assistant/src/session/mod.rs +++ b/crates/code_assistant/src/session/mod.rs @@ -80,29 +80,11 @@ pub struct SessionState { pub model_config: Option, } +#[cfg(test)] impl SessionState { - /// Get the linearized message history for the active path. - /// This returns cloned messages from the tree structure. - #[allow(dead_code)] // Will be used more extensively in Phase 4 - 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() - } - - /// Ensure the messages vec is in sync with the tree. - /// Call this after modifying the tree structure. - #[allow(dead_code)] // Will be used more extensively in Phase 4 - pub fn sync_messages_from_tree(&mut self) { - self.messages = self.get_active_messages(); - } - /// 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. - #[allow(dead_code)] // Used by tests pub fn from_messages( session_id: impl Into, name: impl Into, From 260871a2deacf0e8295d9e7ee63991799e9c3ed2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 11:52:39 +0100 Subject: [PATCH 10/14] Cleanup --- .../src/ui/gpui/branch_switcher.rs | 174 +----------------- crates/code_assistant/src/ui/gpui/messages.rs | 6 +- 2 files changed, 8 insertions(+), 172 deletions(-) diff --git a/crates/code_assistant/src/ui/gpui/branch_switcher.rs b/crates/code_assistant/src/ui/gpui/branch_switcher.rs index 8c7ecf9e..3461dc01 100644 --- a/crates/code_assistant/src/ui/gpui/branch_switcher.rs +++ b/crates/code_assistant/src/ui/gpui/branch_switcher.rs @@ -1,184 +1,20 @@ -use crate::persistence::{BranchInfo, NodeId}; -use gpui::{ - div, prelude::*, px, App, Context, CursorStyle, EventEmitter, FocusHandle, Focusable, - MouseButton, MouseUpEvent, Render, Window, -}; +use crate::persistence::BranchInfo; +use gpui::{div, prelude::*, px, App, CursorStyle, MouseButton, Window}; use gpui_component::ActiveTheme; -/// Events emitted by the BranchSwitcher component -#[derive(Clone, Debug)] -pub enum BranchSwitcherEvent { - /// User clicked to switch to a different branch - #[allow(dead_code)] // Will be used when entity-based BranchSwitcher is used - SwitchToBranch { node_id: NodeId }, -} - -/// A compact branch navigation component that shows "◀ 2/3 ▶" -/// Displayed below user messages where branches exist -#[allow(dead_code)] // Entity-based component - will be used when needed -pub struct BranchSwitcher { - branch_info: BranchInfo, - session_id: String, - focus_handle: FocusHandle, -} - -impl BranchSwitcher { - #[allow(dead_code)] // Will be used when entity-based BranchSwitcher is used - pub fn new(branch_info: BranchInfo, session_id: String, cx: &mut Context) -> Self { - Self { - branch_info, - session_id, - focus_handle: cx.focus_handle(), - } - } - - /// Update the branch info (e.g., after a branch switch) - #[allow(dead_code)] // Will be used when entity-based BranchSwitcher is used - pub fn set_branch_info(&mut self, branch_info: BranchInfo) { - self.branch_info = branch_info; - } - - fn on_prev_click(&mut self, _: &MouseUpEvent, _window: &mut Window, cx: &mut Context) { - if let Some(prev_node_id) = self.get_previous_sibling() { - cx.emit(BranchSwitcherEvent::SwitchToBranch { - node_id: prev_node_id, - }); - } - } - - fn on_next_click(&mut self, _: &MouseUpEvent, _window: &mut Window, cx: &mut Context) { - if let Some(next_node_id) = self.get_next_sibling() { - cx.emit(BranchSwitcherEvent::SwitchToBranch { - node_id: next_node_id, - }); - } - } - - fn get_previous_sibling(&self) -> Option { - if self.branch_info.active_index > 0 { - self.branch_info - .sibling_ids - .get(self.branch_info.active_index - 1) - .copied() - } else { - None - } - } - - fn get_next_sibling(&self) -> Option { - if self.branch_info.active_index + 1 < self.branch_info.sibling_ids.len() { - self.branch_info - .sibling_ids - .get(self.branch_info.active_index + 1) - .copied() - } else { - None - } - } - - fn has_previous(&self) -> bool { - self.branch_info.active_index > 0 - } - - fn has_next(&self) -> bool { - self.branch_info.active_index + 1 < self.branch_info.sibling_ids.len() - } - - /// Get the session ID this switcher belongs to - #[allow(dead_code)] - pub fn session_id(&self) -> &str { - &self.session_id - } -} - -impl EventEmitter for BranchSwitcher {} - -impl Focusable for BranchSwitcher { - fn focus_handle(&self, _: &App) -> FocusHandle { - self.focus_handle.clone() - } -} - -impl Render for BranchSwitcher { - fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { - let has_prev = self.has_previous(); - let has_next = self.has_next(); - let current = self.branch_info.active_index + 1; - let total = self.branch_info.sibling_ids.len(); - - let text_color = cx.theme().muted_foreground; - let active_color = cx.theme().foreground; - let hover_bg = cx.theme().muted; - - div() - .flex() - .flex_row() - .items_center() - .gap_1() - .text_xs() - .children(vec![ - // Previous button - div() - .id("branch-prev") - .px_1() - .py(px(2.)) - .rounded_sm() - .cursor(if has_prev { - CursorStyle::PointingHand - } else { - CursorStyle::OperationNotAllowed - }) - .text_color(if has_prev { active_color } else { text_color }) - .when(has_prev, |el| { - el.hover(|s| s.bg(hover_bg)) - .on_mouse_up(MouseButton::Left, cx.listener(Self::on_prev_click)) - }) - .child("◀") - .into_any_element(), - // Current position display - div() - .px_1() - .text_color(text_color) - .child(format!("{}/{}", current, total)) - .into_any_element(), - // Next button - div() - .id("branch-next") - .px_1() - .py(px(2.)) - .rounded_sm() - .cursor(if has_next { - CursorStyle::PointingHand - } else { - CursorStyle::OperationNotAllowed - }) - .text_color(if has_next { active_color } else { text_color }) - .when(has_next, |el| { - el.hover(|s| s.bg(hover_bg)) - .on_mouse_up(MouseButton::Left, cx.listener(Self::on_next_click)) - }) - .child("▶") - .into_any_element(), - ]) - } -} - -/// A stateless version of BranchSwitcher that can be rendered directly in elements -/// without needing an Entity. Used when rendering message lists. +/// A stateless branch navigation component that shows "◀ 2/3 ▶" +/// Displayed below user messages where branches exist. #[derive(Clone, IntoElement)] pub struct BranchSwitcherElement { branch_info: BranchInfo, session_id: String, - #[allow(dead_code)] // Reserved for future use - node_id: NodeId, } impl BranchSwitcherElement { - pub fn new(branch_info: BranchInfo, session_id: String, node_id: NodeId) -> Self { + pub fn new(branch_info: BranchInfo, session_id: String) -> Self { Self { branch_info, session_id, - node_id, } } } diff --git a/crates/code_assistant/src/ui/gpui/messages.rs b/crates/code_assistant/src/ui/gpui/messages.rs index 9d923096..a3037397 100644 --- a/crates/code_assistant/src/ui/gpui/messages.rs +++ b/crates/code_assistant/src/ui/gpui/messages.rs @@ -242,13 +242,13 @@ impl Render for MessagesView { // Add branch switcher if branch_info is present (only for user messages) if is_user_message { - if let (Some(branch_info), Some(node_id), Some(session_id)) = - (branch_info, node_id, current_session_id.clone()) + 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, node_id)) + .child(BranchSwitcherElement::new(branch_info, session_id)) .into_any_element(); } } From 63bf556c19cb80c1756fe7d5a295909d808e8578 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 12:15:34 +0100 Subject: [PATCH 11/14] Remove dead code --- crates/code_assistant/src/ui/gpui/input_area.rs | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/crates/code_assistant/src/ui/gpui/input_area.rs b/crates/code_assistant/src/ui/gpui/input_area.rs index 293cdf14..e742273e 100644 --- a/crates/code_assistant/src/ui/gpui/input_area.rs +++ b/crates/code_assistant/src/ui/gpui/input_area.rs @@ -140,18 +140,6 @@ impl InputArea { cx.notify(); } - /// Get the current branch parent ID (for editing) - #[allow(dead_code)] - pub fn branch_parent_id(&self) -> Option { - self.branch_parent_id - } - - /// Check if we're currently in edit mode (have a branch parent) - #[allow(dead_code)] // Will be used for edit mode visual indicator - pub fn is_editing(&self) -> bool { - self.branch_parent_id.is_some() - } - /// Clear the edit mode state fn clear_edit_mode(&mut self) { self.branch_parent_id = None; From 45087e6001e260a19f4029935e99435d421996c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 12:26:51 +0100 Subject: [PATCH 12/14] Cleanup --- crates/code_assistant/src/ui/ui_events.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/crates/code_assistant/src/ui/ui_events.rs b/crates/code_assistant/src/ui/ui_events.rs index d582d95f..4530608a 100644 --- a/crates/code_assistant/src/ui/ui_events.rs +++ b/crates/code_assistant/src/ui/ui_events.rs @@ -12,10 +12,8 @@ pub struct MessageData { pub role: MessageRole, pub fragments: Vec, /// Optional node ID for branching support - #[allow(dead_code)] // Will be used in Phase 4 pub node_id: Option, /// Optional branch info if this message is part of a branch - #[allow(dead_code)] // Will be used in Phase 4 pub branch_info: Option, } @@ -138,7 +136,6 @@ pub enum UiEvent { // === Session Branching Events === /// Request to start editing a message (creates a branch point) /// UI should load the message content into the input area - #[allow(dead_code)] // Will be used in Phase 4 StartMessageEdit { session_id: String, /// The node ID of the message being edited @@ -146,7 +143,6 @@ pub enum UiEvent { }, /// Switch to a different branch at a branch point - #[allow(dead_code)] // Will be used in Phase 4 SwitchBranch { session_id: String, /// The node ID to switch to (a sibling of the current node at a branch point) @@ -155,7 +151,6 @@ pub enum UiEvent { /// Response: Message content loaded for editing /// Sent in response to StartMessageEdit - #[allow(dead_code)] // Will be used in Phase 4 MessageEditReady { /// The text content of the message content: String, @@ -170,7 +165,6 @@ pub enum UiEvent { }, /// Response: Branch switch completed, new messages to display - #[allow(dead_code)] // Will be used in Phase 4 BranchSwitched { session_id: String, /// Full message list for the new active path From e35aa48e57d021f658b5818a3b2c8582edd09a83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 12:49:49 +0100 Subject: [PATCH 13/14] Clean up and plan-related fixes --- crates/code_assistant/src/agent/runner.rs | 27 ++++++++++++ crates/code_assistant/src/persistence.rs | 50 ----------------------- 2 files changed, 27 insertions(+), 50 deletions(-) diff --git a/crates/code_assistant/src/agent/runner.rs b/crates/code_assistant/src/agent/runner.rs index f3bb1a3c..0c241790 100644 --- a/crates/code_assistant/src/agent/runner.rs +++ b/crates/code_assistant/src/agent/runner.rs @@ -380,6 +380,27 @@ impl Agent { 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<()> { @@ -1797,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/persistence.rs b/crates/code_assistant/src/persistence.rs index 40532f08..04e124b3 100644 --- a/crates/code_assistant/src/persistence.rs +++ b/crates/code_assistant/src/persistence.rs @@ -44,7 +44,6 @@ pub struct MessageNode { /// Information about a branch point in the conversation (for UI) #[derive(Debug, Clone, PartialEq)] -#[allow(dead_code)] // Will be used by UI in Phase 4 pub struct BranchInfo { /// Node ID where the branch occurs (the node that has multiple children) pub parent_node_id: Option, @@ -56,14 +55,6 @@ pub struct BranchInfo { pub active_index: usize, } -impl BranchInfo { - /// Total number of branches at this point - #[allow(dead_code)] // Will be used by UI in Phase 4 - pub fn total_branches(&self) -> usize { - self.sibling_ids.len() - } -} - /// Model configuration for a session #[derive(Debug, Serialize, Deserialize, Clone)] pub struct SessionModelConfig { @@ -277,7 +268,6 @@ impl ChatSession { } /// Get all direct children of a node. - #[allow(dead_code)] // Will be used by UI in Phase 4 pub fn get_children(&self, parent_id: Option) -> Vec<&MessageNode> { self.message_nodes .values() @@ -286,22 +276,14 @@ impl ChatSession { } /// Get children sorted by creation time (oldest first). - #[allow(dead_code)] // Will be used by UI in Phase 4 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 } - /// Check if a node has multiple children (is a branch point). - #[allow(dead_code)] // Will be used by UI in Phase 4 - 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 (if it's part of a branch). /// Returns None if the node has no siblings (no branching at this point). - #[allow(dead_code)] // Will be used by UI in Phase 4 pub fn get_branch_info(&self, node_id: NodeId) -> Option { let node = self.message_nodes.get(&node_id)?; let siblings: Vec = self @@ -323,18 +305,8 @@ impl ChatSession { }) } - /// Get branch infos for all branch points in the active path. - #[allow(dead_code)] // Will be used by UI in Phase 4 - pub fn get_all_branch_infos(&self) -> Vec<(NodeId, BranchInfo)> { - self.active_path - .iter() - .filter_map(|&node_id| self.get_branch_info(node_id).map(|info| (node_id, info))) - .collect() - } - /// Find the plan state for the active path by walking backwards /// to find the most recent plan_snapshot. - #[allow(dead_code)] // Will be used by UI in Phase 4 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) { @@ -353,7 +325,6 @@ impl ChatSession { /// 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. - #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 pub fn add_message(&mut self, message: Message) -> NodeId { self.add_message_with_parent(message, self.active_path.last().copied()) } @@ -361,7 +332,6 @@ impl ChatSession { /// 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. - #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 pub fn add_message_with_parent( &mut self, message: Message, @@ -400,27 +370,8 @@ impl ChatSession { node_id } - /// Add a message and store a plan snapshot with it. - #[allow(dead_code)] // Will be used by UI in Phase 4 - pub fn add_message_with_plan(&mut self, message: Message, plan: PlanState) -> NodeId { - let node_id = self.add_message(message); - if let Some(node) = self.message_nodes.get_mut(&node_id) { - node.plan_snapshot = Some(plan); - } - node_id - } - - /// Update the plan snapshot for an existing node. - #[allow(dead_code)] // Will be used by UI in Phase 4 - pub fn set_plan_snapshot(&mut self, node_id: NodeId, plan: PlanState) { - if let Some(node) = self.message_nodes.get_mut(&node_id) { - node.plan_snapshot = Some(plan); - } - } - /// Switch to a different branch by making a different sibling node active. /// Updates active_path to follow the new branch to its deepest descendant. - #[allow(dead_code)] // Used by tests, will be used by UI in Phase 4 pub fn switch_branch(&mut self, new_node_id: NodeId) -> Result<()> { let node = self .message_nodes @@ -1377,7 +1328,6 @@ mod tests { 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); - assert_eq!(info_3.total_branches(), 3); let info_5 = session.get_branch_info(5).expect("should have branch info"); assert_eq!(info_5.parent_node_id, Some(2)); From a6721deee39754eecee435b62e223b3ef836cc70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stephan=20A=C3=9Fmus?= Date: Mon, 12 Jan 2026 13:52:08 +0100 Subject: [PATCH 14/14] More elegant branch related UI --- .../src/ui/gpui/branch_switcher.rs | 195 ++++++++++-------- crates/code_assistant/src/ui/gpui/messages.rs | 20 +- 2 files changed, 121 insertions(+), 94 deletions(-) diff --git a/crates/code_assistant/src/ui/gpui/branch_switcher.rs b/crates/code_assistant/src/ui/gpui/branch_switcher.rs index 3461dc01..1f3af163 100644 --- a/crates/code_assistant/src/ui/gpui/branch_switcher.rs +++ b/crates/code_assistant/src/ui/gpui/branch_switcher.rs @@ -1,9 +1,9 @@ use crate::persistence::BranchInfo; -use gpui::{div, prelude::*, px, App, CursorStyle, MouseButton, Window}; -use gpui_component::ActiveTheme; +use gpui::{div, prelude::*, px, App, CursorStyle, MouseButton, SharedString, Window}; +use gpui_component::{ActiveTheme, Icon}; -/// A stateless branch navigation component that shows "◀ 2/3 ▶" -/// Displayed below user messages where branches exist. +/// 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, @@ -26,9 +26,8 @@ impl RenderOnce for BranchSwitcherElement { let current = self.branch_info.active_index + 1; let total = self.branch_info.sibling_ids.len(); - let text_color = cx.theme().muted_foreground; + let muted_color = cx.theme().muted_foreground; let active_color = cx.theme().foreground; - let hover_bg = cx.theme().muted; // Get sibling IDs for navigation let prev_node_id = if has_prev { @@ -52,88 +51,116 @@ impl RenderOnce for BranchSwitcherElement { let session_id = self.session_id.clone(); let session_id_for_next = session_id.clone(); - div() - .flex() - .flex_row() - .items_center() - .gap_1() - .text_xs() - .mt_1() - .children(vec![ - // Previous button - { - let base = div() - .id("branch-prev") - .px_1() - .py(px(2.)) - .rounded_sm() - .cursor(if has_prev { - CursorStyle::PointingHand - } else { - CursorStyle::OperationNotAllowed - }) - .text_color(if has_prev { active_color } else { text_color }) - .child("◀"); + // 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(hover_bg)) - .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { - if let Some(node_id) = prev_node_id { - // Send event through UiEventSender global - 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, - }); + 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(text_color) - .child(format!("{}/{}", current, total)) - .into_any_element(), - // Next button - { - let base = div() - .id("branch-next") - .px_1() - .py(px(2.)) - .rounded_sm() - .cursor(if has_next { - CursorStyle::PointingHand + }) + .into_any_element() } else { - CursorStyle::OperationNotAllowed - }) - .text_color(if has_next { active_color } else { text_color }) - .child("▶"); + 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(hover_bg)) - .on_mouse_up(MouseButton::Left, move |_event, _window, cx| { - if let Some(node_id) = next_node_id { - // Send event through UiEventSender global - 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, - }); + 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() - } - }, - ]) + }) + .into_any_element() + } else { + base.into_any_element() + } + }, + ]), + ) } } diff --git a/crates/code_assistant/src/ui/gpui/messages.rs b/crates/code_assistant/src/ui/gpui/messages.rs index a3037397..b6b7e3b9 100644 --- a/crates/code_assistant/src/ui/gpui/messages.rs +++ b/crates/code_assistant/src/ui/gpui/messages.rs @@ -2,9 +2,9 @@ use super::branch_switcher::BranchSwitcherElement; use super::elements::MessageContainer; use gpui::{ div, prelude::*, px, rgb, App, Context, CursorStyle, Entity, FocusHandle, Focusable, - MouseButton, Window, + MouseButton, SharedString, Window, }; -use gpui_component::{v_flex, ActiveTheme}; +use gpui_component::{v_flex, ActiveTheme, Icon}; use std::sync::{Arc, Mutex}; /// MessagesView - Component responsible for displaying the message history @@ -189,15 +189,10 @@ impl Render for MessagesView { el.child( div() .id("edit-message-btn") - .px_2() - .py_1() + .p_1() .rounded_sm() .cursor(CursorStyle::PointingHand) - .text_xs() - .text_color(cx.theme().muted_foreground) - .hover(|s| { - s.bg(cx.theme().muted).text_color(cx.theme().foreground) - }) + .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) @@ -215,7 +210,12 @@ impl Render for MessagesView { } } }) - .child("Edit"), + .child( + Icon::default() + .path(SharedString::from("icons/pencil.svg")) + .text_color(cx.theme().muted_foreground) + .size_4(), + ), ) });