Make closing the SSH modal equivalent to cancelling the SSH connection task (#19568)

TODO: 
- [x] Fix focus not being on the modal initially

Release Notes:

- N/A

---------

Co-authored-by: Trace <violet.white.batt@gmail.com>

crates/recent_projects/src/remote_servers.rs

@@ -547,7 +547,7 @@ impl RemoteServerProjects {
                         })
                         .ok();
 
-                    let Some(session) = session else {
+                    let Some(Some(session)) = session else {
                         workspace
                             .update(&mut cx, |workspace, cx| {
                                 let weak = cx.view().downgrade();
@@ -1071,7 +1071,7 @@ impl RemoteServerProjects {
                             );
 
                             cx.spawn(|mut cx| async move {
-                                if confirmation.await.ok() == Some(0) {
+                                if confirmation.await.ok() == Some(1) {
                                     remote_servers
                                         .update(&mut cx, |this, cx| {
                                             this.delete_ssh_server(index, cx);

crates/recent_projects/src/ssh_connections.rs

@@ -124,6 +124,7 @@ pub struct SshPrompt {
     nickname: Option<SharedString>,
     status_message: Option<SharedString>,
     prompt: Option<(View<Markdown>, oneshot::Sender<Result<String>>)>,
+    cancellation: Option<oneshot::Sender<()>>,
     editor: View<Editor>,
 }
 
@@ -142,12 +143,17 @@ impl SshPrompt {
         Self {
             connection_string,
             nickname,
-            status_message: None,
-            prompt: None,
             editor: cx.new_view(Editor::single_line),
+            status_message: None,
+            cancellation: None,
+            prompt: None,
         }
     }
 
+    pub fn set_cancellation_tx(&mut self, tx: oneshot::Sender<()>) {
+        self.cancellation = Some(tx);
+    }
+
     pub fn set_prompt(
         &mut self,
         prompt: String,
@@ -270,7 +276,13 @@ impl SshConnectionModal {
     }
 
     fn dismiss(&mut self, _: &menu::Cancel, cx: &mut ViewContext<Self>) {
-        cx.emit(DismissEvent);
+        if let Some(tx) = self
+            .prompt
+            .update(cx, |prompt, _cx| prompt.cancellation.take())
+        {
+            tx.send(()).ok();
+        }
+        self.finished(cx);
     }
 }
 
@@ -322,6 +334,7 @@ impl Render for SshConnectionModal {
             .w(rems(34.))
             .border_1()
             .border_color(theme.colors().border)
+            .key_context("SshConnectionModal")
             .track_focus(&self.focus_handle(cx))
             .on_action(cx.listener(Self::dismiss))
             .on_action(cx.listener(Self::confirm))
@@ -486,7 +499,11 @@ impl SshClientDelegate {
         use smol::process::{Command, Stdio};
 
         async fn run_cmd(command: &mut Command) -> Result<()> {
-            let output = command.stderr(Stdio::inherit()).output().await?;
+            let output = command
+                .kill_on_drop(true)
+                .stderr(Stdio::inherit())
+                .output()
+                .await?;
             if !output.status.success() {
                 Err(anyhow::anyhow!("failed to run command: {:?}", command))?;
             }
@@ -585,13 +602,16 @@ pub fn connect_over_ssh(
     connection_options: SshConnectionOptions,
     ui: View<SshPrompt>,
     cx: &mut WindowContext,
-) -> Task<Result<Model<SshRemoteClient>>> {
+) -> Task<Result<Option<Model<SshRemoteClient>>>> {
     let window = cx.window_handle();
     let known_password = connection_options.password.clone();
+    let (tx, rx) = oneshot::channel();
+    ui.update(cx, |ui, _cx| ui.set_cancellation_tx(tx));
 
     remote::SshRemoteClient::new(
         unique_identifier,
         connection_options,
+        rx,
         Arc::new(SshClientDelegate {
             window,
             ui,
@@ -628,6 +648,7 @@ pub async fn open_ssh_project(
     };
 
     loop {
+        let (cancel_tx, cancel_rx) = oneshot::channel();
         let delegate = window.update(cx, {
             let connection_options = connection_options.clone();
             let nickname = nickname.clone();
@@ -636,26 +657,33 @@ pub async fn open_ssh_project(
                 workspace.toggle_modal(cx, |cx| {
                     SshConnectionModal::new(&connection_options, nickname.clone(), cx)
                 });
+
                 let ui = workspace
-                    .active_modal::<SshConnectionModal>(cx)
-                    .unwrap()
+                    .active_modal::<SshConnectionModal>(cx)?
                     .read(cx)
                     .prompt
                     .clone();
 
-                Arc::new(SshClientDelegate {
+                ui.update(cx, |ui, _cx| {
+                    ui.set_cancellation_tx(cancel_tx);
+                });
+
+                Some(Arc::new(SshClientDelegate {
                     window: cx.window_handle(),
                     ui,
                     known_password: connection_options.password.clone(),
-                })
+                }))
             }
         })?;
 
+        let Some(delegate) = delegate else { break };
+
         let did_open_ssh_project = cx
             .update(|cx| {
                 workspace::open_ssh_project(
                     window,
                     connection_options.clone(),
+                    cancel_rx,
                     delegate.clone(),
                     app_state.clone(),
                     paths.clone(),

crates/remote/src/ssh_session.rs

@@ -14,7 +14,7 @@ use futures::{
         oneshot,
     },
     future::BoxFuture,
-    select_biased, AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,
+    select, select_biased, AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,
 };
 use gpui::{
     AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, SemanticVersion, Task,
@@ -455,54 +455,65 @@ impl SshRemoteClient {
     pub fn new(
         unique_identifier: String,
         connection_options: SshConnectionOptions,
+        cancellation: oneshot::Receiver<()>,
         delegate: Arc<dyn SshClientDelegate>,
         cx: &AppContext,
-    ) -> Task<Result<Model<Self>>> {
+    ) -> Task<Result<Option<Model<Self>>>> {
         cx.spawn(|mut cx| async move {
-            let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
-            let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
-            let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
+            let success = Box::pin(async move {
+                let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
+                let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
+                let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
 
-            let client =
-                cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "client"))?;
-            let this = cx.new_model(|_| Self {
-                client: client.clone(),
-                unique_identifier: unique_identifier.clone(),
-                connection_options: connection_options.clone(),
-                state: Arc::new(Mutex::new(Some(State::Connecting))),
-            })?;
+                let client =
+                    cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "client"))?;
+                let this = cx.new_model(|_| Self {
+                    client: client.clone(),
+                    unique_identifier: unique_identifier.clone(),
+                    connection_options: connection_options.clone(),
+                    state: Arc::new(Mutex::new(Some(State::Connecting))),
+                })?;
 
-            let (ssh_connection, io_task) = Self::establish_connection(
-                unique_identifier,
-                false,
-                connection_options,
-                incoming_tx,
-                outgoing_rx,
-                connection_activity_tx,
-                delegate.clone(),
-                &mut cx,
-            )
-            .await?;
+                let (ssh_connection, io_task) = Self::establish_connection(
+                    unique_identifier,
+                    false,
+                    connection_options,
+                    incoming_tx,
+                    outgoing_rx,
+                    connection_activity_tx,
+                    delegate.clone(),
+                    &mut cx,
+                )
+                .await?;
 
-            let multiplex_task = Self::monitor(this.downgrade(), io_task, &cx);
+                let multiplex_task = Self::monitor(this.downgrade(), io_task, &cx);
 
-            if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
-                log::error!("failed to establish connection: {}", error);
-                return Err(error);
+                if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
+                    log::error!("failed to establish connection: {}", error);
+                    return Err(error);
+                }
+
+                let heartbeat_task =
+                    Self::heartbeat(this.downgrade(), connection_activity_rx, &mut cx);
+
+                this.update(&mut cx, |this, _| {
+                    *this.state.lock() = Some(State::Connected {
+                        ssh_connection,
+                        delegate,
+                        multiplex_task,
+                        heartbeat_task,
+                    });
+                })?;
+
+                Ok(Some(this))
+            });
+
+            select! {
+                _ = cancellation.fuse() => {
+                    Ok(None)
+                }
+                result = success.fuse() =>  result
             }
-
-            let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, &mut cx);
-
-            this.update(&mut cx, |this, _| {
-                *this.state.lock() = Some(State::Connected {
-                    ssh_connection,
-                    delegate,
-                    multiplex_task,
-                    heartbeat_task,
-                });
-            })?;
-
-            Ok(this)
         })
     }
 
@@ -1128,6 +1139,7 @@ impl SshRemoteClient {
 
     #[cfg(any(test, feature = "test-support"))]
     pub async fn fake_client(port: u16, client_cx: &mut gpui::TestAppContext) -> Model<Self> {
+        let (_tx, rx) = oneshot::channel();
         client_cx
             .update(|cx| {
                 Self::new(
@@ -1137,12 +1149,14 @@ impl SshRemoteClient {
                         port: Some(port),
                         ..Default::default()
                     },
+                    rx,
                     Arc::new(fake::Delegate),
                     cx,
                 )
             })
             .await
             .unwrap()
+            .unwrap()
     }
 }
 

crates/workspace/src/workspace.rs

@@ -5534,6 +5534,7 @@ pub fn join_hosted_project(
 pub fn open_ssh_project(
     window: WindowHandle<Workspace>,
     connection_options: SshConnectionOptions,
+    cancel_rx: oneshot::Receiver<()>,
     delegate: Arc<dyn SshClientDelegate>,
     app_state: Arc<AppState>,
     paths: Vec<PathBuf>,
@@ -5555,11 +5556,21 @@ pub fn open_ssh_project(
             workspace_id.0
         );
 
-        let session = cx
+        let session = match cx
             .update(|cx| {
-                remote::SshRemoteClient::new(unique_identifier, connection_options, delegate, cx)
+                remote::SshRemoteClient::new(
+                    unique_identifier,
+                    connection_options,
+                    cancel_rx,
+                    delegate,
+                    cx,
+                )
             })?
-            .await?;
+            .await?
+        {
+            Some(result) => result,
+            None => return Ok(()),
+        };
 
         let project = cx.update(|cx| {
             project::Project::ssh(

crates/zed/src/zed.rs

@@ -274,7 +274,6 @@ pub fn initialize_workspace(
                 workspace.add_panel(channels_panel, cx);
                 workspace.add_panel(chat_panel, cx);
                 workspace.add_panel(notification_panel, cx);
-                cx.focus_self();
             })
         })
         .detach();