From f6405f2079e33d8fdd8e7e5ef35101e6e2b6e728 Mon Sep 17 00:00:00 2001 From: Richard Feldman Date: Tue, 7 Apr 2026 22:30:08 -0400 Subject: [PATCH] Fix worktree archive correctness, cleanup, and error handling - Fix broken cancellation check in archive_worktree by using cancel_rx.is_closed() instead of cancel_rx.try_recv().is_ok() - Fix incomplete worktree linking for multi-worktree threads by matching any thread whose folder_paths contains the archiving root path - Clean up archived git refs and DB records when a thread is permanently deleted, preventing disk space leaks - Replace all silent let _ = on fallible operations with proper error logging via .log_err() or explicit log::error! - Ensure all worktree release tasks complete by logging errors and continuing instead of returning early on first failure --- crates/agent_ui/src/conversation_view.rs | 7 ++ crates/agent_ui/src/thread_metadata_store.rs | 49 ++++++++ .../agent_ui/src/thread_worktree_archive.rs | 107 ++++++++++++++---- crates/agent_ui/src/threads_archive_view.rs | 3 + crates/sidebar/src/sidebar.rs | 19 +--- 5 files changed, 151 insertions(+), 34 deletions(-) diff --git a/crates/agent_ui/src/conversation_view.rs b/crates/agent_ui/src/conversation_view.rs index de02fdc5d384f08c693df848b63f7ef20bdd28f8..1e9407e6b8d0624e3416071a3578bbe763f6fa9f 100644 --- a/crates/agent_ui/src/conversation_view.rs +++ b/crates/agent_ui/src/conversation_view.rs @@ -2584,6 +2584,13 @@ impl ConversationView { if let Some(store) = ThreadMetadataStore::try_global(cx) { store.update(cx, |store, cx| store.delete(session_id.clone(), cx)); } + + let session_id = session_id.clone(); + cx.spawn(async move |_this, cx| { + crate::thread_worktree_archive::cleanup_thread_archived_worktrees(&session_id, cx) + .await; + }) + .detach(); } } diff --git a/crates/agent_ui/src/thread_metadata_store.rs b/crates/agent_ui/src/thread_metadata_store.rs index 1eddd3b9a1eafc3bd63f93a8a3ef637626d18bb2..473c06131084b09ad827ef55280dde15f4b6b2dd 100644 --- a/crates/agent_ui/src/thread_metadata_store.rs +++ b/crates/agent_ui/src/thread_metadata_store.rs @@ -501,6 +501,30 @@ impl ThreadMetadataStore { 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 + }) + } + pub fn all_session_ids_for_path<'a>( &'a self, path_list: &PathList, @@ -921,6 +945,31 @@ impl ThreadMetadataDb { }) .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 { diff --git a/crates/agent_ui/src/thread_worktree_archive.rs b/crates/agent_ui/src/thread_worktree_archive.rs index 05663493f3ea5957720bdac60d8015f12c973d0a..bcf62ec323f48f46c71509508ebc7c7c3a184210 100644 --- a/crates/agent_ui/src/thread_worktree_archive.rs +++ b/crates/agent_ui/src/thread_worktree_archive.rs @@ -12,7 +12,7 @@ use project::{ git_store::{Repository, resolve_git_worktree_to_main_repo}, }; use util::ResultExt; -use workspace::{AppState, MultiWorkspace, PathList, Workspace}; +use workspace::{AppState, MultiWorkspace, Workspace}; use crate::thread_metadata_store::{ArchivedGitWorktree, ThreadMetadataStore}; @@ -144,7 +144,9 @@ async fn remove_root_after_worktree_removal( cx: &mut AsyncApp, ) -> Result<()> { for task in release_tasks { - task.await?; + if let Err(error) = task.await { + log::error!("Failed waiting for worktree release: {error:#}"); + } } let (repo, _temp_project) = find_or_create_repository(&root.main_repo_path, cx).await?; @@ -247,15 +249,11 @@ async fn rollback_root(root: &RootPlan, cx: &mut AsyncApp) { let task = affected.project.update(cx, |project, cx| { project.create_worktree(root.root_path.clone(), true, cx) }); - let _ = task.await; + task.await.log_err(); } } -pub async fn persist_worktree_state( - root: &RootPlan, - folder_paths: &PathList, - cx: &mut AsyncApp, -) -> Result { +pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Result { let worktree_repo = root .worktree_repo .clone() @@ -300,7 +298,7 @@ pub async fn persist_worktree_state( let rx = worktree_repo.update(cx, |repo, cx| { repo.reset("HEAD~1".to_string(), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); return Err(error); } }; @@ -315,7 +313,7 @@ pub async fn persist_worktree_state( let rx = worktree_repo.update(cx, |repo, cx| { repo.reset("HEAD~1".to_string(), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); return Err(error.context("failed to stage all files including untracked")); } @@ -341,7 +339,7 @@ pub async fn persist_worktree_state( let rx = worktree_repo.update(cx, |repo, cx| { repo.reset("HEAD~1".to_string(), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); return Err(error); } @@ -358,7 +356,7 @@ pub async fn persist_worktree_state( let rx = worktree_repo.update(cx, |repo, cx| { repo.reset(format!("{}~1", staged_commit_hash), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); return Err(error); } }; @@ -389,7 +387,7 @@ pub async fn persist_worktree_state( let rx = worktree_repo.update(cx, |repo, cx| { repo.reset(format!("{}~1", staged_commit_hash), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); return Err(error); } }; @@ -397,8 +395,15 @@ pub async fn persist_worktree_state( // Link all threads on this worktree to the archived record let session_ids: Vec = store.read_with(cx, |store, _cx| { store - .all_session_ids_for_path(folder_paths) - .cloned() + .entries() + .filter(|thread| { + thread + .folder_paths + .paths() + .iter() + .any(|p| p.as_path() == root.root_path) + }) + .map(|thread| thread.session_id.clone()) .collect() }); @@ -426,7 +431,7 @@ pub async fn persist_worktree_state( let rx = worktree_repo.update(cx, |repo, cx| { repo.reset(format!("{}~1", staged_commit_hash), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); return Err(error.context("failed to link thread to archived worktree")); } } @@ -474,7 +479,7 @@ pub async fn rollback_persist(outcome: &PersistOutcome, root: &RootPlan, cx: &mu cx, ) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); } // Delete the git ref on main repo @@ -483,7 +488,7 @@ pub async fn rollback_persist(outcome: &PersistOutcome, root: &RootPlan, cx: &mu { let ref_name = archived_worktree_ref_name(outcome.archived_worktree_id); let rx = main_repo.update(cx, |repo, _cx| repo.delete_ref(ref_name)); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); } // Delete the DB record @@ -646,12 +651,12 @@ pub async fn restore_worktree_via_git( let rx = wt_repo.update(cx, |repo, cx| { repo.reset(row.original_commit_hash.clone(), ResetMode::Mixed, cx) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); // Delete the old branch and create fresh let rx = wt_repo.update(cx, |repo, _cx| { repo.create_branch(branch_name.clone(), None) }); - let _ = rx.await; + rx.await.ok().and_then(|r| r.log_err()); } } else { // Branch doesn't exist or can't be switched to — create it. @@ -692,6 +697,68 @@ pub async fn cleanup_archived_worktree_record(row: &ArchivedGitWorktree, cx: &mu .log_err(); } +/// Cleans up all archived worktree data associated with a thread being deleted. +/// +/// This unlinks the thread from all its archived worktrees and, for any +/// archived worktree that is no longer referenced by any other thread, +/// deletes the git ref and DB records. +pub async fn cleanup_thread_archived_worktrees(session_id: &acp::SessionId, cx: &mut AsyncApp) { + let store = cx.update(|cx| ThreadMetadataStore::global(cx)); + + let archived_worktrees = store + .read_with(cx, |store, cx| { + store.get_archived_worktrees_for_thread(session_id.0.to_string(), cx) + }) + .await; + let archived_worktrees = match archived_worktrees { + Ok(rows) => rows, + Err(error) => { + log::error!( + "Failed to fetch archived worktrees for thread {}: {error:#}", + session_id.0 + ); + return; + } + }; + + if archived_worktrees.is_empty() { + return; + } + + if let Err(error) = store + .read_with(cx, |store, cx| { + store.unlink_thread_from_all_archived_worktrees(session_id.0.to_string(), cx) + }) + .await + { + log::error!( + "Failed to unlink thread {} from archived worktrees: {error:#}", + session_id.0 + ); + return; + } + + for row in &archived_worktrees { + let still_referenced = store + .read_with(cx, |store, cx| { + store.is_archived_worktree_referenced(row.id, cx) + }) + .await; + match still_referenced { + Ok(true) => {} + Ok(false) => { + cleanup_archived_worktree_record(row, cx).await; + } + Err(error) => { + log::error!( + "Failed to check if archived worktree {} is still referenced: {error:#}", + row.id + ); + } + } + } +} + pub fn all_open_workspaces(cx: &App) -> Vec> { cx.windows() .into_iter() diff --git a/crates/agent_ui/src/threads_archive_view.rs b/crates/agent_ui/src/threads_archive_view.rs index 7cb8410e5017438b0e8adde673887c13397d9abf..4e9d8b2e0883e6648d729f2cf39832dd6bca41a8 100644 --- a/crates/agent_ui/src/threads_archive_view.rs +++ b/crates/agent_ui/src/threads_archive_view.rs @@ -603,6 +603,9 @@ impl ThreadsArchiveView { .wait_for_connection() }); cx.spawn(async move |_this, cx| { + crate::thread_worktree_archive::cleanup_thread_archived_worktrees(&session_id, cx) + .await; + let state = task.await?; let task = cx.update(|cx| { if let Some(list) = state.connection.session_list(cx) { diff --git a/crates/sidebar/src/sidebar.rs b/crates/sidebar/src/sidebar.rs index 6944590060f778baecb7d87839b1f78b4d9f58df..692929d5801c813c5d11fb7867675722b6cb42ee 100644 --- a/crates/sidebar/src/sidebar.rs +++ b/crates/sidebar/src/sidebar.rs @@ -2493,20 +2493,13 @@ impl Sidebar { } let (cancel_tx, cancel_rx) = smol::channel::bounded(1); - let folder_paths = metadata.folder_paths.clone(); let current_workspace = current_workspace.clone(); let session_id = session_id.clone(); let task = cx.spawn(async move |_this, cx| { - let result = Self::archive_worktree( - roots, - folder_paths, - current_workspace, - window_handle, - cancel_rx, - cx, - ) - .await; + let result = + Self::archive_worktree(roots, current_workspace, window_handle, cancel_rx, cx) + .await; match result { Ok(ArchiveStatus::Success) => { @@ -2652,7 +2645,6 @@ impl Sidebar { async fn archive_worktree( roots: Vec, - folder_paths: PathList, workspace: Option>, window: WindowHandle, cancel_rx: smol::channel::Receiver<()>, @@ -2718,7 +2710,7 @@ impl Sidebar { for root in &roots { // Check for cancellation before each root - if cancel_rx.try_recv().is_ok() { + if cancel_rx.is_closed() { for (outcome, completed_root) in completed_persists.iter().rev() { thread_worktree_archive::rollback_persist(outcome, completed_root, cx).await; } @@ -2727,8 +2719,7 @@ impl Sidebar { // Persist worktree state (git WIP commits + DB record) if root.worktree_repo.is_some() { - match thread_worktree_archive::persist_worktree_state(root, &folder_paths, cx).await - { + match thread_worktree_archive::persist_worktree_state(root, cx).await { Ok(outcome) => { completed_persists.push((outcome, root.clone())); }