use std::{ path::{Path, PathBuf}, sync::Arc, }; use acp_thread::AcpThreadEvent; use agent::{ThreadStore, ZED_AGENT_ID}; use agent_client_protocol as acp; use anyhow::Context as _; use chrono::{DateTime, Utc}; use collections::{HashMap, HashSet}; use db::{ kvp::KeyValueStore, sqlez::{ bindable::Column, domain::Domain, statement::Statement, thread_safe_connection::ThreadSafeConnection, }, sqlez_macros::sql, }; use fs::Fs; use futures::{FutureExt, future::Shared}; use gpui::{AppContext as _, Entity, Global, Subscription, Task}; use project::AgentId; use remote::RemoteConnectionOptions; use ui::{App, Context, SharedString}; use util::ResultExt as _; use workspace::{PathList, SerializedWorkspaceLocation, WorkspaceDb}; use crate::DEFAULT_THREAD_TITLE; const THREAD_REMOTE_CONNECTION_MIGRATION_KEY: &str = "thread-metadata-remote-connection-backfill"; pub fn init(cx: &mut App) { ThreadMetadataStore::init_global(cx); let migration_task = migrate_thread_metadata(cx); migrate_thread_remote_connections(cx, migration_task); } /// Migrate existing thread metadata from native agent thread store to the new metadata storage. /// We skip migrating threads that do not have a project. /// /// TODO: Remove this after N weeks of shipping the sidebar fn migrate_thread_metadata(cx: &mut App) -> Task> { let store = ThreadMetadataStore::global(cx); let db = store.read(cx).db.clone(); cx.spawn(async move |cx| { let existing_entries = db.list_ids()?.into_iter().collect::>(); let is_first_migration = existing_entries.is_empty(); let mut to_migrate = store.read_with(cx, |_store, cx| { ThreadStore::global(cx) .read(cx) .entries() .filter_map(|entry| { if existing_entries.contains(&entry.id.0) { return None; } Some(ThreadMetadata { session_id: entry.id, agent_id: ZED_AGENT_ID.clone(), title: entry.title, updated_at: entry.updated_at, created_at: entry.created_at, worktree_paths: ThreadWorktreePaths::from_folder_paths(&entry.folder_paths), remote_connection: None, archived: true, }) }) .collect::>() }); if to_migrate.is_empty() { return anyhow::Ok(()); } // On the first migration (no entries in DB yet), keep the 5 most // recent threads per project unarchived. if is_first_migration { let mut per_project: HashMap> = HashMap::default(); for entry in &mut to_migrate { if entry.worktree_paths.is_empty() { continue; } per_project .entry(entry.worktree_paths.folder_path_list().clone()) .or_default() .push(entry); } for entries in per_project.values_mut() { entries.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); for entry in entries.iter_mut().take(5) { entry.archived = false; } } } log::info!("Migrating {} thread store entries", to_migrate.len()); // Manually save each entry to the database and call reload, otherwise // we'll end up triggering lots of reloads after each save for entry in to_migrate { db.save(entry).await?; } log::info!("Finished migrating thread store entries"); let _ = store.update(cx, |store, cx| store.reload(cx)); anyhow::Ok(()) }) } fn migrate_thread_remote_connections(cx: &mut App, migration_task: Task>) { let store = ThreadMetadataStore::global(cx); let db = store.read(cx).db.clone(); let kvp = KeyValueStore::global(cx); let workspace_db = WorkspaceDb::global(cx); let fs = ::global(cx); cx.spawn(async move |cx| -> anyhow::Result<()> { migration_task.await?; if kvp .read_kvp(THREAD_REMOTE_CONNECTION_MIGRATION_KEY)? .is_some() { return Ok(()); } let recent_workspaces = workspace_db.recent_workspaces_on_disk(fs.as_ref()).await?; let mut local_path_lists = HashSet::::default(); let mut remote_path_lists = HashMap::::default(); recent_workspaces .iter() .filter(|(_, location, path_list, _)| { !path_list.is_empty() && matches!(location, &SerializedWorkspaceLocation::Local) }) .for_each(|(_, _, path_list, _)| { local_path_lists.insert(path_list.clone()); }); for (_, location, path_list, _) in recent_workspaces { match location { SerializedWorkspaceLocation::Remote(remote_connection) if !local_path_lists.contains(&path_list) => { remote_path_lists .entry(path_list) .or_insert(remote_connection); } _ => {} } } let mut reloaded = false; for metadata in db.list()? { if metadata.remote_connection.is_some() { continue; } if let Some(remote_connection) = remote_path_lists .get(metadata.folder_paths()) .or_else(|| remote_path_lists.get(metadata.main_worktree_paths())) { db.save(ThreadMetadata { remote_connection: Some(remote_connection.clone()), ..metadata }) .await?; reloaded = true; } } let reloaded_task = reloaded .then_some(store.update(cx, |store, cx| store.reload(cx))) .unwrap_or(Task::ready(()).shared()); kvp.write_kvp( THREAD_REMOTE_CONNECTION_MIGRATION_KEY.to_string(), "1".to_string(), ) .await?; reloaded_task.await; Ok(()) }) .detach_and_log_err(cx); } struct GlobalThreadMetadataStore(Entity); impl Global for GlobalThreadMetadataStore {} /// Paired worktree paths for a thread. Each folder path has a corresponding /// main worktree path at the same position. The two lists are always the /// same length and are modified together via `add_path` / `remove_main_path`. /// /// For non-linked worktrees, the main path and folder path are identical. /// For linked worktrees, the main path is the original repo and the folder /// path is the linked worktree location. /// /// Internally stores two `PathList`s with matching insertion order so that /// `ordered_paths()` on both yields positionally-paired results. #[derive(Default, Debug, Clone)] pub struct ThreadWorktreePaths { folder_paths: PathList, main_worktree_paths: PathList, } impl PartialEq for ThreadWorktreePaths { fn eq(&self, other: &Self) -> bool { self.folder_paths == other.folder_paths && self.main_worktree_paths == other.main_worktree_paths } } impl ThreadWorktreePaths { /// Build from a project's current state. Each visible worktree is paired /// with its main repo path (resolved via git), falling back to the /// worktree's own path if no git repo is found. pub fn from_project(project: &project::Project, cx: &App) -> Self { let (mains, folders): (Vec, Vec) = project .visible_worktrees(cx) .map(|worktree| { let snapshot = worktree.read(cx).snapshot(); let folder_path = snapshot.abs_path().to_path_buf(); let main_path = snapshot .root_repo_common_dir() .and_then(|dir| Some(dir.parent()?.to_path_buf())) .unwrap_or_else(|| folder_path.clone()); (main_path, folder_path) }) .unzip(); Self { folder_paths: PathList::new(&folders), main_worktree_paths: PathList::new(&mains), } } /// Build from two parallel `PathList`s that already share the same /// insertion order. Used for deserialization from DB. /// /// Returns an error if the two lists have different lengths, which /// indicates corrupted data from a prior migration bug. pub fn from_path_lists( main_worktree_paths: PathList, folder_paths: PathList, ) -> anyhow::Result { anyhow::ensure!( main_worktree_paths.paths().len() == folder_paths.paths().len(), "main_worktree_paths has {} entries but folder_paths has {}", main_worktree_paths.paths().len(), folder_paths.paths().len(), ); Ok(Self { folder_paths, main_worktree_paths, }) } /// Build for non-linked worktrees where main == folder for every path. pub fn from_folder_paths(folder_paths: &PathList) -> Self { Self { folder_paths: folder_paths.clone(), main_worktree_paths: folder_paths.clone(), } } pub fn is_empty(&self) -> bool { self.folder_paths.is_empty() } /// The folder paths (for workspace matching / `threads_by_paths` index). pub fn folder_path_list(&self) -> &PathList { &self.folder_paths } /// The main worktree paths (for group key / `threads_by_main_paths` index). pub fn main_worktree_path_list(&self) -> &PathList { &self.main_worktree_paths } /// Iterate the (main_worktree_path, folder_path) pairs in insertion order. pub fn ordered_pairs(&self) -> impl Iterator { self.main_worktree_paths .ordered_paths() .zip(self.folder_paths.ordered_paths()) } /// Add a new path pair. If the exact (main, folder) pair already exists, /// this is a no-op. Rebuilds both internal `PathList`s to maintain /// consistent ordering. pub fn add_path(&mut self, main_path: &Path, folder_path: &Path) { let already_exists = self .ordered_pairs() .any(|(m, f)| m.as_path() == main_path && f.as_path() == folder_path); if already_exists { return; } let (mut mains, mut folders): (Vec, Vec) = self .ordered_pairs() .map(|(m, f)| (m.clone(), f.clone())) .unzip(); mains.push(main_path.to_path_buf()); folders.push(folder_path.to_path_buf()); self.main_worktree_paths = PathList::new(&mains); self.folder_paths = PathList::new(&folders); } /// Remove all pairs whose main worktree path matches the given path. /// This removes the corresponding entries from both lists. pub fn remove_main_path(&mut self, main_path: &Path) { let (mains, folders): (Vec, Vec) = self .ordered_pairs() .filter(|(m, _)| m.as_path() != main_path) .map(|(m, f)| (m.clone(), f.clone())) .unzip(); self.main_worktree_paths = PathList::new(&mains); self.folder_paths = PathList::new(&folders); } } /// Lightweight metadata for any thread (native or ACP), enough to populate /// the sidebar list and route to the correct load path when clicked. #[derive(Debug, Clone, PartialEq)] pub struct ThreadMetadata { pub session_id: acp::SessionId, pub agent_id: AgentId, pub title: SharedString, pub updated_at: DateTime, pub created_at: Option>, pub worktree_paths: ThreadWorktreePaths, pub remote_connection: Option, pub archived: bool, } impl ThreadMetadata { pub fn folder_paths(&self) -> &PathList { self.worktree_paths.folder_path_list() } pub fn main_worktree_paths(&self) -> &PathList { self.worktree_paths.main_worktree_path_list() } } impl From<&ThreadMetadata> for acp_thread::AgentSessionInfo { fn from(meta: &ThreadMetadata) -> Self { Self { session_id: meta.session_id.clone(), work_dirs: Some(meta.folder_paths().clone()), title: Some(meta.title.clone()), updated_at: Some(meta.updated_at), created_at: meta.created_at, meta: None, } } } /// Record of a git worktree that was archived (deleted from disk) when its /// last thread was archived. pub struct ArchivedGitWorktree { /// Auto-incrementing primary key. pub id: i64, /// Absolute path to the directory of the worktree before it was deleted. /// Used when restoring, to put the recreated worktree back where it was. /// If the path already exists on disk, the worktree is assumed to be /// already restored and is used as-is. pub worktree_path: PathBuf, /// Absolute path of the main repository ("main worktree") that owned this worktree. /// Used when restoring, to reattach the recreated worktree to the correct main repo. /// If the main repo isn't found on disk, unarchiving fails because we only store /// commit hashes, and without the actual git repo being available, we can't restore /// the files. pub main_repo_path: PathBuf, /// Branch that was checked out in the worktree at archive time. `None` if /// the worktree was in detached HEAD state, which isn't supported in Zed, but /// could happen if the user made a detached one outside of Zed. /// On restore, we try to switch to this branch. If that fails (e.g. it's /// checked out elsewhere), we auto-generate a new one. pub branch_name: Option, /// SHA of the WIP commit that captures files that were staged (but not yet /// committed) at the time of archiving. This commit can be empty if the /// user had no staged files at the time. It sits directly on top of whatever /// the user's last actual commit was. pub staged_commit_hash: String, /// SHA of the WIP commit that captures files that were unstaged (including /// untracked) at the time of archiving. This commit can be empty if the user /// had no unstaged files at the time. It sits on top of `staged_commit_hash`. /// After doing `git reset` past both of these commits, we're back in the state /// we had before archiving, including what was staged, what was unstaged, and /// what was committed. pub unstaged_commit_hash: String, /// SHA of the commit that HEAD pointed at before we created the two WIP /// commits during archival. After resetting past the WIP commits during /// restore, HEAD should land back on this commit. It also serves as a /// pre-restore sanity check (abort if this commit no longer exists in the /// repo) and as a fallback target if the WIP resets fail. pub original_commit_hash: String, } /// The store holds all metadata needed to show threads in the sidebar/the archive. /// /// Automatically listens to AcpThread events and updates metadata if it has changed. pub struct ThreadMetadataStore { db: ThreadMetadataDb, threads: HashMap, threads_by_paths: HashMap>, threads_by_main_paths: HashMap>, reload_task: Option>>, session_subscriptions: HashMap, pending_thread_ops_tx: smol::channel::Sender, in_flight_archives: HashMap, smol::channel::Sender<()>)>, _db_operations_task: Task<()>, } #[derive(Debug, PartialEq)] enum DbOperation { Upsert(ThreadMetadata), Delete(acp::SessionId), } impl DbOperation { fn id(&self) -> &acp::SessionId { match self { DbOperation::Upsert(thread) => &thread.session_id, DbOperation::Delete(session_id) => session_id, } } } impl ThreadMetadataStore { #[cfg(not(any(test, feature = "test-support")))] pub fn init_global(cx: &mut App) { if cx.has_global::() { return; } let db = ThreadMetadataDb::global(cx); let thread_store = cx.new(|cx| Self::new(db, cx)); cx.set_global(GlobalThreadMetadataStore(thread_store)); } #[cfg(any(test, feature = "test-support"))] pub fn init_global(cx: &mut App) { let thread = std::thread::current(); let test_name = thread.name().unwrap_or("unknown_test"); let db_name = format!("THREAD_METADATA_DB_{}", test_name); let db = smol::block_on(db::open_test_db::(&db_name)); let thread_store = cx.new(|cx| Self::new(ThreadMetadataDb(db), cx)); cx.set_global(GlobalThreadMetadataStore(thread_store)); } pub fn try_global(cx: &App) -> Option> { cx.try_global::() .map(|store| store.0.clone()) } pub fn global(cx: &App) -> Entity { cx.global::().0.clone() } pub fn is_empty(&self) -> bool { self.threads.is_empty() } /// Returns all thread IDs. pub fn entry_ids(&self) -> impl Iterator + '_ { self.threads.keys().cloned() } /// Returns the metadata for a specific thread, if it exists. pub fn entry(&self, session_id: &acp::SessionId) -> Option<&ThreadMetadata> { self.threads.get(session_id) } /// Returns all threads. pub fn entries(&self) -> impl Iterator + '_ { self.threads.values() } /// Returns all archived threads. pub fn archived_entries(&self) -> impl Iterator + '_ { self.entries().filter(|t| t.archived) } /// Returns all threads for the given path list, excluding archived threads. pub fn entries_for_path( &self, path_list: &PathList, ) -> impl Iterator + '_ { self.threads_by_paths .get(path_list) .into_iter() .flatten() .filter_map(|s| self.threads.get(s)) .filter(|s| !s.archived) } /// Returns threads whose `main_worktree_paths` matches the given path list, /// excluding archived threads. This finds threads that were opened in a /// linked worktree but are associated with the given main worktree. pub fn entries_for_main_worktree_path( &self, path_list: &PathList, ) -> impl Iterator + '_ { self.threads_by_main_paths .get(path_list) .into_iter() .flatten() .filter_map(|s| self.threads.get(s)) .filter(|s| !s.archived) } fn reload(&mut self, cx: &mut Context) -> Shared> { let db = self.db.clone(); self.reload_task.take(); let list_task = cx .background_spawn(async move { db.list().context("Failed to fetch sidebar metadata") }); let reload_task = cx .spawn(async move |this, cx| { let Some(rows) = list_task.await.log_err() else { return; }; this.update(cx, |this, cx| { this.threads.clear(); this.threads_by_paths.clear(); this.threads_by_main_paths.clear(); for row in rows { this.threads_by_paths .entry(row.folder_paths().clone()) .or_default() .insert(row.session_id.clone()); if !row.main_worktree_paths().is_empty() { this.threads_by_main_paths .entry(row.main_worktree_paths().clone()) .or_default() .insert(row.session_id.clone()); } this.threads.insert(row.session_id.clone(), row); } cx.notify(); }) .ok(); }) .shared(); self.reload_task = Some(reload_task.clone()); reload_task } pub fn save_all(&mut self, metadata: Vec, cx: &mut Context) { for metadata in metadata { self.save_internal(metadata); } cx.notify(); } #[cfg(any(test, feature = "test-support"))] pub fn save_manually(&mut self, metadata: ThreadMetadata, cx: &mut Context) { self.save(metadata, cx) } fn save(&mut self, metadata: ThreadMetadata, cx: &mut Context) { self.save_internal(metadata); cx.notify(); } fn save_internal(&mut self, metadata: ThreadMetadata) { if let Some(thread) = self.threads.get(&metadata.session_id) { if thread.folder_paths() != metadata.folder_paths() { if let Some(session_ids) = self.threads_by_paths.get_mut(thread.folder_paths()) { session_ids.remove(&metadata.session_id); } } if thread.main_worktree_paths() != metadata.main_worktree_paths() && !thread.main_worktree_paths().is_empty() { if let Some(session_ids) = self .threads_by_main_paths .get_mut(thread.main_worktree_paths()) { session_ids.remove(&metadata.session_id); } } } self.threads .insert(metadata.session_id.clone(), metadata.clone()); self.threads_by_paths .entry(metadata.folder_paths().clone()) .or_default() .insert(metadata.session_id.clone()); if !metadata.main_worktree_paths().is_empty() { self.threads_by_main_paths .entry(metadata.main_worktree_paths().clone()) .or_default() .insert(metadata.session_id.clone()); } self.pending_thread_ops_tx .try_send(DbOperation::Upsert(metadata)) .log_err(); } pub fn update_working_directories( &mut self, session_id: &acp::SessionId, work_dirs: PathList, cx: &mut Context, ) { if let Some(thread) = self.threads.get(session_id) { self.save_internal(ThreadMetadata { worktree_paths: ThreadWorktreePaths::from_path_lists( thread.main_worktree_paths().clone(), work_dirs.clone(), ) .unwrap_or_else(|_| ThreadWorktreePaths::from_folder_paths(&work_dirs)), ..thread.clone() }); cx.notify(); } } pub fn archive( &mut self, session_id: &acp::SessionId, archive_job: Option<(Task<()>, smol::channel::Sender<()>)>, cx: &mut Context, ) { self.update_archived(session_id, true, cx); if let Some(job) = archive_job { self.in_flight_archives.insert(session_id.clone(), job); } } pub fn unarchive(&mut self, session_id: &acp::SessionId, cx: &mut Context) { self.update_archived(session_id, false, cx); // Dropping the Sender triggers cancellation in the background task. self.in_flight_archives.remove(session_id); } pub fn cleanup_completed_archive(&mut self, session_id: &acp::SessionId) { self.in_flight_archives.remove(session_id); } /// Updates a thread's `folder_paths` after an archived worktree has been /// restored to disk. The restored worktree may land at a different path /// than it had before archival, so each `(old_path, new_path)` pair in /// `path_replacements` is applied to the thread's stored folder paths. pub fn update_restored_worktree_paths( &mut self, session_id: &acp::SessionId, path_replacements: &[(PathBuf, PathBuf)], cx: &mut Context, ) { if let Some(thread) = self.threads.get(session_id).cloned() { let mut paths: Vec = thread.folder_paths().paths().to_vec(); for (old_path, new_path) in path_replacements { if let Some(pos) = paths.iter().position(|p| p == old_path) { paths[pos] = new_path.clone(); } } let new_folder_paths = PathList::new(&paths); self.save_internal(ThreadMetadata { worktree_paths: ThreadWorktreePaths::from_path_lists( thread.main_worktree_paths().clone(), new_folder_paths.clone(), ) .unwrap_or_else(|_| ThreadWorktreePaths::from_folder_paths(&new_folder_paths)), ..thread }); cx.notify(); } } pub fn complete_worktree_restore( &mut self, session_id: &acp::SessionId, path_replacements: &[(PathBuf, PathBuf)], cx: &mut Context, ) { if let Some(thread) = self.threads.get(session_id).cloned() { let mut paths: Vec = thread.folder_paths().paths().to_vec(); for (old_path, new_path) in path_replacements { for path in &mut paths { if path == old_path { *path = new_path.clone(); } } } let new_folder_paths = PathList::new(&paths); self.save_internal(ThreadMetadata { worktree_paths: ThreadWorktreePaths::from_path_lists( thread.main_worktree_paths().clone(), new_folder_paths.clone(), ) .unwrap_or_else(|_| ThreadWorktreePaths::from_folder_paths(&new_folder_paths)), ..thread }); cx.notify(); } } /// Apply a mutation to the worktree paths of all threads whose current /// `main_worktree_paths` matches `current_main_paths`, then re-index. pub fn change_worktree_paths( &mut self, current_main_paths: &PathList, mutate: impl Fn(&mut ThreadWorktreePaths), cx: &mut Context, ) { let session_ids: Vec<_> = self .threads_by_main_paths .get(current_main_paths) .into_iter() .flatten() .cloned() .collect(); if session_ids.is_empty() { return; } for session_id in &session_ids { if let Some(thread) = self.threads.get_mut(session_id) { if let Some(ids) = self .threads_by_main_paths .get_mut(thread.main_worktree_paths()) { ids.remove(session_id); } if let Some(ids) = self.threads_by_paths.get_mut(thread.folder_paths()) { ids.remove(session_id); } mutate(&mut thread.worktree_paths); self.threads_by_main_paths .entry(thread.main_worktree_paths().clone()) .or_default() .insert(session_id.clone()); self.threads_by_paths .entry(thread.folder_paths().clone()) .or_default() .insert(session_id.clone()); self.pending_thread_ops_tx .try_send(DbOperation::Upsert(thread.clone())) .log_err(); } } cx.notify(); } pub fn create_archived_worktree( &self, worktree_path: String, main_repo_path: String, branch_name: Option, staged_commit_hash: String, unstaged_commit_hash: String, original_commit_hash: String, cx: &App, ) -> Task> { let db = self.db.clone(); cx.background_spawn(async move { db.create_archived_worktree( worktree_path, main_repo_path, branch_name, staged_commit_hash, unstaged_commit_hash, original_commit_hash, ) .await }) } pub fn link_thread_to_archived_worktree( &self, session_id: String, archived_worktree_id: i64, cx: &App, ) -> Task> { let db = self.db.clone(); cx.background_spawn(async move { db.link_thread_to_archived_worktree(session_id, archived_worktree_id) .await }) } pub fn get_archived_worktrees_for_thread( &self, session_id: String, cx: &App, ) -> Task>> { let db = self.db.clone(); cx.background_spawn(async move { db.get_archived_worktrees_for_thread(session_id).await }) } pub fn delete_archived_worktree(&self, id: i64, cx: &App) -> Task> { let db = self.db.clone(); cx.background_spawn(async move { db.delete_archived_worktree(id).await }) } pub fn unlink_thread_from_all_archived_worktrees( &self, session_id: String, cx: &App, ) -> Task> { let db = self.db.clone(); cx.background_spawn(async move { db.unlink_thread_from_all_archived_worktrees(session_id) .await }) } pub fn is_archived_worktree_referenced( &self, archived_worktree_id: i64, cx: &App, ) -> Task> { let db = self.db.clone(); cx.background_spawn(async move { db.is_archived_worktree_referenced(archived_worktree_id) .await }) } fn update_archived( &mut self, session_id: &acp::SessionId, archived: bool, cx: &mut Context, ) { if let Some(thread) = self.threads.get(session_id) { self.save_internal(ThreadMetadata { archived, ..thread.clone() }); cx.notify(); } } pub fn delete(&mut self, session_id: acp::SessionId, cx: &mut Context) { if let Some(thread) = self.threads.get(&session_id) { if let Some(session_ids) = self.threads_by_paths.get_mut(thread.folder_paths()) { session_ids.remove(&session_id); } if !thread.main_worktree_paths().is_empty() { if let Some(session_ids) = self .threads_by_main_paths .get_mut(thread.main_worktree_paths()) { session_ids.remove(&session_id); } } } self.threads.remove(&session_id); self.pending_thread_ops_tx .try_send(DbOperation::Delete(session_id)) .log_err(); cx.notify(); } fn new(db: ThreadMetadataDb, cx: &mut Context) -> Self { let weak_store = cx.weak_entity(); cx.observe_new::(move |thread, _window, cx| { // Don't track subagent threads in the sidebar. if thread.parent_session_id().is_some() { return; } let thread_entity = cx.entity(); cx.on_release({ let weak_store = weak_store.clone(); move |thread, cx| { weak_store .update(cx, |store, _cx| { let session_id = thread.session_id().clone(); store.session_subscriptions.remove(&session_id); }) .ok(); } }) .detach(); weak_store .update(cx, |this, cx| { let subscription = cx.subscribe(&thread_entity, Self::handle_thread_event); this.session_subscriptions .insert(thread.session_id().clone(), subscription); }) .ok(); }) .detach(); let (tx, rx) = smol::channel::unbounded(); let _db_operations_task = cx.background_spawn({ let db = db.clone(); async move { while let Ok(first_update) = rx.recv().await { let mut updates = vec![first_update]; while let Ok(update) = rx.try_recv() { updates.push(update); } let updates = Self::dedup_db_operations(updates); for operation in updates { match operation { DbOperation::Upsert(metadata) => { db.save(metadata).await.log_err(); } DbOperation::Delete(session_id) => { db.delete(session_id).await.log_err(); } } } } } }); let mut this = Self { db, threads: HashMap::default(), threads_by_paths: HashMap::default(), threads_by_main_paths: HashMap::default(), reload_task: None, session_subscriptions: HashMap::default(), pending_thread_ops_tx: tx, in_flight_archives: HashMap::default(), _db_operations_task, }; let _ = this.reload(cx); this } fn dedup_db_operations(operations: Vec) -> Vec { let mut ops = HashMap::default(); for operation in operations.into_iter().rev() { if ops.contains_key(operation.id()) { continue; } ops.insert(operation.id().clone(), operation); } ops.into_values().collect() } fn handle_thread_event( &mut self, thread: Entity, event: &AcpThreadEvent, cx: &mut Context, ) { // Don't track subagent threads in the sidebar. if thread.read(cx).parent_session_id().is_some() { return; } match event { AcpThreadEvent::NewEntry | AcpThreadEvent::TitleUpdated | AcpThreadEvent::EntryUpdated(_) | AcpThreadEvent::EntriesRemoved(_) | AcpThreadEvent::ToolAuthorizationRequested(_) | AcpThreadEvent::ToolAuthorizationReceived(_) | AcpThreadEvent::Retry(_) | AcpThreadEvent::Stopped(_) | AcpThreadEvent::Error | AcpThreadEvent::LoadError(_) | AcpThreadEvent::Refusal | AcpThreadEvent::WorkingDirectoriesUpdated => { let thread_ref = thread.read(cx); if thread_ref.entries().is_empty() { return; } let existing_thread = self.threads.get(thread_ref.session_id()); let session_id = thread_ref.session_id().clone(); let title = thread_ref .title() .unwrap_or_else(|| DEFAULT_THREAD_TITLE.into()); let updated_at = Utc::now(); let created_at = existing_thread .and_then(|t| t.created_at) .unwrap_or_else(|| updated_at); let agent_id = thread_ref.connection().agent_id(); let project = thread_ref.project().read(cx); let worktree_paths = ThreadWorktreePaths::from_project(project, cx); let project_group_key = project.project_group_key(cx); let remote_connection = project_group_key.host(); // Threads without a folder path (e.g. started in an empty // window) are archived by default so they don't get lost, // because they won't show up in the sidebar. Users can reload // them from the archive. let archived = existing_thread .map(|t| t.archived) .unwrap_or(worktree_paths.is_empty()); let metadata = ThreadMetadata { session_id, agent_id, title, created_at: Some(created_at), updated_at, worktree_paths, remote_connection, archived, }; self.save(metadata, cx); } AcpThreadEvent::TokenUsageUpdated | AcpThreadEvent::SubagentSpawned(_) | AcpThreadEvent::PromptCapabilitiesUpdated | AcpThreadEvent::AvailableCommandsUpdated(_) | AcpThreadEvent::ModeUpdated(_) | AcpThreadEvent::ConfigOptionsUpdated(_) => {} } } } impl Global for ThreadMetadataStore {} struct ThreadMetadataDb(ThreadSafeConnection); impl Domain for ThreadMetadataDb { const NAME: &str = stringify!(ThreadMetadataDb); const MIGRATIONS: &[&str] = &[ sql!( CREATE TABLE IF NOT EXISTS sidebar_threads( session_id TEXT PRIMARY KEY, agent_id TEXT, title TEXT NOT NULL, updated_at TEXT NOT NULL, created_at TEXT, folder_paths TEXT, folder_paths_order TEXT ) STRICT; ), sql!(ALTER TABLE sidebar_threads ADD COLUMN archived INTEGER DEFAULT 0), sql!(ALTER TABLE sidebar_threads ADD COLUMN main_worktree_paths TEXT), sql!(ALTER TABLE sidebar_threads ADD COLUMN main_worktree_paths_order TEXT), sql!( CREATE TABLE IF NOT EXISTS archived_git_worktrees( id INTEGER PRIMARY KEY, worktree_path TEXT NOT NULL, main_repo_path TEXT NOT NULL, branch_name TEXT, staged_commit_hash TEXT, unstaged_commit_hash TEXT, original_commit_hash TEXT ) STRICT; CREATE TABLE IF NOT EXISTS thread_archived_worktrees( session_id TEXT NOT NULL, archived_worktree_id INTEGER NOT NULL REFERENCES archived_git_worktrees(id), PRIMARY KEY (session_id, archived_worktree_id) ) STRICT; ), sql!(ALTER TABLE sidebar_threads ADD COLUMN remote_connection TEXT), ]; } db::static_connection!(ThreadMetadataDb, []); impl ThreadMetadataDb { pub fn list_ids(&self) -> anyhow::Result>> { self.select::>( "SELECT session_id FROM sidebar_threads \ ORDER BY updated_at DESC", )?() } /// List all sidebar thread metadata, ordered by updated_at descending. pub fn list(&self) -> anyhow::Result> { self.select::( "SELECT session_id, agent_id, title, updated_at, created_at, folder_paths, folder_paths_order, archived, main_worktree_paths, main_worktree_paths_order, remote_connection \ FROM sidebar_threads \ ORDER BY updated_at DESC" )?() } /// Upsert metadata for a thread. pub async fn save(&self, row: ThreadMetadata) -> anyhow::Result<()> { let id = row.session_id.0.clone(); let agent_id = if row.agent_id.as_ref() == ZED_AGENT_ID.as_ref() { None } else { Some(row.agent_id.to_string()) }; let title = row.title.to_string(); let updated_at = row.updated_at.to_rfc3339(); let created_at = row.created_at.map(|dt| dt.to_rfc3339()); let serialized = row.folder_paths().serialize(); let (folder_paths, folder_paths_order) = if row.folder_paths().is_empty() { (None, None) } else { (Some(serialized.paths), Some(serialized.order)) }; let main_serialized = row.main_worktree_paths().serialize(); let (main_worktree_paths, main_worktree_paths_order) = if row.main_worktree_paths().is_empty() { (None, None) } else { (Some(main_serialized.paths), Some(main_serialized.order)) }; let remote_connection = row .remote_connection .as_ref() .map(serde_json::to_string) .transpose() .context("serialize thread metadata remote connection")?; let archived = row.archived; self.write(move |conn| { let sql = "INSERT INTO sidebar_threads(session_id, agent_id, title, updated_at, created_at, folder_paths, folder_paths_order, archived, main_worktree_paths, main_worktree_paths_order, remote_connection) \ VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) \ ON CONFLICT(session_id) DO UPDATE SET \ agent_id = excluded.agent_id, \ title = excluded.title, \ updated_at = excluded.updated_at, \ created_at = excluded.created_at, \ folder_paths = excluded.folder_paths, \ folder_paths_order = excluded.folder_paths_order, \ archived = excluded.archived, \ main_worktree_paths = excluded.main_worktree_paths, \ main_worktree_paths_order = excluded.main_worktree_paths_order, \ remote_connection = excluded.remote_connection"; let mut stmt = Statement::prepare(conn, sql)?; let mut i = stmt.bind(&id, 1)?; i = stmt.bind(&agent_id, i)?; i = stmt.bind(&title, i)?; i = stmt.bind(&updated_at, i)?; i = stmt.bind(&created_at, i)?; i = stmt.bind(&folder_paths, i)?; i = stmt.bind(&folder_paths_order, i)?; i = stmt.bind(&archived, i)?; i = stmt.bind(&main_worktree_paths, i)?; i = stmt.bind(&main_worktree_paths_order, i)?; stmt.bind(&remote_connection, i)?; stmt.exec() }) .await } /// Delete metadata for a single thread. pub async fn delete(&self, session_id: acp::SessionId) -> anyhow::Result<()> { let id = session_id.0.clone(); self.write(move |conn| { let mut stmt = Statement::prepare(conn, "DELETE FROM sidebar_threads WHERE session_id = ?")?; stmt.bind(&id, 1)?; stmt.exec() }) .await } pub async fn create_archived_worktree( &self, worktree_path: String, main_repo_path: String, branch_name: Option, staged_commit_hash: String, unstaged_commit_hash: String, original_commit_hash: String, ) -> anyhow::Result { self.write(move |conn| { let mut stmt = Statement::prepare( conn, "INSERT INTO archived_git_worktrees(worktree_path, main_repo_path, branch_name, staged_commit_hash, unstaged_commit_hash, original_commit_hash) \ VALUES (?1, ?2, ?3, ?4, ?5, ?6) \ RETURNING id", )?; let mut i = stmt.bind(&worktree_path, 1)?; i = stmt.bind(&main_repo_path, i)?; i = stmt.bind(&branch_name, i)?; i = stmt.bind(&staged_commit_hash, i)?; i = stmt.bind(&unstaged_commit_hash, i)?; stmt.bind(&original_commit_hash, i)?; stmt.maybe_row::()?.context("expected RETURNING id") }) .await } pub async fn link_thread_to_archived_worktree( &self, session_id: String, archived_worktree_id: i64, ) -> anyhow::Result<()> { self.write(move |conn| { let mut stmt = Statement::prepare( conn, "INSERT INTO thread_archived_worktrees(session_id, archived_worktree_id) \ VALUES (?1, ?2)", )?; let i = stmt.bind(&session_id, 1)?; stmt.bind(&archived_worktree_id, i)?; stmt.exec() }) .await } pub async fn get_archived_worktrees_for_thread( &self, session_id: String, ) -> anyhow::Result> { self.select_bound::( "SELECT a.id, a.worktree_path, a.main_repo_path, a.branch_name, a.staged_commit_hash, a.unstaged_commit_hash, a.original_commit_hash \ FROM archived_git_worktrees a \ JOIN thread_archived_worktrees t ON a.id = t.archived_worktree_id \ WHERE t.session_id = ?1", )?(session_id) } pub async fn delete_archived_worktree(&self, id: i64) -> anyhow::Result<()> { self.write(move |conn| { let mut stmt = Statement::prepare( conn, "DELETE FROM thread_archived_worktrees WHERE archived_worktree_id = ?", )?; stmt.bind(&id, 1)?; stmt.exec()?; let mut stmt = Statement::prepare(conn, "DELETE FROM archived_git_worktrees WHERE id = ?")?; stmt.bind(&id, 1)?; stmt.exec() }) .await } pub async fn unlink_thread_from_all_archived_worktrees( &self, session_id: String, ) -> anyhow::Result<()> { self.write(move |conn| { let mut stmt = Statement::prepare( conn, "DELETE FROM thread_archived_worktrees WHERE session_id = ?", )?; stmt.bind(&session_id, 1)?; stmt.exec() }) .await } pub async fn is_archived_worktree_referenced( &self, archived_worktree_id: i64, ) -> anyhow::Result { self.select_row_bound::( "SELECT COUNT(*) FROM thread_archived_worktrees WHERE archived_worktree_id = ?1", )?(archived_worktree_id) .map(|count| count.unwrap_or(0) > 0) } } impl Column for ThreadMetadata { fn column(statement: &mut Statement, start_index: i32) -> anyhow::Result<(Self, i32)> { let (id, next): (Arc, i32) = Column::column(statement, start_index)?; let (agent_id, next): (Option, i32) = Column::column(statement, next)?; let (title, next): (String, i32) = Column::column(statement, next)?; let (updated_at_str, next): (String, i32) = Column::column(statement, next)?; let (created_at_str, next): (Option, i32) = Column::column(statement, next)?; let (folder_paths_str, next): (Option, i32) = Column::column(statement, next)?; let (folder_paths_order_str, next): (Option, i32) = Column::column(statement, next)?; let (archived, next): (bool, i32) = Column::column(statement, next)?; let (main_worktree_paths_str, next): (Option, i32) = Column::column(statement, next)?; let (main_worktree_paths_order_str, next): (Option, i32) = Column::column(statement, next)?; let (remote_connection_json, next): (Option, i32) = Column::column(statement, next)?; let agent_id = agent_id .map(|id| AgentId::new(id)) .unwrap_or(ZED_AGENT_ID.clone()); let updated_at = DateTime::parse_from_rfc3339(&updated_at_str)?.with_timezone(&Utc); let created_at = created_at_str .as_deref() .map(DateTime::parse_from_rfc3339) .transpose()? .map(|dt| dt.with_timezone(&Utc)); let folder_paths = folder_paths_str .map(|paths| { PathList::deserialize(&util::path_list::SerializedPathList { paths, order: folder_paths_order_str.unwrap_or_default(), }) }) .unwrap_or_default(); let main_worktree_paths = main_worktree_paths_str .map(|paths| { PathList::deserialize(&util::path_list::SerializedPathList { paths, order: main_worktree_paths_order_str.unwrap_or_default(), }) }) .unwrap_or_default(); let remote_connection = remote_connection_json .as_deref() .map(serde_json::from_str::) .transpose() .context("deserialize thread metadata remote connection")?; let worktree_paths = ThreadWorktreePaths::from_path_lists(main_worktree_paths, folder_paths) .unwrap_or_else(|_| ThreadWorktreePaths::default()); Ok(( ThreadMetadata { session_id: acp::SessionId::new(id), agent_id, title: title.into(), updated_at, created_at, worktree_paths, remote_connection, archived, }, next, )) } } impl Column for ArchivedGitWorktree { fn column(statement: &mut Statement, start_index: i32) -> anyhow::Result<(Self, i32)> { let (id, next): (i64, i32) = Column::column(statement, start_index)?; let (worktree_path_str, next): (String, i32) = Column::column(statement, next)?; let (main_repo_path_str, next): (String, i32) = Column::column(statement, next)?; let (branch_name, next): (Option, i32) = Column::column(statement, next)?; let (staged_commit_hash, next): (String, i32) = Column::column(statement, next)?; let (unstaged_commit_hash, next): (String, i32) = Column::column(statement, next)?; let (original_commit_hash, next): (String, i32) = Column::column(statement, next)?; Ok(( ArchivedGitWorktree { id, worktree_path: PathBuf::from(worktree_path_str), main_repo_path: PathBuf::from(main_repo_path_str), branch_name, staged_commit_hash, unstaged_commit_hash, original_commit_hash, }, next, )) } } #[cfg(test)] mod tests { use super::*; use acp_thread::{AgentConnection, StubAgentConnection}; use action_log::ActionLog; use agent::DbThread; use agent_client_protocol as acp; use gpui::TestAppContext; use project::FakeFs; use project::Project; use remote::WslConnectionOptions; use std::path::Path; use std::rc::Rc; fn make_db_thread(title: &str, updated_at: DateTime) -> DbThread { DbThread { title: title.to_string().into(), messages: Vec::new(), updated_at, detailed_summary: None, initial_project_snapshot: None, cumulative_token_usage: Default::default(), request_token_usage: Default::default(), model: None, profile: None, imported: false, subagent_context: None, speed: None, thinking_enabled: false, thinking_effort: None, draft_prompt: None, ui_scroll_position: None, } } fn make_metadata( session_id: &str, title: &str, updated_at: DateTime, folder_paths: PathList, ) -> ThreadMetadata { ThreadMetadata { archived: false, session_id: acp::SessionId::new(session_id), agent_id: agent::ZED_AGENT_ID.clone(), title: title.to_string().into(), updated_at, created_at: Some(updated_at), worktree_paths: ThreadWorktreePaths::from_folder_paths(&folder_paths), remote_connection: None, } } fn init_test(cx: &mut TestAppContext) { let fs = FakeFs::new(cx.executor()); cx.update(|cx| { let settings_store = settings::SettingsStore::test(cx); cx.set_global(settings_store); ::set_global(fs, cx); ThreadMetadataStore::init_global(cx); ThreadStore::init_global(cx); }); cx.run_until_parked(); } fn clear_thread_metadata_remote_connection_backfill(cx: &mut TestAppContext) { let kvp = cx.update(|cx| KeyValueStore::global(cx)); smol::block_on(kvp.delete_kvp("thread-metadata-remote-connection-backfill".to_string())) .unwrap(); } fn run_thread_metadata_migrations(cx: &mut TestAppContext) { clear_thread_metadata_remote_connection_backfill(cx); cx.update(|cx| { let migration_task = migrate_thread_metadata(cx); migrate_thread_remote_connections(cx, migration_task); }); cx.run_until_parked(); } #[gpui::test] async fn test_store_initializes_cache_from_database(cx: &mut TestAppContext) { let first_paths = PathList::new(&[Path::new("/project-a")]); let second_paths = PathList::new(&[Path::new("/project-b")]); let now = Utc::now(); let older = now - chrono::Duration::seconds(1); let thread = std::thread::current(); let test_name = thread.name().unwrap_or("unknown_test"); let db_name = format!("THREAD_METADATA_DB_{}", test_name); let db = ThreadMetadataDb(smol::block_on(db::open_test_db::( &db_name, ))); db.save(make_metadata( "session-1", "First Thread", now, first_paths.clone(), )) .await .unwrap(); db.save(make_metadata( "session-2", "Second Thread", older, second_paths.clone(), )) .await .unwrap(); cx.update(|cx| { let settings_store = settings::SettingsStore::test(cx); cx.set_global(settings_store); ThreadMetadataStore::init_global(cx); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let entry_ids = store .entry_ids() .map(|session_id| session_id.0.to_string()) .collect::>(); assert_eq!(entry_ids.len(), 2); assert!(entry_ids.contains(&"session-1".to_string())); assert!(entry_ids.contains(&"session-2".to_string())); let first_path_entries = store .entries_for_path(&first_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert_eq!(first_path_entries, vec!["session-1"]); let second_path_entries = store .entries_for_path(&second_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert_eq!(second_path_entries, vec!["session-2"]); }); } #[gpui::test] async fn test_store_cache_updates_after_save_and_delete(cx: &mut TestAppContext) { init_test(cx); let first_paths = PathList::new(&[Path::new("/project-a")]); let second_paths = PathList::new(&[Path::new("/project-b")]); let initial_time = Utc::now(); let updated_time = initial_time + chrono::Duration::seconds(1); let initial_metadata = make_metadata( "session-1", "First Thread", initial_time, first_paths.clone(), ); let second_metadata = make_metadata( "session-2", "Second Thread", initial_time, second_paths.clone(), ); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(initial_metadata, cx); store.save(second_metadata, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let first_path_entries = store .entries_for_path(&first_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert_eq!(first_path_entries, vec!["session-1"]); let second_path_entries = store .entries_for_path(&second_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert_eq!(second_path_entries, vec!["session-2"]); }); let moved_metadata = make_metadata( "session-1", "First Thread", updated_time, second_paths.clone(), ); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(moved_metadata, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let entry_ids = store .entry_ids() .map(|session_id| session_id.0.to_string()) .collect::>(); assert_eq!(entry_ids.len(), 2); assert!(entry_ids.contains(&"session-1".to_string())); assert!(entry_ids.contains(&"session-2".to_string())); let first_path_entries = store .entries_for_path(&first_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert!(first_path_entries.is_empty()); let second_path_entries = store .entries_for_path(&second_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert_eq!(second_path_entries.len(), 2); assert!(second_path_entries.contains(&"session-1".to_string())); assert!(second_path_entries.contains(&"session-2".to_string())); }); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.delete(acp::SessionId::new("session-2"), cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let entry_ids = store .entry_ids() .map(|session_id| session_id.0.to_string()) .collect::>(); assert_eq!(entry_ids, vec!["session-1"]); let second_path_entries = store .entries_for_path(&second_paths) .map(|entry| entry.session_id.0.to_string()) .collect::>(); assert_eq!(second_path_entries, vec!["session-1"]); }); } #[gpui::test] async fn test_migrate_thread_metadata_migrates_only_missing_threads(cx: &mut TestAppContext) { init_test(cx); let project_a_paths = PathList::new(&[Path::new("/project-a")]); let project_b_paths = PathList::new(&[Path::new("/project-b")]); let now = Utc::now(); let existing_metadata = ThreadMetadata { session_id: acp::SessionId::new("a-session-0"), agent_id: agent::ZED_AGENT_ID.clone(), title: "Existing Metadata".into(), updated_at: now - chrono::Duration::seconds(10), created_at: Some(now - chrono::Duration::seconds(10)), worktree_paths: ThreadWorktreePaths::from_folder_paths(&project_a_paths), remote_connection: None, archived: false, }; cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(existing_metadata, cx); }); }); cx.run_until_parked(); let threads_to_save = vec![ ( "a-session-0", "Thread A0 From Native Store", project_a_paths.clone(), now, ), ( "a-session-1", "Thread A1", project_a_paths.clone(), now + chrono::Duration::seconds(1), ), ( "b-session-0", "Thread B0", project_b_paths.clone(), now + chrono::Duration::seconds(2), ), ( "projectless", "Projectless", PathList::default(), now + chrono::Duration::seconds(3), ), ]; for (session_id, title, paths, updated_at) in &threads_to_save { let save_task = cx.update(|cx| { let thread_store = ThreadStore::global(cx); let session_id = session_id.to_string(); let title = title.to_string(); let paths = paths.clone(); thread_store.update(cx, |store, cx| { store.save_thread( acp::SessionId::new(session_id), make_db_thread(&title, *updated_at), paths, cx, ) }) }); save_task.await.unwrap(); cx.run_until_parked(); } run_thread_metadata_migrations(cx); let list = cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.read(cx).entries().cloned().collect::>() }); assert_eq!(list.len(), 4); assert!( list.iter() .all(|metadata| metadata.agent_id.as_ref() == agent::ZED_AGENT_ID.as_ref()) ); let existing_metadata = list .iter() .find(|metadata| metadata.session_id.0.as_ref() == "a-session-0") .unwrap(); assert_eq!(existing_metadata.title.as_ref(), "Existing Metadata"); assert!(!existing_metadata.archived); let migrated_session_ids = list .iter() .map(|metadata| metadata.session_id.0.as_ref()) .collect::>(); assert!(migrated_session_ids.contains(&"a-session-1")); assert!(migrated_session_ids.contains(&"b-session-0")); assert!(migrated_session_ids.contains(&"projectless")); let migrated_entries = list .iter() .filter(|metadata| metadata.session_id.0.as_ref() != "a-session-0") .collect::>(); assert!(migrated_entries.iter().all(|metadata| metadata.archived)); } #[gpui::test] async fn test_migrate_thread_metadata_noops_when_all_threads_already_exist( cx: &mut TestAppContext, ) { init_test(cx); let project_paths = PathList::new(&[Path::new("/project-a")]); let existing_updated_at = Utc::now(); let existing_metadata = ThreadMetadata { session_id: acp::SessionId::new("existing-session"), agent_id: agent::ZED_AGENT_ID.clone(), title: "Existing Metadata".into(), updated_at: existing_updated_at, created_at: Some(existing_updated_at), worktree_paths: ThreadWorktreePaths::from_folder_paths(&project_paths), remote_connection: None, archived: false, }; cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(existing_metadata, cx); }); }); cx.run_until_parked(); let save_task = cx.update(|cx| { let thread_store = ThreadStore::global(cx); thread_store.update(cx, |store, cx| { store.save_thread( acp::SessionId::new("existing-session"), make_db_thread( "Updated Native Thread Title", existing_updated_at + chrono::Duration::seconds(1), ), project_paths.clone(), cx, ) }) }); save_task.await.unwrap(); cx.run_until_parked(); run_thread_metadata_migrations(cx); let list = cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.read(cx).entries().cloned().collect::>() }); assert_eq!(list.len(), 1); assert_eq!(list[0].session_id.0.as_ref(), "existing-session"); } #[gpui::test] async fn test_migrate_thread_remote_connections_backfills_from_workspace_db( cx: &mut TestAppContext, ) { init_test(cx); let folder_paths = PathList::new(&[Path::new("/remote-project")]); let updated_at = Utc::now(); let metadata = make_metadata( "remote-session", "Remote Thread", updated_at, folder_paths.clone(), ); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(metadata, cx); }); }); cx.run_until_parked(); let workspace_db = cx.update(|cx| WorkspaceDb::global(cx)); let workspace_id = workspace_db.next_id().await.unwrap(); let serialized_paths = folder_paths.serialize(); let remote_connection_id = 1_i64; workspace_db .write(move |conn| { let mut stmt = Statement::prepare( conn, "INSERT INTO remote_connections(id, kind, user, distro) VALUES (?1, ?2, ?3, ?4)", )?; let mut next_index = stmt.bind(&remote_connection_id, 1)?; next_index = stmt.bind(&"wsl", next_index)?; next_index = stmt.bind(&Some("anth".to_string()), next_index)?; stmt.bind(&Some("Ubuntu".to_string()), next_index)?; stmt.exec()?; let mut stmt = Statement::prepare( conn, "UPDATE workspaces SET paths = ?2, paths_order = ?3, remote_connection_id = ?4, timestamp = CURRENT_TIMESTAMP WHERE workspace_id = ?1", )?; let mut next_index = stmt.bind(&workspace_id, 1)?; next_index = stmt.bind(&serialized_paths.paths, next_index)?; next_index = stmt.bind(&serialized_paths.order, next_index)?; stmt.bind(&Some(remote_connection_id as i32), next_index)?; stmt.exec() }) .await .unwrap(); clear_thread_metadata_remote_connection_backfill(cx); cx.update(|cx| { migrate_thread_remote_connections(cx, Task::ready(Ok(()))); }); cx.run_until_parked(); let metadata = cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store .read(cx) .entry(&acp::SessionId::new("remote-session")) .cloned() .expect("expected migrated metadata row") }); assert_eq!( metadata.remote_connection, Some(RemoteConnectionOptions::Wsl(WslConnectionOptions { distro_name: "Ubuntu".to_string(), user: Some("anth".to_string()), })) ); } #[gpui::test] async fn test_migrate_thread_metadata_archives_beyond_five_most_recent_per_project( cx: &mut TestAppContext, ) { init_test(cx); let project_a_paths = PathList::new(&[Path::new("/project-a")]); let project_b_paths = PathList::new(&[Path::new("/project-b")]); let now = Utc::now(); // Create 7 threads for project A and 3 for project B let mut threads_to_save = Vec::new(); for i in 0..7 { threads_to_save.push(( format!("a-session-{i}"), format!("Thread A{i}"), project_a_paths.clone(), now + chrono::Duration::seconds(i as i64), )); } for i in 0..3 { threads_to_save.push(( format!("b-session-{i}"), format!("Thread B{i}"), project_b_paths.clone(), now + chrono::Duration::seconds(i as i64), )); } for (session_id, title, paths, updated_at) in &threads_to_save { let save_task = cx.update(|cx| { let thread_store = ThreadStore::global(cx); let session_id = session_id.to_string(); let title = title.to_string(); let paths = paths.clone(); thread_store.update(cx, |store, cx| { store.save_thread( acp::SessionId::new(session_id), make_db_thread(&title, *updated_at), paths, cx, ) }) }); save_task.await.unwrap(); cx.run_until_parked(); } run_thread_metadata_migrations(cx); let list = cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.read(cx).entries().cloned().collect::>() }); assert_eq!(list.len(), 10); // Project A: 5 most recent should be unarchived, 2 oldest should be archived let mut project_a_entries: Vec<_> = list .iter() .filter(|m| *m.folder_paths() == project_a_paths) .collect(); assert_eq!(project_a_entries.len(), 7); project_a_entries.sort_by(|a, b| b.updated_at.cmp(&a.updated_at)); for entry in &project_a_entries[..5] { assert!( !entry.archived, "Expected {} to be unarchived (top 5 most recent)", entry.session_id.0 ); } for entry in &project_a_entries[5..] { assert!( entry.archived, "Expected {} to be archived (older than top 5)", entry.session_id.0 ); } // Project B: all 3 should be unarchived (under the limit) let project_b_entries: Vec<_> = list .iter() .filter(|m| *m.folder_paths() == project_b_paths) .collect(); assert_eq!(project_b_entries.len(), 3); assert!(project_b_entries.iter().all(|m| !m.archived)); } #[gpui::test] async fn test_empty_thread_events_do_not_create_metadata(cx: &mut TestAppContext) { init_test(cx); let fs = FakeFs::new(cx.executor()); let project = Project::test(fs, None::<&Path>, cx).await; let connection = Rc::new(StubAgentConnection::new()); let thread = cx .update(|cx| { connection .clone() .new_session(project.clone(), PathList::default(), cx) }) .await .unwrap(); let session_id = cx.read(|cx| thread.read(cx).session_id().clone()); cx.update(|cx| { thread.update(cx, |thread, cx| { thread.set_title("Draft Thread".into(), cx).detach(); }); }); cx.run_until_parked(); let metadata_ids = cx.update(|cx| { ThreadMetadataStore::global(cx) .read(cx) .entry_ids() .collect::>() }); assert!( metadata_ids.is_empty(), "expected empty draft thread title updates to be ignored" ); cx.update(|cx| { thread.update(cx, |thread, cx| { thread.push_user_content_block(None, "Hello".into(), cx); }); }); cx.run_until_parked(); let metadata_ids = cx.update(|cx| { ThreadMetadataStore::global(cx) .read(cx) .entry_ids() .collect::>() }); assert_eq!(metadata_ids, vec![session_id]); } #[gpui::test] async fn test_nonempty_thread_metadata_preserved_when_thread_released(cx: &mut TestAppContext) { init_test(cx); let fs = FakeFs::new(cx.executor()); let project = Project::test(fs, None::<&Path>, cx).await; let connection = Rc::new(StubAgentConnection::new()); let thread = cx .update(|cx| { connection .clone() .new_session(project.clone(), PathList::default(), cx) }) .await .unwrap(); let session_id = cx.read(|cx| thread.read(cx).session_id().clone()); cx.update(|cx| { thread.update(cx, |thread, cx| { thread.push_user_content_block(None, "Hello".into(), cx); }); }); cx.run_until_parked(); let metadata_ids = cx.update(|cx| { ThreadMetadataStore::global(cx) .read(cx) .entry_ids() .collect::>() }); assert_eq!(metadata_ids, vec![session_id.clone()]); drop(thread); cx.update(|_| {}); cx.run_until_parked(); let metadata_ids = cx.update(|cx| { ThreadMetadataStore::global(cx) .read(cx) .entry_ids() .collect::>() }); assert_eq!(metadata_ids, vec![session_id]); } #[gpui::test] async fn test_threads_without_project_association_are_archived_by_default( cx: &mut TestAppContext, ) { init_test(cx); let fs = FakeFs::new(cx.executor()); let project_without_worktree = Project::test(fs.clone(), None::<&Path>, cx).await; let project_with_worktree = Project::test(fs, [Path::new("/project-a")], cx).await; let connection = Rc::new(StubAgentConnection::new()); let thread_without_worktree = cx .update(|cx| { connection.clone().new_session( project_without_worktree.clone(), PathList::default(), cx, ) }) .await .unwrap(); let session_without_worktree = cx.read(|cx| thread_without_worktree.read(cx).session_id().clone()); cx.update(|cx| { thread_without_worktree.update(cx, |thread, cx| { thread.push_user_content_block(None, "content".into(), cx); thread.set_title("No Project Thread".into(), cx).detach(); }); }); cx.run_until_parked(); let thread_with_worktree = cx .update(|cx| { connection.clone().new_session( project_with_worktree.clone(), PathList::default(), cx, ) }) .await .unwrap(); let session_with_worktree = cx.read(|cx| thread_with_worktree.read(cx).session_id().clone()); cx.update(|cx| { thread_with_worktree.update(cx, |thread, cx| { thread.push_user_content_block(None, "content".into(), cx); thread.set_title("Project Thread".into(), cx).detach(); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let without_worktree = store .entry(&session_without_worktree) .expect("missing metadata for thread without project association"); assert!(without_worktree.folder_paths().is_empty()); assert!( without_worktree.archived, "expected thread without project association to be archived" ); let with_worktree = store .entry(&session_with_worktree) .expect("missing metadata for thread with project association"); assert_eq!( *with_worktree.folder_paths(), PathList::new(&[Path::new("/project-a")]) ); assert!( !with_worktree.archived, "expected thread with project association to remain unarchived" ); }); } #[gpui::test] async fn test_subagent_threads_excluded_from_sidebar_metadata(cx: &mut TestAppContext) { init_test(cx); let fs = FakeFs::new(cx.executor()); let project = Project::test(fs, None::<&Path>, cx).await; let connection = Rc::new(StubAgentConnection::new()); // Create a regular (non-subagent) AcpThread. let regular_thread = cx .update(|cx| { connection .clone() .new_session(project.clone(), PathList::default(), cx) }) .await .unwrap(); let regular_session_id = cx.read(|cx| regular_thread.read(cx).session_id().clone()); // Set a title on the regular thread to trigger a save via handle_thread_update. cx.update(|cx| { regular_thread.update(cx, |thread, cx| { thread.push_user_content_block(None, "content".into(), cx); thread.set_title("Regular Thread".into(), cx).detach(); }); }); cx.run_until_parked(); // Create a subagent AcpThread let subagent_session_id = acp::SessionId::new("subagent-session"); let subagent_thread = cx.update(|cx| { let action_log = cx.new(|_| ActionLog::new(project.clone())); cx.new(|cx| { acp_thread::AcpThread::new( Some(regular_session_id.clone()), Some("Subagent Thread".into()), None, connection.clone(), project.clone(), action_log, subagent_session_id.clone(), watch::Receiver::constant(acp::PromptCapabilities::new()), cx, ) }) }); // Set a title on the subagent thread to trigger handle_thread_update. cx.update(|cx| { subagent_thread.update(cx, |thread, cx| { thread .set_title("Subagent Thread Title".into(), cx) .detach(); }); }); cx.run_until_parked(); // List all metadata from the store cache. let list = cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.read(cx).entries().cloned().collect::>() }); // The subagent thread should NOT appear in the sidebar metadata. // Only the regular thread should be listed. assert_eq!( list.len(), 1, "Expected only the regular thread in sidebar metadata, \ but found {} entries (subagent threads are leaking into the sidebar)", list.len(), ); assert_eq!(list[0].session_id, regular_session_id); assert_eq!(list[0].title.as_ref(), "Regular Thread"); } #[test] fn test_dedup_db_operations_keeps_latest_operation_for_session() { let now = Utc::now(); let operations = vec![ DbOperation::Upsert(make_metadata( "session-1", "First Thread", now, PathList::default(), )), DbOperation::Delete(acp::SessionId::new("session-1")), ]; let deduped = ThreadMetadataStore::dedup_db_operations(operations); assert_eq!(deduped.len(), 1); assert_eq!( deduped[0], DbOperation::Delete(acp::SessionId::new("session-1")) ); } #[test] fn test_dedup_db_operations_keeps_latest_insert_for_same_session() { let now = Utc::now(); let later = now + chrono::Duration::seconds(1); let old_metadata = make_metadata("session-1", "Old Title", now, PathList::default()); let new_metadata = make_metadata("session-1", "New Title", later, PathList::default()); let deduped = ThreadMetadataStore::dedup_db_operations(vec![ DbOperation::Upsert(old_metadata), DbOperation::Upsert(new_metadata.clone()), ]); assert_eq!(deduped.len(), 1); assert_eq!(deduped[0], DbOperation::Upsert(new_metadata)); } #[test] fn test_dedup_db_operations_preserves_distinct_sessions() { let now = Utc::now(); let metadata1 = make_metadata("session-1", "First Thread", now, PathList::default()); let metadata2 = make_metadata("session-2", "Second Thread", now, PathList::default()); let deduped = ThreadMetadataStore::dedup_db_operations(vec![ DbOperation::Upsert(metadata1.clone()), DbOperation::Upsert(metadata2.clone()), ]); assert_eq!(deduped.len(), 2); assert!(deduped.contains(&DbOperation::Upsert(metadata1))); assert!(deduped.contains(&DbOperation::Upsert(metadata2))); } #[gpui::test] async fn test_archive_and_unarchive_thread(cx: &mut TestAppContext) { init_test(cx); let paths = PathList::new(&[Path::new("/project-a")]); let now = Utc::now(); let metadata = make_metadata("session-1", "Thread 1", now, paths.clone()); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(metadata, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let path_entries = store .entries_for_path(&paths) .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(path_entries, vec!["session-1"]); let archived = store .archived_entries() .map(|e| e.session_id.0.to_string()) .collect::>(); assert!(archived.is_empty()); }); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.archive(&acp::SessionId::new("session-1"), None, cx); }); }); // Thread 1 should now be archived cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let path_entries = store .entries_for_path(&paths) .map(|e| e.session_id.0.to_string()) .collect::>(); assert!(path_entries.is_empty()); let archived = store.archived_entries().collect::>(); assert_eq!(archived.len(), 1); assert_eq!(archived[0].session_id.0.as_ref(), "session-1"); assert!(archived[0].archived); }); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.unarchive(&acp::SessionId::new("session-1"), cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let path_entries = store .entries_for_path(&paths) .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(path_entries, vec!["session-1"]); let archived = store .archived_entries() .map(|e| e.session_id.0.to_string()) .collect::>(); assert!(archived.is_empty()); }); } #[gpui::test] async fn test_entries_for_path_excludes_archived(cx: &mut TestAppContext) { init_test(cx); let paths = PathList::new(&[Path::new("/project-a")]); let now = Utc::now(); let metadata1 = make_metadata("session-1", "Active Thread", now, paths.clone()); let metadata2 = make_metadata( "session-2", "Archived Thread", now - chrono::Duration::seconds(1), paths.clone(), ); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(metadata1, cx); store.save(metadata2, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.archive(&acp::SessionId::new("session-2"), None, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let path_entries = store .entries_for_path(&paths) .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(path_entries, vec!["session-1"]); let all_entries = store .entries() .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(all_entries.len(), 2); assert!(all_entries.contains(&"session-1".to_string())); assert!(all_entries.contains(&"session-2".to_string())); let archived = store .archived_entries() .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(archived, vec!["session-2"]); }); } #[gpui::test] async fn test_save_all_persists_multiple_threads(cx: &mut TestAppContext) { init_test(cx); let paths = PathList::new(&[Path::new("/project-a")]); let now = Utc::now(); let m1 = make_metadata("session-1", "Thread One", now, paths.clone()); let m2 = make_metadata( "session-2", "Thread Two", now - chrono::Duration::seconds(1), paths.clone(), ); let m3 = make_metadata( "session-3", "Thread Three", now - chrono::Duration::seconds(2), paths, ); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save_all(vec![m1, m2, m3], cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let all_entries = store .entries() .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(all_entries.len(), 3); assert!(all_entries.contains(&"session-1".to_string())); assert!(all_entries.contains(&"session-2".to_string())); assert!(all_entries.contains(&"session-3".to_string())); let entry_ids = store.entry_ids().collect::>(); assert_eq!(entry_ids.len(), 3); }); } #[gpui::test] async fn test_archived_flag_persists_across_reload(cx: &mut TestAppContext) { init_test(cx); let paths = PathList::new(&[Path::new("/project-a")]); let now = Utc::now(); let metadata = make_metadata("session-1", "Thread 1", now, paths.clone()); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(metadata, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.archive(&acp::SessionId::new("session-1"), None, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { let _ = store.reload(cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let thread = store .entries() .find(|e| e.session_id.0.as_ref() == "session-1") .expect("thread should exist after reload"); assert!(thread.archived); let path_entries = store .entries_for_path(&paths) .map(|e| e.session_id.0.to_string()) .collect::>(); assert!(path_entries.is_empty()); let archived = store .archived_entries() .map(|e| e.session_id.0.to_string()) .collect::>(); assert_eq!(archived, vec!["session-1"]); }); } #[gpui::test] async fn test_archive_nonexistent_thread_is_noop(cx: &mut TestAppContext) { init_test(cx); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.archive(&acp::SessionId::new("nonexistent"), None, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); assert!(store.is_empty()); assert_eq!(store.entries().count(), 0); assert_eq!(store.archived_entries().count(), 0); }); } #[gpui::test] async fn test_save_followed_by_archiving_without_parking(cx: &mut TestAppContext) { init_test(cx); let paths = PathList::new(&[Path::new("/project-a")]); let now = Utc::now(); let metadata = make_metadata("session-1", "Thread 1", now, paths); let session_id = metadata.session_id.clone(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); store.update(cx, |store, cx| { store.save(metadata.clone(), cx); store.archive(&session_id, None, cx); }); }); cx.run_until_parked(); cx.update(|cx| { let store = ThreadMetadataStore::global(cx); let store = store.read(cx); let entries: Vec = store.entries().cloned().collect(); pretty_assertions::assert_eq!( entries, vec![ThreadMetadata { archived: true, ..metadata }] ); }); } #[gpui::test] async fn test_create_and_retrieve_archived_worktree(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let id = store .read_with(cx, |store, cx| { store.create_archived_worktree( "/tmp/worktree".to_string(), "/home/user/repo".to_string(), Some("feature-branch".to_string()), "staged_aaa".to_string(), "unstaged_bbb".to_string(), "original_000".to_string(), cx, ) }) .await .unwrap(); store .read_with(cx, |store, cx| { store.link_thread_to_archived_worktree("session-1".to_string(), id, cx) }) .await .unwrap(); let worktrees = store .read_with(cx, |store, cx| { store.get_archived_worktrees_for_thread("session-1".to_string(), cx) }) .await .unwrap(); assert_eq!(worktrees.len(), 1); let wt = &worktrees[0]; assert_eq!(wt.id, id); assert_eq!(wt.worktree_path, PathBuf::from("/tmp/worktree")); assert_eq!(wt.main_repo_path, PathBuf::from("/home/user/repo")); assert_eq!(wt.branch_name.as_deref(), Some("feature-branch")); assert_eq!(wt.staged_commit_hash, "staged_aaa"); assert_eq!(wt.unstaged_commit_hash, "unstaged_bbb"); assert_eq!(wt.original_commit_hash, "original_000"); } #[gpui::test] async fn test_delete_archived_worktree(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let id = store .read_with(cx, |store, cx| { store.create_archived_worktree( "/tmp/worktree".to_string(), "/home/user/repo".to_string(), Some("main".to_string()), "deadbeef".to_string(), "deadbeef".to_string(), "original_000".to_string(), cx, ) }) .await .unwrap(); store .read_with(cx, |store, cx| { store.link_thread_to_archived_worktree("session-1".to_string(), id, cx) }) .await .unwrap(); store .read_with(cx, |store, cx| store.delete_archived_worktree(id, cx)) .await .unwrap(); let worktrees = store .read_with(cx, |store, cx| { store.get_archived_worktrees_for_thread("session-1".to_string(), cx) }) .await .unwrap(); assert!(worktrees.is_empty()); } #[gpui::test] async fn test_link_multiple_threads_to_archived_worktree(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let id = store .read_with(cx, |store, cx| { store.create_archived_worktree( "/tmp/worktree".to_string(), "/home/user/repo".to_string(), None, "abc123".to_string(), "abc123".to_string(), "original_000".to_string(), cx, ) }) .await .unwrap(); store .read_with(cx, |store, cx| { store.link_thread_to_archived_worktree("session-1".to_string(), id, cx) }) .await .unwrap(); store .read_with(cx, |store, cx| { store.link_thread_to_archived_worktree("session-2".to_string(), id, cx) }) .await .unwrap(); let wt1 = store .read_with(cx, |store, cx| { store.get_archived_worktrees_for_thread("session-1".to_string(), cx) }) .await .unwrap(); let wt2 = store .read_with(cx, |store, cx| { store.get_archived_worktrees_for_thread("session-2".to_string(), cx) }) .await .unwrap(); assert_eq!(wt1.len(), 1); assert_eq!(wt2.len(), 1); assert_eq!(wt1[0].id, wt2[0].id); } #[gpui::test] async fn test_complete_worktree_restore_multiple_paths(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let original_paths = PathList::new(&[ Path::new("/projects/worktree-a"), Path::new("/projects/worktree-b"), Path::new("/other/unrelated"), ]); let meta = make_metadata("session-multi", "Multi Thread", Utc::now(), original_paths); store.update(cx, |store, cx| { store.save_manually(meta, cx); }); let replacements = vec![ ( PathBuf::from("/projects/worktree-a"), PathBuf::from("/restored/worktree-a"), ), ( PathBuf::from("/projects/worktree-b"), PathBuf::from("/restored/worktree-b"), ), ]; store.update(cx, |store, cx| { store.complete_worktree_restore( &acp::SessionId::new("session-multi"), &replacements, cx, ); }); let entry = store.read_with(cx, |store, _cx| { store.entry(&acp::SessionId::new("session-multi")).cloned() }); let entry = entry.unwrap(); let paths = entry.folder_paths().paths(); assert_eq!(paths.len(), 3); assert!(paths.contains(&PathBuf::from("/restored/worktree-a"))); assert!(paths.contains(&PathBuf::from("/restored/worktree-b"))); assert!(paths.contains(&PathBuf::from("/other/unrelated"))); } #[gpui::test] async fn test_complete_worktree_restore_preserves_unmatched_paths(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let original_paths = PathList::new(&[Path::new("/projects/worktree-a"), Path::new("/other/path")]); let meta = make_metadata("session-partial", "Partial", Utc::now(), original_paths); store.update(cx, |store, cx| { store.save_manually(meta, cx); }); let replacements = vec![ ( PathBuf::from("/projects/worktree-a"), PathBuf::from("/new/worktree-a"), ), ( PathBuf::from("/nonexistent/path"), PathBuf::from("/should/not/appear"), ), ]; store.update(cx, |store, cx| { store.complete_worktree_restore( &acp::SessionId::new("session-partial"), &replacements, cx, ); }); let entry = store.read_with(cx, |store, _cx| { store .entry(&acp::SessionId::new("session-partial")) .cloned() }); let entry = entry.unwrap(); let paths = entry.folder_paths().paths(); assert_eq!(paths.len(), 2); assert!(paths.contains(&PathBuf::from("/new/worktree-a"))); assert!(paths.contains(&PathBuf::from("/other/path"))); assert!(!paths.contains(&PathBuf::from("/should/not/appear"))); } #[gpui::test] async fn test_update_restored_worktree_paths_multiple(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let original_paths = PathList::new(&[ Path::new("/projects/worktree-a"), Path::new("/projects/worktree-b"), Path::new("/other/unrelated"), ]); let meta = make_metadata("session-multi", "Multi Thread", Utc::now(), original_paths); store.update(cx, |store, cx| { store.save_manually(meta, cx); }); let replacements = vec![ ( PathBuf::from("/projects/worktree-a"), PathBuf::from("/restored/worktree-a"), ), ( PathBuf::from("/projects/worktree-b"), PathBuf::from("/restored/worktree-b"), ), ]; store.update(cx, |store, cx| { store.update_restored_worktree_paths( &acp::SessionId::new("session-multi"), &replacements, cx, ); }); let entry = store.read_with(cx, |store, _cx| { store.entry(&acp::SessionId::new("session-multi")).cloned() }); let entry = entry.unwrap(); let paths = entry.folder_paths().paths(); assert_eq!(paths.len(), 3); assert!(paths.contains(&PathBuf::from("/restored/worktree-a"))); assert!(paths.contains(&PathBuf::from("/restored/worktree-b"))); assert!(paths.contains(&PathBuf::from("/other/unrelated"))); } #[gpui::test] async fn test_update_restored_worktree_paths_preserves_unmatched(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let original_paths = PathList::new(&[Path::new("/projects/worktree-a"), Path::new("/other/path")]); let meta = make_metadata("session-partial", "Partial", Utc::now(), original_paths); store.update(cx, |store, cx| { store.save_manually(meta, cx); }); let replacements = vec![ ( PathBuf::from("/projects/worktree-a"), PathBuf::from("/new/worktree-a"), ), ( PathBuf::from("/nonexistent/path"), PathBuf::from("/should/not/appear"), ), ]; store.update(cx, |store, cx| { store.update_restored_worktree_paths( &acp::SessionId::new("session-partial"), &replacements, cx, ); }); let entry = store.read_with(cx, |store, _cx| { store .entry(&acp::SessionId::new("session-partial")) .cloned() }); let entry = entry.unwrap(); let paths = entry.folder_paths().paths(); assert_eq!(paths.len(), 2); assert!(paths.contains(&PathBuf::from("/new/worktree-a"))); assert!(paths.contains(&PathBuf::from("/other/path"))); assert!(!paths.contains(&PathBuf::from("/should/not/appear"))); } #[gpui::test] async fn test_multiple_archived_worktrees_per_thread(cx: &mut TestAppContext) { init_test(cx); let store = cx.update(|cx| ThreadMetadataStore::global(cx)); let id1 = store .read_with(cx, |store, cx| { store.create_archived_worktree( "/projects/worktree-a".to_string(), "/home/user/repo".to_string(), Some("branch-a".to_string()), "staged_a".to_string(), "unstaged_a".to_string(), "original_000".to_string(), cx, ) }) .await .unwrap(); let id2 = store .read_with(cx, |store, cx| { store.create_archived_worktree( "/projects/worktree-b".to_string(), "/home/user/repo".to_string(), Some("branch-b".to_string()), "staged_b".to_string(), "unstaged_b".to_string(), "original_000".to_string(), cx, ) }) .await .unwrap(); store .read_with(cx, |store, cx| { store.link_thread_to_archived_worktree("session-1".to_string(), id1, cx) }) .await .unwrap(); store .read_with(cx, |store, cx| { store.link_thread_to_archived_worktree("session-1".to_string(), id2, cx) }) .await .unwrap(); let worktrees = store .read_with(cx, |store, cx| { store.get_archived_worktrees_for_thread("session-1".to_string(), cx) }) .await .unwrap(); assert_eq!(worktrees.len(), 2); let paths: Vec<&Path> = worktrees .iter() .map(|w| w.worktree_path.as_path()) .collect(); assert!(paths.contains(&Path::new("/projects/worktree-a"))); assert!(paths.contains(&Path::new("/projects/worktree-b"))); } // ── ThreadWorktreePaths tests ────────────────────────────────────── /// Helper to build a `ThreadWorktreePaths` from (main, folder) pairs. fn make_worktree_paths(pairs: &[(&str, &str)]) -> ThreadWorktreePaths { let (mains, folders): (Vec<&Path>, Vec<&Path>) = pairs .iter() .map(|(m, f)| (Path::new(*m), Path::new(*f))) .unzip(); ThreadWorktreePaths::from_path_lists(PathList::new(&mains), PathList::new(&folders)) .unwrap() } #[test] fn test_thread_worktree_paths_full_add_then_remove_cycle() { // Full scenario from the issue: // 1. Start with linked worktree selectric → zed // 2. Add cloud // 3. Remove zed let mut paths = make_worktree_paths(&[("/projects/zed", "/worktrees/selectric/zed")]); // Step 2: add cloud paths.add_path(Path::new("/projects/cloud"), Path::new("/projects/cloud")); assert_eq!(paths.ordered_pairs().count(), 2); assert_eq!( paths.folder_path_list(), &PathList::new(&[ Path::new("/worktrees/selectric/zed"), Path::new("/projects/cloud"), ]) ); assert_eq!( paths.main_worktree_path_list(), &PathList::new(&[Path::new("/projects/zed"), Path::new("/projects/cloud"),]) ); // Step 3: remove zed paths.remove_main_path(Path::new("/projects/zed")); assert_eq!(paths.ordered_pairs().count(), 1); assert_eq!( paths.folder_path_list(), &PathList::new(&[Path::new("/projects/cloud")]) ); assert_eq!( paths.main_worktree_path_list(), &PathList::new(&[Path::new("/projects/cloud")]) ); } #[test] fn test_thread_worktree_paths_add_is_idempotent() { let mut paths = make_worktree_paths(&[("/projects/zed", "/projects/zed")]); paths.add_path(Path::new("/projects/zed"), Path::new("/projects/zed")); assert_eq!(paths.ordered_pairs().count(), 1); } #[test] fn test_thread_worktree_paths_remove_nonexistent_is_noop() { let mut paths = make_worktree_paths(&[("/projects/zed", "/worktrees/selectric/zed")]); paths.remove_main_path(Path::new("/projects/nonexistent")); assert_eq!(paths.ordered_pairs().count(), 1); } #[test] fn test_thread_worktree_paths_from_path_lists_preserves_association() { let folder = PathList::new(&[ Path::new("/worktrees/selectric/zed"), Path::new("/projects/cloud"), ]); let main = PathList::new(&[Path::new("/projects/zed"), Path::new("/projects/cloud")]); let paths = ThreadWorktreePaths::from_path_lists(main, folder).unwrap(); let pairs: Vec<_> = paths .ordered_pairs() .map(|(m, f)| (m.clone(), f.clone())) .collect(); assert_eq!(pairs.len(), 2); assert!(pairs.contains(&( PathBuf::from("/projects/zed"), PathBuf::from("/worktrees/selectric/zed") ))); assert!(pairs.contains(&( PathBuf::from("/projects/cloud"), PathBuf::from("/projects/cloud") ))); } #[test] fn test_thread_worktree_paths_main_deduplicates_linked_worktrees() { // Two linked worktrees of the same main repo: the main_worktree_path_list // deduplicates because PathList stores unique sorted paths, but // ordered_pairs still has both entries. let paths = make_worktree_paths(&[ ("/projects/zed", "/worktrees/selectric/zed"), ("/projects/zed", "/worktrees/feature/zed"), ]); // main_worktree_path_list has the duplicate main path twice // (PathList keeps all entries from its input) assert_eq!(paths.ordered_pairs().count(), 2); assert_eq!( paths.folder_path_list(), &PathList::new(&[ Path::new("/worktrees/selectric/zed"), Path::new("/worktrees/feature/zed"), ]) ); assert_eq!( paths.main_worktree_path_list(), &PathList::new(&[Path::new("/projects/zed"), Path::new("/projects/zed"),]) ); } #[test] fn test_thread_worktree_paths_mismatched_lengths_returns_error() { let folder = PathList::new(&[ Path::new("/worktrees/selectric/zed"), Path::new("/projects/cloud"), ]); let main = PathList::new(&[Path::new("/projects/zed")]); let result = ThreadWorktreePaths::from_path_lists(main, folder); assert!(result.is_err()); } }