Add cancel button and restore spinner for worktree restores

Richard Feldman created

- worktree_restore_tasks HashMap on Sidebar tracks in-flight restore
  operations keyed by SessionId
- Restore task is now stored (not detached) so it can be tracked and
  the pending state cleaned up when it completes
- cancel_worktree_restore drops the task from the map and notifies
- ThreadItem wired up with pending_worktree_restore bool and
  on_cancel_restore callback from the Sidebar's render_thread
- Action slot (archive/stop buttons) hidden while a restore is pending
- Reconciled with existing PR4 restore spinner in worktree labels

Change summary

crates/agent_ui/src/thread_worktree_archive.rs | 26 ++---------
crates/sidebar/src/sidebar.rs                  | 44 +++++++++++++++++--
crates/ui/src/components/ai/thread_item.rs     | 17 ++++++-
3 files changed, 59 insertions(+), 28 deletions(-)

Detailed changes

crates/agent_ui/src/thread_worktree_archive.rs 🔗

@@ -968,10 +968,10 @@ pub async fn restore_worktree_via_git(
 
     let commit_exists = main_repo
         .update(cx, |repo, _cx| {
-            repo.commit_exists(row.original_commit_hash.clone())
+            repo.resolve_commit(row.original_commit_hash.clone())
         })
         .await
-        .map_err(|_| anyhow!("commit_exists check was canceled"))?
+        .map_err(|_| anyhow!("resolve_commit was canceled"))?
         .context("failed to check if original commit exists")?;
 
     if !commit_exists {
@@ -993,9 +993,7 @@ pub async fn restore_worktree_via_git(
         // a `.git` file (worktrees have a `.git` file, not a directory).
         let dot_git_path = worktree_path.join(".git");
         let dot_git_metadata = app_state.fs.metadata(&dot_git_path).await?;
-        let is_git_worktree = dot_git_metadata
-            .as_ref()
-            .is_some_and(|meta| !meta.is_dir);
+        let is_git_worktree = dot_git_metadata.as_ref().is_some_and(|meta| !meta.is_dir);
 
         if is_git_worktree {
             // Already a git worktree — another thread on the same worktree
@@ -1047,11 +1045,7 @@ pub async fn restore_worktree_via_git(
 
     let soft_reset_ok = if mixed_reset_ok {
         let rx = wt_repo.update(cx, |repo, cx| {
-            repo.reset(
-                row.original_commit_hash.clone(),
-                ResetMode::Soft,
-                cx,
-            )
+            repo.reset(row.original_commit_hash.clone(), ResetMode::Soft, cx)
         });
         match rx.await {
             Ok(Ok(())) => true,
@@ -1077,11 +1071,7 @@ pub async fn restore_worktree_via_git(
             row.original_commit_hash
         );
         let rx = wt_repo.update(cx, |repo, cx| {
-            repo.reset(
-                row.original_commit_hash.clone(),
-                ResetMode::Mixed,
-                cx,
-            )
+            repo.reset(row.original_commit_hash.clone(), ResetMode::Mixed, cx)
         });
         match rx.await {
             Ok(Ok(())) => {}
@@ -1142,11 +1132,7 @@ pub async fn restore_worktree_via_git(
                     row.original_commit_hash
                 );
                 let rx = wt_repo.update(cx, |repo, cx| {
-                    repo.reset(
-                        row.original_commit_hash.clone(),
-                        ResetMode::Mixed,
-                        cx,
-                    )
+                    repo.reset(row.original_commit_hash.clone(), ResetMode::Mixed, cx)
                 });
                 let _ = rx.await;
                 // Delete the old branch and create fresh

crates/sidebar/src/sidebar.rs 🔗

@@ -18,8 +18,8 @@ use editor::Editor;
 use feature_flags::{AgentV2FeatureFlag, FeatureFlagViewExt as _};
 use gpui::{
     Action as _, AnyElement, App, ClickEvent, Context, Entity, FocusHandle, Focusable, KeyContext,
-    ListState, Pixels, Render, SharedString, WeakEntity, Window, WindowHandle, linear_color_stop,
-    linear_gradient, list, prelude::*, px,
+    ListState, Pixels, Render, SharedString, Task, WeakEntity, Window, WindowHandle,
+    linear_color_stop, linear_gradient, list, prelude::*, px,
 };
 use menu::{
     Cancel, Confirm, SelectChild, SelectFirst, SelectLast, SelectNext, SelectParent, SelectPrevious,
@@ -372,6 +372,7 @@ pub struct Sidebar {
     view: SidebarView,
     recent_projects_popover_handle: PopoverMenuHandle<SidebarRecentProjects>,
     project_header_menu_ix: Option<usize>,
+    worktree_restore_tasks: HashMap<acp::SessionId, Task<()>>,
     _subscriptions: Vec<gpui::Subscription>,
     pending_worktree_restores: HashSet<acp::SessionId>,
     _draft_observation: Option<gpui::Subscription>,
@@ -464,6 +465,7 @@ impl Sidebar {
             view: SidebarView::default(),
             recent_projects_popover_handle: PopoverMenuHandle::default(),
             project_header_menu_ix: None,
+            worktree_restore_tasks: HashMap::new(),
             _subscriptions: Vec::new(),
             pending_worktree_restores: HashSet::default(),
             _draft_observation: None,
@@ -2209,8 +2211,10 @@ impl Sidebar {
             .read(cx)
             .get_archived_worktrees_for_thread(session_id.0.to_string(), cx);
         let path_list = metadata.folder_paths.clone();
+        let cleanup_session_id = session_id.clone();
+        let insert_session_id = session_id.clone();
 
-        cx.spawn_in(window, async move |this, cx| {
+        let restore_task = cx.spawn_in(window, async move |this, cx| {
             let archived_worktrees = task.await?;
 
             if archived_worktrees.is_empty() {
@@ -2307,8 +2311,27 @@ impl Sidebar {
             }
 
             anyhow::Ok(())
-        })
-        .detach_and_log_err(cx);
+        });
+
+        let monitor_task = cx.spawn(async move |this, _cx| {
+            let result = restore_task.await;
+            if let Err(err) = &result {
+                log::error!("{err:#}");
+            }
+            let _ = this.update(_cx, |this, cx| {
+                this.worktree_restore_tasks.remove(&cleanup_session_id);
+                cx.notify();
+            });
+        });
+
+        self.worktree_restore_tasks
+            .insert(insert_session_id, monitor_task);
+        cx.notify();
+    }
+
+    fn cancel_worktree_restore(&mut self, session_id: &acp::SessionId, cx: &mut Context<Self>) {
+        self.worktree_restore_tasks.remove(session_id);
+        cx.notify();
     }
 
     fn expand_selected_entry(
@@ -2893,6 +2916,10 @@ impl Sidebar {
         let session_id_for_delete = thread.metadata.session_id.clone();
         let focus_handle = self.focus_handle.clone();
 
+        let is_restoring_worktree = self
+            .worktree_restore_tasks
+            .contains_key(&thread.metadata.session_id);
+
         let id = SharedString::from(format!("thread-entry-{}", ix));
 
         let color = cx.theme().colors();
@@ -2956,6 +2983,13 @@ impl Sidebar {
             .selected(is_selected)
             .focused(is_focused)
             .hovered(is_hovered)
+            .pending_worktree_restore(is_restoring_worktree)
+            .when(is_restoring_worktree, |this| {
+                let session_id = session_id_for_delete.clone();
+                this.on_cancel_restore(cx.listener(move |this, _, _window, cx| {
+                    this.cancel_worktree_restore(&session_id, cx);
+                }))
+            })
             .on_hover(cx.listener(move |this, is_hovered: &bool, _window, cx| {
                 if *is_hovered {
                     this.hovered_thread_index = Some(ix);

crates/ui/src/components/ai/thread_item.rs 🔗

@@ -229,6 +229,8 @@ impl ThreadItem {
 
 impl RenderOnce for ThreadItem {
     fn render(mut self, _: &mut Window, cx: &mut App) -> impl IntoElement {
+        let pending_worktree_restore = self.pending_worktree_restore;
+        let action_slot = self.action_slot.take();
         let color = cx.theme().colors();
         let sidebar_base_bg = color
             .title_bar_background
@@ -311,7 +313,16 @@ impl RenderOnce for ThreadItem {
             (None, None)
         };
 
-        let icon = if self.status == AgentThreadStatus::Running {
+        let icon = if self.pending_worktree_restore {
+            icon_container()
+                .child(
+                    Icon::new(IconName::LoadCircle)
+                        .size(IconSize::Small)
+                        .color(Color::Muted)
+                        .with_rotate_animation(2),
+                )
+                .into_any_element()
+        } else if self.status == AgentThreadStatus::Running {
             icon_container()
                 .child(
                     Icon::new(IconName::LoadCircle)
@@ -413,8 +424,8 @@ impl RenderOnce for ThreadItem {
                             .when_some(self.tooltip, |this, tooltip| this.tooltip(tooltip)),
                     )
                     .child(gradient_overlay)
-                    .when(self.hovered, |this| {
-                        this.when_some(self.action_slot, |this, slot| {
+                    .when(!pending_worktree_restore && self.hovered, |this| {
+                        this.when_some(action_slot, |this, slot| {
                             let overlay = GradientFade::new(base_bg, hover_bg, hover_bg)
                                 .width(px(64.0))
                                 .right(px(6.))