debugger: Fix connections over SSH (#32834)

Before this change, we would see "connection reset" when sending the
initialize
request over SSH in the case that the debug adapter was slow to boot.

(Although we'd have successfully created a connection to the local SSH
port,
trying to read/write from it would not work until the remote end of the
connection had been established)

Fixes  #32575

Release Notes:

- debugger: Fix connecting to a Python debugger over SSH

crates/dap/src/client.rs

@@ -8,8 +8,7 @@ use dap_types::{
     requests::Request,
 };
 use futures::channel::oneshot;
-use gpui::{AppContext, AsyncApp};
-use smol::channel::{Receiver, Sender};
+use gpui::AsyncApp;
 use std::{
     hash::Hash,
     sync::atomic::{AtomicU64, Ordering},
@@ -44,99 +43,56 @@ impl DebugAdapterClient {
         id: SessionId,
         binary: DebugAdapterBinary,
         message_handler: DapMessageHandler,
-        cx: AsyncApp,
+        cx: &mut AsyncApp,
     ) -> Result<Self> {
-        let ((server_rx, server_tx), transport_delegate) =
-            TransportDelegate::start(&binary, cx.clone()).await?;
+        let transport_delegate = TransportDelegate::start(&binary, cx).await?;
         let this = Self {
             id,
             binary,
             transport_delegate,
             sequence_count: AtomicU64::new(1),
         };
-        log::info!("Successfully connected to debug adapter");
-
-        let client_id = this.id;
-
-        // start handling events/reverse requests
-        cx.background_spawn(Self::handle_receive_messages(
-            client_id,
-            server_rx,
-            server_tx.clone(),
-            message_handler,
-        ))
-        .detach();
+        this.connect(message_handler, cx).await?;
 
         Ok(this)
     }
 
-    pub async fn reconnect(
+    pub fn should_reconnect_for_ssh(&self) -> bool {
+        self.transport_delegate.tcp_arguments().is_some()
+            && self.binary.command.as_deref() == Some("ssh")
+    }
+
+    pub async fn connect(
+        &self,
+        message_handler: DapMessageHandler,
+        cx: &mut AsyncApp,
+    ) -> Result<()> {
+        self.transport_delegate.connect(message_handler, cx).await
+    }
+
+    pub async fn create_child_connection(
         &self,
         session_id: SessionId,
         binary: DebugAdapterBinary,
         message_handler: DapMessageHandler,
-        cx: AsyncApp,
+        cx: &mut AsyncApp,
     ) -> Result<Self> {
-        let binary = match self.transport_delegate.transport() {
-            crate::transport::Transport::Tcp(tcp_transport) => DebugAdapterBinary {
+        let binary = if let Some(connection) = self.transport_delegate.tcp_arguments() {
+            DebugAdapterBinary {
                 command: None,
                 arguments: Default::default(),
                 envs: Default::default(),
                 cwd: Default::default(),
-                connection: Some(crate::adapters::TcpArguments {
-                    host: tcp_transport.host,
-                    port: tcp_transport.port,
-                    timeout: Some(tcp_transport.timeout),
-                }),
+                connection: Some(connection),
                 request_args: binary.request_args,
-            },
-            _ => self.binary.clone(),
+            }
+        } else {
+            self.binary.clone()
         };
 
         Self::start(session_id, binary, message_handler, cx).await
     }
 
-    async fn handle_receive_messages(
-        client_id: SessionId,
-        server_rx: Receiver<Message>,
-        client_tx: Sender<Message>,
-        mut message_handler: DapMessageHandler,
-    ) -> Result<()> {
-        let result = loop {
-            let message = match server_rx.recv().await {
-                Ok(message) => message,
-                Err(e) => break Err(e.into()),
-            };
-            match message {
-                Message::Event(ev) => {
-                    log::debug!("Client {} received event `{}`", client_id.0, &ev);
-
-                    message_handler(Message::Event(ev))
-                }
-                Message::Request(req) => {
-                    log::debug!(
-                        "Client {} received reverse request `{}`",
-                        client_id.0,
-                        &req.command
-                    );
-
-                    message_handler(Message::Request(req))
-                }
-                Message::Response(response) => {
-                    log::debug!("Received response after request timeout: {:#?}", response);
-                }
-            }
-
-            smol::future::yield_now().await;
-        };
-
-        drop(client_tx);
-
-        log::debug!("Handle receive messages dropped");
-
-        result
-    }
-
     /// Send a request to an adapter and get a response back
     /// Note: This function will block until a response is sent back from the adapter
     pub async fn request<R: Request>(&self, arguments: R::Arguments) -> Result<R::Response> {
@@ -152,8 +108,7 @@ impl DebugAdapterClient {
             arguments: Some(serialized_arguments),
         };
         self.transport_delegate
-            .add_pending_request(sequence_id, callback_tx)
-            .await;
+            .add_pending_request(sequence_id, callback_tx);
 
         log::debug!(
             "Client {} send `{}` request with sequence_id: {}",
@@ -230,8 +185,11 @@ impl DebugAdapterClient {
             + Send
             + FnMut(u64, R::Arguments) -> Result<R::Response, dap_types::ErrorResponse>,
     {
-        let transport = self.transport_delegate.transport().as_fake();
-        transport.on_request::<R, F>(handler);
+        self.transport_delegate
+            .transport
+            .lock()
+            .as_fake()
+            .on_request::<R, F>(handler);
     }
 
     #[cfg(any(test, feature = "test-support"))]
@@ -250,8 +208,11 @@ impl DebugAdapterClient {
     where
         F: 'static + Send + Fn(Response),
     {
-        let transport = self.transport_delegate.transport().as_fake();
-        transport.on_response::<R, F>(handler).await;
+        self.transport_delegate
+            .transport
+            .lock()
+            .as_fake()
+            .on_response::<R, F>(handler);
     }
 
     #[cfg(any(test, feature = "test-support"))]
@@ -308,7 +269,7 @@ mod tests {
                 },
             },
             Box::new(|_| panic!("Did not expect to hit this code path")),
-            cx.to_async(),
+            &mut cx.to_async(),
         )
         .await
         .unwrap();
@@ -390,7 +351,7 @@ mod tests {
                     );
                 }
             }),
-            cx.to_async(),
+            &mut cx.to_async(),
         )
         .await
         .unwrap();
@@ -448,7 +409,7 @@ mod tests {
                     );
                 }
             }),
-            cx.to_async(),
+            &mut cx.to_async(),
         )
         .await
         .unwrap();