Fix worktree archive correctness, cleanup, and error handling

Richard Feldman created

- 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

Change summary

crates/agent_ui/src/conversation_view.rs       |   7 +
crates/agent_ui/src/thread_metadata_store.rs   |  49 +++++++++
crates/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(-)

Detailed changes

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();
     }
 }
 

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<anyhow::Result<()>> {
+        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<anyhow::Result<bool>> {
+        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<bool> {
+        self.select_row_bound::<i64, i64>(
+            "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 {

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<PersistOutcome> {
+pub async fn persist_worktree_state(root: &RootPlan, cx: &mut AsyncApp) -> Result<PersistOutcome> {
     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<acp::SessionId> = 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<Entity<Workspace>> {
     cx.windows()
         .into_iter()

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) {

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<thread_worktree_archive::RootPlan>,
-        folder_paths: PathList,
         workspace: Option<Entity<Workspace>>,
         window: WindowHandle<MultiWorkspace>,
         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()));
                     }