1 files changed, 95 insertions, 9 deletions

diff --git a/src/app/api/worktrees/deferred.rs b/src/app/api/worktrees/deferred.rs
index 110f00c9..9adbd70c 100644
--- a/src/app/api/worktrees/deferred.rs
+++ b/src/app/api/worktrees/deferred.rs
@@ -1,5 +1,5 @@
 use std::path::Path;
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use crate::api::schema::{
     EventData, EventEnvelope, EventKind, Request, ResponseResult, WorktreeCreateParams,
@@ -289,6 +289,14 @@ impl App {
             || self
                 .pending_api_worktree_creates
                 .contains_key(&checkout_key)
+            || self
+                .state
+                .pane_ids_for_workspace(ws_idx)
+                .iter()
+                .any(|pane_id| {
+                    self.pending_worktree_remove_runtime_restores
+                        .contains_key(pane_id)
+                })
         {
             Self::send_api_response(
                 respond_to,
@@ -301,9 +309,12 @@ impl App {
             return;
         }
 
-        if Self::should_shutdown_workspace_terminal_runtimes_for_worktree_remove(params.force) {
-            self.shutdown_workspace_terminal_runtimes_for_worktree_remove(ws_idx);
-        }
+        let shutdown_panes =
+            if Self::should_shutdown_workspace_terminal_runtimes_for_worktree_remove(params.force) {
+                self.shutdown_workspace_terminal_runtimes_for_worktree_remove(ws_idx)
+            } else {
+                Vec::new()
+            };
 
         let operation_id = self.next_api_worktree_operation_id();
         self.pending_api_worktree_removes
@@ -322,6 +333,7 @@ impl App {
             id,
             operation_id,
             checkout_key,
+            shutdown_panes,
             respond_to,
         };
         let repo_root = space.repo_root;
@@ -468,9 +480,9 @@ impl App {
     pub(crate) fn handle_api_worktree_remove_finished(
         &mut self,
         mut result: crate::events::WorktreeRemoveResult,
-    ) {
+    ) -> Vec<crate::app::actions::PaneStateUpdate> {
         let Some(api) = result.api_request.take() else {
-            return;
+            return Vec::new();
         };
         let operation_matches = self
             .pending_api_worktree_removes
@@ -489,7 +501,7 @@ impl App {
                     "worktree remove completed after the operation was superseded",
                 ),
             );
-            return;
+            return Vec::new();
         }
         self.pending_api_worktree_removes
             .remove(&result.workspace_id);
@@ -497,6 +509,11 @@ impl App {
             .remove(&api.checkout_key);
 
         if let Err(message) = result.result {
+            let pane_updates = self.restore_shutdown_worktree_panes(
+                &api.shutdown_panes,
+                api.operation_id,
+                &result.path,
+            );
             let code =
                 if !result.forced && crate::worktree::is_dirty_worktree_remove_error(&message) {
                     "dirty_worktree_requires_force"
@@ -504,7 +521,7 @@ impl App {
                     "worktree_remove_failed"
                 };
             Self::send_api_response(api.respond_to, encode_error(api.id, code, message));
-            return;
+            return pane_updates;
         }
 
         let mut workspace_id = result.workspace_id.clone();
@@ -546,6 +563,11 @@ impl App {
         } else if let Some(snapshot) = workspace_snapshot.as_ref() {
             workspace_id = snapshot.workspace_id.clone();
         }
+        let pane_updates = self.restore_shutdown_worktree_panes(
+            &api.shutdown_panes,
+            api.operation_id,
+            &result.path,
+        );
 
         let Some(worktree) = worktree else {
             Self::send_api_response(
@@ -556,7 +578,7 @@ impl App {
                     "removed worktree but lost worktree snapshot",
                 ),
             );
-            return;
+            return pane_updates;
         };
         self.emit_worktree_removed_event(
             workspace_id.clone(),
@@ -573,5 +595,69 @@ impl App {
             },
         );
         Self::send_api_response(api.respond_to, response);
+        pane_updates
+    }
+
+    fn restore_shutdown_worktree_panes(
+        &mut self,
+        shutdown_panes: &[crate::layout::PaneId],
+        operation_id: u64,
+        removed_checkout: &std::path::Path,
+    ) -> Vec<crate::app::actions::PaneStateUpdate> {
+        let mut pane_updates = Vec::new();
+        let removed_checkout = crate::worktree::canonical_or_original(removed_checkout);
+        for &pane_id in shutdown_panes {
+            let Some((ws_idx, terminal_id)) = self
+                .find_pane(pane_id)
+                .map(|(ws_idx, pane)| (ws_idx, pane.attached_terminal_id.clone()))
+            else {
+                self.pending_worktree_remove_runtime_exits.remove(&pane_id);
+                self.pending_worktree_remove_runtime_restores
+                    .remove(&pane_id);
+                continue;
+            };
+            let runtime_missing = self.terminal_runtimes.get(&terminal_id).is_none();
+            if runtime_missing {
+                let workspace = &self.state.workspaces[ws_idx];
+                let current_checkout = workspace
+                    .worktree_space()
+                    .map(|space| &space.checkout_path)
+                    .unwrap_or(&workspace.identity_cwd);
+                if crate::worktree::canonical_or_original(current_checkout) != removed_checkout {
+                    if let Some(terminal) = self.state.terminals.get_mut(&terminal_id) {
+                        terminal.cwd = current_checkout.clone();
+                    }
+                }
+                if self
+                    .pending_worktree_remove_runtime_exits
+                    .contains_key(&pane_id)
+                {
+                    if self
+                        .pending_worktree_remove_runtime_restores
+                        .insert(pane_id, operation_id)
+                        .is_none()
+                    {
+                        let event_tx = self.event_tx.clone();
+                        tokio::spawn(async move {
+                            tokio::time::sleep(Duration::from_secs(1)).await;
+                            let _ = event_tx
+                                .send(AppEvent::WorktreeRuntimeRestoreFailed {
+                                    pane_id,
+                                    operation_id,
+                                })
+                                .await;
+                        });
+                    }
+                } else {
+                    pane_updates.extend(self.publish_worktree_runtime_agent_release(pane_id));
+                    if !self.respawn_shell_for_launch_pane(pane_id, false) {
+                        self.pending_worktree_remove_runtime_restores
+                            .insert(pane_id, operation_id);
+                        self.queue_worktree_runtime_restore_failed(pane_id, operation_id);
+                    }
+                }
+            }
+        }
+        pane_updates
     }
 }