remote: Fix not being able to cancel in connecting state (#46789)

Lukas Wirth created

Release Notes:

- Fixed being unable to cancel remote connections while still connecting
to the remote server

Change summary

crates/recent_projects/src/remote_connections.rs | 147 ++++++++++-------
crates/remote/src/remote_client.rs               |  32 ++-
crates/workspace/src/modal_layer.rs              |  28 ++
crates/workspace/src/workspace.rs                |   7 
4 files changed, 135 insertions(+), 79 deletions(-)

Detailed changes

crates/recent_projects/src/remote_connections.rs 🔗

@@ -8,7 +8,7 @@ use askpass::EncryptedPassword;
 use auto_update::AutoUpdater;
 use editor::Editor;
 use extension_host::ExtensionStore;
-use futures::channel::oneshot;
+use futures::{FutureExt as _, channel::oneshot, select};
 use gpui::{
     AnyWindowHandle, App, AsyncApp, DismissEvent, Entity, EventEmitter, Focusable, FontFeatures,
     ParentElement as _, PromptLevel, Render, SharedString, Task, TextStyleRefinement, WeakEntity,
@@ -142,6 +142,7 @@ pub struct RemoteConnectionPrompt {
 impl Drop for RemoteConnectionPrompt {
     fn drop(&mut self) {
         if let Some(cancel) = self.cancellation.take() {
+            log::debug!("cancelling remote connection");
             cancel.send(()).ok();
         }
     }
@@ -345,6 +346,7 @@ impl RemoteConnectionModal {
             .prompt
             .update(cx, |prompt, _cx| prompt.cancellation.take())
         {
+            log::debug!("cancelling remote connection");
             tx.send(()).ok();
         }
         self.finished(cx);
@@ -579,13 +581,13 @@ impl remote::RemoteClientDelegate for RemoteClientDelegate {
 
 impl RemoteClientDelegate {
     fn update_status(&self, status: Option<&str>, cx: &mut AsyncApp) {
-        self.window
-            .update(cx, |_, _, cx| {
-                self.ui.update(cx, |modal, cx| {
+        cx.update(|cx| {
+            self.ui
+                .update(cx, |modal, cx| {
                     modal.set_status(status.map(|s| s.to_string()), cx);
                 })
-            })
-            .ok();
+                .ok()
+        });
     }
 }
 
@@ -604,7 +606,7 @@ pub fn connect(
             .and_then(|pw| pw.try_into().ok()),
         _ => None,
     };
-    let (tx, rx) = oneshot::channel();
+    let (tx, mut rx) = oneshot::channel();
     ui.update(cx, |ui, _cx| ui.set_cancellation_tx(tx));
 
     let delegate = Arc::new(RemoteClientDelegate {
@@ -614,7 +616,12 @@ pub fn connect(
     });
 
     cx.spawn(async move |cx| {
-        let connection = remote::connect(connection_options, delegate.clone(), cx).await?;
+        let connection = remote::connect(connection_options, delegate.clone(), cx);
+        let connection = select! {
+            _ = rx => return Ok(None),
+            result = connection.fuse() => result,
+        }?;
+
         cx.update(|cx| remote::RemoteClient::new(unique_identifier, connection, rx, delegate, cx))
             .await
     })
@@ -663,12 +670,13 @@ pub async fn open_remote_project(
     };
 
     loop {
-        let (cancel_tx, cancel_rx) = oneshot::channel();
+        let (cancel_tx, mut cancel_rx) = oneshot::channel();
         let delegate = window.update(cx, {
             let paths = paths.clone();
             let connection_options = connection_options.clone();
             move |workspace, window, cx| {
                 window.activate_window();
+                workspace.hide_modal(window, cx);
                 workspace.toggle_modal(window, cx, |window, cx| {
                     RemoteConnectionModal::new(&connection_options, paths, window, cx)
                 });
@@ -702,52 +710,66 @@ pub async fn open_remote_project(
 
         let Some(delegate) = delegate else { break };
 
-        let remote_connection =
-            match remote::connect(connection_options.clone(), delegate.clone(), cx).await {
-                Ok(connection) => connection,
-                Err(e) => {
+        let connection = remote::connect(connection_options.clone(), delegate.clone(), cx);
+        let connection = select! {
+            _ = cancel_rx => {
+                window
+                    .update(cx, |workspace, _, cx| {
+                        if let Some(ui) = workspace.active_modal::<RemoteConnectionModal>(cx) {
+                            ui.update(cx, |modal, cx| modal.finished(cx))
+                        }
+                    })
+                    .ok();
+
+                break;
+            },
+            result = connection.fuse() => result,
+        };
+        let remote_connection = match connection {
+            Ok(connection) => connection,
+            Err(e) => {
+                window
+                    .update(cx, |workspace, _, cx| {
+                        if let Some(ui) = workspace.active_modal::<RemoteConnectionModal>(cx) {
+                            ui.update(cx, |modal, cx| modal.finished(cx))
+                        }
+                    })
+                    .ok();
+                log::error!("Failed to open project: {e:#}");
+                let response = window
+                    .update(cx, |_, window, cx| {
+                        window.prompt(
+                            PromptLevel::Critical,
+                            match connection_options {
+                                RemoteConnectionOptions::Ssh(_) => "Failed to connect over SSH",
+                                RemoteConnectionOptions::Wsl(_) => "Failed to connect to WSL",
+                                RemoteConnectionOptions::Docker(_) => {
+                                    "Failed to connect to Dev Container"
+                                }
+                                #[cfg(any(test, feature = "test-support"))]
+                                RemoteConnectionOptions::Mock(_) => {
+                                    "Failed to connect to mock server"
+                                }
+                            },
+                            Some(&format!("{e:#}")),
+                            &["Retry", "Cancel"],
+                            cx,
+                        )
+                    })?
+                    .await;
+
+                if response == Ok(0) {
+                    continue;
+                }
+
+                if created_new_window {
                     window
-                        .update(cx, |workspace, _, cx| {
-                            if let Some(ui) = workspace.active_modal::<RemoteConnectionModal>(cx) {
-                                ui.update(cx, |modal, cx| modal.finished(cx))
-                            }
-                        })
+                        .update(cx, |_, window, _| window.remove_window())
                         .ok();
-                    log::error!("Failed to open project: {e:#}");
-                    let response = window
-                        .update(cx, |_, window, cx| {
-                            window.prompt(
-                                PromptLevel::Critical,
-                                match connection_options {
-                                    RemoteConnectionOptions::Ssh(_) => "Failed to connect over SSH",
-                                    RemoteConnectionOptions::Wsl(_) => "Failed to connect to WSL",
-                                    RemoteConnectionOptions::Docker(_) => {
-                                        "Failed to connect to Dev Container"
-                                    }
-                                    #[cfg(any(test, feature = "test-support"))]
-                                    RemoteConnectionOptions::Mock(_) => {
-                                        "Failed to connect to mock server"
-                                    }
-                                },
-                                Some(&format!("{e:#}")),
-                                &["Retry", "Cancel"],
-                                cx,
-                            )
-                        })?
-                        .await;
-
-                    if response == Ok(0) {
-                        continue;
-                    }
-
-                    if created_new_window {
-                        window
-                            .update(cx, |_, window, _| window.remove_window())
-                            .ok();
-                    }
-                    break;
                 }
-            };
+                return Ok(());
+            }
+        };
 
         let (paths, paths_with_positions) =
             determine_paths_with_positions(&remote_connection, paths.clone()).await;
@@ -845,20 +867,19 @@ pub async fn open_remote_project(
             }
         }
 
-        window
-            .update(cx, |workspace, _, cx| {
-                if let Some(client) = workspace.project().read(cx).remote_client() {
-                    if let Some(extension_store) = ExtensionStore::try_global(cx) {
-                        extension_store
-                            .update(cx, |store, cx| store.register_remote_client(client, cx));
-                    }
-                }
-            })
-            .ok();
-
         break;
     }
 
+    window
+        .update(cx, |workspace, _, cx| {
+            if let Some(client) = workspace.project().read(cx).remote_client() {
+                if let Some(extension_store) = ExtensionStore::try_global(cx) {
+                    extension_store
+                        .update(cx, |store, cx| store.register_remote_client(client, cx));
+                }
+            }
+        })
+        .ok();
     // Already showed the error to the user
     Ok(())
 }

crates/remote/src/remote_client.rs 🔗

@@ -20,7 +20,7 @@ use futures::{
         mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},
         oneshot,
     },
-    future::{BoxFuture, Shared},
+    future::{BoxFuture, Shared, WeakShared},
     select, select_biased,
 };
 use gpui::{
@@ -1131,7 +1131,7 @@ impl RemoteClient {
 }
 
 enum ConnectionPoolEntry {
-    Connecting(Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>>),
+    Connecting(WeakShared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>>),
     Connected(Weak<dyn RemoteConnection>),
 }
 
@@ -1152,21 +1152,30 @@ impl ConnectionPool {
         let connection = self.connections.get(&opts);
         match connection {
             Some(ConnectionPoolEntry::Connecting(task)) => {
-                delegate.set_status(
-                    Some("Waiting for existing connection attempt"),
-                    &mut cx.to_async(),
-                );
-                return task.clone();
+                if let Some(task) = task.upgrade() {
+                    log::debug!("Connecting task is still alive");
+                    cx.spawn(async move |cx| {
+                        delegate.set_status(Some("Waiting for existing connection attempt"), cx)
+                    })
+                    .detach();
+                    return task;
+                }
+                log::debug!("Connecting task is dead, removing it and restarting a connection");
+                self.connections.remove(&opts);
             }
             Some(ConnectionPoolEntry::Connected(remote)) => {
                 if let Some(remote) = remote.upgrade()
                     && !remote.has_been_killed()
                 {
+                    log::debug!("Connection is still alive");
                     return Task::ready(Ok(remote)).shared();
                 }
+                log::debug!("Connection is dead, removing it and restarting a connection");
                 self.connections.remove(&opts);
             }
-            None => {}
+            None => {
+                log::debug!("No existing connection found, starting a new one");
+            }
         }
 
         let task = cx
@@ -1224,9 +1233,10 @@ impl ConnectionPool {
                 }
             })
             .shared();
-
-        self.connections
-            .insert(opts.clone(), ConnectionPoolEntry::Connecting(task.clone()));
+        if let Some(task) = task.downgrade() {
+            self.connections
+                .insert(opts.clone(), ConnectionPoolEntry::Connecting(task));
+        }
         task
     }
 }

crates/workspace/src/modal_layer.rs 🔗

@@ -83,15 +83,22 @@ impl ModalLayer {
         }
     }
 
+    /// Toggles a modal of type `V`. If a modal of the same type is currently active,
+    /// it will be hidden. If a different modal is active, it will be replaced with the new one.
+    /// If no modal is active, the new modal will be shown.
+    ///
+    /// If closing the current modal fails (e.g., due to `on_before_dismiss` returning
+    /// `DismissDecision::Dismiss(false)` or `DismissDecision::Pending`), the new modal
+    /// will not be shown.
     pub fn toggle_modal<V, B>(&mut self, window: &mut Window, cx: &mut Context<Self>, build_view: B)
     where
         V: ModalView,
         B: FnOnce(&mut Window, &mut Context<V>) -> V,
     {
         if let Some(active_modal) = &self.active_modal {
-            let is_close = active_modal.modal.view().downcast::<V>().is_ok();
+            let should_close = active_modal.modal.view().downcast::<V>().is_ok();
             let did_close = self.hide_modal(window, cx);
-            if is_close || !did_close {
+            if should_close || !did_close {
                 return;
             }
         }
@@ -100,6 +107,8 @@ impl ModalLayer {
         cx.emit(ModalOpenedEvent);
     }
 
+    /// Shows a modal and sets up subscriptions for dismiss events and focus tracking.
+    /// The modal is automatically focused after being shown.
     fn show_modal<V>(&mut self, new_modal: Entity<V>, window: &mut Window, cx: &mut Context<Self>)
     where
         V: ModalView,
@@ -130,6 +139,13 @@ impl ModalLayer {
         cx.notify();
     }
 
+    /// Attempts to hide the currently active modal.
+    ///
+    /// The modal's `on_before_dismiss` method is called to determine if dismissal should proceed.
+    /// If dismissal is allowed, the modal is removed and focus is restored to the previously
+    /// focused element.
+    ///
+    /// Returns `true` if the modal was successfully hidden, `false` otherwise.
     pub fn hide_modal(&mut self, window: &mut Window, cx: &mut Context<Self>) -> bool {
         let Some(active_modal) = self.active_modal.as_mut() else {
             self.dismiss_on_focus_lost = false;
@@ -137,9 +153,9 @@ impl ModalLayer {
         };
 
         match active_modal.modal.on_before_dismiss(window, cx) {
-            DismissDecision::Dismiss(dismiss) => {
-                self.dismiss_on_focus_lost = !dismiss;
-                if !dismiss {
+            DismissDecision::Dismiss(should_dismiss) => {
+                if !should_dismiss {
+                    self.dismiss_on_focus_lost = !should_dismiss;
                     return false;
                 }
             }
@@ -157,9 +173,11 @@ impl ModalLayer {
             }
             cx.notify();
         }
+        self.dismiss_on_focus_lost = false;
         true
     }
 
+    /// Returns the currently active modal if it is of type `V`.
     pub fn active_modal<V>(&self) -> Option<Entity<V>>
     where
         V: 'static,

crates/workspace/src/workspace.rs 🔗

@@ -6423,6 +6423,13 @@ impl Workspace {
         self.modal_layer.read(cx).active_modal()
     }
 
+    /// Toggles a modal of type `V`. If a modal of the same type is currently active,
+    /// it will be hidden. If a different modal is active, it will be replaced with the new one.
+    /// If no modal is active, the new modal will be shown.
+    ///
+    /// If closing the current modal fails (e.g., due to `on_before_dismiss` returning
+    /// `DismissDecision::Dismiss(false)` or `DismissDecision::Pending`), the new modal
+    /// will not be shown.
     pub fn toggle_modal<V: ModalView, B>(&mut self, window: &mut Window, cx: &mut App, build: B)
     where
         B: FnOnce(&mut Window, &mut Context<V>) -> V,