tenant.openagents/omega

1#[cfg(any(test, feature = "test-support"))]
2use crate::transport::mock::ConnectGuard;
3use crate::{
4    SshConnectionOptions,
5    protocol::MessageId,
6    proxy::ProxyLaunchError,
7    transport::{
8        docker::{DockerConnectionOptions, DockerExecConnection},
9        ssh::SshRemoteConnection,
10        wsl::{WslConnectionOptions, WslRemoteConnection},
11    },
12};
13use anyhow::{Context as _, Result, anyhow};
14use askpass::EncryptedPassword;
15use async_trait::async_trait;
16use collections::HashMap;
17use futures::{
18    Future, FutureExt as _, StreamExt as _,
19    channel::{
20        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},
21        oneshot,
22    },
23    future::{BoxFuture, Shared, WeakShared},
24    select, select_biased,
25    stream::BoxStream,
26};
27use gpui::{
28    App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,
29    EventEmitter, FutureExt, Global, Task, TaskExt, WeakEntity,
30};
31use parking_lot::Mutex;
32
33use release_channel::ReleaseChannel;
34use rpc::{
35    AnyProtoClient, ErrorExt, ProtoClient, ProtoMessageHandlerSet, RpcError,
36    proto::{self, Envelope, EnvelopedMessage, PeerId, RequestMessage, build_typed_envelope},
37};
38use semver::Version;
39use std::{
40    collections::VecDeque,
41    fmt,
42    ops::ControlFlow,
43    path::PathBuf,
44    sync::{
45        Arc, Weak,
46        atomic::{AtomicU32, AtomicU64, Ordering::SeqCst},
47    },
48    time::{Duration, Instant},
49};
50use util::{
51    ResultExt,
52    paths::{PathStyle, RemotePathBuf},
53};
54
55#[derive(Copy, Clone, Debug, PartialEq, Eq)]
56pub enum RemoteOs {
57    Linux,
58    MacOs,
59    Windows,
60}
61
62impl RemoteOs {
63    pub fn as_str(&self) -> &'static str {
64        match self {
65            RemoteOs::Linux => "linux",
66            RemoteOs::MacOs => "macos",
67            RemoteOs::Windows => "windows",
68        }
69    }
70
71    pub fn is_windows(&self) -> bool {
72        matches!(self, RemoteOs::Windows)
73    }
74
75    /// A human-readable OS name for telemetry. Matches `client::telemetry::os_name`
76    /// ignoring the compositor (as we run headless on remotes).
77    pub fn display_name(&self) -> &'static str {
78        match self {
79            RemoteOs::Linux => "Linux",
80            RemoteOs::MacOs => "macOS",
81            RemoteOs::Windows => "Windows",
82        }
83    }
84}
85
86impl std::fmt::Display for RemoteOs {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.write_str(self.as_str())
89    }
90}
91
92#[derive(Copy, Clone, Debug, PartialEq, Eq)]
93pub enum RemoteArch {
94    X86_64,
95    Aarch64,
96}
97
98impl RemoteArch {
99    pub fn as_str(&self) -> &'static str {
100        match self {
101            RemoteArch::X86_64 => "x86_64",
102            RemoteArch::Aarch64 => "aarch64",
103        }
104    }
105}
106
107impl std::fmt::Display for RemoteArch {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        f.write_str(self.as_str())
110    }
111}
112
113#[derive(Copy, Clone, Debug)]
114pub struct RemotePlatform {
115    pub os: RemoteOs,
116    pub arch: RemoteArch,
117}
118
119#[derive(Clone, Debug)]
120pub struct CommandTemplate {
121    pub program: String,
122    pub args: Vec<String>,
123    pub env: HashMap<String, String>,
124}
125
126/// Whether a command should be run with TTY allocation for interactive use.
127#[derive(Clone, Copy, Debug, PartialEq, Eq)]
128pub enum Interactive {
129    /// Allocate a pseudo-TTY for interactive terminal use.
130    Yes,
131    /// Do not allocate a TTY - for commands that communicate via piped stdio.
132    No,
133}
134
135pub trait RemoteClientDelegate: Send + Sync {
136    fn ask_password(
137        &self,
138        prompt: String,
139        tx: oneshot::Sender<EncryptedPassword>,
140        cx: &mut AsyncApp,
141    );
142    fn get_download_url(
143        &self,
144        platform: RemotePlatform,
145        release_channel: ReleaseChannel,
146        version: Option<Version>,
147        cx: &mut AsyncApp,
148    ) -> Task<Result<Option<String>>>;
149    fn download_server_binary_locally(
150        &self,
151        platform: RemotePlatform,
152        release_channel: ReleaseChannel,
153        version: Option<Version>,
154        cx: &mut AsyncApp,
155    ) -> Task<Result<PathBuf>>;
156    fn set_status(&self, status: Option<&str>, cx: &mut AsyncApp);
157}
158
159const MAX_MISSED_HEARTBEATS: usize = 5;
160const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
161const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
162const INITIAL_CONNECTION_TIMEOUT: Duration =
163    Duration::from_secs(if cfg!(debug_assertions) { 5 } else { 60 });
164
165pub const MAX_RECONNECT_ATTEMPTS: usize = 3;
166
167enum State {
168    Connecting,
169    Connected {
170        remote_connection: Arc<dyn RemoteConnection>,
171        delegate: Arc<dyn RemoteClientDelegate>,
172
173        multiplex_task: Task<Result<()>>,
174        heartbeat_task: Task<Result<()>>,
175    },
176    HeartbeatMissed {
177        missed_heartbeats: usize,
178
179        remote_connection: Arc<dyn RemoteConnection>,
180        delegate: Arc<dyn RemoteClientDelegate>,
181
182        multiplex_task: Task<Result<()>>,
183        heartbeat_task: Task<Result<()>>,
184    },
185    Reconnecting,
186    ReconnectFailed {
187        remote_connection: Arc<dyn RemoteConnection>,
188        delegate: Arc<dyn RemoteClientDelegate>,
189
190        error: anyhow::Error,
191        attempts: usize,
192    },
193    ReconnectExhausted,
194    ServerNotRunning,
195}
196
197impl fmt::Display for State {
198    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
199        match self {
200            Self::Connecting => write!(f, "connecting"),
201            Self::Connected { .. } => write!(f, "connected"),
202            Self::Reconnecting => write!(f, "reconnecting"),
203            Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
204            Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
205            Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
206            Self::ServerNotRunning { .. } => write!(f, "server not running"),
207        }
208    }
209}
210
211impl State {
212    fn remote_connection(&self) -> Option<Arc<dyn RemoteConnection>> {
213        match self {
214            Self::Connected {
215                remote_connection, ..
216            } => Some(remote_connection.clone()),
217            Self::HeartbeatMissed {
218                remote_connection, ..
219            } => Some(remote_connection.clone()),
220            Self::ReconnectFailed {
221                remote_connection, ..
222            } => Some(remote_connection.clone()),
223            _ => None,
224        }
225    }
226
227    fn can_reconnect(&self) -> bool {
228        match self {
229            Self::Connected { .. }
230            | Self::HeartbeatMissed { .. }
231            | Self::ReconnectFailed { .. } => true,
232            State::Connecting
233            | State::Reconnecting
234            | State::ReconnectExhausted
235            | State::ServerNotRunning => false,
236        }
237    }
238
239    fn is_reconnect_failed(&self) -> bool {
240        matches!(self, Self::ReconnectFailed { .. })
241    }
242
243    fn is_reconnect_exhausted(&self) -> bool {
244        matches!(self, Self::ReconnectExhausted { .. })
245    }
246
247    fn is_server_not_running(&self) -> bool {
248        matches!(self, Self::ServerNotRunning)
249    }
250
251    fn is_reconnecting(&self) -> bool {
252        matches!(self, Self::Reconnecting { .. })
253    }
254
255    fn heartbeat_recovered(self) -> Self {
256        match self {
257            Self::HeartbeatMissed {
258                remote_connection,
259                delegate,
260                multiplex_task,
261                heartbeat_task,
262                ..
263            } => Self::Connected {
264                remote_connection,
265                delegate,
266                multiplex_task,
267                heartbeat_task,
268            },
269            _ => self,
270        }
271    }
272
273    fn heartbeat_missed(self) -> Self {
274        match self {
275            Self::Connected {
276                remote_connection,
277                delegate,
278                multiplex_task,
279                heartbeat_task,
280            } => Self::HeartbeatMissed {
281                missed_heartbeats: 1,
282                remote_connection,
283                delegate,
284                multiplex_task,
285                heartbeat_task,
286            },
287            Self::HeartbeatMissed {
288                missed_heartbeats,
289                remote_connection,
290                delegate,
291                multiplex_task,
292                heartbeat_task,
293            } => Self::HeartbeatMissed {
294                missed_heartbeats: missed_heartbeats + 1,
295                remote_connection,
296                delegate,
297                multiplex_task,
298                heartbeat_task,
299            },
300            _ => self,
301        }
302    }
303}
304
305/// The state of the ssh connection.
306#[derive(Clone, Copy, Debug, PartialEq, Eq)]
307pub enum ConnectionState {
308    Connecting,
309    Connected,
310    HeartbeatMissed,
311    Reconnecting,
312    Disconnected,
313}
314
315impl From<&State> for ConnectionState {
316    fn from(value: &State) -> Self {
317        match value {
318            State::Connecting => Self::Connecting,
319            State::Connected { .. } => Self::Connected,
320            State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
321            State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
322            State::ReconnectExhausted => Self::Disconnected,
323            State::ServerNotRunning => Self::Disconnected,
324        }
325    }
326}
327
328pub struct RemoteClient {
329    client: Arc<ChannelClient>,
330    unique_identifier: String,
331    connection_options: RemoteConnectionOptions,
332    path_style: PathStyle,
333    platform: RemotePlatform,
334    os_version: Option<String>,
335    state: Option<State>,
336}
337
338#[derive(Debug)]
339pub enum RemoteClientEvent {
340    Disconnected { server_not_running: bool },
341}
342
343impl EventEmitter<RemoteClientEvent> for RemoteClient {}
344
345/// Identifies the socket on the remote server so that reconnects
346/// can re-join the same project.
347pub enum ConnectionIdentifier {
348    Setup(u64),
349    Workspace(i64),
350}
351
352static NEXT_ID: AtomicU64 = AtomicU64::new(1);
353
354impl ConnectionIdentifier {
355    pub fn setup() -> Self {
356        Self::Setup(NEXT_ID.fetch_add(1, SeqCst))
357    }
358
359    // This string gets used in a socket name, and so must be relatively short.
360    // The total length of:
361    //   /home/{username}/.local/share/zed/server_state/{name}/stdout.sock
362    // Must be less than about 100 characters
363    //   https://unix.stackexchange.com/questions/367008/why-is-socket-path-length-limited-to-a-hundred-chars
364    // So our strings should be at most 20 characters or so.
365    fn to_string(&self, cx: &App) -> String {
366        let identifier_prefix = match ReleaseChannel::global(cx) {
367            ReleaseChannel::Stable => "".to_string(),
368            release_channel => format!("{}-", release_channel.dev_name()),
369        };
370        match self {
371            Self::Setup(setup_id) => format!("{identifier_prefix}setup-{setup_id}"),
372            Self::Workspace(workspace_id) => {
373                format!("{identifier_prefix}workspace-{workspace_id}",)
374            }
375        }
376    }
377}
378
379pub async fn connect(
380    connection_options: RemoteConnectionOptions,
381    delegate: Arc<dyn RemoteClientDelegate>,
382    cx: &mut AsyncApp,
383) -> Result<Arc<dyn RemoteConnection>> {
384    cx.update(|cx| {
385        cx.update_default_global(|pool: &mut ConnectionPool, cx| {
386            pool.connect(connection_options.clone(), delegate.clone(), cx)
387        })
388    })
389    .await
390    .map_err(|e| e.cloned())
391}
392
393/// Returns `true` if the global [`ConnectionPool`] already has a live
394/// connection for the given options. Callers can use this to decide
395/// whether to show interactive UI (e.g., a password modal) before
396/// connecting.
397pub fn has_active_connection(opts: &RemoteConnectionOptions, cx: &App) -> bool {
398    cx.try_global::<ConnectionPool>().is_some_and(|pool| {
399        matches!(
400            pool.connections.get(opts),
401            Some(ConnectionPoolEntry::Connected(remote))
402                if remote.upgrade().is_some_and(|r| !r.has_been_killed())
403        )
404    })
405}
406
407impl RemoteClient {
408    pub fn new(
409        unique_identifier: ConnectionIdentifier,
410        remote_connection: Arc<dyn RemoteConnection>,
411        cancellation: oneshot::Receiver<()>,
412        delegate: Arc<dyn RemoteClientDelegate>,
413        cx: &mut App,
414    ) -> Task<Result<Option<Entity<Self>>>> {
415        let unique_identifier = unique_identifier.to_string(cx);
416        cx.spawn(async move |cx| {
417            let success = Box::pin(async move {
418                let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
419                let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
420                let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
421
422                let client = cx.update(|cx| {
423                    ChannelClient::new(
424                        incoming_rx,
425                        outgoing_tx,
426                        cx,
427                        "client",
428                        remote_connection.has_wsl_interop(),
429                    )
430                });
431
432                let path_style = remote_connection.path_style();
433                let platform = remote_connection.remote_platform();
434                let os_version = remote_connection.remote_os_version();
435                let connection_options = remote_connection.connection_options();
436                let connection_type = connection_options.connection_type();
437                let this = cx.new(|_| Self {
438                    client: client.clone(),
439                    unique_identifier: unique_identifier.clone(),
440                    connection_options,
441                    path_style,
442                    platform,
443                    os_version: os_version.clone(),
444                    state: Some(State::Connecting),
445                });
446
447                let io_task = remote_connection.start_proxy(
448                    unique_identifier,
449                    false,
450                    incoming_tx,
451                    outgoing_rx,
452                    connection_activity_tx,
453                    delegate.clone(),
454                    cx,
455                );
456
457                let ready = client
458                    .wait_for_remote_started()
459                    .with_timeout(INITIAL_CONNECTION_TIMEOUT, cx.background_executor())
460                    .await;
461                match ready {
462                    Ok(Some(_)) => {}
463                    Ok(None) => {
464                        let mut error = "remote client exited before becoming ready".to_owned();
465                        if let Some(status) = io_task.now_or_never() {
466                            match status {
467                                Ok(exit_code) => {
468                                    error.push_str(&format!(", exit_code={exit_code:?}"))
469                                }
470                                Err(e) => error.push_str(&format!(", error={e:?}")),
471                            }
472                        }
473                        let error = anyhow::anyhow!("{error}");
474                        log::error!("failed to establish connection: {}", error);
475                        return Err(error);
476                    }
477                    Err(_) => {
478                        let mut error = String::new();
479                        if let Some(status) = io_task.now_or_never() {
480                            error.push_str("Client exited with ");
481                            match status {
482                                Ok(exit_code) => {
483                                    error.push_str(&format!("exit_code {exit_code:?}"))
484                                }
485                                Err(e) => error.push_str(&format!("error {e:?}")),
486                            }
487                        } else {
488                            error.push_str("client did not become ready within the timeout");
489                        }
490                        let error = anyhow::anyhow!("{error}");
491                        log::error!("failed to establish connection: {error}");
492                        return Err(error);
493                    }
494                }
495                let multiplex_task = Self::monitor(this.downgrade(), io_task, cx);
496                if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
497                    log::error!("failed to establish connection: {}", error);
498                    return Err(error);
499                }
500
501                let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, cx);
502
503                this.update(cx, |this, _| {
504                    this.state = Some(State::Connected {
505                        remote_connection,
506                        delegate,
507                        multiplex_task,
508                        heartbeat_task,
509                    });
510                });
511
512                // Use the same `remote_*` property schema as the forwarded
513                // remote events (see `client::telemetry::report_remote_event`)
514                // so all remote-origin telemetry can be queried uniformly.
515                telemetry::event!(
516                    "Remote Connection Established",
517                    remote = true,
518                    remote_connection_type = connection_type,
519                    remote_os_name = platform.os.display_name(),
520                    remote_os_version = os_version,
521                    remote_architecture = platform.arch.as_str(),
522                );
523
524                Ok(Some(this))
525            });
526
527            select! {
528                _ = cancellation.fuse() => {
529                    Ok(None)
530                }
531                result = success.fuse() =>  result
532            }
533        })
534    }
535
536    pub fn proto_client_from_channels(
537        incoming_rx: mpsc::UnboundedReceiver<Envelope>,
538        outgoing_tx: mpsc::UnboundedSender<Envelope>,
539        cx: &App,
540        name: &'static str,
541        has_wsl_interop: bool,
542    ) -> AnyProtoClient {
543        ChannelClient::new(incoming_rx, outgoing_tx, cx, name, has_wsl_interop).into()
544    }
545
546    pub fn shutdown_processes<T: RequestMessage>(
547        &mut self,
548        shutdown_request: Option<T>,
549        executor: BackgroundExecutor,
550    ) -> Option<impl Future<Output = ()> + use<T>> {
551        let state = self.state.take()?;
552        log::info!("shutting down remote processes");
553
554        let State::Connected {
555            multiplex_task,
556            heartbeat_task,
557            remote_connection,
558            delegate,
559        } = state
560        else {
561            return None;
562        };
563
564        let client = self.client.clone();
565
566        Some(async move {
567            if let Some(shutdown_request) = shutdown_request {
568                client.send(shutdown_request).log_err();
569                // We wait 50ms instead of waiting for a response, because
570                // waiting for a response would require us to wait on the main thread
571                // which we want to avoid in an `on_app_quit` callback.
572                executor.timer(Duration::from_millis(50)).await;
573            }
574
575            // Drop `multiplex_task` because it owns our remote_connection_proxy_process, which is a
576            // child of master_process.
577            drop(multiplex_task);
578            // Now drop the rest of state, which kills master process.
579            drop(heartbeat_task);
580            drop(remote_connection);
581            drop(delegate);
582        })
583    }
584
585    fn reconnect(&mut self, cx: &mut Context<Self>) -> Result<()> {
586        let can_reconnect = self
587            .state
588            .as_ref()
589            .map(|state| state.can_reconnect())
590            .unwrap_or(false);
591        if !can_reconnect {
592            let state = if let Some(state) = self.state.as_ref() {
593                state.to_string()
594            } else {
595                "no state set".to_string()
596            };
597            log::info!(
598                "aborting reconnect, because not in state that allows reconnecting: {state}"
599            );
600            anyhow::bail!(
601                "aborting reconnect, because not in state that allows reconnecting: {state}"
602            );
603        }
604
605        let state = self.state.take().unwrap();
606        let (attempts, remote_connection, delegate) = match state {
607            State::Connected {
608                remote_connection,
609                delegate,
610                multiplex_task,
611                heartbeat_task,
612            }
613            | State::HeartbeatMissed {
614                remote_connection,
615                delegate,
616                multiplex_task,
617                heartbeat_task,
618                ..
619            } => {
620                drop(multiplex_task);
621                drop(heartbeat_task);
622                (0, remote_connection, delegate)
623            }
624            State::ReconnectFailed {
625                attempts,
626                remote_connection,
627                delegate,
628                ..
629            } => (attempts, remote_connection, delegate),
630            State::Connecting
631            | State::Reconnecting
632            | State::ReconnectExhausted
633            | State::ServerNotRunning => unreachable!(),
634        };
635
636        let attempts = attempts + 1;
637        if attempts > MAX_RECONNECT_ATTEMPTS {
638            log::error!(
639                "Failed to reconnect to after {} attempts, giving up",
640                MAX_RECONNECT_ATTEMPTS
641            );
642            self.set_state(State::ReconnectExhausted, cx);
643            return Ok(());
644        }
645
646        self.set_state(State::Reconnecting, cx);
647
648        log::info!(
649            "Trying to reconnect to remote server... Attempt {}",
650            attempts
651        );
652
653        let unique_identifier = self.unique_identifier.clone();
654        let client = self.client.clone();
655        let reconnect_task = cx.spawn(async move |this, cx| {
656            macro_rules! failed {
657                ($error:expr, $attempts:expr, $remote_connection:expr, $delegate:expr) => {
658                    delegate.set_status(Some(&format!("{error:#}", error = $error)), cx);
659                    return State::ReconnectFailed {
660                        error: anyhow!($error),
661                        attempts: $attempts,
662                        remote_connection: $remote_connection,
663                        delegate: $delegate,
664                    };
665                };
666            }
667
668            if let Err(error) = remote_connection
669                .kill()
670                .await
671                .context("Failed to kill remote_connection process")
672            {
673                failed!(error, attempts, remote_connection, delegate);
674            };
675
676            let connection_options = remote_connection.connection_options();
677
678            let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
679            let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
680            let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
681
682            let (remote_connection, io_task) = match async {
683                let remote_connection = cx
684                    .update_global(|pool: &mut ConnectionPool, cx| {
685                        pool.connect(connection_options, delegate.clone(), cx)
686                    })
687                    .await
688                    .map_err(|error| error.cloned())?;
689
690                let io_task = remote_connection.start_proxy(
691                    unique_identifier,
692                    true,
693                    incoming_tx,
694                    outgoing_rx,
695                    connection_activity_tx,
696                    delegate.clone(),
697                    cx,
698                );
699                anyhow::Ok((remote_connection, io_task))
700            }
701            .await
702            {
703                Ok((remote_connection, io_task)) => (remote_connection, io_task),
704                Err(error) => {
705                    failed!(error, attempts, remote_connection, delegate);
706                }
707            };
708
709            let multiplex_task = Self::monitor(this.clone(), io_task, cx);
710            client.reconnect(incoming_rx, outgoing_tx, cx);
711
712            if let Err(error) = client.resync(HEARTBEAT_TIMEOUT).await {
713                failed!(error, attempts, remote_connection, delegate);
714            };
715
716            State::Connected {
717                remote_connection,
718                delegate,
719                multiplex_task,
720                heartbeat_task: Self::heartbeat(this.clone(), connection_activity_rx, cx),
721            }
722        });
723
724        cx.spawn(async move |this, cx| {
725            let new_state = reconnect_task.await;
726            this.update(cx, |this, cx| {
727                this.try_set_state(cx, |old_state| {
728                    if old_state.is_reconnecting() {
729                        match &new_state {
730                            State::Connecting
731                            | State::Reconnecting
732                            | State::HeartbeatMissed { .. }
733                            | State::ServerNotRunning => {}
734                            State::Connected { .. } => {
735                                log::info!("Successfully reconnected");
736                            }
737                            State::ReconnectFailed {
738                                error, attempts, ..
739                            } => {
740                                log::error!(
741                                    "Reconnect attempt {} failed: {:?}. Starting new attempt...",
742                                    attempts,
743                                    error
744                                );
745                            }
746                            State::ReconnectExhausted => {
747                                log::error!("Reconnect attempt failed and all attempts exhausted");
748                            }
749                        }
750                        Some(new_state)
751                    } else {
752                        None
753                    }
754                });
755
756                if this.state_is(State::is_reconnect_failed) {
757                    this.reconnect(cx)
758                } else if this.state_is(State::is_reconnect_exhausted) {
759                    Ok(())
760                } else {
761                    log::debug!("State has transition from Reconnecting into new state while attempting reconnect.");
762                    Ok(())
763                }
764            })
765        })
766        .detach_and_log_err(cx);
767
768        Ok(())
769    }
770
771    fn heartbeat(
772        this: WeakEntity<Self>,
773        mut connection_activity_rx: mpsc::Receiver<()>,
774        cx: &mut AsyncApp,
775    ) -> Task<Result<()>> {
776        let Ok(client) = this.read_with(cx, |this, _| this.client.clone()) else {
777            return Task::ready(Err(anyhow!("remote_connectionRemoteClient lost")));
778        };
779
780        cx.spawn(async move |cx| {
781            let mut missed_heartbeats = 0;
782
783            let keepalive_timer = cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse();
784            futures::pin_mut!(keepalive_timer);
785
786            loop {
787                select_biased! {
788                    result = connection_activity_rx.next().fuse() => {
789                        if result.is_none() {
790                            log::warn!("remote heartbeat: connection activity channel has been dropped. stopping.");
791                            return Ok(());
792                        }
793
794                        if missed_heartbeats != 0 {
795                            missed_heartbeats = 0;
796                            let _ =this.update(cx, |this, cx| {
797                                this.handle_heartbeat_result(missed_heartbeats, cx)
798                            })?;
799                        }
800                    }
801                    _ = keepalive_timer => {
802                        log::debug!("Sending heartbeat to server...");
803
804                        let result = select_biased! {
805                            _ = connection_activity_rx.next().fuse() => {
806                                Ok(())
807                            }
808                            ping_result = client.ping(HEARTBEAT_TIMEOUT).fuse() => {
809                                ping_result
810                            }
811                        };
812
813                        if result.is_err() {
814                            missed_heartbeats += 1;
815                            log::warn!(
816                                "No heartbeat from server after {:?}. Missed heartbeat {} out of {}.",
817                                HEARTBEAT_TIMEOUT,
818                                missed_heartbeats,
819                                MAX_MISSED_HEARTBEATS
820                            );
821                        } else if missed_heartbeats != 0 {
822                            missed_heartbeats = 0;
823                        } else {
824                            continue;
825                        }
826
827                        let result = this.update(cx, |this, cx| {
828                            this.handle_heartbeat_result(missed_heartbeats, cx)
829                        })?;
830                        if result.is_break() {
831                            return Ok(());
832                        }
833                    }
834                }
835
836                keepalive_timer.set(cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse());
837            }
838        })
839    }
840
841    fn handle_heartbeat_result(
842        &mut self,
843        missed_heartbeats: usize,
844        cx: &mut Context<Self>,
845    ) -> ControlFlow<()> {
846        let state = self.state.take().unwrap();
847        let next_state = if missed_heartbeats > 0 {
848            state.heartbeat_missed()
849        } else {
850            state.heartbeat_recovered()
851        };
852
853        self.set_state(next_state, cx);
854
855        if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
856            log::error!(
857                "Missed last {} heartbeats. Reconnecting...",
858                missed_heartbeats
859            );
860
861            self.reconnect(cx)
862                .context("failed to start reconnect process after missing heartbeats")
863                .log_err();
864            ControlFlow::Break(())
865        } else {
866            ControlFlow::Continue(())
867        }
868    }
869
870    fn monitor(
871        this: WeakEntity<Self>,
872        io_task: Task<Result<i32>>,
873        cx: &AsyncApp,
874    ) -> Task<Result<()>> {
875        cx.spawn(async move |cx| {
876            let result = io_task.await;
877
878            match result {
879                Ok(exit_code) => {
880                    if let Some(error) = ProxyLaunchError::from_exit_code(exit_code) {
881                        match error {
882                            ProxyLaunchError::ServerNotRunning => {
883                                log::error!("failed to reconnect because server is not running");
884                                this.update(cx, |this, cx| {
885                                    this.set_state(State::ServerNotRunning, cx);
886                                })?;
887                            }
888                        }
889                    } else {
890                        log::error!("proxy process terminated unexpectedly: {exit_code}");
891                        this.update(cx, |this, cx| {
892                            this.reconnect(cx).ok();
893                        })?;
894                    }
895                }
896                Err(error) => {
897                    log::warn!(
898                        "remote io task died with error: {:?}. reconnecting...",
899                        error
900                    );
901                    this.update(cx, |this, cx| {
902                        this.reconnect(cx).ok();
903                    })?;
904                }
905            }
906
907            Ok(())
908        })
909    }
910
911    fn state_is(&self, check: impl FnOnce(&State) -> bool) -> bool {
912        self.state.as_ref().is_some_and(check)
913    }
914
915    fn try_set_state(&mut self, cx: &mut Context<Self>, map: impl FnOnce(&State) -> Option<State>) {
916        let new_state = self.state.as_ref().and_then(map);
917        if let Some(new_state) = new_state {
918            self.state.replace(new_state);
919            cx.notify();
920        }
921    }
922
923    fn set_state(&mut self, state: State, cx: &mut Context<Self>) {
924        log::info!("setting state to '{}'", &state);
925
926        let is_reconnect_exhausted = state.is_reconnect_exhausted();
927        let is_server_not_running = state.is_server_not_running();
928        self.state.replace(state);
929
930        if is_reconnect_exhausted || is_server_not_running {
931            cx.emit(RemoteClientEvent::Disconnected {
932                server_not_running: is_server_not_running,
933            });
934        }
935        cx.notify();
936    }
937
938    pub fn shell(&self) -> Option<String> {
939        Some(self.remote_connection()?.shell())
940    }
941
942    pub fn default_system_shell(&self) -> Option<String> {
943        Some(self.remote_connection()?.default_system_shell())
944    }
945
946    pub fn shares_network_interface(&self) -> bool {
947        self.remote_connection()
948            .map_or(false, |connection| connection.shares_network_interface())
949    }
950
951    pub fn has_wsl_interop(&self) -> bool {
952        self.remote_connection()
953            .map_or(false, |connection| connection.has_wsl_interop())
954    }
955
956    pub fn build_command(
957        &self,
958        program: Option<String>,
959        args: &[String],
960        env: &HashMap<String, String>,
961        working_dir: Option<String>,
962        port_forward: Option<(u16, String, u16)>,
963        interactive: Interactive,
964    ) -> Result<CommandTemplate> {
965        let Some(connection) = self.remote_connection() else {
966            return Err(anyhow!("no remote connection"));
967        };
968        connection.build_command(program, args, env, working_dir, port_forward, interactive)
969    }
970
971    pub fn build_forward_ports_command(
972        &self,
973        forwards: Vec<(u16, String, u16)>,
974    ) -> Result<CommandTemplate> {
975        let Some(connection) = self.remote_connection() else {
976            return Err(anyhow!("no remote connection"));
977        };
978        connection.build_forward_ports_command(forwards)
979    }
980
981    pub fn upload_directory(
982        &self,
983        src_path: PathBuf,
984        dest_path: RemotePathBuf,
985        cx: &App,
986    ) -> Task<Result<()>> {
987        let Some(connection) = self.remote_connection() else {
988            return Task::ready(Err(anyhow!("no remote connection")));
989        };
990        connection.upload_directory(src_path, dest_path, cx)
991    }
992
993    pub fn proto_client(&self) -> AnyProtoClient {
994        self.client.clone().into()
995    }
996
997    pub fn connection_options(&self) -> RemoteConnectionOptions {
998        self.connection_options.clone()
999    }
1000
1001    pub fn connection(&self) -> Option<Arc<dyn RemoteConnection>> {
1002        if let State::Connected {
1003            remote_connection, ..
1004        } = self.state.as_ref()?
1005        {
1006            Some(remote_connection.clone())
1007        } else {
1008            None
1009        }
1010    }
1011
1012    pub fn connection_state(&self) -> ConnectionState {
1013        self.state
1014            .as_ref()
1015            .map(ConnectionState::from)
1016            .unwrap_or(ConnectionState::Disconnected)
1017    }
1018
1019    pub fn is_disconnected(&self) -> bool {
1020        self.connection_state() == ConnectionState::Disconnected
1021    }
1022
1023    pub fn path_style(&self) -> PathStyle {
1024        self.path_style
1025    }
1026
1027    /// The platform (OS and architecture) of the remote host, detected during
1028    /// connection setup.
1029    pub fn remote_platform(&self) -> RemotePlatform {
1030        self.platform
1031    }
1032
1033    /// The OS version of the remote host (e.g. `"ubuntu 24.04"`), detected
1034    /// during connection setup. `None` if it could not be determined.
1035    pub fn remote_os_version(&self) -> Option<String> {
1036        self.os_version.clone()
1037    }
1038
1039    /// A stable identifier for the kind of remote connection (e.g. `"ssh"`,
1040    /// `"wsl"`, `"docker"`, `"podman"`).
1041    pub fn connection_type(&self) -> &'static str {
1042        self.connection_options.connection_type()
1043    }
1044
1045    /// Forcibly disconnects from the remote server by killing the underlying connection.
1046    /// This will trigger the reconnection logic if reconnection attempts remain.
1047    /// Useful for testing reconnection behavior in real environments.
1048    pub fn force_disconnect(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
1049        let Some(connection) = self.remote_connection() else {
1050            return Task::ready(Err(anyhow!("no active remote connection to disconnect")));
1051        };
1052
1053        log::info!("force_disconnect: killing remote connection");
1054
1055        cx.spawn(async move |_, _| {
1056            connection.kill().await?;
1057            Ok(())
1058        })
1059    }
1060
1061    /// Simulates a timeout by pausing heartbeat responses.
1062    /// This will cause heartbeat failures and eventually trigger reconnection
1063    /// after MAX_MISSED_HEARTBEATS are missed.
1064    /// Useful for testing timeout behavior in real environments.
1065    pub fn force_heartbeat_timeout(&mut self, attempts: usize, cx: &mut Context<Self>) {
1066        log::info!("force_heartbeat_timeout: triggering heartbeat failure state");
1067
1068        if let Some(State::Connected {
1069            remote_connection,
1070            delegate,
1071            multiplex_task,
1072            heartbeat_task,
1073        }) = self.state.take()
1074        {
1075            self.set_state(
1076                if attempts == 0 {
1077                    State::HeartbeatMissed {
1078                        missed_heartbeats: MAX_MISSED_HEARTBEATS,
1079                        remote_connection,
1080                        delegate,
1081                        multiplex_task,
1082                        heartbeat_task,
1083                    }
1084                } else {
1085                    State::ReconnectFailed {
1086                        remote_connection,
1087                        delegate,
1088                        error: anyhow!("forced heartbeat timeout"),
1089                        attempts,
1090                    }
1091                },
1092                cx,
1093            );
1094
1095            self.reconnect(cx)
1096                .context("failed to start reconnect after forced timeout")
1097                .log_err();
1098        } else {
1099            log::warn!("force_heartbeat_timeout: not in Connected state, ignoring");
1100        }
1101    }
1102
1103    #[cfg(any(test, feature = "test-support"))]
1104    pub fn force_server_not_running(&mut self, cx: &mut Context<Self>) {
1105        self.set_state(State::ServerNotRunning, cx);
1106    }
1107
1108    #[cfg(any(test, feature = "test-support"))]
1109    pub fn simulate_disconnect(&self, client_cx: &mut App) -> Task<()> {
1110        let opts = self.connection_options();
1111        client_cx.spawn(async move |cx| {
1112            let connection = cx.update_global(|c: &mut ConnectionPool, _| {
1113                if let Some(ConnectionPoolEntry::Connected(c)) = c.connections.get(&opts) {
1114                    if let Some(connection) = c.upgrade() {
1115                        connection
1116                    } else {
1117                        panic!("connection was dropped")
1118                    }
1119                } else {
1120                    panic!("missing test connection")
1121                }
1122            });
1123
1124            connection.simulate_disconnect(cx);
1125        })
1126    }
1127
1128    /// Creates a mock connection pair for testing.
1129    ///
1130    /// This is the recommended way to create mock remote connections for tests.
1131    /// It returns the `MockConnectionOptions` (which can be passed to create a
1132    /// `HeadlessProject`), an `AnyProtoClient` for the server side and a
1133    /// `ConnectGuard` for the client side which blocks the connection from
1134    /// being established until dropped.
1135    ///
1136    /// # Example
1137    /// ```ignore
1138    /// let (opts, server_session, connect_guard) = RemoteClient::fake_server(cx, server_cx);
1139    /// // Set up HeadlessProject with server_session...
1140    /// drop(connect_guard);
1141    /// let client = RemoteClient::fake_client(opts, cx).await;
1142    /// ```
1143    #[cfg(any(test, feature = "test-support"))]
1144    pub fn fake_server(
1145        client_cx: &mut gpui::TestAppContext,
1146        server_cx: &mut gpui::TestAppContext,
1147    ) -> (RemoteConnectionOptions, AnyProtoClient, ConnectGuard) {
1148        use crate::transport::mock::MockConnection;
1149        let (opts, server_client, connect_guard) = MockConnection::new(client_cx, server_cx);
1150        (opts.into(), server_client, connect_guard)
1151    }
1152
1153    /// Registers a new mock server for existing connection options.
1154    ///
1155    /// Use this to simulate reconnection: after forcing a disconnect, register
1156    /// a new server so the next `connect()` call succeeds.
1157    #[cfg(any(test, feature = "test-support"))]
1158    pub fn fake_server_with_opts(
1159        opts: &RemoteConnectionOptions,
1160        client_cx: &mut gpui::TestAppContext,
1161        server_cx: &mut gpui::TestAppContext,
1162    ) -> (AnyProtoClient, ConnectGuard) {
1163        use crate::transport::mock::MockConnection;
1164        let mock_opts = match opts {
1165            RemoteConnectionOptions::Mock(mock_opts) => mock_opts.clone(),
1166            _ => panic!("fake_server_with_opts requires Mock connection options"),
1167        };
1168        MockConnection::new_with_opts(mock_opts, client_cx, server_cx)
1169    }
1170
1171    /// Creates a `RemoteClient` connected to a mock server.
1172    ///
1173    /// Call `fake_server` first to get the connection options, set up the
1174    /// `HeadlessProject` with the server session, then call this method
1175    /// to create the client.
1176    #[cfg(any(test, feature = "test-support"))]
1177    pub async fn connect_mock(
1178        opts: RemoteConnectionOptions,
1179        client_cx: &mut gpui::TestAppContext,
1180    ) -> Entity<Self> {
1181        assert!(matches!(opts, RemoteConnectionOptions::Mock(..)));
1182        use crate::transport::mock::MockDelegate;
1183        let (_tx, rx) = oneshot::channel();
1184        let mut cx = client_cx.to_async();
1185        let connection = connect(opts, Arc::new(MockDelegate), &mut cx)
1186            .await
1187            .unwrap();
1188        client_cx
1189            .update(|cx| {
1190                Self::new(
1191                    ConnectionIdentifier::setup(),
1192                    connection,
1193                    rx,
1194                    Arc::new(MockDelegate),
1195                    cx,
1196                )
1197            })
1198            .await
1199            .unwrap()
1200            .unwrap()
1201    }
1202
1203    pub fn remote_connection(&self) -> Option<Arc<dyn RemoteConnection>> {
1204        self.state
1205            .as_ref()
1206            .and_then(|state| state.remote_connection())
1207    }
1208}
1209
1210enum ConnectionPoolEntry {
1211    Connecting(WeakShared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>>),
1212    Connected(Weak<dyn RemoteConnection>),
1213}
1214
1215#[derive(Default)]
1216struct ConnectionPool {
1217    connections: HashMap<RemoteConnectionOptions, ConnectionPoolEntry>,
1218}
1219
1220impl Global for ConnectionPool {}
1221
1222impl ConnectionPool {
1223    fn connect(
1224        &mut self,
1225        opts: RemoteConnectionOptions,
1226        delegate: Arc<dyn RemoteClientDelegate>,
1227        cx: &mut App,
1228    ) -> Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>> {
1229        let connection = self.connections.get(&opts);
1230        match connection {
1231            Some(ConnectionPoolEntry::Connecting(task)) => {
1232                if let Some(task) = task.upgrade() {
1233                    log::debug!("Connecting task is still alive");
1234                    cx.spawn(async move |cx| {
1235                        delegate.set_status(Some("Waiting for existing connection attempt"), cx)
1236                    })
1237                    .detach();
1238                    return task;
1239                }
1240                log::debug!("Connecting task is dead, removing it and restarting a connection");
1241                self.connections.remove(&opts);
1242            }
1243            Some(ConnectionPoolEntry::Connected(remote)) => {
1244                if let Some(remote) = remote.upgrade()
1245                    && !remote.has_been_killed()
1246                {
1247                    log::debug!("Connection is still alive");
1248                    return Task::ready(Ok(remote)).shared();
1249                }
1250                log::debug!("Connection is dead, removing it and restarting a connection");
1251                self.connections.remove(&opts);
1252            }
1253            None => {
1254                log::debug!("No existing connection found, starting a new one");
1255            }
1256        }
1257
1258        let task = cx
1259            .spawn({
1260                let opts = opts.clone();
1261                let delegate = delegate.clone();
1262                async move |cx| {
1263                    let connection = match opts.clone() {
1264                        RemoteConnectionOptions::Ssh(opts) => {
1265                            SshRemoteConnection::new(opts, delegate, cx)
1266                                .await
1267                                .map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>)
1268                        }
1269                        RemoteConnectionOptions::Wsl(opts) => {
1270                            WslRemoteConnection::new(opts, delegate, cx)
1271                                .await
1272                                .map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>)
1273                        }
1274                        RemoteConnectionOptions::Docker(opts) => {
1275                            DockerExecConnection::new(opts, delegate, cx)
1276                                .await
1277                                .map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>)
1278                        }
1279                        #[cfg(any(test, feature = "test-support"))]
1280                        RemoteConnectionOptions::Mock(opts) => match cx.update(|cx| {
1281                            cx.default_global::<crate::transport::mock::MockConnectionRegistry>()
1282                                .take(&opts)
1283                        }) {
1284                            Some(connection) => Ok(connection.await as Arc<dyn RemoteConnection>),
1285                            None => Err(anyhow!(
1286                                "Mock connection not found. Call MockConnection::new() first."
1287                            )),
1288                        },
1289                    };
1290
1291                    cx.update_global(|pool: &mut Self, _| {
1292                        debug_assert!(matches!(
1293                            pool.connections.get(&opts),
1294                            Some(ConnectionPoolEntry::Connecting(_))
1295                        ));
1296                        match connection {
1297                            Ok(connection) => {
1298                                pool.connections.insert(
1299                                    opts.clone(),
1300                                    ConnectionPoolEntry::Connected(Arc::downgrade(&connection)),
1301                                );
1302                                Ok(connection)
1303                            }
1304                            Err(error) => {
1305                                pool.connections.remove(&opts);
1306                                Err(Arc::new(error))
1307                            }
1308                        }
1309                    })
1310                }
1311            })
1312            .shared();
1313        if let Some(task) = task.downgrade() {
1314            self.connections
1315                .insert(opts.clone(), ConnectionPoolEntry::Connecting(task));
1316        }
1317        task
1318    }
1319}
1320
1321#[derive(Debug, Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
1322pub enum RemoteConnectionOptions {
1323    Ssh(SshConnectionOptions),
1324    Wsl(WslConnectionOptions),
1325    Docker(DockerConnectionOptions),
1326    #[cfg(any(test, feature = "test-support"))]
1327    Mock(crate::transport::mock::MockConnectionOptions),
1328}
1329
1330impl RemoteConnectionOptions {
1331    pub fn display_name(&self) -> String {
1332        match self {
1333            RemoteConnectionOptions::Ssh(opts) => opts
1334                .nickname
1335                .clone()
1336                .unwrap_or_else(|| opts.host.to_string()),
1337            RemoteConnectionOptions::Wsl(opts) => opts.distro_name.clone(),
1338            RemoteConnectionOptions::Docker(opts) => {
1339                if opts.use_podman {
1340                    format!("[podman] {}", opts.name)
1341                } else {
1342                    opts.name.clone()
1343                }
1344            }
1345            #[cfg(any(test, feature = "test-support"))]
1346            RemoteConnectionOptions::Mock(opts) => format!("mock-{}", opts.id),
1347        }
1348    }
1349
1350    /// A stable identifier for the kind of remote connection, suitable for
1351    /// telemetry (e.g. `"ssh"`, `"wsl"`, `"docker"`, `"podman"`).
1352    pub fn connection_type(&self) -> &'static str {
1353        match self {
1354            RemoteConnectionOptions::Ssh(_) => "ssh",
1355            RemoteConnectionOptions::Wsl(_) => "wsl",
1356            RemoteConnectionOptions::Docker(opts) => {
1357                if opts.use_podman {
1358                    "podman"
1359                } else {
1360                    "docker"
1361                }
1362            }
1363            #[cfg(any(test, feature = "test-support"))]
1364            RemoteConnectionOptions::Mock(_) => "mock",
1365        }
1366    }
1367}
1368
1369#[cfg(test)]
1370mod tests {
1371    use super::*;
1372    use gpui::TestAppContext;
1373    use rpc::{ErrorCodeExt, proto::ErrorCode};
1374
1375    #[test]
1376    fn test_ssh_display_name_prefers_nickname() {
1377        let options = RemoteConnectionOptions::Ssh(SshConnectionOptions {
1378            host: "1.2.3.4".into(),
1379            nickname: Some("My Cool Project".to_string()),
1380            ..Default::default()
1381        });
1382
1383        assert_eq!(options.display_name(), "My Cool Project");
1384    }
1385
1386    #[test]
1387    fn test_ssh_display_name_falls_back_to_host() {
1388        let options = RemoteConnectionOptions::Ssh(SshConnectionOptions {
1389            host: "1.2.3.4".into(),
1390            ..Default::default()
1391        });
1392
1393        assert_eq!(options.display_name(), "1.2.3.4");
1394    }
1395
1396    #[test]
1397    fn test_connection_type() {
1398        assert_eq!(
1399            RemoteConnectionOptions::Ssh(SshConnectionOptions::default()).connection_type(),
1400            "ssh"
1401        );
1402        assert_eq!(
1403            RemoteConnectionOptions::Wsl(WslConnectionOptions {
1404                distro_name: "Ubuntu".to_string(),
1405                user: None,
1406            })
1407            .connection_type(),
1408            "wsl"
1409        );
1410        assert_eq!(
1411            RemoteConnectionOptions::Docker(DockerConnectionOptions {
1412                use_podman: false,
1413                ..Default::default()
1414            })
1415            .connection_type(),
1416            "docker"
1417        );
1418        assert_eq!(
1419            RemoteConnectionOptions::Docker(DockerConnectionOptions {
1420                use_podman: true,
1421                ..Default::default()
1422            })
1423            .connection_type(),
1424            "podman"
1425        );
1426    }
1427
1428    #[gpui::test]
1429    async fn test_channel_client_request_stream_terminates_on_error(cx: &mut TestAppContext) {
1430        let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
1431        let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
1432
1433        let client =
1434            cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
1435
1436        // The client sends RemoteStarted on startup; drain the outgoing channel
1437        // so it doesn't block.
1438        let _drain_outgoing = cx
1439            .executor()
1440            .spawn(async move { while outgoing_rx.next().await.is_some() {} });
1441
1442        let mut stream = client
1443            .request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
1444            .await
1445            .unwrap();
1446
1447        let request_id = 0;
1448
1449        incoming_tx
1450            .unbounded_send(proto::Test { id: 1 }.into_envelope(100, Some(request_id), None))
1451            .unwrap();
1452
1453        let first = stream.next().await.unwrap().unwrap();
1454        assert_eq!(
1455            proto::Test::from_envelope(first).unwrap(),
1456            proto::Test { id: 1 }
1457        );
1458
1459        // Send an Error without a trailing EndStream. The Error alone should
1460        // terminate the stream.
1461        incoming_tx
1462            .unbounded_send(
1463                ErrorCode::Internal
1464                    .message("boom".to_string())
1465                    .to_proto()
1466                    .into_envelope(101, Some(request_id), None),
1467            )
1468            .unwrap();
1469
1470        let second = stream.next().await.unwrap();
1471        let error = second.unwrap_err();
1472        assert!(
1473            format!("{error}").contains("boom"),
1474            "expected error to surface server message, got: {error}"
1475        );
1476
1477        assert!(stream.next().await.is_none());
1478        assert_eq!(client.stream_response_channels.lock().len(), 0);
1479    }
1480
1481    #[gpui::test]
1482    async fn test_channel_client_dropping_stream_request_before_response_cleans_up_channel(
1483        cx: &mut TestAppContext,
1484    ) {
1485        let (_incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
1486        let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
1487
1488        let client =
1489            cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
1490
1491        let _drain_outgoing = cx
1492            .executor()
1493            .spawn(async move { while outgoing_rx.next().await.is_some() {} });
1494
1495        let stream = client
1496            .request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
1497            .await
1498            .unwrap();
1499
1500        assert_eq!(client.stream_response_channels.lock().len(), 1);
1501
1502        drop(stream);
1503        cx.run_until_parked();
1504
1505        assert_eq!(
1506            client.stream_response_channels.lock().len(),
1507            0,
1508            "dropping a stream before any responses arrive should remove response channel bookkeeping"
1509        );
1510    }
1511
1512    #[gpui::test]
1513    async fn test_channel_client_dropping_stream_request_before_completion(
1514        cx: &mut TestAppContext,
1515    ) {
1516        let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
1517        let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
1518
1519        let client =
1520            cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
1521
1522        let _drain_outgoing = cx
1523            .executor()
1524            .spawn(async move { while outgoing_rx.next().await.is_some() {} });
1525
1526        let mut stream = client
1527            .request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
1528            .await
1529            .unwrap();
1530
1531        let request_id = 0;
1532
1533        incoming_tx
1534            .unbounded_send(proto::Test { id: 1 }.into_envelope(100, Some(request_id), None))
1535            .unwrap();
1536        let _ = stream.next().await.unwrap().unwrap();
1537
1538        assert_eq!(client.stream_response_channels.lock().len(), 1);
1539
1540        drop(stream);
1541
1542        // Inject an orphaned non-terminal response. The read loop should detect
1543        // that the consumer has been dropped and clean up its bookkeeping (no
1544        // EndStream sent here on purpose, otherwise the cleanup would happen
1545        // via the terminal-response path and mask the bug under test).
1546        incoming_tx
1547            .unbounded_send(proto::Test { id: 2 }.into_envelope(101, Some(request_id), None))
1548            .unwrap();
1549
1550        cx.run_until_parked();
1551
1552        assert_eq!(
1553            client.stream_response_channels.lock().len(),
1554            0,
1555            "stream channel should be removed once the consumer has dropped the stream"
1556        );
1557    }
1558}
1559
1560impl From<SshConnectionOptions> for RemoteConnectionOptions {
1561    fn from(opts: SshConnectionOptions) -> Self {
1562        RemoteConnectionOptions::Ssh(opts)
1563    }
1564}
1565
1566impl From<WslConnectionOptions> for RemoteConnectionOptions {
1567    fn from(opts: WslConnectionOptions) -> Self {
1568        RemoteConnectionOptions::Wsl(opts)
1569    }
1570}
1571
1572#[cfg(any(test, feature = "test-support"))]
1573impl From<crate::transport::mock::MockConnectionOptions> for RemoteConnectionOptions {
1574    fn from(opts: crate::transport::mock::MockConnectionOptions) -> Self {
1575        RemoteConnectionOptions::Mock(opts)
1576    }
1577}
1578
1579#[cfg(target_os = "windows")]
1580/// Open a wsl path (\\wsl.localhost\<distro>\path)
1581#[derive(Debug, Clone, PartialEq, Eq, gpui::Action)]
1582#[action(namespace = workspace, no_json, no_register)]
1583pub struct OpenWslPath {
1584    pub distro: WslConnectionOptions,
1585    pub paths: Vec<PathBuf>,
1586}
1587
1588#[async_trait(?Send)]
1589pub trait RemoteConnection: Send + Sync {
1590    fn start_proxy(
1591        &self,
1592        unique_identifier: String,
1593        reconnect: bool,
1594        incoming_tx: UnboundedSender<Envelope>,
1595        outgoing_rx: UnboundedReceiver<Envelope>,
1596        connection_activity_tx: Sender<()>,
1597        delegate: Arc<dyn RemoteClientDelegate>,
1598        cx: &mut AsyncApp,
1599    ) -> Task<Result<i32>>;
1600    fn upload_directory(
1601        &self,
1602        src_path: PathBuf,
1603        dest_path: RemotePathBuf,
1604        cx: &App,
1605    ) -> Task<Result<()>>;
1606    async fn kill(&self) -> Result<()>;
1607    fn has_been_killed(&self) -> bool;
1608    fn shares_network_interface(&self) -> bool {
1609        false
1610    }
1611    fn build_command(
1612        &self,
1613        program: Option<String>,
1614        args: &[String],
1615        env: &HashMap<String, String>,
1616        working_dir: Option<String>,
1617        port_forward: Option<(u16, String, u16)>,
1618        interactive: Interactive,
1619    ) -> Result<CommandTemplate>;
1620    fn build_forward_ports_command(
1621        &self,
1622        forwards: Vec<(u16, String, u16)>,
1623    ) -> Result<CommandTemplate>;
1624    fn connection_options(&self) -> RemoteConnectionOptions;
1625    fn path_style(&self) -> PathStyle;
1626    /// The remote platform (OS and architecture), detected during connection setup.
1627    fn remote_platform(&self) -> RemotePlatform;
1628    /// The remote host's OS version (e.g. `"ubuntu 24.04"` or `"15.6.1"`),
1629    /// detected during connection setup. `None` if it could not be determined.
1630    fn remote_os_version(&self) -> Option<String>;
1631    fn shell(&self) -> String;
1632    fn default_system_shell(&self) -> String;
1633    fn has_wsl_interop(&self) -> bool;
1634
1635    #[cfg(any(test, feature = "test-support"))]
1636    fn simulate_disconnect(&self, _: &AsyncApp) {}
1637}
1638
1639type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
1640type StreamResponseChannels =
1641    Arc<Mutex<HashMap<MessageId, UnboundedSender<(Result<Envelope>, oneshot::Sender<()>)>>>>;
1642
1643struct Signal<T: 'static> {
1644    tx: Mutex<Option<oneshot::Sender<T>>>,
1645    rx: Shared<Task<Option<T>>>,
1646}
1647
1648impl<T: Send + Clone + 'static> Signal<T> {
1649    pub fn new(cx: &App) -> Self {
1650        let (tx, rx) = oneshot::channel();
1651
1652        let task = cx
1653            .background_executor()
1654            .spawn(async move { rx.await.ok() })
1655            .shared();
1656
1657        Self {
1658            tx: Mutex::new(Some(tx)),
1659            rx: task,
1660        }
1661    }
1662
1663    fn set(&self, value: T) {
1664        if let Some(tx) = self.tx.lock().take() {
1665            let _ = tx.send(value);
1666        }
1667    }
1668
1669    fn wait(&self) -> Shared<Task<Option<T>>> {
1670        self.rx.clone()
1671    }
1672}
1673
1674pub(crate) struct ChannelClient {
1675    next_message_id: AtomicU32,
1676    outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
1677    buffer: Mutex<VecDeque<Envelope>>,
1678    response_channels: ResponseChannels,
1679    stream_response_channels: StreamResponseChannels,
1680    message_handlers: Mutex<ProtoMessageHandlerSet>,
1681    max_received: AtomicU32,
1682    name: &'static str,
1683    task: Mutex<Task<Result<()>>>,
1684    remote_started: Signal<()>,
1685    has_wsl_interop: bool,
1686    executor: BackgroundExecutor,
1687}
1688
1689impl ChannelClient {
1690    pub(crate) fn new(
1691        incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1692        outgoing_tx: mpsc::UnboundedSender<Envelope>,
1693        cx: &App,
1694        name: &'static str,
1695        has_wsl_interop: bool,
1696    ) -> Arc<Self> {
1697        Arc::new_cyclic(|this| Self {
1698            outgoing_tx: Mutex::new(outgoing_tx),
1699            next_message_id: AtomicU32::new(0),
1700            max_received: AtomicU32::new(0),
1701            response_channels: ResponseChannels::default(),
1702            stream_response_channels: StreamResponseChannels::default(),
1703            message_handlers: Default::default(),
1704            buffer: Mutex::new(VecDeque::new()),
1705            name,
1706            executor: cx.background_executor().clone(),
1707            task: Mutex::new(Self::start_handling_messages(
1708                this.clone(),
1709                incoming_rx,
1710                &cx.to_async(),
1711            )),
1712            remote_started: Signal::new(cx),
1713            has_wsl_interop,
1714        })
1715    }
1716
1717    fn wait_for_remote_started(&self) -> Shared<Task<Option<()>>> {
1718        self.remote_started.wait()
1719    }
1720
1721    fn start_handling_messages(
1722        this: Weak<Self>,
1723        mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1724        cx: &AsyncApp,
1725    ) -> Task<Result<()>> {
1726        cx.spawn(async move |cx| {
1727            if let Some(this) = this.upgrade() {
1728                let envelope = proto::RemoteStarted {}.into_envelope(0, None, None);
1729                this.outgoing_tx.lock().unbounded_send(envelope).ok();
1730            };
1731
1732            let peer_id = PeerId { owner_id: 0, id: 0 };
1733            while let Some(incoming) = incoming_rx.next().await {
1734                let Some(this) = this.upgrade() else {
1735                    return anyhow::Ok(());
1736                };
1737                if let Some(ack_id) = incoming.ack_id {
1738                    let mut buffer = this.buffer.lock();
1739                    while buffer.front().is_some_and(|msg| msg.id <= ack_id) {
1740                        buffer.pop_front();
1741                    }
1742                }
1743                if let Some(proto::envelope::Payload::FlushBufferedMessages(_)) = &incoming.payload
1744                {
1745                    log::debug!(
1746                        "{}:remote message received. name:FlushBufferedMessages",
1747                        this.name
1748                    );
1749                    {
1750                        let buffer = this.buffer.lock();
1751                        for envelope in buffer.iter() {
1752                            this.outgoing_tx
1753                                .lock()
1754                                .unbounded_send(envelope.clone())
1755                                .ok();
1756                        }
1757                    }
1758                    let mut envelope = proto::Ack {}.into_envelope(0, Some(incoming.id), None);
1759                    envelope.id = this.next_message_id.fetch_add(1, SeqCst);
1760                    this.outgoing_tx.lock().unbounded_send(envelope).ok();
1761                    continue;
1762                }
1763
1764                if let Some(proto::envelope::Payload::RemoteStarted(_)) = &incoming.payload {
1765                    this.remote_started.set(());
1766                    let mut envelope = proto::Ack {}.into_envelope(0, Some(incoming.id), None);
1767                    envelope.id = this.next_message_id.fetch_add(1, SeqCst);
1768                    this.outgoing_tx.lock().unbounded_send(envelope).ok();
1769                    continue;
1770                }
1771
1772                this.max_received.store(incoming.id, SeqCst);
1773
1774                if let Some(request_id) = incoming.responding_to {
1775                    let request_id = MessageId(request_id);
1776                    // An incoming response with no payload is malformed; drop
1777                    // it. The request future and any stream consumers will
1778                    // remain pending until either a real response arrives or
1779                    // the connection is torn down.
1780                    if incoming.payload.is_none() {
1781                        continue;
1782                    }
1783                    let sender = this.response_channels.lock().remove(&request_id);
1784                    if let Some(sender) = sender {
1785                        let (tx, rx) = oneshot::channel();
1786                        sender.send((incoming, tx)).ok();
1787                        rx.await.ok();
1788                    } else {
1789                        let terminal_stream_response = matches!(
1790                            &incoming.payload,
1791                            Some(proto::envelope::Payload::Error(_))
1792                                | Some(proto::envelope::Payload::EndStream(_))
1793                        );
1794                        let sender = if terminal_stream_response {
1795                            this.stream_response_channels.lock().remove(&request_id)
1796                        } else {
1797                            this.stream_response_channels
1798                                .lock()
1799                                .get(&request_id)
1800                                .cloned()
1801                        };
1802                        if let Some(sender) = sender {
1803                            let (tx, rx) = oneshot::channel();
1804                            if sender.unbounded_send((Ok(incoming), tx)).is_err() {
1805                                this.stream_response_channels.lock().remove(&request_id);
1806                                continue;
1807                            }
1808                            rx.await.ok();
1809                        }
1810                    }
1811                } else if let Some(envelope) =
1812                    build_typed_envelope(peer_id, Instant::now(), incoming)
1813                {
1814                    let type_name = envelope.payload_type_name();
1815                    let message_id = envelope.message_id();
1816                    if let Some(future) = ProtoMessageHandlerSet::handle_message(
1817                        &this.message_handlers,
1818                        envelope,
1819                        this.clone().into(),
1820                        cx.clone(),
1821                    ) {
1822                        log::debug!("{}:remote message received. name:{type_name}", this.name);
1823                        cx.foreground_executor()
1824                            .spawn(async move {
1825                                match future.await {
1826                                    Ok(_) => {
1827                                        log::debug!(
1828                                            "{}:remote message handled. name:{type_name}",
1829                                            this.name
1830                                        );
1831                                    }
1832                                    Err(error) => {
1833                                        log::error!(
1834                                            "{}:error handling message. type:{}, error:{:#}",
1835                                            this.name,
1836                                            type_name,
1837                                            format!("{error:#}").lines().fold(
1838                                                String::new(),
1839                                                |mut message, line| {
1840                                                    if !message.is_empty() {
1841                                                        message.push(' ');
1842                                                    }
1843                                                    message.push_str(line);
1844                                                    message
1845                                                }
1846                                            )
1847                                        );
1848                                    }
1849                                }
1850                            })
1851                            .detach()
1852                    } else {
1853                        log::error!("{}:unhandled remote message name:{type_name}", this.name);
1854                        if let Err(e) = AnyProtoClient::from(this.clone()).send_response(
1855                            message_id,
1856                            anyhow::anyhow!("no handler registered for {type_name}").to_proto(),
1857                        ) {
1858                            log::error!(
1859                                "{}:error sending error response for {type_name}:{e:#}",
1860                                this.name
1861                            );
1862                        }
1863                    }
1864                }
1865            }
1866            anyhow::Ok(())
1867        })
1868    }
1869
1870    pub(crate) fn reconnect(
1871        self: &Arc<Self>,
1872        incoming_rx: UnboundedReceiver<Envelope>,
1873        outgoing_tx: UnboundedSender<Envelope>,
1874        cx: &AsyncApp,
1875    ) {
1876        *self.outgoing_tx.lock() = outgoing_tx;
1877        *self.task.lock() = Self::start_handling_messages(Arc::downgrade(self), incoming_rx, cx);
1878    }
1879
1880    fn request<T: RequestMessage>(
1881        &self,
1882        payload: T,
1883    ) -> impl 'static + Future<Output = Result<T::Response>> {
1884        self.request_internal(payload, true)
1885    }
1886
1887    fn request_internal<T: RequestMessage>(
1888        &self,
1889        payload: T,
1890        use_buffer: bool,
1891    ) -> impl 'static + Future<Output = Result<T::Response>> {
1892        log::debug!("remote request start. name:{}", T::NAME);
1893        let response =
1894            self.request_dynamic(payload.into_envelope(0, None, None), T::NAME, use_buffer);
1895        async move {
1896            let response = response.await?;
1897            log::debug!("remote request finish. name:{}", T::NAME);
1898            T::Response::from_envelope(response).context("received a response of the wrong type")
1899        }
1900    }
1901
1902    async fn resync(&self, timeout: Duration) -> Result<()> {
1903        smol::future::or(
1904            async {
1905                self.request_internal(proto::FlushBufferedMessages {}, false)
1906                    .await?;
1907
1908                for envelope in self.buffer.lock().iter() {
1909                    self.outgoing_tx
1910                        .lock()
1911                        .unbounded_send(envelope.clone())
1912                        .ok();
1913                }
1914                Ok(())
1915            },
1916            async {
1917                self.executor.timer(timeout).await;
1918                anyhow::bail!("Timed out resyncing remote client")
1919            },
1920        )
1921        .await
1922    }
1923
1924    async fn ping(&self, timeout: Duration) -> Result<()> {
1925        smol::future::or(
1926            async {
1927                self.request(proto::Ping {}).await?;
1928                Ok(())
1929            },
1930            async {
1931                self.executor.timer(timeout).await;
1932                anyhow::bail!("Timed out pinging remote client")
1933            },
1934        )
1935        .await
1936    }
1937
1938    fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
1939        log::debug!("remote send name:{}", T::NAME);
1940        self.send_dynamic(payload.into_envelope(0, None, None))
1941    }
1942
1943    fn request_dynamic(
1944        &self,
1945        mut envelope: proto::Envelope,
1946        type_name: &'static str,
1947        use_buffer: bool,
1948    ) -> impl 'static + Future<Output = Result<proto::Envelope>> {
1949        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1950        let (tx, rx) = oneshot::channel();
1951        let mut response_channels_lock = self.response_channels.lock();
1952        response_channels_lock.insert(MessageId(envelope.id), tx);
1953        drop(response_channels_lock);
1954
1955        let result = if use_buffer {
1956            self.send_buffered(envelope)
1957        } else {
1958            self.send_unbuffered(envelope)
1959        };
1960        async move {
1961            if let Err(error) = &result {
1962                log::error!("failed to send message: {error}");
1963                anyhow::bail!("failed to send message: {error}");
1964            }
1965
1966            let response = rx.await.context("connection lost")?.0;
1967            if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
1968                return Err(RpcError::from_proto(error, type_name));
1969            }
1970            Ok(response)
1971        }
1972    }
1973
1974    fn request_stream_dynamic(
1975        &self,
1976        mut envelope: proto::Envelope,
1977        type_name: &'static str,
1978    ) -> impl 'static + Future<Output = Result<BoxStream<'static, Result<proto::Envelope>>>> {
1979        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1980        let message_id = MessageId(envelope.id);
1981        let (tx, rx) = mpsc::unbounded();
1982        let stream_response_channels = self.stream_response_channels.clone();
1983        stream_response_channels.lock().insert(message_id, tx);
1984
1985        let result = self.send_buffered(envelope);
1986        async move {
1987            if let Err(error) = &result {
1988                log::error!("failed to send message: {error}");
1989                anyhow::bail!("failed to send message: {error}");
1990            }
1991
1992            let cleanup_stream_response_channel = util::defer({
1993                let stream_response_channels = stream_response_channels.clone();
1994                move || {
1995                    stream_response_channels.lock().remove(&message_id);
1996                }
1997            });
1998
1999            Ok(rx
2000                .filter_map(move |(response, _barrier)| {
2001                    // Keep the cleanup guard alive until the returned stream is dropped.
2002                    let _keep_cleanup_guard_alive = &cleanup_stream_response_channel;
2003                    futures::future::ready(match response {
2004                        Ok(response) => {
2005                            if let Some(proto::envelope::Payload::Error(error)) = &response.payload
2006                            {
2007                                Some(Err(RpcError::from_proto(error, type_name)))
2008                            } else if let Some(proto::envelope::Payload::EndStream(_)) =
2009                                &response.payload
2010                            {
2011                                None
2012                            } else {
2013                                Some(Ok(response))
2014                            }
2015                        }
2016                        Err(error) => Some(Err(error)),
2017                    })
2018                })
2019                .boxed())
2020        }
2021    }
2022
2023    pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
2024        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
2025        self.send_buffered(envelope)
2026    }
2027
2028    fn send_buffered(&self, mut envelope: proto::Envelope) -> Result<()> {
2029        envelope.ack_id = Some(self.max_received.load(SeqCst));
2030        self.buffer.lock().push_back(envelope.clone());
2031        // ignore errors on send (happen while we're reconnecting)
2032        // assume that the global "disconnected" overlay is sufficient.
2033        self.outgoing_tx.lock().unbounded_send(envelope).ok();
2034        Ok(())
2035    }
2036
2037    fn send_unbuffered(&self, mut envelope: proto::Envelope) -> Result<()> {
2038        envelope.ack_id = Some(self.max_received.load(SeqCst));
2039        self.outgoing_tx.lock().unbounded_send(envelope).ok();
2040        Ok(())
2041    }
2042}
2043
2044impl ProtoClient for ChannelClient {
2045    fn request(
2046        &self,
2047        envelope: proto::Envelope,
2048        request_type: &'static str,
2049    ) -> BoxFuture<'static, Result<proto::Envelope>> {
2050        self.request_dynamic(envelope, request_type, true).boxed()
2051    }
2052
2053    fn request_stream(
2054        &self,
2055        envelope: proto::Envelope,
2056        request_type: &'static str,
2057    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<proto::Envelope>>>> {
2058        self.request_stream_dynamic(envelope, request_type).boxed()
2059    }
2060
2061    fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
2062        self.send_dynamic(envelope)
2063    }
2064
2065    fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
2066        self.send_dynamic(envelope)
2067    }
2068
2069    fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
2070        &self.message_handlers
2071    }
2072
2073    fn is_via_collab(&self) -> bool {
2074        false
2075    }
2076
2077    fn has_wsl_interop(&self) -> bool {
2078        self.has_wsl_interop
2079    }
2080}
2081