pi/interactive/
tree_ui.rs

1use super::*;
2
3impl PiApp {
4    #[allow(clippy::too_many_lines)]
5    pub(super) fn handle_tree_ui_key(&mut self, key: &KeyMsg) -> Option<Cmd> {
6        let tree_ui = self.tree_ui.take()?;
7
8        match tree_ui {
9            TreeUiState::Selector(mut selector) => {
10                match key.key_type {
11                    KeyType::Up => selector.move_selection(-1),
12                    KeyType::Down => selector.move_selection(1),
13                    KeyType::CtrlU => {
14                        selector.user_only = !selector.user_only;
15                        if let Ok(session_guard) = self.session.try_lock() {
16                            selector.rebuild(&session_guard);
17                        }
18                    }
19                    KeyType::CtrlO => {
20                        selector.show_all = !selector.show_all;
21                        if let Ok(session_guard) = self.session.try_lock() {
22                            selector.rebuild(&session_guard);
23                        }
24                    }
25                    KeyType::Esc | KeyType::CtrlC => {
26                        self.status_message = Some("Tree navigation cancelled".to_string());
27                        self.tree_ui = None;
28                        return None;
29                    }
30                    KeyType::Enter => {
31                        if selector.rows.is_empty() {
32                            self.tree_ui = None;
33                            return None;
34                        }
35
36                        let selected = selector.rows[selector.selected].clone();
37                        selector.last_selected_id = Some(selected.id.clone());
38
39                        let (new_leaf_id, editor_text) = if let Some(text) = selected.resubmit_text
40                        {
41                            (selected.parent_id.clone(), Some(text))
42                        } else {
43                            (Some(selected.id.clone()), None)
44                        };
45
46                        // No-op if already at target leaf.
47                        if selector.current_leaf_id.as_deref() == new_leaf_id.as_deref() {
48                            self.status_message = Some("Already on that branch".to_string());
49                            self.tree_ui = None;
50                            return None;
51                        }
52
53                        let Ok(session_guard) = self.session.try_lock() else {
54                            self.status_message = Some("Session busy; try again".to_string());
55                            self.tree_ui = None;
56                            return None;
57                        };
58
59                        let old_leaf_id = session_guard.leaf_id.clone();
60                        let (entries_to_summarize, summary_from_id) = collect_tree_branch_entries(
61                            &session_guard,
62                            old_leaf_id.as_deref(),
63                            new_leaf_id.as_deref(),
64                        );
65                        let session_id = session_guard.header.id.clone();
66                        drop(session_guard);
67
68                        let api_key_present = self.agent.try_lock().is_ok_and(|agent_guard| {
69                            agent_guard.stream_options().api_key.is_some()
70                        });
71
72                        let pending = PendingTreeNavigation {
73                            session_id,
74                            old_leaf_id,
75                            selected_entry_id: selected.id,
76                            new_leaf_id,
77                            editor_text,
78                            entries_to_summarize,
79                            summary_from_id,
80                            api_key_present,
81                        };
82
83                        if pending.entries_to_summarize.is_empty() {
84                            // Nothing to summarize; switch immediately.
85                            if !self.start_tree_navigation(
86                                pending,
87                                TreeSummaryChoice::NoSummary,
88                                None,
89                            ) {
90                                self.tree_ui = Some(TreeUiState::Selector(selector));
91                            }
92                            return None;
93                        }
94
95                        self.tree_ui = Some(TreeUiState::SummaryPrompt(TreeSummaryPromptState {
96                            pending,
97                            selected: 0,
98                        }));
99                        return None;
100                    }
101                    _ => {}
102                }
103
104                self.tree_ui = Some(TreeUiState::Selector(selector));
105            }
106            TreeUiState::SummaryPrompt(mut prompt) => {
107                match key.key_type {
108                    KeyType::Up if prompt.selected > 0 => {
109                        prompt.selected -= 1;
110                    }
111                    KeyType::Down
112                        if prompt.selected < TreeSummaryChoice::all().len().saturating_sub(1) =>
113                    {
114                        prompt.selected += 1;
115                    }
116                    KeyType::Esc | KeyType::CtrlC => {
117                        self.status_message = Some("Tree navigation cancelled".to_string());
118                        self.tree_ui = None;
119                        return None;
120                    }
121                    KeyType::Enter => {
122                        let choice = TreeSummaryChoice::all()[prompt.selected];
123                        match choice {
124                            TreeSummaryChoice::NoSummary | TreeSummaryChoice::Summarize => {
125                                let pending = prompt.pending.clone();
126                                if !self.start_tree_navigation(pending, choice, None) {
127                                    self.tree_ui = Some(TreeUiState::SummaryPrompt(prompt));
128                                }
129                                return None;
130                            }
131                            TreeSummaryChoice::SummarizeWithCustomPrompt => {
132                                self.tree_ui =
133                                    Some(TreeUiState::CustomPrompt(TreeCustomPromptState {
134                                        pending: prompt.pending,
135                                        instructions: String::new(),
136                                    }));
137                                return None;
138                            }
139                        }
140                    }
141                    _ => {}
142                }
143                self.tree_ui = Some(TreeUiState::SummaryPrompt(prompt));
144            }
145            TreeUiState::CustomPrompt(mut custom) => {
146                match key.key_type {
147                    KeyType::Esc | KeyType::CtrlC => {
148                        self.tree_ui = Some(TreeUiState::SummaryPrompt(TreeSummaryPromptState {
149                            pending: custom.pending,
150                            selected: 2,
151                        }));
152                        return None;
153                    }
154                    KeyType::Backspace => {
155                        custom.instructions.pop();
156                    }
157                    KeyType::Enter => {
158                        let pending = custom.pending.clone();
159                        let instructions = if custom.instructions.trim().is_empty() {
160                            None
161                        } else {
162                            Some(custom.instructions.clone())
163                        };
164                        if !self.start_tree_navigation(
165                            pending,
166                            TreeSummaryChoice::SummarizeWithCustomPrompt,
167                            instructions,
168                        ) {
169                            self.tree_ui = Some(TreeUiState::CustomPrompt(custom));
170                        }
171                        return None;
172                    }
173                    KeyType::Runes => {
174                        for ch in key.runes.iter().copied() {
175                            custom.instructions.push(ch);
176                        }
177                    }
178                    _ => {}
179                }
180                self.tree_ui = Some(TreeUiState::CustomPrompt(custom));
181            }
182        }
183
184        None
185    }
186
187    /// Handle keyboard input when the branch picker overlay is active.
188    pub fn handle_branch_picker_key(&mut self, key: &KeyMsg) -> Option<Cmd> {
189        let picker = self.branch_picker.as_mut()?;
190
191        match key.key_type {
192            KeyType::Up => picker.select_prev(),
193            KeyType::Down => picker.select_next(),
194            KeyType::PgUp => picker.select_page_up(),
195            KeyType::PgDown => picker.select_page_down(),
196            KeyType::Runes if key.runes == ['k'] => picker.select_prev(),
197            KeyType::Runes if key.runes == ['j'] => picker.select_next(),
198            KeyType::Enter => {
199                if let Some(branch) = picker.selected_branch().cloned() {
200                    if self.switch_to_branch_leaf(&branch.leaf_id) {
201                        self.branch_picker = None;
202                    }
203                    return None;
204                }
205                self.branch_picker = None;
206            }
207            KeyType::Esc | KeyType::CtrlC => {
208                self.branch_picker = None;
209                self.status_message = Some("Branch picker cancelled".to_string());
210            }
211            KeyType::Runes if key.runes == ['q'] => {
212                self.branch_picker = None;
213            }
214            _ => {} // consume all other input while picker is open
215        }
216
217        None
218    }
219
220    /// Switch the active branch to a different leaf. Reloads the conversation.
221    fn switch_to_branch_leaf(&mut self, leaf_id: &str) -> bool {
222        let Ok(session_guard) = self.session.try_lock() else {
223            self.status_message = Some("Session busy; try again".to_string());
224            return false;
225        };
226        let session_id = session_guard.header.id.clone();
227        let old_leaf_id = session_guard.leaf_id.clone();
228        drop(session_guard);
229
230        let pending = PendingTreeNavigation {
231            session_id,
232            old_leaf_id,
233            selected_entry_id: leaf_id.to_string(),
234            new_leaf_id: Some(leaf_id.to_string()),
235            editor_text: None,
236            entries_to_summarize: Vec::new(),
237            summary_from_id: String::new(),
238            api_key_present: false,
239        };
240        self.start_tree_navigation(pending, TreeSummaryChoice::NoSummary, None)
241    }
242
243    /// Open the branch picker if the session has sibling branches.
244    pub fn open_branch_picker(&mut self) {
245        if self.agent_state != AgentState::Idle {
246            self.status_message = Some("Cannot switch branches while processing".to_string());
247            return;
248        }
249
250        let Ok(session_guard) = self.session.try_lock() else {
251            self.status_message = Some("Session busy; try again".to_string());
252            return;
253        };
254        let branches = session_guard.sibling_branches().map(|(_, b)| b);
255        drop(session_guard);
256
257        match branches {
258            Some(branches) if branches.len() > 1 => {
259                let mut picker = BranchPickerOverlay::new(branches);
260                picker.max_visible = super::overlay_max_visible(self.term_height);
261                self.branch_picker = Some(picker);
262            }
263            _ => {
264                self.status_message =
265                    Some("No branches to pick (use /fork to create one)".to_string());
266            }
267        }
268    }
269
270    /// Cycle to the next or previous sibling branch (Ctrl+Right / Ctrl+Left).
271    pub fn cycle_sibling_branch(&mut self, forward: bool) {
272        if self.agent_state != AgentState::Idle {
273            self.status_message = Some("Cannot switch branches while processing".to_string());
274            return;
275        }
276
277        let Ok(session_guard) = self.session.try_lock() else {
278            self.status_message = Some("Session busy; try again".to_string());
279            return;
280        };
281        let target = session_guard.sibling_branches().and_then(|(_, branches)| {
282            if branches.len() <= 1 {
283                return None;
284            }
285            let current_idx = branches.iter().position(|b| b.is_current)?;
286            let next_idx = if forward {
287                (current_idx + 1) % branches.len()
288            } else {
289                current_idx.checked_sub(1).unwrap_or(branches.len() - 1)
290            };
291            Some(branches[next_idx].leaf_id.clone())
292        });
293        drop(session_guard);
294
295        if let Some(leaf_id) = target {
296            self.switch_to_branch_leaf(&leaf_id);
297        } else {
298            self.status_message = Some("No sibling branches (use /fork to create one)".to_string());
299        }
300    }
301
302    #[allow(clippy::too_many_lines)]
303    pub(super) fn start_tree_navigation(
304        &mut self,
305        pending: PendingTreeNavigation,
306        choice: TreeSummaryChoice,
307        custom_instructions: Option<String>,
308    ) -> bool {
309        let summary_requested = matches!(
310            choice,
311            TreeSummaryChoice::Summarize | TreeSummaryChoice::SummarizeWithCustomPrompt
312        );
313
314        // Fast path: no summary + no extensions. Keep it synchronous so unit tests can drive it
315        // without running the async runtime.
316        if !summary_requested && self.extensions.is_none() {
317            let Ok(mut session_guard) = self.session.try_lock() else {
318                self.status_message = Some("Session busy; try again".to_string());
319                return false;
320            };
321
322            if let Some(target_id) = &pending.new_leaf_id {
323                if !session_guard.navigate_to(target_id) {
324                    self.status_message = Some(format!("Branch target not found: {target_id}"));
325                    return false;
326                }
327            } else {
328                session_guard.reset_leaf();
329            }
330
331            let (messages, usage) = conversation_from_session(&session_guard);
332            let agent_messages = session_guard.to_messages_for_current_path();
333            let status_leaf = pending
334                .new_leaf_id
335                .clone()
336                .unwrap_or_else(|| "root".to_string());
337            drop(session_guard);
338
339            self.spawn_save_session();
340
341            if let Ok(mut agent_guard) = self.agent.try_lock() {
342                agent_guard.replace_messages(agent_messages);
343            }
344
345            self.messages = messages;
346            self.message_render_cache.clear();
347            self.total_usage = usage;
348            self.current_response.clear();
349            self.current_thinking.clear();
350            self.agent_state = AgentState::Idle;
351            self.current_tool = None;
352            self.abort_handle = None;
353            self.status_message = Some(format!("Switched to {status_leaf}"));
354            if let Err(message) = self.sync_runtime_selection_from_session_header() {
355                self.status_message = Some(message);
356            }
357            self.scroll_to_bottom();
358
359            if let Some(text) = pending.editor_text {
360                self.input.set_value(&text);
361            }
362            self.input.focus();
363
364            return true;
365        }
366
367        let event_tx = self.event_tx.clone();
368        let session = Arc::clone(&self.session);
369        let agent = Arc::clone(&self.agent);
370        let extensions = self.extensions.clone();
371        let reserve_tokens = self.config.branch_summary_reserve_tokens();
372        let runtime_handle = self.runtime_handle.clone();
373
374        let Ok(agent_guard) = self.agent.try_lock() else {
375            self.status_message = Some("Agent busy; try again".to_string());
376            self.agent_state = AgentState::Idle;
377            return false;
378        };
379        let provider = agent_guard.provider();
380        let key_opt = agent_guard.stream_options().api_key.clone();
381
382        self.tree_ui = None;
383        self.agent_state = AgentState::Processing;
384        self.status_message = Some("Switching branches...".to_string());
385
386        runtime_handle.spawn(async move {
387            let cx = Cx::for_request();
388
389            let from_id_for_event = pending
390                .old_leaf_id
391                .clone()
392                .unwrap_or_else(|| "root".to_string());
393            let to_id_for_event = pending
394                .new_leaf_id
395                .clone()
396                .unwrap_or_else(|| "root".to_string());
397
398            if let Some(manager) = extensions.clone() {
399                let cancelled = manager
400                    .dispatch_cancellable_event(
401                        ExtensionEventName::SessionBeforeSwitch,
402                        Some(json!({
403                            "fromId": from_id_for_event.clone(),
404                            "toId": to_id_for_event.clone(),
405                            "sessionId": pending.session_id.clone(),
406                        })),
407                        EXTENSION_EVENT_TIMEOUT_MS,
408                    )
409                    .await
410                    .unwrap_or(false);
411                if cancelled {
412                    let _ = crate::interactive::enqueue_pi_event(
413                        &event_tx,
414                        &asupersync::Cx::current().unwrap_or_else(asupersync::Cx::for_request),
415                        PiMsg::System("Session switch cancelled by extension".to_string()),
416                    )
417                    .await;
418                    return;
419                }
420            }
421
422            let summary_skipped =
423                summary_requested && key_opt.is_none() && !pending.entries_to_summarize.is_empty();
424            let summary_text = if !summary_requested || pending.entries_to_summarize.is_empty() {
425                None
426            } else if let Some(key) = key_opt.as_deref() {
427                match crate::compaction::summarize_entries(
428                    &pending.entries_to_summarize,
429                    provider,
430                    key,
431                    reserve_tokens,
432                    custom_instructions.as_deref(),
433                )
434                .await
435                {
436                    Ok(summary) => summary,
437                    Err(err) => {
438                        let _ = crate::interactive::enqueue_pi_event(
439                            &event_tx,
440                            &cx,
441                            PiMsg::AgentError(format!("Branch summary failed: {err}")),
442                        )
443                        .await;
444                        return;
445                    }
446                }
447            } else {
448                None
449            };
450
451            let mut summary_entry_payload: Option<Value> = None;
452            let mut summary_entry_id: Option<String> = None;
453
454            let messages_for_agent = {
455                let mut guard = match session.lock(&cx).await {
456                    Ok(guard) => guard,
457                    Err(err) => {
458                        let _ = crate::interactive::enqueue_pi_event(
459                            &event_tx,
460                            &cx,
461                            PiMsg::AgentError(format!("Failed to lock session: {err}")),
462                        )
463                        .await;
464                        return;
465                    }
466                };
467
468                if let Some(target_id) = &pending.new_leaf_id {
469                    if !guard.navigate_to(target_id) {
470                        let _ = crate::interactive::enqueue_pi_event(
471                            &event_tx,
472                            &asupersync::Cx::current().unwrap_or_else(asupersync::Cx::for_request),
473                            PiMsg::AgentError(format!("Branch target not found: {target_id}")),
474                        )
475                        .await;
476                        return;
477                    }
478                } else {
479                    guard.reset_leaf();
480                }
481
482                if let Some(summary_text) = summary_text {
483                    let summary_clone = summary_text.clone();
484                    guard.append_branch_summary(
485                        pending.summary_from_id.clone(),
486                        summary_text,
487                        None,
488                        None,
489                    );
490                    summary_entry_id = guard.leaf_id.clone();
491                    let mut summary_entry = serde_json::Map::new();
492                    summary_entry.insert(
493                        "type".to_string(),
494                        Value::String("branch_summary".to_string()),
495                    );
496                    summary_entry.insert(
497                        "fromId".to_string(),
498                        Value::String(pending.summary_from_id.clone()),
499                    );
500                    summary_entry.insert("summary".to_string(), Value::String(summary_clone));
501                    summary_entry.insert("fromHook".to_string(), Value::Bool(false));
502                    summary_entry_payload = Some(Value::Object(summary_entry));
503                }
504
505                let _ = guard.save().await;
506                guard.to_messages_for_current_path()
507            };
508
509            {
510                let mut agent_guard = match agent.lock(&cx).await {
511                    Ok(guard) => guard,
512                    Err(err) => {
513                        let _ = crate::interactive::enqueue_pi_event(
514                            &event_tx,
515                            &cx,
516                            PiMsg::AgentError(format!("Failed to lock agent: {err}")),
517                        )
518                        .await;
519                        return;
520                    }
521                };
522                agent_guard.replace_messages(messages_for_agent);
523            }
524
525            let (messages, usage) = {
526                let guard = match session.lock(&cx).await {
527                    Ok(guard) => guard,
528                    Err(err) => {
529                        let _ = crate::interactive::enqueue_pi_event(
530                            &event_tx,
531                            &cx,
532                            PiMsg::AgentError(format!("Failed to lock session: {err}")),
533                        )
534                        .await;
535                        return;
536                    }
537                };
538                conversation_from_session(&guard)
539            };
540
541            let status = if summary_skipped {
542                Some(format!(
543                    "Switched to {to_id_for_event} (no summary: missing API key)"
544                ))
545            } else {
546                Some(format!("Switched to {to_id_for_event}"))
547            };
548
549            let _ = crate::interactive::enqueue_pi_event(
550                &event_tx,
551                &asupersync::Cx::current().unwrap_or_else(asupersync::Cx::for_request),
552                PiMsg::ConversationReset {
553                    messages,
554                    usage,
555                    status,
556                },
557            )
558            .await;
559
560            if let Some(text) = pending.editor_text {
561                let _ = crate::interactive::enqueue_pi_event(
562                    &event_tx,
563                    &asupersync::Cx::current().unwrap_or_else(asupersync::Cx::for_request),
564                    PiMsg::SetEditorText(text),
565                )
566                .await;
567            }
568
569            if let Some(manager) = extensions {
570                let new_leaf_id = summary_entry_id
571                    .clone()
572                    .or_else(|| pending.new_leaf_id.clone());
573                let old_leaf_value = pending
574                    .old_leaf_id
575                    .clone()
576                    .map_or(Value::Null, Value::String);
577                let new_leaf_value = new_leaf_id.clone().map_or(Value::Null, Value::String);
578                let mut tree_payload = serde_json::Map::new();
579                tree_payload.insert("newLeafId".to_string(), new_leaf_value);
580                tree_payload.insert("oldLeafId".to_string(), old_leaf_value);
581                if let Some(summary_entry) = summary_entry_payload {
582                    tree_payload.insert("summaryEntry".to_string(), summary_entry);
583                }
584
585                let _ = manager
586                    .dispatch_event(
587                        ExtensionEventName::SessionSwitch,
588                        Some(json!({
589                            "fromId": from_id_for_event,
590                            "toId": to_id_for_event,
591                            "sessionId": pending.session_id,
592                        })),
593                    )
594                    .await;
595                let _ = manager
596                    .dispatch_event(
597                        ExtensionEventName::SessionTree,
598                        Some(Value::Object(tree_payload)),
599                    )
600                    .await;
601            }
602        });
603        true
604    }
605}
pi/interactive/tree_ui.rs

pi/interactive/
tree_ui.rs
