[APP-3801] implement remote environments auth (#9331)

app/src/lib.rs

@@ -565,12 +565,12 @@ pub fn run() -> Result<()> {
                 }
             }
             #[cfg(not(target_family = "wasm"))]
-            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerProxy) => {
-                return crate::remote_server::run_proxy();
+            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerProxy(args)) => {
+                return crate::remote_server::run_proxy(args.identity_key.clone());
             }
             #[cfg(not(target_family = "wasm"))]
-            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerDaemon) => {
-                return crate::remote_server::run_daemon();
+            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerDaemon(args)) => {
+                return crate::remote_server::run_daemon(args.identity_key.clone());
             }
             #[cfg(not(target_family = "wasm"))]
             warp_cli::Command::Worker(warp_cli::WorkerCommand::RipgrepSearch {
@@ -1253,6 +1253,8 @@ fn initialize_app(
     ctx.add_singleton_model(|_ctx| SyncedInputState::new());
 
     ctx.add_singleton_model(remote_server::manager::RemoteServerManager::new);
+    #[cfg(not(target_family = "wasm"))]
+    remote_server::wire_auth_token_rotation(ctx);
 
     log::info!(
         "Starting warp with channel state {} and version {:?}",

app/src/remote_server/auth_context.rs

@@ -0,0 +1,49 @@
+use std::sync::Arc;
+
+use remote_server::auth::RemoteServerAuthContext;
+use warpui::r#async::BoxFuture;
+
+use crate::auth::auth_state::AuthState;
+use crate::server::server_api::auth::AuthClient;
+
+/// Builds the app-wide auth context used by remote-server connections.
+pub fn server_api_auth_context(
+    auth_state: Arc<AuthState>,
+    auth_client: Arc<dyn AuthClient>,
+) -> RemoteServerAuthContext {
+    let token_auth_state = auth_state.clone();
+    let token_auth_client = auth_client;
+    let identity_auth_state = auth_state;
+
+    RemoteServerAuthContext::new(
+        move || -> BoxFuture<'static, Option<String>> {
+            if !use_authenticated_user_identity(&token_auth_state) {
+                return Box::pin(async { None });
+            }
+
+            let auth_client = token_auth_client.clone();
+            Box::pin(async move {
+                match auth_client.get_or_refresh_access_token().await {
+                    Ok(token) => token.bearer_token(),
+                    Err(_) => None,
+                }
+            })
+        },
+        move || remote_server_identity_key(&identity_auth_state),
+    )
+}
+
+fn use_authenticated_user_identity(auth_state: &AuthState) -> bool {
+    auth_state.is_logged_in() && !auth_state.is_user_anonymous().unwrap_or(true)
+}
+
+fn remote_server_identity_key(auth_state: &AuthState) -> String {
+    if use_authenticated_user_identity(auth_state) {
+        auth_state
+            .user_id()
+            .map(|uid| uid.as_string())
+            .unwrap_or_else(|| auth_state.anonymous_id())
+    } else {
+        auth_state.anonymous_id()
+    }
+}

app/src/remote_server/mod.rs

@@ -1,7 +1,15 @@
+#[cfg(not(target_family = "wasm"))]
+use crate::server::server_api::{ServerApiEvent, ServerApiProvider};
+#[cfg(not(target_family = "wasm"))]
+use remote_server::manager::RemoteServerManager;
+#[cfg(not(target_family = "wasm"))]
+use warpui::SingletonEntity;
 // Re-export everything from the `remote_server` crate so existing
 // `crate::remote_server::*` imports in `app` continue to work.
 pub use remote_server::*;
 
+#[cfg(not(target_family = "wasm"))]
+pub mod auth_context;
 #[cfg(not(target_family = "wasm"))]
 pub mod server_model;
 #[cfg(not(target_family = "wasm"))]
@@ -11,23 +19,23 @@ pub mod unix;
 
 /// Run the `remote-server-proxy` subcommand.
 #[cfg(unix)]
-pub fn run_proxy() -> anyhow::Result<()> {
-    unix::run_proxy()
+pub fn run_proxy(identity_key: String) -> anyhow::Result<()> {
+    unix::run_proxy(identity_key)
 }
 
 #[cfg(not(unix))]
-pub fn run_proxy() -> anyhow::Result<()> {
+pub fn run_proxy(_identity_key: String) -> anyhow::Result<()> {
     anyhow::bail!("remote-server-proxy is not supported on this platform")
 }
 
 /// Run the `remote-server-daemon` subcommand.
 #[cfg(unix)]
-pub fn run_daemon() -> anyhow::Result<()> {
-    unix::run_daemon()
+pub fn run_daemon(identity_key: String) -> anyhow::Result<()> {
+    unix::run_daemon(identity_key)
 }
 
 #[cfg(not(unix))]
-pub fn run_daemon() -> anyhow::Result<()> {
+pub fn run_daemon(_identity_key: String) -> anyhow::Result<()> {
     anyhow::bail!("remote-server-daemon is not supported on this platform")
 }
 
@@ -79,3 +87,17 @@ pub(super) fn run_daemon_app(
     })?;
     Ok(())
 }
+
+/// Forwards app auth-token rotation events to the remote-server manager.
+#[cfg(not(target_family = "wasm"))]
+pub fn wire_auth_token_rotation(ctx: &mut warpui::AppContext) {
+    let server_api = ServerApiProvider::handle(ctx);
+    let manager = RemoteServerManager::handle(ctx);
+    ctx.subscribe_to_model(&server_api, move |_, event, ctx| {
+        if let ServerApiEvent::AccessTokenRefreshed { token } = event {
+            manager.update(ctx, |manager, _| {
+                manager.rotate_auth_token(token.clone());
+            });
+        }
+    });
+}

app/src/remote_server/server_model.rs

@@ -18,12 +18,12 @@ use warp_util::file::FileId;
 
 use super::proto::{
     client_message, delete_file_response, run_command_response, server_message,
-    write_file_response, Abort, ClientMessage, DeleteFile, DeleteFileResponse, DeleteFileSuccess,
-    ErrorCode, ErrorResponse, FailedFileRead, FileContextProto, FileOperationError,
-    InitializeResponse, NavigatedToDirectory, NavigatedToDirectoryResponse,
-    ReadFileContextResponse, RunCommandError, RunCommandErrorCode, RunCommandRequest,
-    RunCommandResponse, RunCommandSuccess, ServerMessage, SessionBootstrapped, WriteFile,
-    WriteFileResponse, WriteFileSuccess,
+    write_file_response, Abort, Authenticate, ClientMessage, DeleteFile, DeleteFileResponse,
+    DeleteFileSuccess, ErrorCode, ErrorResponse, FailedFileRead, FileContextProto,
+    FileOperationError, Initialize, InitializeResponse, NavigatedToDirectory,
+    NavigatedToDirectoryResponse, ReadFileContextResponse, RunCommandError, RunCommandErrorCode,
+    RunCommandRequest, RunCommandResponse, RunCommandSuccess, ServerMessage, SessionBootstrapped,
+    WriteFile, WriteFileResponse, WriteFileSuccess,
 };
 
 /// How long the daemon waits with no connections before exiting.
@@ -164,6 +164,13 @@ pub struct ServerModel {
     executors: HashMap<SessionId, Arc<LocalCommandExecutor>>,
     /// Tracks in-flight file write/delete operations and handles cleanup.
     pending_file_ops: PendingFileOps,
+    /// Daemon-wide bearer credential for the identity-scoped daemon.
+    ///
+    /// The token is written by Initialize when the client supplies a
+    /// non-empty credential, or by Authenticate during token rotation. It is
+    /// intentionally retained across proxy connection teardown and cleared
+    /// only by daemon process exit.
+    auth_token: Option<String>,
 }
 
 impl Entity for ServerModel {
@@ -188,6 +195,7 @@ impl ServerModel {
             host_id,
             executors: HashMap::new(),
             pending_file_ops: PendingFileOps::new(),
+            auth_token: None,
         };
         // Subscribe to FileModel and RepoMetadataModel events
         // file operation results and repo metadata pushes are forwarded to all
@@ -377,7 +385,13 @@ impl ServerModel {
         let request_id = RequestId::from(msg.request_id);
 
         let outcome = match msg.message {
-            Some(client_message::Message::Initialize(_)) => self.handle_initialize(&request_id),
+            Some(client_message::Message::Initialize(msg)) => {
+                self.handle_initialize(msg, &request_id)
+            }
+            Some(client_message::Message::Authenticate(msg)) => {
+                self.handle_authenticate(msg);
+                return;
+            }
             Some(client_message::Message::SessionBootstrapped(msg)) => {
                 self.handle_session_bootstrapped(msg);
                 return;
@@ -497,8 +511,11 @@ impl ServerModel {
     }
 
     /// Handles `Initialize` by returning the server version and host id.
-    fn handle_initialize(&self, request_id: &RequestId) -> HandlerOutcome {
+    fn handle_initialize(&mut self, msg: Initialize, request_id: &RequestId) -> HandlerOutcome {
         log::info!("Handling Initialize (request_id={request_id})");
+        if !msg.auth_token.is_empty() {
+            self.auth_token = Some(msg.auth_token);
+        }
         let server_version = ChannelState::app_version()
             .unwrap_or(env!("CARGO_PKG_VERSION"))
             .to_string();
@@ -510,6 +527,20 @@ impl ServerModel {
         ))
     }
 
+    /// Handles `Authenticate` by replacing the daemon-wide credential.
+    /// This is a notification — no response is sent.
+    fn handle_authenticate(&mut self, msg: Authenticate) {
+        if msg.auth_token.is_empty() {
+            log::warn!("Received Authenticate notification with empty auth token; ignoring");
+            return;
+        }
+        self.auth_token = Some(msg.auth_token);
+    }
+
+    pub fn auth_token(&self) -> Option<&str> {
+        self.auth_token.as_deref()
+    }
+
     /// Handles `Abort` by cancelling the in-progress request it targets.
     /// This is a notification — no response is sent.
     fn handle_abort(&mut self, abort: Abort, request_id: &RequestId) {
@@ -1074,3 +1105,7 @@ fn file_context_result_to_proto(result: ReadFileContextResult) -> ReadFileContex
         failed_files,
     }
 }
+
+#[cfg(test)]
+#[path = "server_model_tests.rs"]
+mod tests;

app/src/remote_server/server_model_tests.rs

@@ -0,0 +1,97 @@
+use std::collections::HashMap;
+
+use super::super::proto::{Authenticate, Initialize};
+use super::super::protocol::RequestId;
+use super::{PendingFileOps, ServerModel};
+
+fn test_model() -> ServerModel {
+    ServerModel {
+        connection_senders: HashMap::new(),
+        snapshot_sent_roots_by_connection: HashMap::new(),
+        grace_timer_cancel: None,
+        in_progress: HashMap::new(),
+        host_id: "test-host-id".to_string(),
+        executors: HashMap::new(),
+        pending_file_ops: PendingFileOps::new(),
+        auth_token: None,
+    }
+}
+
+fn request_id() -> RequestId {
+    RequestId::from("test-request".to_string())
+}
+
+#[test]
+fn fresh_model_starts_without_auth_token() {
+    let model = test_model();
+
+    assert_eq!(model.auth_token(), None);
+}
+
+#[test]
+fn initialize_with_auth_token_stores_token() {
+    let mut model = test_model();
+
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    assert_eq!(model.auth_token(), Some("initial-token"));
+}
+
+#[test]
+fn empty_initialize_preserves_existing_auth_token() {
+    let mut model = test_model();
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    model.handle_initialize(
+        Initialize {
+            auth_token: String::new(),
+        },
+        &request_id(),
+    );
+
+    assert_eq!(model.auth_token(), Some("initial-token"));
+}
+
+#[test]
+fn authenticate_with_auth_token_replaces_auth_token() {
+    let mut model = test_model();
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    model.handle_authenticate(Authenticate {
+        auth_token: "rotated-token".to_string(),
+    });
+
+    assert_eq!(model.auth_token(), Some("rotated-token"));
+}
+
+#[test]
+fn empty_authenticate_preserves_existing_auth_token() {
+    let mut model = test_model();
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    model.handle_authenticate(Authenticate {
+        auth_token: String::new(),
+    });
+
+    assert_eq!(model.auth_token(), Some("initial-token"));
+}

app/src/remote_server/ssh_transport.rs

@@ -3,6 +3,7 @@
 //! [`SshTransport`] uses an existing SSH ControlMaster socket to check/install
 //! the remote server binary and to launch the `remote-server-proxy` process
 //! whose stdin/stdout become the protocol channel.
+use std::fmt;
 use std::future::Future;
 use std::path::PathBuf;
 use std::pin::Pin;
@@ -12,8 +13,11 @@ use std::sync::Arc;
 use anyhow::Result;
 use warpui::r#async::executor;
 
+use remote_server::auth::RemoteServerAuthContext;
 use remote_server::client::RemoteServerClient;
-use remote_server::setup::{self, RemotePlatform, CHECK_TIMEOUT, INSTALL_TIMEOUT};
+use remote_server::setup::{
+    self, remote_server_daemon_dir, RemotePlatform, CHECK_TIMEOUT, INSTALL_TIMEOUT,
+};
 use remote_server::ssh::{run_ssh_command, run_ssh_script, ssh_args};
 use remote_server::transport::{Connection, RemoteTransport};
 
@@ -23,16 +27,54 @@ use remote_server::transport::{Connection, RemoteTransport};
 /// process (`ssh -N -o ControlMaster=yes -o ControlPath=<path>`). All SSH
 /// commands (binary check, install, proxy launch) are multiplexed through
 /// this socket without re-authenticating.
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct SshTransport {
     socket_path: PathBuf,
+    auth_context: Arc<RemoteServerAuthContext>,
+}
+
+impl fmt::Debug for SshTransport {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("SshTransport")
+            .field("socket_path", &self.socket_path)
+            .finish_non_exhaustive()
+    }
 }
 
 impl SshTransport {
-    pub fn new(socket_path: PathBuf) -> Self {
-        Self { socket_path }
+    pub fn new(socket_path: PathBuf, auth_context: Arc<RemoteServerAuthContext>) -> Self {
+        Self {
+            socket_path,
+            auth_context,
+        }
+    }
+
+    pub fn socket_path(&self) -> &PathBuf {
+        &self.socket_path
+    }
+
+    pub fn remote_daemon_socket_path(&self) -> String {
+        format!(
+            "{}/server.sock",
+            remote_server_daemon_dir(&self.auth_context.remote_server_identity_key())
+        )
+    }
+
+    pub fn remote_daemon_pid_path(&self) -> String {
+        format!(
+            "{}/server.pid",
+            remote_server_daemon_dir(&self.auth_context.remote_server_identity_key())
+        )
+    }
+
+    fn remote_proxy_command(&self) -> String {
+        let binary = remote_server::setup::remote_server_binary();
+        let identity_key = self.auth_context.remote_server_identity_key();
+        let quoted_identity_key = shell_words::quote(&identity_key);
+        format!("{binary} remote-server-proxy --identity-key {quoted_identity_key}")
     }
 }
+
 impl RemoteTransport for SshTransport {
     fn detect_platform(
         &self,
@@ -100,10 +142,10 @@ impl RemoteTransport for SshTransport {
         executor: Arc<executor::Background>,
     ) -> Pin<Box<dyn Future<Output = Result<Connection>> + Send>> {
         let socket_path = self.socket_path.clone();
+        let remote_proxy_command = self.remote_proxy_command();
         Box::pin(async move {
-            let binary = setup::remote_server_binary();
             let mut args = ssh_args(&socket_path);
-            args.push(format!("{binary} remote-server-proxy"));
+            args.push(remote_proxy_command);
 
             // `kill_on_drop(true)` pairs with ownership of the `Child` being
             // returned in the [`Connection`] below: the
@@ -142,3 +184,28 @@ impl RemoteTransport for SshTransport {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use warpui::r#async::BoxFuture;
+    fn static_auth_context() -> Arc<RemoteServerAuthContext> {
+        Arc::new(RemoteServerAuthContext::new(
+            || -> BoxFuture<'static, Option<String>> { Box::pin(async { None }) },
+            || "user id/with spaces".to_string(),
+        ))
+    }
+
+    #[test]
+    fn remote_proxy_command_quotes_identity_key() {
+        let transport = SshTransport::new(
+            PathBuf::from("/tmp/control-master.sock"),
+            static_auth_context(),
+        );
+
+        let command = transport.remote_proxy_command();
+
+        assert!(command.contains("remote-server-proxy --identity-key"));
+        assert!(command.contains("'user id/with spaces'"));
+    }
+}

app/src/remote_server/unix/mod.rs

@@ -14,6 +14,8 @@
 mod proxy;
 
 use super::server_model::{ConnectionId, ServerModel};
+use std::fs::Permissions;
+use std::os::unix::fs::PermissionsExt;
 use warpui::r#async::executor;
 
 /// Run the `remote-server-proxy` subcommand.
@@ -21,11 +23,11 @@ use warpui::r#async::executor;
 /// Ensures the daemon is running (starting it if necessary), then bridges
 /// this process's stdin/stdout to the daemon's Unix socket for the lifetime
 /// of the SSH session.
-pub fn run_proxy() -> anyhow::Result<()> {
+pub fn run_proxy(identity_key: String) -> anyhow::Result<()> {
     env_logger::Builder::from_default_env()
         .target(env_logger::Target::Stderr)
         .init();
-    proxy::run()
+    proxy::run(&identity_key)
 }
 
 /// Run the `remote-server-daemon` subcommand.
@@ -33,7 +35,7 @@ pub fn run_proxy() -> anyhow::Result<()> {
 /// Binds a Unix domain socket and writes a PID file, then delegates the
 /// WarpUI app startup to [`super::run_daemon_app`] with the Unix-specific
 /// `ServerModel` constructor.
-pub fn run_daemon() -> anyhow::Result<()> {
+pub fn run_daemon(identity_key: String) -> anyhow::Result<()> {
     // Log to a rotating file so daemon output is preserved across invocations.
     // The file is written to the same directory as client logs (~/Library/Logs
     // on macOS, ~/.local/share/warp-terminal on Linux). Since the daemon runs
@@ -43,19 +45,19 @@ pub fn run_daemon() -> anyhow::Result<()> {
         log_destination: Some(warp_logging::LogDestination::File),
     })?;
 
-    // socket_path: ~/.warp[-channel]/remote-server/server.sock
+    // socket_path: ~/.warp[-channel]/remote-server/{identity_key}/server.sock
     //   The Unix domain socket the daemon binds on.  Proxy processes connect
     //   to it and bridge their SSH stdio channel through it.
     //
-    // pid_path:    ~/.warp[-channel]/remote-server/server.pid
+    // pid_path:    ~/.warp[-channel]/remote-server/{identity_key}/server.pid
     //   Contains the daemon's PID.  Proxy processes read it and use
     //   kill(pid, 0) to detect whether the daemon is still alive before
     //   deciding whether to start a new one.
-    let socket_path = proxy::socket_path();
-    let pid_path = proxy::pid_path();
+    let socket_path = proxy::socket_path(&identity_key);
+    let pid_path = proxy::pid_path(&identity_key);
 
     if let Some(parent) = socket_path.parent() {
-        std::fs::create_dir_all(parent)?;
+        proxy::ensure_private_daemon_dir(parent)?;
     }
     if socket_path.exists() {
         std::fs::remove_file(&socket_path)?;
@@ -64,6 +66,7 @@ pub fn run_daemon() -> anyhow::Result<()> {
     // Bind with std (no async runtime needed yet); converted to
     // async_io::Async inside the closure where the executor is active.
     let listener = std::os::unix::net::UnixListener::bind(&socket_path)?;
+    std::fs::set_permissions(&socket_path, Permissions::from_mode(0o600))?;
     // async_io::Async::new() requires non-blocking mode.
     listener.set_nonblocking(true)?;
     log::info!("Daemon bound to {}", socket_path.display());

app/src/remote_server/unix/proxy.rs

@@ -10,6 +10,8 @@
 //! 4. Connect to `server.sock` and bridge stdin/stdout to the socket using
 //!    the existing 4-byte length-prefixed frame format.
 
+use std::fs::Permissions;
+use std::os::unix::fs::PermissionsExt;
 use std::os::unix::io::AsRawFd;
 use std::path::PathBuf;
 use std::process::Stdio;
@@ -18,30 +20,37 @@ use std::time::Duration;
 use super::super::setup;
 
 /// Path to the daemon's Unix domain socket.
-pub(super) fn socket_path() -> PathBuf {
-    let dir = setup::remote_server_dir();
+pub(super) fn socket_path(identity_key: &str) -> PathBuf {
+    let dir = setup::remote_server_daemon_dir(identity_key);
     let expanded = shellexpand::tilde(&dir).into_owned();
     PathBuf::from(expanded).join("server.sock")
 }
 
 /// Path to the daemon's PID file (also used as the flock target).
-pub(super) fn pid_path() -> PathBuf {
-    let dir = setup::remote_server_dir();
+pub(super) fn pid_path(identity_key: &str) -> PathBuf {
+    let dir = setup::remote_server_daemon_dir(identity_key);
     let expanded = shellexpand::tilde(&dir).into_owned();
     PathBuf::from(expanded).join("server.pid")
 }
 
+/// Ensures the daemon directory exists with owner-only permissions.
+pub(super) fn ensure_private_daemon_dir(path: &std::path::Path) -> anyhow::Result<()> {
+    std::fs::create_dir_all(path)?;
+    std::fs::set_permissions(path, Permissions::from_mode(0o700))?;
+    Ok(())
+}
+
 /// Entry point for `remote-server-proxy`.
 ///
 /// Ensures the daemon is running, then bridges stdin/stdout to the daemon's
 /// Unix socket for the lifetime of this SSH session.
-pub fn run() -> anyhow::Result<()> {
-    let socket_path = socket_path();
-    let pid_path = pid_path();
+pub fn run(identity_key: &str) -> anyhow::Result<()> {
+    let socket_path = socket_path(identity_key);
+    let pid_path = pid_path(identity_key);
 
     // Ensure the parent directory exists.
     if let Some(parent) = socket_path.parent() {
-        std::fs::create_dir_all(parent)?;
+        ensure_private_daemon_dir(parent)?;
     }
 
     // ---- Acquire exclusive flock on the PID file --------------------------------
@@ -83,6 +92,8 @@ pub fn run() -> anyhow::Result<()> {
         let exe = std::env::current_exe()?;
         let mut cmd = command::blocking::Command::new(&exe);
         cmd.arg("remote-server-daemon")
+            .arg("--identity-key")
+            .arg(identity_key)
             .stdin(Stdio::null())
             .stdout(Stdio::null())
             .stderr(Stdio::null());

app/src/server/server_api.rs

@@ -54,6 +54,7 @@ use parking_lot::{Mutex, RwLock};
 use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
 use std::borrow::Cow;
+use std::fmt;
 use std::path::Path;
 use std::sync::Arc;
 use std::time::Duration;
@@ -355,7 +356,7 @@ cfg_if::cfg_if! {
 /// An event related to the server API itself (and not a particular API call).
 /// Most errors should be handled in callbacks to individual APIs, rather than sent over the
 /// server API channel.
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub enum ServerApiEvent {
     /// We made a staging API call that was blocked, which may indicate a firewall misconfiguration.
     StagingAccessBlocked,
@@ -364,6 +365,25 @@ pub enum ServerApiEvent {
     NeedsReauth,
     /// The user's account has been disabled.
     UserAccountDisabled,
+    /// The current bearer token was refreshed.
+    AccessTokenRefreshed {
+        #[cfg_attr(target_family = "wasm", allow(dead_code))]
+        token: String,
+    },
+}
+
+impl fmt::Debug for ServerApiEvent {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::StagingAccessBlocked => f.write_str("StagingAccessBlocked"),
+            Self::NeedsReauth => f.write_str("NeedsReauth"),
+            Self::UserAccountDisabled => f.write_str("UserAccountDisabled"),
+            Self::AccessTokenRefreshed { .. } => f
+                .debug_struct("AccessTokenRefreshed")
+                .field("token", &"<redacted>")
+                .finish(),
+        }
+    }
 }
 
 /// An API wrapper struct with methods to requests to warp-server.

app/src/server/server_api/auth.rs

@@ -268,6 +268,12 @@ impl AuthClient for ServerApi {
                     let new_firebase_token_info = result?;
                     self.auth_state
                         .update_firebase_tokens(new_firebase_token_info.clone());
+                    let _ = self
+                        .event_sender
+                        .send(ServerApiEvent::AccessTokenRefreshed {
+                            token: new_firebase_token_info.id_token.clone(),
+                        })
+                        .await;
                     return Ok(AuthToken::Firebase(new_firebase_token_info.id_token));
                 }
 

app/src/terminal/writeable_pty/remote_server_controller.rs

@@ -1,5 +1,10 @@
+use crate::auth::auth_state::AuthStateProvider;
+use crate::remote_server::auth_context::server_api_auth_context;
 use instant::Instant;
+use remote_server::auth::RemoteServerAuthContext;
 use settings::Setting;
+use std::path::PathBuf;
+use std::sync::Arc;
 use warp_core::SessionId;
 use warpui::{Entity, ModelContext, ModelHandle, SingletonEntity, WeakModelHandle};
 
@@ -7,6 +12,7 @@ use crate::terminal::warpify::settings::SshExtensionInstallMode;
 
 use crate::remote_server::manager::{RemoteServerManager, RemoteServerManagerEvent};
 use crate::remote_server::ssh_transport::SshTransport;
+use crate::server::server_api::ServerApiProvider;
 use crate::terminal::model::session::{IsLegacySSHSession, SessionInfo};
 use crate::terminal::model_events::{ModelEvent, ModelEventDispatcher};
 use crate::terminal::warpify::settings::WarpifySettings;
@@ -59,6 +65,7 @@ enum SshInitState {
 pub struct RemoteServerController<T: EventLoopSender> {
     pty_controller: WeakModelHandle<PtyController<T>>,
     model_event_dispatcher: ModelHandle<ModelEventDispatcher>,
+    auth_context: Arc<RemoteServerAuthContext>,
     state: SshInitState,
     /// Whether the binary was installed during this setup flow.
     did_install: bool,
@@ -76,6 +83,10 @@ impl<T: EventLoopSender> RemoteServerController<T> {
         model_event_dispatcher: ModelHandle<ModelEventDispatcher>,
         ctx: &mut ModelContext<Self>,
     ) -> Self {
+        let auth_context = Arc::new(server_api_auth_context(
+            AuthStateProvider::as_ref(ctx).get().clone(),
+            ServerApiProvider::as_ref(ctx).get_auth_client(),
+        ));
         ctx.subscribe_to_model(&model_event_dispatcher, |me, event, ctx| {
             if let ModelEvent::SshInitShell {
                 pending_session_info,
@@ -104,12 +115,25 @@ impl<T: EventLoopSender> RemoteServerController<T> {
             RemoteServerManagerEvent::SessionConnectionFailed { session_id, .. } => {
                 me.on_session_connection_failed(*session_id, ctx);
             }
-            _ => {}
+            RemoteServerManagerEvent::SessionConnecting { .. }
+            | RemoteServerManagerEvent::SessionDisconnected { .. }
+            | RemoteServerManagerEvent::SessionReconnected { .. }
+            | RemoteServerManagerEvent::SessionDeregistered { .. }
+            | RemoteServerManagerEvent::HostConnected { .. }
+            | RemoteServerManagerEvent::HostDisconnected { .. }
+            | RemoteServerManagerEvent::NavigatedToDirectory { .. }
+            | RemoteServerManagerEvent::RepoMetadataSnapshot { .. }
+            | RemoteServerManagerEvent::RepoMetadataUpdated { .. }
+            | RemoteServerManagerEvent::RepoMetadataDirectoryLoaded { .. }
+            | RemoteServerManagerEvent::SetupStateChanged { .. }
+            | RemoteServerManagerEvent::ClientRequestFailed { .. }
+            | RemoteServerManagerEvent::ServerMessageDecodingError { .. } => {}
         });
 
         Self {
             pty_controller,
             model_event_dispatcher,
+            auth_context,
             state: SshInitState::Idle,
             did_install: false,
             remote_platform: None,
@@ -135,7 +159,6 @@ impl<T: EventLoopSender> RemoteServerController<T> {
         };
         let session_id = info.session_id;
         let socket_path = socket_path.clone();
-
         debug_assert!(matches!(self.state, SshInitState::Idle));
         match std::mem::replace(&mut self.state, SshInitState::Idle) {
             SshInitState::Idle => {}
@@ -158,8 +181,7 @@ impl<T: EventLoopSender> RemoteServerController<T> {
                 self.flush_stashed_bootstrap(old_info, ctx);
             }
         }
-
-        let transport = SshTransport::new(socket_path);
+        let transport = SshTransport::new(socket_path, self.auth_context.clone());
         self.did_install = false;
         self.remote_platform = None;
         self.state = SshInitState::AwaitingCheck {
@@ -199,14 +221,13 @@ impl<T: EventLoopSender> RemoteServerController<T> {
 
         match result {
             Ok(true) => {
+                let socket_path = transport.socket_path().clone();
                 self.state = SshInitState::AwaitingConnect {
                     session_id,
                     session_info,
                     setup_start,
                 };
-                RemoteServerManager::handle(ctx).update(ctx, |mgr, ctx| {
-                    mgr.connect_session(session_id, transport, ctx);
-                });
+                self.connect_session_for_current_identity(session_id, socket_path, ctx);
             }
             Ok(false) => {
                 let install_mode = *WarpifySettings::as_ref(ctx)
@@ -404,14 +425,13 @@ impl<T: EventLoopSender> RemoteServerController<T> {
         };
         match result {
             Ok(()) => {
+                let socket_path = transport.socket_path().clone();
                 self.state = SshInitState::AwaitingConnect {
                     session_id,
                     session_info,
                     setup_start,
                 };
-                RemoteServerManager::handle(ctx).update(ctx, |mgr, ctx| {
-                    mgr.connect_session(session_id, transport, ctx);
-                });
+                self.connect_session_for_current_identity(session_id, socket_path, ctx);
             }
             Err(err) => {
                 log::error!("Binary install failed for {session_id:?}: {err}");
@@ -419,4 +439,17 @@ impl<T: EventLoopSender> RemoteServerController<T> {
             }
         }
     }
+
+    fn connect_session_for_current_identity(
+        &mut self,
+        session_id: SessionId,
+        socket_path: PathBuf,
+        ctx: &mut ModelContext<Self>,
+    ) {
+        let transport = SshTransport::new(socket_path, self.auth_context.clone());
+        let auth_context = self.auth_context.clone();
+        RemoteServerManager::handle(ctx).update(ctx, |mgr, ctx| {
+            mgr.connect_session(session_id, transport, auth_context, ctx);
+        });
+    }
 }