From aa5c1ff84e9f7e8920dee5750c1f1e2b24d29cf3 Mon Sep 17 00:00:00 2001 From: Mikayla Maki Date: Mon, 9 Mar 2026 10:21:35 -0700 Subject: [PATCH] Optimize update_entries (#51122) Before you mark this PR as ready for review, make sure that you have: - [x] Added a solid test coverage and/or screenshots from doing manual testing - [x] Done a self-review taking into account security and performance aspects - [x] Aligned any UI changes with the [UI checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) Release Notes: - N/A --- crates/agent/src/thread_store.rs | 32 ++- crates/sidebar/src/sidebar.rs | 370 ++++++++++++++----------------- 2 files changed, 189 insertions(+), 213 deletions(-) diff --git a/crates/agent/src/thread_store.rs b/crates/agent/src/thread_store.rs index 961be1da4c09890691adbd5448d7678b2808fe7b..dd1f650de2f59a0e681e15e7eae3fad1a49ccc41 100644 --- a/crates/agent/src/thread_store.rs +++ b/crates/agent/src/thread_store.rs @@ -2,6 +2,7 @@ use crate::{DbThread, DbThreadMetadata, ThreadsDatabase}; use agent_client_protocol as acp; use anyhow::{Result, anyhow}; use gpui::{App, Context, Entity, Global, Task, prelude::*}; +use std::collections::HashMap; use util::path_list::PathList; struct GlobalThreadStore(Entity); @@ -10,6 +11,7 @@ impl Global for GlobalThreadStore {} pub struct ThreadStore { threads: Vec, + threads_by_paths: HashMap>, } impl ThreadStore { @@ -29,6 +31,7 @@ impl ThreadStore { pub fn new(cx: &mut Context) -> Self { let this = Self { threads: Vec::new(), + threads_by_paths: HashMap::default(), }; this.reload(cx); this @@ -91,14 +94,21 @@ impl ThreadStore { let database_connection = ThreadsDatabase::connect(cx); cx.spawn(async move |this, cx| { let database = database_connection.await.map_err(|err| anyhow!(err))?; - let threads = database - .list_threads() - .await? - .into_iter() - .filter(|thread| thread.parent_session_id.is_none()) - .collect::>(); + let all_threads = database.list_threads().await?; this.update(cx, |this, cx| { - this.threads = threads; + this.threads.clear(); + this.threads_by_paths.clear(); + for thread in all_threads { + if thread.parent_session_id.is_some() { + continue; + } + let index = this.threads.len(); + this.threads_by_paths + .entry(thread.folder_paths.clone()) + .or_default() + .push(index); + this.threads.push(thread); + } cx.notify(); }) }) @@ -114,10 +124,12 @@ impl ThreadStore { } /// Returns threads whose folder_paths match the given paths exactly. + /// Uses a cached index for O(1) lookup per path list. pub fn threads_for_paths(&self, paths: &PathList) -> impl Iterator { - self.threads - .iter() - .filter(move |thread| &thread.folder_paths == paths) + self.threads_by_paths + .get(paths) + .into_iter() + .flat_map(|indices| indices.iter().map(|&index| &self.threads[index])) } } diff --git a/crates/sidebar/src/sidebar.rs b/crates/sidebar/src/sidebar.rs index 4b56efb81a90f30ab75cb567ab07e28deef424a2..1e50a75e2841fb471b2d630b71c2df59200c5bea 100644 --- a/crates/sidebar/src/sidebar.rs +++ b/crates/sidebar/src/sidebar.rs @@ -65,8 +65,19 @@ impl From<&ActiveThreadInfo> for acp_thread::AgentSessionInfo { } } -#[derive(Clone, Debug)] -#[allow(dead_code)] +#[derive(Clone)] +struct ThreadEntry { + session_info: acp_thread::AgentSessionInfo, + icon: IconName, + icon_from_external_svg: Option, + status: AgentThreadStatus, + workspace: Entity, + is_live: bool, + is_background: bool, + highlight_positions: Vec, +} + +#[derive(Clone)] enum ListEntry { ProjectHeader { path_list: PathList, @@ -75,17 +86,7 @@ enum ListEntry { highlight_positions: Vec, has_threads: bool, }, - Thread { - session_info: acp_thread::AgentSessionInfo, - icon: IconName, - icon_from_external_svg: Option, - status: AgentThreadStatus, - diff_stats: Option<(usize, usize)>, - workspace: Entity, - is_live: bool, - is_background: bool, - highlight_positions: Vec, - }, + Thread(ThreadEntry), ViewMore { path_list: PathList, remaining_count: usize, @@ -97,6 +98,12 @@ enum ListEntry { }, } +impl From for ListEntry { + fn from(thread: ThreadEntry) -> Self { + ListEntry::Thread(thread) + } +} + #[derive(Default)] struct SidebarContents { entries: Vec, @@ -227,7 +234,7 @@ impl Sidebar { .contents .entries .iter() - .position(|entry| matches!(entry, ListEntry::Thread { .. })) + .position(|entry| matches!(entry, ListEntry::Thread(_))) .or_else(|| { if this.contents.entries.is_empty() { None @@ -416,18 +423,20 @@ impl Sidebar { .entries .iter() .filter_map(|entry| match entry { - ListEntry::Thread { - session_info, - status, - is_live: true, - .. - } => Some((session_info.session_id.clone(), *status)), + ListEntry::Thread(thread) if thread.is_live => { + Some((thread.session_info.session_id.clone(), thread.status)) + } _ => None, }) .collect(); let mut entries = Vec::new(); let mut notified_threads = previous.notified_threads; + // Track all session IDs we add to entries so we can prune stale + // notifications without a separate pass at the end. + let mut current_session_ids: HashSet = HashSet::new(); + // Compute active_entry_index inline during the build pass. + let mut active_entry_index: Option = None; for workspace in workspaces.iter() { let (path_list, label) = workspace_path_list_and_label(workspace, cx); @@ -435,17 +444,16 @@ impl Sidebar { let is_collapsed = self.collapsed_groups.contains(&path_list); let should_load_threads = !is_collapsed || !query.is_empty(); - let mut threads: Vec = Vec::new(); + let mut threads: Vec = Vec::new(); if should_load_threads { if let Some(ref thread_store) = thread_store { for meta in thread_store.read(cx).threads_for_paths(&path_list) { - threads.push(ListEntry::Thread { + threads.push(ThreadEntry { session_info: meta.into(), icon: IconName::ZedAgent, icon_from_external_svg: None, status: AgentThreadStatus::default(), - diff_stats: None, workspace: workspace.clone(), is_live: false, is_background: false, @@ -456,76 +464,50 @@ impl Sidebar { let live_infos = Self::all_thread_infos_for_workspace(workspace, cx); - for info in &live_infos { - let Some(existing) = threads.iter_mut().find(|t| { - matches!(t, ListEntry::Thread { session_info, .. } if session_info.session_id == info.session_id) - }) else { - continue; - }; - - if let ListEntry::Thread { - session_info, - status, - icon, - icon_from_external_svg, - workspace: _, - is_live, - is_background, - .. - } = existing - { - session_info.title = Some(info.title.clone()); - *status = info.status; - *icon = info.icon; - *icon_from_external_svg = info.icon_from_external_svg.clone(); - *is_live = true; - *is_background = info.is_background; + if !live_infos.is_empty() { + let thread_index_by_session: HashMap = threads + .iter() + .enumerate() + .map(|(i, t)| (t.session_info.session_id.clone(), i)) + .collect(); + + for info in &live_infos { + let Some(&idx) = thread_index_by_session.get(&info.session_id) else { + continue; + }; + + let thread = &mut threads[idx]; + thread.session_info.title = Some(info.title.clone()); + thread.status = info.status; + thread.icon = info.icon; + thread.icon_from_external_svg = info.icon_from_external_svg.clone(); + thread.is_live = true; + thread.is_background = info.is_background; } } - // Update notification state for live threads. + // Update notification state for live threads in the same pass. + let is_active_workspace = active_workspace + .as_ref() + .is_some_and(|active| active == workspace); + for thread in &threads { - if let ListEntry::Thread { - workspace: thread_workspace, - session_info, - status, - is_background, - .. - } = thread + let session_id = &thread.session_info.session_id; + if thread.is_background && thread.status == AgentThreadStatus::Completed { + notified_threads.insert(session_id.clone()); + } else if thread.status == AgentThreadStatus::Completed + && !is_active_workspace + && old_statuses.get(session_id) == Some(&AgentThreadStatus::Running) { - let session_id = &session_info.session_id; - if *is_background && *status == AgentThreadStatus::Completed { - notified_threads.insert(session_id.clone()); - } else if *status == AgentThreadStatus::Completed - && active_workspace - .as_ref() - .is_none_or(|active| active != thread_workspace) - && old_statuses.get(session_id) == Some(&AgentThreadStatus::Running) - { - notified_threads.insert(session_id.clone()); - } + notified_threads.insert(session_id.clone()); + } - if active_workspace - .as_ref() - .is_some_and(|active| active == thread_workspace) - && !*is_background - { - notified_threads.remove(session_id); - } + if is_active_workspace && !thread.is_background { + notified_threads.remove(session_id); } } - threads.sort_by(|a, b| { - let a_time = match a { - ListEntry::Thread { session_info, .. } => session_info.updated_at, - _ => unreachable!(), - }; - let b_time = match b { - ListEntry::Thread { session_info, .. } => session_info.updated_at, - _ => unreachable!(), - }; - b_time.cmp(&a_time) - }); + threads.sort_by(|a, b| b.session_info.updated_at.cmp(&a.session_info.updated_at)); } if !query.is_empty() { @@ -535,25 +517,19 @@ impl Sidebar { fuzzy_match_positions(&query, &label).unwrap_or_default(); let workspace_matched = !workspace_highlight_positions.is_empty(); - let mut matched_threads = Vec::new(); + let mut matched_threads: Vec = Vec::new(); for mut thread in threads { - if let ListEntry::Thread { - session_info, - highlight_positions, - .. - } = &mut thread - { - let title = session_info - .title - .as_ref() - .map(|s| s.as_ref()) - .unwrap_or(""); - if let Some(positions) = fuzzy_match_positions(&query, title) { - *highlight_positions = positions; - } - if workspace_matched || !highlight_positions.is_empty() { - matched_threads.push(thread); - } + let title = thread + .session_info + .title + .as_ref() + .map(|s| s.as_ref()) + .unwrap_or(""); + if let Some(positions) = fuzzy_match_positions(&query, title) { + thread.highlight_positions = positions; + } + if workspace_matched || !thread.highlight_positions.is_empty() { + matched_threads.push(thread); } } @@ -561,6 +537,15 @@ impl Sidebar { continue; } + if active_entry_index.is_none() + && self.focused_thread.is_none() + && active_workspace + .as_ref() + .is_some_and(|active| active == workspace) + { + active_entry_index = Some(entries.len()); + } + entries.push(ListEntry::ProjectHeader { path_list: path_list.clone(), label, @@ -568,9 +553,33 @@ impl Sidebar { highlight_positions: workspace_highlight_positions, has_threads, }); - entries.extend(matched_threads); + + // Track session IDs and compute active_entry_index as we add + // thread entries. + for thread in matched_threads { + current_session_ids.insert(thread.session_info.session_id.clone()); + if active_entry_index.is_none() { + if let Some(focused) = &self.focused_thread { + if &thread.session_info.session_id == focused { + active_entry_index = Some(entries.len()); + } + } + } + entries.push(thread.into()); + } } else { let has_threads = !threads.is_empty(); + + // Check if this header is the active entry before pushing it. + if active_entry_index.is_none() + && self.focused_thread.is_none() + && active_workspace + .as_ref() + .is_some_and(|active| active == workspace) + { + active_entry_index = Some(entries.len()); + } + entries.push(ListEntry::ProjectHeader { path_list: path_list.clone(), label, @@ -591,7 +600,19 @@ impl Sidebar { let count = threads_to_show.min(total); let is_fully_expanded = count >= total; - entries.extend(threads.into_iter().take(count)); + // Track session IDs and compute active_entry_index as we add + // thread entries. + for thread in threads.into_iter().take(count) { + current_session_ids.insert(thread.session_info.session_id.clone()); + if active_entry_index.is_none() { + if let Some(focused) = &self.focused_thread { + if &thread.session_info.session_id == focused { + active_entry_index = Some(entries.len()); + } + } + } + entries.push(thread.into()); + } if total > DEFAULT_THREADS_SHOWN { entries.push(ListEntry::ViewMore { @@ -610,16 +631,11 @@ impl Sidebar { } } - // Prune stale entries from notified_threads. - let current_session_ids: HashSet<&acp::SessionId> = entries - .iter() - .filter_map(|e| match e { - ListEntry::Thread { session_info, .. } => Some(&session_info.session_id), - _ => None, - }) - .collect(); + // Prune stale notifications using the session IDs we collected during + // the build pass (no extra scan needed). notified_threads.retain(|id| current_session_ids.contains(id)); + self.active_entry_index = active_entry_index; self.contents = SidebarContents { entries, notified_threads, @@ -639,7 +655,6 @@ impl Sidebar { let scroll_position = self.list_state.logical_scroll_top(); self.rebuild_contents(cx); - self.recompute_active_entry_index(cx); self.list_state.reset(self.contents.entries.len()); self.list_state.scroll_to(scroll_position); @@ -653,24 +668,6 @@ impl Sidebar { cx.notify(); } - fn recompute_active_entry_index(&mut self, cx: &App) { - self.active_entry_index = if let Some(session_id) = &self.focused_thread { - self.contents.entries.iter().position(|entry| { - matches!(entry, ListEntry::Thread { session_info, .. } if &session_info.session_id == session_id) - }) - } else { - let active_workspace = self - .multi_workspace - .upgrade() - .map(|mw| mw.read(cx).workspace().clone()); - active_workspace.and_then(|active| { - self.contents.entries.iter().position(|entry| { - matches!(entry, ListEntry::ProjectHeader { workspace, .. } if workspace == &active) - }) - }) - }; - } - fn render_list_entry( &mut self, ix: usize, @@ -705,25 +702,7 @@ impl Sidebar { is_selected, cx, ), - ListEntry::Thread { - session_info, - icon, - icon_from_external_svg, - status, - workspace, - highlight_positions, - .. - } => self.render_thread( - ix, - session_info, - *icon, - icon_from_external_svg.clone(), - *status, - workspace, - highlight_positions, - is_selected, - cx, - ), + ListEntry::Thread(thread) => self.render_thread(ix, thread, is_selected, cx), ListEntry::ViewMore { path_list, remaining_count, @@ -975,8 +954,8 @@ impl Sidebar { }) } - fn filter_query(&self, cx: &App) -> String { - self.filter_editor.read(cx).text(cx) + fn has_filter_query(&self, cx: &App) -> bool { + self.filter_editor.read(cx).buffer().read(cx).is_empty() } fn editor_move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context) { @@ -1041,13 +1020,9 @@ impl Sidebar { let workspace = workspace.clone(); self.activate_workspace(&workspace, window, cx); } - ListEntry::Thread { - session_info, - workspace, - .. - } => { - let session_info = session_info.clone(); - let workspace = workspace.clone(); + ListEntry::Thread(thread) => { + let session_info = thread.session_info.clone(); + let workspace = thread.workspace.clone(); self.activate_thread(session_info, &workspace, window, cx); } ListEntry::ViewMore { @@ -1144,7 +1119,7 @@ impl Sidebar { } } Some( - ListEntry::Thread { .. } | ListEntry::ViewMore { .. } | ListEntry::NewThread { .. }, + ListEntry::Thread(_) | ListEntry::ViewMore { .. } | ListEntry::NewThread { .. }, ) => { for i in (0..ix).rev() { if let Some(ListEntry::ProjectHeader { path_list, .. }) = @@ -1165,32 +1140,30 @@ impl Sidebar { fn render_thread( &self, ix: usize, - session_info: &acp_thread::AgentSessionInfo, - icon: IconName, - icon_from_external_svg: Option, - status: AgentThreadStatus, - workspace: &Entity, - highlight_positions: &[usize], + thread: &ThreadEntry, is_selected: bool, cx: &mut Context, ) -> AnyElement { - let has_notification = self.contents.is_thread_notified(&session_info.session_id); + let has_notification = self + .contents + .is_thread_notified(&thread.session_info.session_id); - let title: SharedString = session_info + let title: SharedString = thread + .session_info .title .clone() .unwrap_or_else(|| "Untitled".into()); - let session_info = session_info.clone(); - let workspace = workspace.clone(); + let session_info = thread.session_info.clone(); + let workspace = thread.workspace.clone(); let id = SharedString::from(format!("thread-entry-{}", ix)); ThreadItem::new(id, title) - .icon(icon) - .when_some(icon_from_external_svg, |this, svg| { + .icon(thread.icon) + .when_some(thread.icon_from_external_svg.clone(), |this, svg| { this.custom_icon_from_external_svg(svg) }) - .highlight_positions(highlight_positions.to_vec()) - .status(status) + .highlight_positions(thread.highlight_positions.to_vec()) + .status(thread.status) .notified(has_notification) .selected(self.focused_thread.as_ref() == Some(&session_info.session_id)) .focused(is_selected) @@ -1356,7 +1329,7 @@ impl Render for Sidebar { let ui_font = theme::setup_ui_font(window, cx); let is_focused = self.focus_handle.is_focused(window) || self.filter_editor.focus_handle(cx).is_focused(window); - let has_query = !self.filter_query(cx).is_empty(); + let has_query = self.has_filter_query(cx); let focus_tooltip_label = if is_focused { "Focus Workspace" @@ -1666,19 +1639,15 @@ mod tests { }; format!("{} [{}]{}", icon, label, selected) } - ListEntry::Thread { - session_info, - status, - is_live, - .. - } => { - let title = session_info + ListEntry::Thread(thread) => { + let title = thread + .session_info .title .as_ref() .map(|s| s.as_ref()) .unwrap_or("Untitled"); - let active = if *is_live { " *" } else { "" }; - let status_str = match status { + let active = if thread.is_live { " *" } else { "" }; + let status_str = match thread.status { AgentThreadStatus::Running => " (running)", AgentThreadStatus::Error => " (error)", AgentThreadStatus::WaitingForConfirmation => " (waiting)", @@ -1686,7 +1655,7 @@ mod tests { }; let notified = if sidebar .contents - .is_thread_notified(&session_info.session_id) + .is_thread_notified(&thread.session_info.session_id) { " (!)" } else { @@ -2007,7 +1976,7 @@ mod tests { has_threads: true, }, // Thread with default (Completed) status, not active - ListEntry::Thread { + ListEntry::Thread(ThreadEntry { session_info: acp_thread::AgentSessionInfo { session_id: acp::SessionId::new(Arc::from("t-1")), cwd: None, @@ -2018,14 +1987,13 @@ mod tests { icon: IconName::ZedAgent, icon_from_external_svg: None, status: AgentThreadStatus::Completed, - diff_stats: None, workspace: workspace.clone(), is_live: false, is_background: false, highlight_positions: Vec::new(), - }, + }), // Active thread with Running status - ListEntry::Thread { + ListEntry::Thread(ThreadEntry { session_info: acp_thread::AgentSessionInfo { session_id: acp::SessionId::new(Arc::from("t-2")), cwd: None, @@ -2036,14 +2004,13 @@ mod tests { icon: IconName::ZedAgent, icon_from_external_svg: None, status: AgentThreadStatus::Running, - diff_stats: None, workspace: workspace.clone(), is_live: true, is_background: false, highlight_positions: Vec::new(), - }, + }), // Active thread with Error status - ListEntry::Thread { + ListEntry::Thread(ThreadEntry { session_info: acp_thread::AgentSessionInfo { session_id: acp::SessionId::new(Arc::from("t-3")), cwd: None, @@ -2054,14 +2021,13 @@ mod tests { icon: IconName::ZedAgent, icon_from_external_svg: None, status: AgentThreadStatus::Error, - diff_stats: None, workspace: workspace.clone(), is_live: true, is_background: false, highlight_positions: Vec::new(), - }, + }), // Thread with WaitingForConfirmation status, not active - ListEntry::Thread { + ListEntry::Thread(ThreadEntry { session_info: acp_thread::AgentSessionInfo { session_id: acp::SessionId::new(Arc::from("t-4")), cwd: None, @@ -2072,14 +2038,13 @@ mod tests { icon: IconName::ZedAgent, icon_from_external_svg: None, status: AgentThreadStatus::WaitingForConfirmation, - diff_stats: None, workspace: workspace.clone(), is_live: false, is_background: false, highlight_positions: Vec::new(), - }, + }), // Background thread that completed (should show notification) - ListEntry::Thread { + ListEntry::Thread(ThreadEntry { session_info: acp_thread::AgentSessionInfo { session_id: acp::SessionId::new(Arc::from("t-5")), cwd: None, @@ -2090,12 +2055,11 @@ mod tests { icon: IconName::ZedAgent, icon_from_external_svg: None, status: AgentThreadStatus::Completed, - diff_stats: None, workspace: workspace.clone(), is_live: true, is_background: true, highlight_positions: Vec::new(), - }, + }), // View More entry ListEntry::ViewMore { path_list: expanded_path.clone(), @@ -3475,7 +3439,7 @@ mod tests { let active_entry = sidebar.active_entry_index .and_then(|ix| sidebar.contents.entries.get(ix)); assert!( - matches!(active_entry, Some(ListEntry::Thread { session_info, .. }) if session_info.session_id == session_id_a), + matches!(active_entry, Some(ListEntry::Thread(thread)) if thread.session_info.session_id == session_id_a), "Active entry should be the clicked thread" ); }); @@ -3531,7 +3495,7 @@ mod tests { .active_entry_index .and_then(|ix| sidebar.contents.entries.get(ix)); assert!( - matches!(active_entry, Some(ListEntry::Thread { session_info, .. }) if session_info.session_id == session_id_b), + matches!(active_entry, Some(ListEntry::Thread(thread)) if thread.session_info.session_id == session_id_b), "Active entry should be the cross-workspace thread" ); }); @@ -3626,7 +3590,7 @@ mod tests { .active_entry_index .and_then(|ix| sidebar.contents.entries.get(ix)); assert!( - matches!(active_entry, Some(ListEntry::Thread { session_info, .. }) if session_info.session_id == session_id_b2), + matches!(active_entry, Some(ListEntry::Thread(thread)) if thread.session_info.session_id == session_id_b2), "Active entry should be the focused thread" ); });