fix: focus parent after worktree removal (#1075)

refs #1004

Co-authored-by: akbash-bot <akbash-bot@users.noreply.github.com>

2 files changed, 88 insertions, 4 deletions

diff --git a/src/app/api/worktrees/deferred.rs b/src/app/api/worktrees/deferred.rs
index 516e9455..3cd5e9f1 100644
--- a/src/app/api/worktrees/deferred.rs
+++ b/src/app/api/worktrees/deferred.rs
@@ -548,8 +548,7 @@ impl App {
                         .cloned()
                         .map(|space| self.worktree_info_for_membership(&space, None));
                 }
-                self.state.selected = ws_idx;
-                self.state.close_selected_workspace();
+                self.close_removed_linked_worktree_workspace(ws_idx);
                 self.shutdown_detached_terminal_runtimes();
                 self.emit_event(EventEnvelope {
                     event: EventKind::WorkspaceClosed,
diff --git a/src/app/worktrees.rs b/src/app/worktrees.rs
index ecaa01a4..a61a78ab 100644
--- a/src/app/worktrees.rs
+++ b/src/app/worktrees.rs
@@ -916,8 +916,7 @@ impl App {
                             space.is_linked_worktree && space.checkout_path == result.path
                         });
                     if still_same_linked_worktree {
-                        self.state.selected = ws_idx;
-                        self.state.close_selected_workspace();
+                        self.close_removed_linked_worktree_workspace(ws_idx);
                         self.shutdown_detached_terminal_runtimes();
                         self.emit_event(crate::api::schema::EventEnvelope {
                             event: crate::api::schema::EventKind::WorkspaceClosed,
@@ -969,6 +968,31 @@ impl App {
         force || cfg!(windows)
     }
 
+    pub(crate) fn close_removed_linked_worktree_workspace(&mut self, ws_idx: usize) {
+        let parent_key = self
+            .state
+            .workspaces
+            .get(ws_idx)
+            .and_then(|workspace| workspace.worktree_space())
+            .filter(|space| space.is_linked_worktree)
+            .map(|space| space.key.clone());
+
+        self.state.selected = ws_idx;
+        self.state.close_selected_workspace();
+
+        let Some(parent_key) = parent_key else {
+            return;
+        };
+        let Some(parent_idx) = self.state.workspaces.iter().position(|workspace| {
+            workspace
+                .worktree_space()
+                .is_some_and(|space| !space.is_linked_worktree && space.key == parent_key)
+        }) else {
+            return;
+        };
+        self.state.switch_workspace(parent_idx);
+    }
+
     pub(crate) fn shutdown_workspace_terminal_runtimes_for_worktree_remove(
         &mut self,
         ws_idx: usize,
@@ -2144,6 +2168,67 @@ mod tests {
     }
 
     #[test]
+    fn worktree_remove_finished_focuses_parent_workspace() {
+        let mut app = app_for_worktree_tests();
+        let checkout = std::path::PathBuf::from("/repo/herdr-issue");
+        app.state.workspaces = vec![
+            crate::workspace::Workspace::test_new("parent"),
+            crate::workspace::Workspace::test_new("issue"),
+            crate::workspace::Workspace::test_new("sibling"),
+        ];
+        app.state.workspaces[0].worktree_space = Some(crate::workspace::WorktreeSpaceMembership {
+            key: "repo-key".into(),
+            label: "herdr".into(),
+            repo_root: "/repo/herdr".into(),
+            checkout_path: "/repo/herdr".into(),
+            is_linked_worktree: false,
+        });
+        app.state.workspaces[1].worktree_space = Some(crate::workspace::WorktreeSpaceMembership {
+            key: "repo-key".into(),
+            label: "herdr".into(),
+            repo_root: "/repo/herdr".into(),
+            checkout_path: checkout.clone(),
+            is_linked_worktree: true,
+        });
+        app.state.workspaces[2].worktree_space = Some(crate::workspace::WorktreeSpaceMembership {
+            key: "repo-key".into(),
+            label: "herdr".into(),
+            repo_root: "/repo/herdr".into(),
+            checkout_path: "/repo/herdr-sibling".into(),
+            is_linked_worktree: true,
+        });
+        let child_id = app.state.workspaces[1].id.clone();
+        let parent_id = app.state.workspaces[0].id.clone();
+        app.state.active = Some(1);
+        app.state.selected = 1;
+        app.state.worktree_remove = Some(WorktreeRemoveState {
+            workspace_id: child_id.clone(),
+            repo_root: std::path::PathBuf::from("/repo/herdr"),
+            path: checkout.clone(),
+            error: None,
+            removing: true,
+            force_confirmation: false,
+        });
+
+        app.handle_worktree_remove_finished(WorktreeRemoveResult {
+            workspace_id: child_id,
+            path: checkout,
+            workspace: None,
+            worktree: None,
+            forced: false,
+            api_request: None,
+            result: Ok(()),
+        });
+
+        assert_eq!(app.state.workspaces.len(), 2);
+        assert_eq!(app.state.active, Some(0));
+        assert_eq!(app.state.selected, 0);
+        assert_eq!(app.state.workspaces[0].id, parent_id);
+        assert_eq!(app.state.workspaces[1].display_name(), "sibling");
+        assert!(app.state.worktree_remove.is_none());
+    }
+
+    #[test]
     fn worktree_remove_finished_emits_removed_event_from_snapshot_after_workspace_closed() {
         let event_hub = crate::api::EventHub::default();
         let mut app = app_for_worktree_tests_with_event_hub(event_hub.clone());