use super::*;
impl PiApp {
#[allow(clippy::too_many_lines)]
pub(super) fn handle_tree_ui_key(&mut self, key: &KeyMsg) -> Option<Cmd> {
let tree_ui = self.tree_ui.take()?;
match tree_ui {
TreeUiState::Selector(mut selector) => {
match key.key_type {
KeyType::Up => selector.move_selection(-1),
KeyType::Down => selector.move_selection(1),
KeyType::CtrlU => {
selector.user_only = !selector.user_only;
if let Ok(session_guard) = self.session.try_lock() {
selector.rebuild(&session_guard);
}
}
KeyType::CtrlO => {
selector.show_all = !selector.show_all;
if let Ok(session_guard) = self.session.try_lock() {
selector.rebuild(&session_guard);
}
}
KeyType::Esc | KeyType::CtrlC => {
self.status_message = Some("Tree navigation cancelled".to_string());
self.tree_ui = None;
return None;
}
KeyType::Enter => {
if selector.rows.is_empty() {
self.tree_ui = None;
return None;
}
let selected = selector.rows[selector.selected].clone();
selector.last_selected_id = Some(selected.id.clone());
let (new_leaf_id, editor_text) = if let Some(text) = selected.resubmit_text
{
(selected.parent_id.clone(), Some(text))
} else {
(Some(selected.id.clone()), None)
};
if selector.current_leaf_id.as_deref() == new_leaf_id.as_deref() {
self.status_message = Some("Already on that branch".to_string());
self.tree_ui = None;
return None;
}
let Ok(session_guard) = self.session.try_lock() else {
self.status_message = Some("Session busy; try again".to_string());
self.tree_ui = None;
return None;
};
let old_leaf_id = session_guard.leaf_id.clone();
let (entries_to_summarize, summary_from_id) = collect_tree_branch_entries(
&session_guard,
old_leaf_id.as_deref(),
new_leaf_id.as_deref(),
);
let session_id = session_guard.header.id.clone();
drop(session_guard);
let api_key_present = self.agent.try_lock().is_ok_and(|agent_guard| {
agent_guard.stream_options().api_key.is_some()
});
let pending = PendingTreeNavigation {
session_id,
old_leaf_id,
selected_entry_id: selected.id,
new_leaf_id,
editor_text,
entries_to_summarize,
summary_from_id,
api_key_present,
};
if pending.entries_to_summarize.is_empty() {
self.start_tree_navigation(pending, TreeSummaryChoice::NoSummary, None);
return None;
}
self.tree_ui = Some(TreeUiState::SummaryPrompt(TreeSummaryPromptState {
pending,
selected: 0,
}));
return None;
}
_ => {}
}
self.tree_ui = Some(TreeUiState::Selector(selector));
}
TreeUiState::SummaryPrompt(mut prompt) => {
match key.key_type {
KeyType::Up => {
if prompt.selected > 0 {
prompt.selected -= 1;
}
}
KeyType::Down => {
if prompt.selected < TreeSummaryChoice::all().len().saturating_sub(1) {
prompt.selected += 1;
}
}
KeyType::Esc | KeyType::CtrlC => {
self.status_message = Some("Tree navigation cancelled".to_string());
self.tree_ui = None;
return None;
}
KeyType::Enter => {
let choice = TreeSummaryChoice::all()[prompt.selected];
match choice {
TreeSummaryChoice::NoSummary | TreeSummaryChoice::Summarize => {
let pending = prompt.pending;
self.start_tree_navigation(pending, choice, None);
return None;
}
TreeSummaryChoice::SummarizeWithCustomPrompt => {
self.tree_ui =
Some(TreeUiState::CustomPrompt(TreeCustomPromptState {
pending: prompt.pending,
instructions: String::new(),
}));
return None;
}
}
}
_ => {}
}
self.tree_ui = Some(TreeUiState::SummaryPrompt(prompt));
}
TreeUiState::CustomPrompt(mut custom) => {
match key.key_type {
KeyType::Esc | KeyType::CtrlC => {
self.tree_ui = Some(TreeUiState::SummaryPrompt(TreeSummaryPromptState {
pending: custom.pending,
selected: 2,
}));
return None;
}
KeyType::Backspace => {
custom.instructions.pop();
}
KeyType::Enter => {
let pending = custom.pending;
let instructions = if custom.instructions.trim().is_empty() {
None
} else {
Some(custom.instructions)
};
self.start_tree_navigation(
pending,
TreeSummaryChoice::SummarizeWithCustomPrompt,
instructions,
);
return None;
}
KeyType::Runes => {
for ch in key.runes.iter().copied() {
custom.instructions.push(ch);
}
}
_ => {}
}
self.tree_ui = Some(TreeUiState::CustomPrompt(custom));
}
}
None
}
pub fn handle_branch_picker_key(&mut self, key: &KeyMsg) -> Option<Cmd> {
let picker = self.branch_picker.as_mut()?;
match key.key_type {
KeyType::Up => picker.select_prev(),
KeyType::Down => picker.select_next(),
KeyType::Runes if key.runes == ['k'] => picker.select_prev(),
KeyType::Runes if key.runes == ['j'] => picker.select_next(),
KeyType::Enter => {
if let Some(branch) = picker.selected_branch().cloned() {
self.branch_picker = None;
return self.switch_to_branch_leaf(&branch.leaf_id);
}
self.branch_picker = None;
}
KeyType::Esc | KeyType::CtrlC => {
self.branch_picker = None;
self.status_message = Some("Branch picker cancelled".to_string());
}
KeyType::Runes if key.runes == ['q'] => {
self.branch_picker = None;
}
_ => {} }
None
}
fn switch_to_branch_leaf(&mut self, leaf_id: &str) -> Option<Cmd> {
let (session_id, old_leaf_id) = self
.session
.try_lock()
.ok()
.map(|g| (g.header.id.clone(), g.leaf_id.clone()))
.unwrap_or_default();
let pending = PendingTreeNavigation {
session_id,
old_leaf_id,
selected_entry_id: leaf_id.to_string(),
new_leaf_id: Some(leaf_id.to_string()),
editor_text: None,
entries_to_summarize: Vec::new(),
summary_from_id: String::new(),
api_key_present: false,
};
self.start_tree_navigation(pending, TreeSummaryChoice::NoSummary, None);
None
}
pub fn open_branch_picker(&mut self) {
if self.agent_state != AgentState::Idle {
self.status_message = Some("Cannot switch branches while processing".to_string());
return;
}
let branches = self
.session
.try_lock()
.ok()
.and_then(|guard| guard.sibling_branches().map(|(_, b)| b));
match branches {
Some(branches) if branches.len() > 1 => {
self.branch_picker = Some(BranchPickerOverlay::new(branches));
}
_ => {
self.status_message =
Some("No branches to pick (use /fork to create one)".to_string());
}
}
}
pub fn cycle_sibling_branch(&mut self, forward: bool) {
if self.agent_state != AgentState::Idle {
self.status_message = Some("Cannot switch branches while processing".to_string());
return;
}
let target = self.session.try_lock().ok().and_then(|guard| {
let (_, branches) = guard.sibling_branches()?;
if branches.len() <= 1 {
return None;
}
let current_idx = branches.iter().position(|b| b.is_current)?;
let next_idx = if forward {
(current_idx + 1) % branches.len()
} else {
current_idx.checked_sub(1).unwrap_or(branches.len() - 1)
};
Some(branches[next_idx].leaf_id.clone())
});
if let Some(leaf_id) = target {
self.switch_to_branch_leaf(&leaf_id);
} else {
self.status_message = Some("No sibling branches (use /fork to create one)".to_string());
}
}
#[allow(clippy::too_many_lines)]
pub(super) fn start_tree_navigation(
&mut self,
pending: PendingTreeNavigation,
choice: TreeSummaryChoice,
custom_instructions: Option<String>,
) {
let summary_requested = matches!(
choice,
TreeSummaryChoice::Summarize | TreeSummaryChoice::SummarizeWithCustomPrompt
);
if !summary_requested && self.extensions.is_none() {
let Ok(mut session_guard) = self.session.try_lock() else {
self.status_message = Some("Session busy; try again".to_string());
return;
};
if let Some(target_id) = &pending.new_leaf_id {
if !session_guard.navigate_to(target_id) {
self.status_message = Some(format!("Branch target not found: {target_id}"));
return;
}
} else {
session_guard.reset_leaf();
}
let (messages, usage) = conversation_from_session(&session_guard);
let agent_messages = session_guard.to_messages_for_current_path();
let status_leaf = pending
.new_leaf_id
.clone()
.unwrap_or_else(|| "root".to_string());
drop(session_guard);
self.spawn_save_session();
if let Ok(mut agent_guard) = self.agent.try_lock() {
agent_guard.replace_messages(agent_messages);
}
self.messages = messages;
self.message_render_cache.clear();
self.total_usage = usage;
self.current_response.clear();
self.current_thinking.clear();
self.agent_state = AgentState::Idle;
self.current_tool = None;
self.abort_handle = None;
self.status_message = Some(format!("Switched to {status_leaf}"));
self.scroll_to_bottom();
if let Some(text) = pending.editor_text {
self.input.set_value(&text);
}
self.input.focus();
return;
}
let event_tx = self.event_tx.clone();
let session = Arc::clone(&self.session);
let agent = Arc::clone(&self.agent);
let extensions = self.extensions.clone();
let reserve_tokens = self.config.branch_summary_reserve_tokens();
let runtime_handle = self.runtime_handle.clone();
let Ok(agent_guard) = self.agent.try_lock() else {
self.status_message = Some("Agent busy; try again".to_string());
self.agent_state = AgentState::Idle;
return;
};
let provider = agent_guard.provider();
let key_opt = agent_guard.stream_options().api_key.clone();
self.tree_ui = None;
self.agent_state = AgentState::Processing;
self.status_message = Some("Switching branches...".to_string());
runtime_handle.spawn(async move {
let cx = Cx::for_request();
let from_id_for_event = pending
.old_leaf_id
.clone()
.unwrap_or_else(|| "root".to_string());
let to_id_for_event = pending
.new_leaf_id
.clone()
.unwrap_or_else(|| "root".to_string());
if let Some(manager) = extensions.clone() {
let cancelled = manager
.dispatch_cancellable_event(
ExtensionEventName::SessionBeforeSwitch,
Some(json!({
"fromId": from_id_for_event.clone(),
"toId": to_id_for_event.clone(),
"sessionId": pending.session_id.clone(),
})),
EXTENSION_EVENT_TIMEOUT_MS,
)
.await
.unwrap_or(false);
if cancelled {
let _ = event_tx.try_send(PiMsg::System(
"Session switch cancelled by extension".to_string(),
));
return;
}
}
let summary_skipped =
summary_requested && key_opt.is_none() && !pending.entries_to_summarize.is_empty();
let summary_text = if !summary_requested || pending.entries_to_summarize.is_empty() {
None
} else if let Some(key) = key_opt.as_deref() {
match crate::compaction::summarize_entries(
&pending.entries_to_summarize,
provider,
key,
reserve_tokens,
custom_instructions.as_deref(),
)
.await
{
Ok(summary) => summary,
Err(err) => {
let _ = event_tx
.try_send(PiMsg::AgentError(format!("Branch summary failed: {err}")));
return;
}
}
} else {
None
};
let messages_for_agent = {
let mut guard = match session.lock(&cx).await {
Ok(guard) => guard,
Err(err) => {
let _ = event_tx
.try_send(PiMsg::AgentError(format!("Failed to lock session: {err}")));
return;
}
};
if let Some(target_id) = &pending.new_leaf_id {
if !guard.navigate_to(target_id) {
let _ = event_tx.try_send(PiMsg::AgentError(format!(
"Branch target not found: {target_id}"
)));
return;
}
} else {
guard.reset_leaf();
}
if let Some(summary_text) = summary_text {
guard.append_branch_summary(
pending.summary_from_id.clone(),
summary_text,
None,
None,
);
}
let _ = guard.save().await;
guard.to_messages_for_current_path()
};
{
let mut agent_guard = match agent.lock(&cx).await {
Ok(guard) => guard,
Err(err) => {
let _ = event_tx
.try_send(PiMsg::AgentError(format!("Failed to lock agent: {err}")));
return;
}
};
agent_guard.replace_messages(messages_for_agent);
}
let (messages, usage) = {
let guard = match session.lock(&cx).await {
Ok(guard) => guard,
Err(err) => {
let _ = event_tx
.try_send(PiMsg::AgentError(format!("Failed to lock session: {err}")));
return;
}
};
conversation_from_session(&guard)
};
let status = if summary_skipped {
Some(format!(
"Switched to {to_id_for_event} (no summary: missing API key)"
))
} else {
Some(format!("Switched to {to_id_for_event}"))
};
let _ = event_tx.try_send(PiMsg::ConversationReset {
messages,
usage,
status,
});
if let Some(text) = pending.editor_text {
let _ = event_tx.try_send(PiMsg::SetEditorText(text));
}
if let Some(manager) = extensions {
let _ = manager
.dispatch_event(
ExtensionEventName::SessionSwitch,
Some(json!({
"fromId": from_id_for_event,
"toId": to_id_for_event,
"sessionId": pending.session_id,
})),
)
.await;
}
});
}
}