client.rs

   1#[cfg(any(test, feature = "test-support"))]
   2pub mod test;
   3
   4mod proxy;
   5pub mod telemetry;
   6pub mod user;
   7pub mod zed_urls;
   8
   9use anyhow::{Context as _, Result, anyhow};
  10use async_tungstenite::tungstenite::{
  11    client::IntoClientRequest,
  12    error::Error as WebsocketError,
  13    http::{HeaderValue, Request, StatusCode},
  14};
  15use clock::SystemClock;
  16use cloud_api_client::CloudApiClient;
  17use cloud_api_client::websocket_protocol::MessageToClient;
  18use credentials_provider::CredentialsProvider;
  19use feature_flags::FeatureFlagAppExt as _;
  20use futures::{
  21    AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,
  22    channel::oneshot, future::BoxFuture,
  23};
  24use gpui::{App, AsyncApp, Entity, Global, Task, WeakEntity, actions};
  25use http_client::{HttpClient, HttpClientWithUrl, http, read_proxy_from_env};
  26use parking_lot::RwLock;
  27use postage::watch;
  28use proxy::connect_proxy_stream;
  29use rand::prelude::*;
  30use release_channel::{AppVersion, ReleaseChannel};
  31use rpc::proto::{AnyTypedEnvelope, EnvelopedMessage, PeerId, RequestMessage};
  32use serde::{Deserialize, Serialize};
  33use settings::{Settings, SettingsContent};
  34use std::{
  35    any::TypeId,
  36    convert::TryFrom,
  37    fmt::Write as _,
  38    future::Future,
  39    marker::PhantomData,
  40    path::PathBuf,
  41    sync::{
  42        Arc, LazyLock, Weak,
  43        atomic::{AtomicU64, Ordering},
  44    },
  45    time::{Duration, Instant},
  46};
  47use std::{cmp, pin::Pin};
  48use telemetry::Telemetry;
  49use thiserror::Error;
  50use tokio::net::TcpStream;
  51use url::Url;
  52use util::{ConnectionResult, ResultExt};
  53
  54pub use rpc::*;
  55pub use telemetry_events::Event;
  56pub use user::*;
  57
  58static ZED_SERVER_URL: LazyLock<Option<String>> =
  59    LazyLock::new(|| std::env::var("ZED_SERVER_URL").ok());
  60static ZED_RPC_URL: LazyLock<Option<String>> = LazyLock::new(|| std::env::var("ZED_RPC_URL").ok());
  61
  62pub static IMPERSONATE_LOGIN: LazyLock<Option<String>> = LazyLock::new(|| {
  63    std::env::var("ZED_IMPERSONATE")
  64        .ok()
  65        .and_then(|s| if s.is_empty() { None } else { Some(s) })
  66});
  67
  68pub static USE_WEB_LOGIN: LazyLock<bool> = LazyLock::new(|| std::env::var("ZED_WEB_LOGIN").is_ok());
  69
  70pub static ADMIN_API_TOKEN: LazyLock<Option<String>> = LazyLock::new(|| {
  71    std::env::var("ZED_ADMIN_API_TOKEN")
  72        .ok()
  73        .and_then(|s| if s.is_empty() { None } else { Some(s) })
  74});
  75
  76pub static ZED_APP_PATH: LazyLock<Option<PathBuf>> =
  77    LazyLock::new(|| std::env::var("ZED_APP_PATH").ok().map(PathBuf::from));
  78
  79pub static ZED_ALWAYS_ACTIVE: LazyLock<bool> =
  80    LazyLock::new(|| std::env::var("ZED_ALWAYS_ACTIVE").is_ok_and(|e| !e.is_empty()));
  81
  82pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(500);
  83pub const MAX_RECONNECTION_DELAY: Duration = Duration::from_secs(30);
  84pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(20);
  85
  86actions!(
  87    client,
  88    [
  89        /// Signs in to Zed account.
  90        SignIn,
  91        /// Signs out of Zed account.
  92        SignOut,
  93        /// Reconnects to the collaboration server.
  94        Reconnect
  95    ]
  96);
  97
  98#[derive(Deserialize)]
  99pub struct ClientSettings {
 100    pub server_url: String,
 101}
 102
 103impl Settings for ClientSettings {
 104    fn from_settings(content: &settings::SettingsContent, _cx: &mut App) -> Self {
 105        if let Some(server_url) = &*ZED_SERVER_URL {
 106            return Self {
 107                server_url: server_url.clone(),
 108            };
 109        }
 110        Self {
 111            server_url: content.server_url.clone().unwrap(),
 112        }
 113    }
 114}
 115
 116#[derive(Deserialize, Default)]
 117pub struct ProxySettings {
 118    pub proxy: Option<String>,
 119}
 120
 121impl ProxySettings {
 122    pub fn proxy_url(&self) -> Option<Url> {
 123        self.proxy
 124            .as_ref()
 125            .and_then(|input| {
 126                input
 127                    .parse::<Url>()
 128                    .inspect_err(|e| log::error!("Error parsing proxy settings: {}", e))
 129                    .ok()
 130            })
 131            .or_else(read_proxy_from_env)
 132    }
 133}
 134
 135impl Settings for ProxySettings {
 136    fn from_settings(content: &settings::SettingsContent, _cx: &mut App) -> Self {
 137        Self {
 138            proxy: content.proxy.clone(),
 139        }
 140    }
 141
 142    fn import_from_vscode(vscode: &settings::VsCodeSettings, current: &mut SettingsContent) {
 143        vscode.string_setting("http.proxy", &mut current.proxy);
 144    }
 145}
 146
 147pub fn init_settings(cx: &mut App) {
 148    TelemetrySettings::register(cx);
 149    ClientSettings::register(cx);
 150    ProxySettings::register(cx);
 151}
 152
 153#[profiling::function]
 154pub fn init(client: &Arc<Client>, cx: &mut App) {
 155    let client = Arc::downgrade(client);
 156    cx.on_action({
 157        let client = client.clone();
 158        move |_: &SignIn, cx| {
 159            if let Some(client) = client.upgrade() {
 160                cx.spawn(async move |cx| client.sign_in_with_optional_connect(true, cx).await)
 161                    .detach_and_log_err(cx);
 162            }
 163        }
 164    });
 165
 166    cx.on_action({
 167        let client = client.clone();
 168        move |_: &SignOut, cx| {
 169            if let Some(client) = client.upgrade() {
 170                cx.spawn(async move |cx| {
 171                    client.sign_out(cx).await;
 172                })
 173                .detach();
 174            }
 175        }
 176    });
 177
 178    cx.on_action({
 179        let client = client;
 180        move |_: &Reconnect, cx| {
 181            if let Some(client) = client.upgrade() {
 182                cx.spawn(async move |cx| {
 183                    client.reconnect(cx);
 184                })
 185                .detach();
 186            }
 187        }
 188    });
 189}
 190
 191pub type MessageToClientHandler = Box<dyn Fn(&MessageToClient, &mut App) + Send + Sync + 'static>;
 192
 193struct GlobalClient(Arc<Client>);
 194
 195impl Global for GlobalClient {}
 196
 197pub struct Client {
 198    id: AtomicU64,
 199    peer: Arc<Peer>,
 200    http: Arc<HttpClientWithUrl>,
 201    cloud_client: Arc<CloudApiClient>,
 202    telemetry: Arc<Telemetry>,
 203    credentials_provider: ClientCredentialsProvider,
 204    state: RwLock<ClientState>,
 205    handler_set: parking_lot::Mutex<ProtoMessageHandlerSet>,
 206    message_to_client_handlers: parking_lot::Mutex<Vec<MessageToClientHandler>>,
 207
 208    #[allow(clippy::type_complexity)]
 209    #[cfg(any(test, feature = "test-support"))]
 210    authenticate:
 211        RwLock<Option<Box<dyn 'static + Send + Sync + Fn(&AsyncApp) -> Task<Result<Credentials>>>>>,
 212
 213    #[allow(clippy::type_complexity)]
 214    #[cfg(any(test, feature = "test-support"))]
 215    establish_connection: RwLock<
 216        Option<
 217            Box<
 218                dyn 'static
 219                    + Send
 220                    + Sync
 221                    + Fn(
 222                        &Credentials,
 223                        &AsyncApp,
 224                    ) -> Task<Result<Connection, EstablishConnectionError>>,
 225            >,
 226        >,
 227    >,
 228
 229    #[cfg(any(test, feature = "test-support"))]
 230    rpc_url: RwLock<Option<Url>>,
 231}
 232
 233#[derive(Error, Debug)]
 234pub enum EstablishConnectionError {
 235    #[error("upgrade required")]
 236    UpgradeRequired,
 237    #[error("unauthorized")]
 238    Unauthorized,
 239    #[error("{0}")]
 240    Other(#[from] anyhow::Error),
 241    #[error("{0}")]
 242    InvalidHeaderValue(#[from] async_tungstenite::tungstenite::http::header::InvalidHeaderValue),
 243    #[error("{0}")]
 244    Io(#[from] std::io::Error),
 245    #[error("{0}")]
 246    Websocket(#[from] async_tungstenite::tungstenite::http::Error),
 247}
 248
 249impl From<WebsocketError> for EstablishConnectionError {
 250    fn from(error: WebsocketError) -> Self {
 251        if let WebsocketError::Http(response) = &error {
 252            match response.status() {
 253                StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,
 254                StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,
 255                _ => {}
 256            }
 257        }
 258        EstablishConnectionError::Other(error.into())
 259    }
 260}
 261
 262impl EstablishConnectionError {
 263    pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {
 264        Self::Other(error.into())
 265    }
 266}
 267
 268#[derive(Copy, Clone, Debug, PartialEq)]
 269pub enum Status {
 270    SignedOut,
 271    UpgradeRequired,
 272    Authenticating,
 273    Authenticated,
 274    AuthenticationError,
 275    Connecting,
 276    ConnectionError,
 277    Connected {
 278        peer_id: PeerId,
 279        connection_id: ConnectionId,
 280    },
 281    ConnectionLost,
 282    Reauthenticating,
 283    Reauthenticated,
 284    Reconnecting,
 285    ReconnectionError {
 286        next_reconnection: Instant,
 287    },
 288}
 289
 290impl Status {
 291    pub fn is_connected(&self) -> bool {
 292        matches!(self, Self::Connected { .. })
 293    }
 294
 295    pub fn was_connected(&self) -> bool {
 296        matches!(
 297            self,
 298            Self::ConnectionLost
 299                | Self::Reauthenticating
 300                | Self::Reauthenticated
 301                | Self::Reconnecting
 302        )
 303    }
 304
 305    /// Returns whether the client is currently connected or was connected at some point.
 306    pub fn is_or_was_connected(&self) -> bool {
 307        self.is_connected() || self.was_connected()
 308    }
 309
 310    pub fn is_signing_in(&self) -> bool {
 311        matches!(
 312            self,
 313            Self::Authenticating | Self::Reauthenticating | Self::Connecting | Self::Reconnecting
 314        )
 315    }
 316
 317    pub fn is_signed_out(&self) -> bool {
 318        matches!(self, Self::SignedOut | Self::UpgradeRequired)
 319    }
 320}
 321
 322struct ClientState {
 323    credentials: Option<Credentials>,
 324    status: (watch::Sender<Status>, watch::Receiver<Status>),
 325    _reconnect_task: Option<Task<()>>,
 326}
 327
 328#[derive(Clone, Debug, Eq, PartialEq)]
 329pub struct Credentials {
 330    pub user_id: u64,
 331    pub access_token: String,
 332}
 333
 334impl Credentials {
 335    pub fn authorization_header(&self) -> String {
 336        format!("{} {}", self.user_id, self.access_token)
 337    }
 338}
 339
 340pub struct ClientCredentialsProvider {
 341    provider: Arc<dyn CredentialsProvider>,
 342}
 343
 344impl ClientCredentialsProvider {
 345    pub fn new(cx: &App) -> Self {
 346        Self {
 347            provider: <dyn CredentialsProvider>::global(cx),
 348        }
 349    }
 350
 351    fn server_url(&self, cx: &AsyncApp) -> Result<String> {
 352        cx.update(|cx| ClientSettings::get_global(cx).server_url.clone())
 353    }
 354
 355    /// Reads the credentials from the provider.
 356    fn read_credentials<'a>(
 357        &'a self,
 358        cx: &'a AsyncApp,
 359    ) -> Pin<Box<dyn Future<Output = Option<Credentials>> + 'a>> {
 360        async move {
 361            if IMPERSONATE_LOGIN.is_some() {
 362                return None;
 363            }
 364
 365            let server_url = self.server_url(cx).ok()?;
 366            let (user_id, access_token) = self
 367                .provider
 368                .read_credentials(&server_url, cx)
 369                .await
 370                .log_err()
 371                .flatten()?;
 372
 373            Some(Credentials {
 374                user_id: user_id.parse().ok()?,
 375                access_token: String::from_utf8(access_token).ok()?,
 376            })
 377        }
 378        .boxed_local()
 379    }
 380
 381    /// Writes the credentials to the provider.
 382    fn write_credentials<'a>(
 383        &'a self,
 384        user_id: u64,
 385        access_token: String,
 386        cx: &'a AsyncApp,
 387    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
 388        async move {
 389            let server_url = self.server_url(cx)?;
 390            self.provider
 391                .write_credentials(
 392                    &server_url,
 393                    &user_id.to_string(),
 394                    access_token.as_bytes(),
 395                    cx,
 396                )
 397                .await
 398        }
 399        .boxed_local()
 400    }
 401
 402    /// Deletes the credentials from the provider.
 403    fn delete_credentials<'a>(
 404        &'a self,
 405        cx: &'a AsyncApp,
 406    ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
 407        async move {
 408            let server_url = self.server_url(cx)?;
 409            self.provider.delete_credentials(&server_url, cx).await
 410        }
 411        .boxed_local()
 412    }
 413}
 414
 415impl Default for ClientState {
 416    fn default() -> Self {
 417        Self {
 418            credentials: None,
 419            status: watch::channel_with(Status::SignedOut),
 420            _reconnect_task: None,
 421        }
 422    }
 423}
 424
 425pub enum Subscription {
 426    Entity {
 427        client: Weak<Client>,
 428        id: (TypeId, u64),
 429    },
 430    Message {
 431        client: Weak<Client>,
 432        id: TypeId,
 433    },
 434}
 435
 436impl Drop for Subscription {
 437    fn drop(&mut self) {
 438        match self {
 439            Subscription::Entity { client, id } => {
 440                if let Some(client) = client.upgrade() {
 441                    let mut state = client.handler_set.lock();
 442                    let _ = state.entities_by_type_and_remote_id.remove(id);
 443                }
 444            }
 445            Subscription::Message { client, id } => {
 446                if let Some(client) = client.upgrade() {
 447                    let mut state = client.handler_set.lock();
 448                    let _ = state.entity_types_by_message_type.remove(id);
 449                    let _ = state.message_handlers.remove(id);
 450                }
 451            }
 452        }
 453    }
 454}
 455
 456pub struct PendingEntitySubscription<T: 'static> {
 457    client: Arc<Client>,
 458    remote_id: u64,
 459    _entity_type: PhantomData<T>,
 460    consumed: bool,
 461}
 462
 463impl<T: 'static> PendingEntitySubscription<T> {
 464    pub fn set_entity(mut self, entity: &Entity<T>, cx: &AsyncApp) -> Subscription {
 465        self.consumed = true;
 466        let mut handlers = self.client.handler_set.lock();
 467        let id = (TypeId::of::<T>(), self.remote_id);
 468        let Some(EntityMessageSubscriber::Pending(messages)) =
 469            handlers.entities_by_type_and_remote_id.remove(&id)
 470        else {
 471            unreachable!()
 472        };
 473
 474        handlers.entities_by_type_and_remote_id.insert(
 475            id,
 476            EntityMessageSubscriber::Entity {
 477                handle: entity.downgrade().into(),
 478            },
 479        );
 480        drop(handlers);
 481        for message in messages {
 482            let client_id = self.client.id();
 483            let type_name = message.payload_type_name();
 484            let sender_id = message.original_sender_id();
 485            log::debug!(
 486                "handling queued rpc message. client_id:{}, sender_id:{:?}, type:{}",
 487                client_id,
 488                sender_id,
 489                type_name
 490            );
 491            self.client.handle_message(message, cx);
 492        }
 493        Subscription::Entity {
 494            client: Arc::downgrade(&self.client),
 495            id,
 496        }
 497    }
 498}
 499
 500impl<T: 'static> Drop for PendingEntitySubscription<T> {
 501    fn drop(&mut self) {
 502        if !self.consumed {
 503            let mut state = self.client.handler_set.lock();
 504            if let Some(EntityMessageSubscriber::Pending(messages)) = state
 505                .entities_by_type_and_remote_id
 506                .remove(&(TypeId::of::<T>(), self.remote_id))
 507            {
 508                for message in messages {
 509                    log::info!("unhandled message {}", message.payload_type_name());
 510                }
 511            }
 512        }
 513    }
 514}
 515
 516#[derive(Copy, Clone, Deserialize, Debug)]
 517pub struct TelemetrySettings {
 518    pub diagnostics: bool,
 519    pub metrics: bool,
 520}
 521
 522impl settings::Settings for TelemetrySettings {
 523    fn from_settings(content: &SettingsContent, _cx: &mut App) -> Self {
 524        Self {
 525            diagnostics: content.telemetry.as_ref().unwrap().diagnostics.unwrap(),
 526            metrics: content.telemetry.as_ref().unwrap().metrics.unwrap(),
 527        }
 528    }
 529
 530    fn import_from_vscode(vscode: &settings::VsCodeSettings, current: &mut SettingsContent) {
 531        let mut telemetry = settings::TelemetrySettingsContent::default();
 532        vscode.enum_setting("telemetry.telemetryLevel", &mut telemetry.metrics, |s| {
 533            Some(s == "all")
 534        });
 535        vscode.enum_setting(
 536            "telemetry.telemetryLevel",
 537            &mut telemetry.diagnostics,
 538            |s| Some(matches!(s, "all" | "error" | "crash")),
 539        );
 540        // we could translate telemetry.telemetryLevel, but just because users didn't want
 541        // to send microsoft telemetry doesn't mean they don't want to send it to zed. their
 542        // all/error/crash/off correspond to combinations of our "diagnostics" and "metrics".
 543        if let Some(diagnostics) = telemetry.diagnostics {
 544            current.telemetry.get_or_insert_default().diagnostics = Some(diagnostics)
 545        }
 546        if let Some(metrics) = telemetry.metrics {
 547            current.telemetry.get_or_insert_default().metrics = Some(metrics)
 548        }
 549    }
 550}
 551
 552impl Client {
 553    pub fn new(
 554        clock: Arc<dyn SystemClock>,
 555        http: Arc<HttpClientWithUrl>,
 556        cx: &mut App,
 557    ) -> Arc<Self> {
 558        Arc::new(Self {
 559            id: AtomicU64::new(0),
 560            peer: Peer::new(0),
 561            telemetry: Telemetry::new(clock, http.clone(), cx),
 562            cloud_client: Arc::new(CloudApiClient::new(http.clone())),
 563            http,
 564            credentials_provider: ClientCredentialsProvider::new(cx),
 565            state: Default::default(),
 566            handler_set: Default::default(),
 567            message_to_client_handlers: parking_lot::Mutex::new(Vec::new()),
 568
 569            #[cfg(any(test, feature = "test-support"))]
 570            authenticate: Default::default(),
 571            #[cfg(any(test, feature = "test-support"))]
 572            establish_connection: Default::default(),
 573            #[cfg(any(test, feature = "test-support"))]
 574            rpc_url: RwLock::default(),
 575        })
 576    }
 577
 578    pub fn production(cx: &mut App) -> Arc<Self> {
 579        let clock = Arc::new(clock::RealSystemClock);
 580        let http = Arc::new(HttpClientWithUrl::new_url(
 581            cx.http_client(),
 582            &ClientSettings::get_global(cx).server_url,
 583            cx.http_client().proxy().cloned(),
 584        ));
 585        Self::new(clock, http, cx)
 586    }
 587
 588    pub fn id(&self) -> u64 {
 589        self.id.load(Ordering::SeqCst)
 590    }
 591
 592    pub fn http_client(&self) -> Arc<HttpClientWithUrl> {
 593        self.http.clone()
 594    }
 595
 596    pub fn cloud_client(&self) -> Arc<CloudApiClient> {
 597        self.cloud_client.clone()
 598    }
 599
 600    pub fn set_id(&self, id: u64) -> &Self {
 601        self.id.store(id, Ordering::SeqCst);
 602        self
 603    }
 604
 605    #[cfg(any(test, feature = "test-support"))]
 606    pub fn teardown(&self) {
 607        let mut state = self.state.write();
 608        state._reconnect_task.take();
 609        self.handler_set.lock().clear();
 610        self.peer.teardown();
 611    }
 612
 613    #[cfg(any(test, feature = "test-support"))]
 614    pub fn override_authenticate<F>(&self, authenticate: F) -> &Self
 615    where
 616        F: 'static + Send + Sync + Fn(&AsyncApp) -> Task<Result<Credentials>>,
 617    {
 618        *self.authenticate.write() = Some(Box::new(authenticate));
 619        self
 620    }
 621
 622    #[cfg(any(test, feature = "test-support"))]
 623    pub fn override_establish_connection<F>(&self, connect: F) -> &Self
 624    where
 625        F: 'static
 626            + Send
 627            + Sync
 628            + Fn(&Credentials, &AsyncApp) -> Task<Result<Connection, EstablishConnectionError>>,
 629    {
 630        *self.establish_connection.write() = Some(Box::new(connect));
 631        self
 632    }
 633
 634    #[cfg(any(test, feature = "test-support"))]
 635    pub fn override_rpc_url(&self, url: Url) -> &Self {
 636        *self.rpc_url.write() = Some(url);
 637        self
 638    }
 639
 640    pub fn global(cx: &App) -> Arc<Self> {
 641        cx.global::<GlobalClient>().0.clone()
 642    }
 643    pub fn set_global(client: Arc<Client>, cx: &mut App) {
 644        cx.set_global(GlobalClient(client))
 645    }
 646
 647    pub fn user_id(&self) -> Option<u64> {
 648        self.state
 649            .read()
 650            .credentials
 651            .as_ref()
 652            .map(|credentials| credentials.user_id)
 653    }
 654
 655    pub fn peer_id(&self) -> Option<PeerId> {
 656        if let Status::Connected { peer_id, .. } = &*self.status().borrow() {
 657            Some(*peer_id)
 658        } else {
 659            None
 660        }
 661    }
 662
 663    pub fn status(&self) -> watch::Receiver<Status> {
 664        self.state.read().status.1.clone()
 665    }
 666
 667    fn set_status(self: &Arc<Self>, status: Status, cx: &AsyncApp) {
 668        log::info!("set status on client {}: {:?}", self.id(), status);
 669        let mut state = self.state.write();
 670        *state.status.0.borrow_mut() = status;
 671
 672        match status {
 673            Status::Connected { .. } => {
 674                state._reconnect_task = None;
 675            }
 676            Status::ConnectionLost => {
 677                let client = self.clone();
 678                state._reconnect_task = Some(cx.spawn(async move |cx| {
 679                    #[cfg(any(test, feature = "test-support"))]
 680                    let mut rng = StdRng::seed_from_u64(0);
 681                    #[cfg(not(any(test, feature = "test-support")))]
 682                    let mut rng = StdRng::from_os_rng();
 683
 684                    let mut delay = INITIAL_RECONNECTION_DELAY;
 685                    loop {
 686                        match client.connect(true, cx).await {
 687                            ConnectionResult::Timeout => {
 688                                log::error!("client connect attempt timed out")
 689                            }
 690                            ConnectionResult::ConnectionReset => {
 691                                log::error!("client connect attempt reset")
 692                            }
 693                            ConnectionResult::Result(r) => {
 694                                if let Err(error) = r {
 695                                    log::error!("failed to connect: {error}");
 696                                } else {
 697                                    break;
 698                                }
 699                            }
 700                        }
 701
 702                        if matches!(
 703                            *client.status().borrow(),
 704                            Status::AuthenticationError | Status::ConnectionError
 705                        ) {
 706                            client.set_status(
 707                                Status::ReconnectionError {
 708                                    next_reconnection: Instant::now() + delay,
 709                                },
 710                                cx,
 711                            );
 712                            let jitter = Duration::from_millis(
 713                                rng.random_range(0..delay.as_millis() as u64),
 714                            );
 715                            cx.background_executor().timer(delay + jitter).await;
 716                            delay = cmp::min(delay * 2, MAX_RECONNECTION_DELAY);
 717                        } else {
 718                            break;
 719                        }
 720                    }
 721                }));
 722            }
 723            Status::SignedOut | Status::UpgradeRequired => {
 724                self.telemetry.set_authenticated_user_info(None, false);
 725                state._reconnect_task.take();
 726            }
 727            _ => {}
 728        }
 729    }
 730
 731    pub fn subscribe_to_entity<T>(
 732        self: &Arc<Self>,
 733        remote_id: u64,
 734    ) -> Result<PendingEntitySubscription<T>>
 735    where
 736        T: 'static,
 737    {
 738        let id = (TypeId::of::<T>(), remote_id);
 739
 740        let mut state = self.handler_set.lock();
 741        anyhow::ensure!(
 742            !state.entities_by_type_and_remote_id.contains_key(&id),
 743            "already subscribed to entity"
 744        );
 745
 746        state
 747            .entities_by_type_and_remote_id
 748            .insert(id, EntityMessageSubscriber::Pending(Default::default()));
 749
 750        Ok(PendingEntitySubscription {
 751            client: self.clone(),
 752            remote_id,
 753            consumed: false,
 754            _entity_type: PhantomData,
 755        })
 756    }
 757
 758    #[track_caller]
 759    pub fn add_message_handler<M, E, H, F>(
 760        self: &Arc<Self>,
 761        entity: WeakEntity<E>,
 762        handler: H,
 763    ) -> Subscription
 764    where
 765        M: EnvelopedMessage,
 766        E: 'static,
 767        H: 'static + Sync + Fn(Entity<E>, TypedEnvelope<M>, AsyncApp) -> F + Send + Sync,
 768        F: 'static + Future<Output = Result<()>>,
 769    {
 770        self.add_message_handler_impl(entity, move |entity, message, _, cx| {
 771            handler(entity, message, cx)
 772        })
 773    }
 774
 775    fn add_message_handler_impl<M, E, H, F>(
 776        self: &Arc<Self>,
 777        entity: WeakEntity<E>,
 778        handler: H,
 779    ) -> Subscription
 780    where
 781        M: EnvelopedMessage,
 782        E: 'static,
 783        H: 'static
 784            + Sync
 785            + Fn(Entity<E>, TypedEnvelope<M>, AnyProtoClient, AsyncApp) -> F
 786            + Send
 787            + Sync,
 788        F: 'static + Future<Output = Result<()>>,
 789    {
 790        let message_type_id = TypeId::of::<M>();
 791        let mut state = self.handler_set.lock();
 792        state
 793            .entities_by_message_type
 794            .insert(message_type_id, entity.into());
 795
 796        let prev_handler = state.message_handlers.insert(
 797            message_type_id,
 798            Arc::new(move |subscriber, envelope, client, cx| {
 799                let subscriber = subscriber.downcast::<E>().unwrap();
 800                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
 801                handler(subscriber, *envelope, client, cx).boxed_local()
 802            }),
 803        );
 804        if prev_handler.is_some() {
 805            let location = std::panic::Location::caller();
 806            panic!(
 807                "{}:{} registered handler for the same message {} twice",
 808                location.file(),
 809                location.line(),
 810                std::any::type_name::<M>()
 811            );
 812        }
 813
 814        Subscription::Message {
 815            client: Arc::downgrade(self),
 816            id: message_type_id,
 817        }
 818    }
 819
 820    pub fn add_request_handler<M, E, H, F>(
 821        self: &Arc<Self>,
 822        entity: WeakEntity<E>,
 823        handler: H,
 824    ) -> Subscription
 825    where
 826        M: RequestMessage,
 827        E: 'static,
 828        H: 'static + Sync + Fn(Entity<E>, TypedEnvelope<M>, AsyncApp) -> F + Send + Sync,
 829        F: 'static + Future<Output = Result<M::Response>>,
 830    {
 831        self.add_message_handler_impl(entity, move |handle, envelope, this, cx| {
 832            Self::respond_to_request(envelope.receipt(), handler(handle, envelope, cx), this)
 833        })
 834    }
 835
 836    async fn respond_to_request<T: RequestMessage, F: Future<Output = Result<T::Response>>>(
 837        receipt: Receipt<T>,
 838        response: F,
 839        client: AnyProtoClient,
 840    ) -> Result<()> {
 841        match response.await {
 842            Ok(response) => {
 843                client.send_response(receipt.message_id, response)?;
 844                Ok(())
 845            }
 846            Err(error) => {
 847                client.send_response(receipt.message_id, error.to_proto())?;
 848                Err(error)
 849            }
 850        }
 851    }
 852
 853    pub async fn has_credentials(&self, cx: &AsyncApp) -> bool {
 854        self.credentials_provider
 855            .read_credentials(cx)
 856            .await
 857            .is_some()
 858    }
 859
 860    pub async fn sign_in(
 861        self: &Arc<Self>,
 862        try_provider: bool,
 863        cx: &AsyncApp,
 864    ) -> Result<Credentials> {
 865        let is_reauthenticating = if self.status().borrow().is_signed_out() {
 866            self.set_status(Status::Authenticating, cx);
 867            false
 868        } else {
 869            self.set_status(Status::Reauthenticating, cx);
 870            true
 871        };
 872
 873        let mut credentials = None;
 874
 875        let old_credentials = self.state.read().credentials.clone();
 876        if let Some(old_credentials) = old_credentials
 877            && self.validate_credentials(&old_credentials, cx).await?
 878        {
 879            credentials = Some(old_credentials);
 880        }
 881
 882        if credentials.is_none()
 883            && try_provider
 884            && let Some(stored_credentials) = self.credentials_provider.read_credentials(cx).await
 885        {
 886            if self.validate_credentials(&stored_credentials, cx).await? {
 887                credentials = Some(stored_credentials);
 888            } else {
 889                self.credentials_provider
 890                    .delete_credentials(cx)
 891                    .await
 892                    .log_err();
 893            }
 894        }
 895
 896        if credentials.is_none() {
 897            let mut status_rx = self.status();
 898            let _ = status_rx.next().await;
 899            futures::select_biased! {
 900                authenticate = self.authenticate(cx).fuse() => {
 901                    match authenticate {
 902                        Ok(creds) => {
 903                            if IMPERSONATE_LOGIN.is_none() {
 904                                self.credentials_provider
 905                                    .write_credentials(creds.user_id, creds.access_token.clone(), cx)
 906                                    .await
 907                                    .log_err();
 908                            }
 909
 910                            credentials = Some(creds);
 911                        },
 912                        Err(err) => {
 913                            self.set_status(Status::AuthenticationError, cx);
 914                            return Err(err);
 915                        }
 916                    }
 917                }
 918                _ = status_rx.next().fuse() => {
 919                    return Err(anyhow!("authentication canceled"));
 920                }
 921            }
 922        }
 923
 924        let credentials = credentials.unwrap();
 925        self.set_id(credentials.user_id);
 926        self.cloud_client
 927            .set_credentials(credentials.user_id as u32, credentials.access_token.clone());
 928        self.state.write().credentials = Some(credentials.clone());
 929        self.set_status(
 930            if is_reauthenticating {
 931                Status::Reauthenticated
 932            } else {
 933                Status::Authenticated
 934            },
 935            cx,
 936        );
 937
 938        Ok(credentials)
 939    }
 940
 941    async fn validate_credentials(
 942        self: &Arc<Self>,
 943        credentials: &Credentials,
 944        cx: &AsyncApp,
 945    ) -> Result<bool> {
 946        match self
 947            .cloud_client
 948            .validate_credentials(credentials.user_id as u32, &credentials.access_token)
 949            .await
 950        {
 951            Ok(valid) => Ok(valid),
 952            Err(err) => {
 953                self.set_status(Status::AuthenticationError, cx);
 954                Err(anyhow!("failed to validate credentials: {}", err))
 955            }
 956        }
 957    }
 958
 959    /// Establishes a WebSocket connection with Cloud for receiving updates from the server.
 960    async fn connect_to_cloud(self: &Arc<Self>, cx: &AsyncApp) -> Result<()> {
 961        let connect_task = cx.update({
 962            let cloud_client = self.cloud_client.clone();
 963            move |cx| cloud_client.connect(cx)
 964        })??;
 965        let connection = connect_task.await?;
 966
 967        let (mut messages, task) = cx.update(|cx| connection.spawn(cx))?;
 968        task.detach();
 969
 970        cx.spawn({
 971            let this = self.clone();
 972            async move |cx| {
 973                while let Some(message) = messages.next().await {
 974                    if let Some(message) = message.log_err() {
 975                        this.handle_message_to_client(message, cx);
 976                    }
 977                }
 978            }
 979        })
 980        .detach();
 981
 982        Ok(())
 983    }
 984
 985    /// Performs a sign-in and also (optionally) connects to Collab.
 986    ///
 987    /// Only Zed staff automatically connect to Collab.
 988    pub async fn sign_in_with_optional_connect(
 989        self: &Arc<Self>,
 990        try_provider: bool,
 991        cx: &AsyncApp,
 992    ) -> Result<()> {
 993        // Don't try to sign in again if we're already connected to Collab, as it will temporarily disconnect us.
 994        if self.status().borrow().is_connected() {
 995            return Ok(());
 996        }
 997
 998        let (is_staff_tx, is_staff_rx) = oneshot::channel::<bool>();
 999        let mut is_staff_tx = Some(is_staff_tx);
1000        cx.update(|cx| {
1001            cx.on_flags_ready(move |state, _cx| {
1002                if let Some(is_staff_tx) = is_staff_tx.take() {
1003                    is_staff_tx.send(state.is_staff).log_err();
1004                }
1005            })
1006            .detach();
1007        })
1008        .log_err();
1009
1010        let credentials = self.sign_in(try_provider, cx).await?;
1011
1012        self.connect_to_cloud(cx).await.log_err();
1013
1014        cx.update(move |cx| {
1015            cx.spawn({
1016                let client = self.clone();
1017                async move |cx| {
1018                    let is_staff = is_staff_rx.await?;
1019                    if is_staff {
1020                        match client.connect_with_credentials(credentials, cx).await {
1021                            ConnectionResult::Timeout => Err(anyhow!("connection timed out")),
1022                            ConnectionResult::ConnectionReset => Err(anyhow!("connection reset")),
1023                            ConnectionResult::Result(result) => {
1024                                result.context("client auth and connect")
1025                            }
1026                        }
1027                    } else {
1028                        Ok(())
1029                    }
1030                }
1031            })
1032            .detach_and_log_err(cx);
1033        })
1034        .log_err();
1035
1036        Ok(())
1037    }
1038
1039    pub async fn connect(
1040        self: &Arc<Self>,
1041        try_provider: bool,
1042        cx: &AsyncApp,
1043    ) -> ConnectionResult<()> {
1044        let was_disconnected = match *self.status().borrow() {
1045            Status::SignedOut | Status::Authenticated => true,
1046            Status::ConnectionError
1047            | Status::ConnectionLost
1048            | Status::Authenticating
1049            | Status::AuthenticationError
1050            | Status::Reauthenticating
1051            | Status::Reauthenticated
1052            | Status::ReconnectionError { .. } => false,
1053            Status::Connected { .. } | Status::Connecting | Status::Reconnecting => {
1054                return ConnectionResult::Result(Ok(()));
1055            }
1056            Status::UpgradeRequired => {
1057                return ConnectionResult::Result(
1058                    Err(EstablishConnectionError::UpgradeRequired)
1059                        .context("client auth and connect"),
1060                );
1061            }
1062        };
1063        let credentials = match self.sign_in(try_provider, cx).await {
1064            Ok(credentials) => credentials,
1065            Err(err) => return ConnectionResult::Result(Err(err)),
1066        };
1067
1068        if was_disconnected {
1069            self.set_status(Status::Connecting, cx);
1070        } else {
1071            self.set_status(Status::Reconnecting, cx);
1072        }
1073
1074        self.connect_with_credentials(credentials, cx).await
1075    }
1076
1077    async fn connect_with_credentials(
1078        self: &Arc<Self>,
1079        credentials: Credentials,
1080        cx: &AsyncApp,
1081    ) -> ConnectionResult<()> {
1082        let mut timeout =
1083            futures::FutureExt::fuse(cx.background_executor().timer(CONNECTION_TIMEOUT));
1084        futures::select_biased! {
1085            connection = self.establish_connection(&credentials, cx).fuse() => {
1086                match connection {
1087                    Ok(conn) => {
1088                        futures::select_biased! {
1089                            result = self.set_connection(conn, cx).fuse() => {
1090                                match result.context("client auth and connect") {
1091                                    Ok(()) => ConnectionResult::Result(Ok(())),
1092                                    Err(err) => {
1093                                        self.set_status(Status::ConnectionError, cx);
1094                                        ConnectionResult::Result(Err(err))
1095                                    },
1096                                }
1097                            },
1098                            _ = timeout => {
1099                                self.set_status(Status::ConnectionError, cx);
1100                                ConnectionResult::Timeout
1101                            }
1102                        }
1103                    }
1104                    Err(EstablishConnectionError::Unauthorized) => {
1105                        self.set_status(Status::ConnectionError, cx);
1106                        ConnectionResult::Result(Err(EstablishConnectionError::Unauthorized).context("client auth and connect"))
1107                    }
1108                    Err(EstablishConnectionError::UpgradeRequired) => {
1109                        self.set_status(Status::UpgradeRequired, cx);
1110                        ConnectionResult::Result(Err(EstablishConnectionError::UpgradeRequired).context("client auth and connect"))
1111                    }
1112                    Err(error) => {
1113                        self.set_status(Status::ConnectionError, cx);
1114                        ConnectionResult::Result(Err(error).context("client auth and connect"))
1115                    }
1116                }
1117            }
1118            _ = &mut timeout => {
1119                self.set_status(Status::ConnectionError, cx);
1120                ConnectionResult::Timeout
1121            }
1122        }
1123    }
1124
1125    async fn set_connection(self: &Arc<Self>, conn: Connection, cx: &AsyncApp) -> Result<()> {
1126        let executor = cx.background_executor();
1127        log::debug!("add connection to peer");
1128        let (connection_id, handle_io, mut incoming) = self.peer.add_connection(conn, {
1129            let executor = executor.clone();
1130            move |duration| executor.timer(duration)
1131        });
1132        let handle_io = executor.spawn(handle_io);
1133
1134        let peer_id = async {
1135            log::debug!("waiting for server hello");
1136            let message = incoming.next().await.context("no hello message received")?;
1137            log::debug!("got server hello");
1138            let hello_message_type_name = message.payload_type_name().to_string();
1139            let hello = message
1140                .into_any()
1141                .downcast::<TypedEnvelope<proto::Hello>>()
1142                .map_err(|_| {
1143                    anyhow!(
1144                        "invalid hello message received: {:?}",
1145                        hello_message_type_name
1146                    )
1147                })?;
1148            let peer_id = hello.payload.peer_id.context("invalid peer id")?;
1149            Ok(peer_id)
1150        };
1151
1152        let peer_id = match peer_id.await {
1153            Ok(peer_id) => peer_id,
1154            Err(error) => {
1155                self.peer.disconnect(connection_id);
1156                return Err(error);
1157            }
1158        };
1159
1160        log::debug!(
1161            "set status to connected (connection id: {:?}, peer id: {:?})",
1162            connection_id,
1163            peer_id
1164        );
1165        self.set_status(
1166            Status::Connected {
1167                peer_id,
1168                connection_id,
1169            },
1170            cx,
1171        );
1172
1173        cx.spawn({
1174            let this = self.clone();
1175            async move |cx| {
1176                while let Some(message) = incoming.next().await {
1177                    this.handle_message(message, cx);
1178                    // Don't starve the main thread when receiving lots of messages at once.
1179                    smol::future::yield_now().await;
1180                }
1181            }
1182        })
1183        .detach();
1184
1185        cx.spawn({
1186            let this = self.clone();
1187            async move |cx| match handle_io.await {
1188                Ok(()) => {
1189                    if *this.status().borrow()
1190                        == (Status::Connected {
1191                            connection_id,
1192                            peer_id,
1193                        })
1194                    {
1195                        this.set_status(Status::SignedOut, cx);
1196                    }
1197                }
1198                Err(err) => {
1199                    log::error!("connection error: {:?}", err);
1200                    this.set_status(Status::ConnectionLost, cx);
1201                }
1202            }
1203        })
1204        .detach();
1205
1206        Ok(())
1207    }
1208
1209    fn authenticate(self: &Arc<Self>, cx: &AsyncApp) -> Task<Result<Credentials>> {
1210        #[cfg(any(test, feature = "test-support"))]
1211        if let Some(callback) = self.authenticate.read().as_ref() {
1212            return callback(cx);
1213        }
1214
1215        self.authenticate_with_browser(cx)
1216    }
1217
1218    fn establish_connection(
1219        self: &Arc<Self>,
1220        credentials: &Credentials,
1221        cx: &AsyncApp,
1222    ) -> Task<Result<Connection, EstablishConnectionError>> {
1223        #[cfg(any(test, feature = "test-support"))]
1224        if let Some(callback) = self.establish_connection.read().as_ref() {
1225            return callback(credentials, cx);
1226        }
1227
1228        self.establish_websocket_connection(credentials, cx)
1229    }
1230
1231    fn rpc_url(
1232        &self,
1233        http: Arc<HttpClientWithUrl>,
1234        release_channel: Option<ReleaseChannel>,
1235    ) -> impl Future<Output = Result<url::Url>> + use<> {
1236        #[cfg(any(test, feature = "test-support"))]
1237        let url_override = self.rpc_url.read().clone();
1238
1239        async move {
1240            #[cfg(any(test, feature = "test-support"))]
1241            if let Some(url) = url_override {
1242                return Ok(url);
1243            }
1244
1245            if let Some(url) = &*ZED_RPC_URL {
1246                return Url::parse(url).context("invalid rpc url");
1247            }
1248
1249            let mut url = http.build_url("/rpc");
1250            if let Some(preview_param) =
1251                release_channel.and_then(|channel| channel.release_query_param())
1252            {
1253                url += "?";
1254                url += preview_param;
1255            }
1256
1257            let response = http.get(&url, Default::default(), false).await?;
1258            anyhow::ensure!(
1259                response.status().is_redirection(),
1260                "unexpected /rpc response status {}",
1261                response.status()
1262            );
1263            let collab_url = response
1264                .headers()
1265                .get("Location")
1266                .context("missing location header in /rpc response")?
1267                .to_str()
1268                .map_err(EstablishConnectionError::other)?
1269                .to_string();
1270            Url::parse(&collab_url).with_context(|| format!("parsing collab rpc url {collab_url}"))
1271        }
1272    }
1273
1274    fn establish_websocket_connection(
1275        self: &Arc<Self>,
1276        credentials: &Credentials,
1277        cx: &AsyncApp,
1278    ) -> Task<Result<Connection, EstablishConnectionError>> {
1279        let release_channel = cx
1280            .update(|cx| ReleaseChannel::try_global(cx))
1281            .ok()
1282            .flatten();
1283        let app_version = cx
1284            .update(|cx| AppVersion::global(cx).to_string())
1285            .ok()
1286            .unwrap_or_default();
1287
1288        let http = self.http.clone();
1289        let proxy = http.proxy().cloned();
1290        let user_agent = http.user_agent().cloned();
1291        let credentials = credentials.clone();
1292        let rpc_url = self.rpc_url(http, release_channel);
1293        let system_id = self.telemetry.system_id();
1294        let metrics_id = self.telemetry.metrics_id();
1295        cx.spawn(async move |cx| {
1296            use HttpOrHttps::*;
1297
1298            #[derive(Debug)]
1299            enum HttpOrHttps {
1300                Http,
1301                Https,
1302            }
1303
1304            let mut rpc_url = rpc_url.await?;
1305            let url_scheme = match rpc_url.scheme() {
1306                "https" => Https,
1307                "http" => Http,
1308                _ => Err(anyhow!("invalid rpc url: {}", rpc_url))?,
1309            };
1310
1311            let stream = gpui_tokio::Tokio::spawn_result(cx, {
1312                let rpc_url = rpc_url.clone();
1313                async move {
1314                    let rpc_host = rpc_url
1315                        .host_str()
1316                        .zip(rpc_url.port_or_known_default())
1317                        .context("missing host in rpc url")?;
1318                    Ok(match proxy {
1319                        Some(proxy) => connect_proxy_stream(&proxy, rpc_host).await?,
1320                        None => Box::new(TcpStream::connect(rpc_host).await?),
1321                    })
1322                }
1323            })?
1324            .await?;
1325
1326            log::info!("connected to rpc endpoint {}", rpc_url);
1327
1328            rpc_url
1329                .set_scheme(match url_scheme {
1330                    Https => "wss",
1331                    Http => "ws",
1332                })
1333                .unwrap();
1334
1335            // We call `into_client_request` to let `tungstenite` construct the WebSocket request
1336            // for us from the RPC URL.
1337            //
1338            // Among other things, it will generate and set a `Sec-WebSocket-Key` header for us.
1339            let mut request = IntoClientRequest::into_client_request(rpc_url.as_str())?;
1340
1341            // We then modify the request to add our desired headers.
1342            let request_headers = request.headers_mut();
1343            request_headers.insert(
1344                http::header::AUTHORIZATION,
1345                HeaderValue::from_str(&credentials.authorization_header())?,
1346            );
1347            request_headers.insert(
1348                "x-zed-protocol-version",
1349                HeaderValue::from_str(&rpc::PROTOCOL_VERSION.to_string())?,
1350            );
1351            request_headers.insert("x-zed-app-version", HeaderValue::from_str(&app_version)?);
1352            request_headers.insert(
1353                "x-zed-release-channel",
1354                HeaderValue::from_str(release_channel.map(|r| r.dev_name()).unwrap_or("unknown"))?,
1355            );
1356            if let Some(user_agent) = user_agent {
1357                request_headers.insert(http::header::USER_AGENT, user_agent);
1358            }
1359            if let Some(system_id) = system_id {
1360                request_headers.insert("x-zed-system-id", HeaderValue::from_str(&system_id)?);
1361            }
1362            if let Some(metrics_id) = metrics_id {
1363                request_headers.insert("x-zed-metrics-id", HeaderValue::from_str(&metrics_id)?);
1364            }
1365
1366            let (stream, _) = async_tungstenite::tokio::client_async_tls_with_connector_and_config(
1367                request,
1368                stream,
1369                Some(Arc::new(http_client_tls::tls_config()).into()),
1370                None,
1371            )
1372            .await?;
1373
1374            Ok(Connection::new(
1375                stream
1376                    .map_err(|error| anyhow!(error))
1377                    .sink_map_err(|error| anyhow!(error)),
1378            ))
1379        })
1380    }
1381
1382    pub fn authenticate_with_browser(self: &Arc<Self>, cx: &AsyncApp) -> Task<Result<Credentials>> {
1383        let http = self.http.clone();
1384        let this = self.clone();
1385        cx.spawn(async move |cx| {
1386            let background = cx.background_executor().clone();
1387
1388            let (open_url_tx, open_url_rx) = oneshot::channel::<String>();
1389            cx.update(|cx| {
1390                cx.spawn(async move |cx| {
1391                    let url = open_url_rx.await?;
1392                    cx.update(|cx| cx.open_url(&url))
1393                })
1394                .detach_and_log_err(cx);
1395            })
1396            .log_err();
1397
1398            let credentials = background
1399                .clone()
1400                .spawn(async move {
1401                    // Generate a pair of asymmetric encryption keys. The public key will be used by the
1402                    // zed server to encrypt the user's access token, so that it can'be intercepted by
1403                    // any other app running on the user's device.
1404                    let (public_key, private_key) =
1405                        rpc::auth::keypair().expect("failed to generate keypair for auth");
1406                    let public_key_string = String::try_from(public_key)
1407                        .expect("failed to serialize public key for auth");
1408
1409                    if let Some((login, token)) =
1410                        IMPERSONATE_LOGIN.as_ref().zip(ADMIN_API_TOKEN.as_ref())
1411                    {
1412                        if !*USE_WEB_LOGIN {
1413                            eprintln!("authenticate as admin {login}, {token}");
1414
1415                            return this
1416                                .authenticate_as_admin(http, login.clone(), token.clone())
1417                                .await;
1418                        }
1419                    }
1420
1421                    // Start an HTTP server to receive the redirect from Zed's sign-in page.
1422                    let server =
1423                        tiny_http::Server::http("127.0.0.1:0").expect("failed to find open port");
1424                    let port = server.server_addr().port();
1425
1426                    // Open the Zed sign-in page in the user's browser, with query parameters that indicate
1427                    // that the user is signing in from a Zed app running on the same device.
1428                    let mut url = http.build_url(&format!(
1429                        "/native_app_signin?native_app_port={}&native_app_public_key={}",
1430                        port, public_key_string
1431                    ));
1432
1433                    if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
1434                        log::info!("impersonating user @{}", impersonate_login);
1435                        write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
1436                    }
1437
1438                    open_url_tx.send(url).log_err();
1439
1440                    #[derive(Deserialize)]
1441                    struct CallbackParams {
1442                        pub user_id: String,
1443                        pub access_token: String,
1444                    }
1445
1446                    // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
1447                    // access token from the query params.
1448                    //
1449                    // TODO - Avoid ever starting more than one HTTP server. Maybe switch to using a
1450                    // custom URL scheme instead of this local HTTP server.
1451                    let (user_id, access_token) = background
1452                        .spawn(async move {
1453                            for _ in 0..100 {
1454                                if let Some(req) = server.recv_timeout(Duration::from_secs(1))? {
1455                                    let path = req.url();
1456                                    let url = Url::parse(&format!("http://example.com{}", path))
1457                                        .context("failed to parse login notification url")?;
1458                                    let callback_params: CallbackParams =
1459                                        serde_urlencoded::from_str(url.query().unwrap_or_default())
1460                                            .context(
1461                                                "failed to parse sign-in callback query parameters",
1462                                            )?;
1463
1464                                    let post_auth_url =
1465                                        http.build_url("/native_app_signin_succeeded");
1466                                    req.respond(
1467                                        tiny_http::Response::empty(302).with_header(
1468                                            tiny_http::Header::from_bytes(
1469                                                &b"Location"[..],
1470                                                post_auth_url.as_bytes(),
1471                                            )
1472                                            .unwrap(),
1473                                        ),
1474                                    )
1475                                    .context("failed to respond to login http request")?;
1476                                    return Ok((
1477                                        callback_params.user_id,
1478                                        callback_params.access_token,
1479                                    ));
1480                                }
1481                            }
1482
1483                            anyhow::bail!("didn't receive login redirect");
1484                        })
1485                        .await?;
1486
1487                    let access_token = private_key
1488                        .decrypt_string(&access_token)
1489                        .context("failed to decrypt access token")?;
1490
1491                    Ok(Credentials {
1492                        user_id: user_id.parse()?,
1493                        access_token,
1494                    })
1495                })
1496                .await?;
1497
1498            cx.update(|cx| cx.activate(true))?;
1499            Ok(credentials)
1500        })
1501    }
1502
1503    async fn authenticate_as_admin(
1504        self: &Arc<Self>,
1505        http: Arc<HttpClientWithUrl>,
1506        login: String,
1507        api_token: String,
1508    ) -> Result<Credentials> {
1509        #[derive(Serialize)]
1510        struct ImpersonateUserBody {
1511            github_login: String,
1512        }
1513
1514        #[derive(Deserialize)]
1515        struct ImpersonateUserResponse {
1516            user_id: u64,
1517            access_token: String,
1518        }
1519
1520        let url = self
1521            .http
1522            .build_zed_cloud_url("/internal/users/impersonate", &[])?;
1523        let request = Request::post(url.as_str())
1524            .header("Content-Type", "application/json")
1525            .header("Authorization", format!("Bearer {api_token}"))
1526            .body(
1527                serde_json::to_string(&ImpersonateUserBody {
1528                    github_login: login,
1529                })?
1530                .into(),
1531            )?;
1532
1533        let mut response = http.send(request).await?;
1534        let mut body = String::new();
1535        response.body_mut().read_to_string(&mut body).await?;
1536        anyhow::ensure!(
1537            response.status().is_success(),
1538            "admin user request failed {} - {}",
1539            response.status().as_u16(),
1540            body,
1541        );
1542        let response: ImpersonateUserResponse = serde_json::from_str(&body)?;
1543
1544        Ok(Credentials {
1545            user_id: response.user_id,
1546            access_token: response.access_token,
1547        })
1548    }
1549
1550    pub async fn sign_out(self: &Arc<Self>, cx: &AsyncApp) {
1551        self.state.write().credentials = None;
1552        self.cloud_client.clear_credentials();
1553        self.disconnect(cx);
1554
1555        if self.has_credentials(cx).await {
1556            self.credentials_provider
1557                .delete_credentials(cx)
1558                .await
1559                .log_err();
1560        }
1561    }
1562
1563    pub fn disconnect(self: &Arc<Self>, cx: &AsyncApp) {
1564        self.peer.teardown();
1565        self.set_status(Status::SignedOut, cx);
1566    }
1567
1568    pub fn reconnect(self: &Arc<Self>, cx: &AsyncApp) {
1569        self.peer.teardown();
1570        self.set_status(Status::ConnectionLost, cx);
1571    }
1572
1573    fn connection_id(&self) -> Result<ConnectionId> {
1574        if let Status::Connected { connection_id, .. } = *self.status().borrow() {
1575            Ok(connection_id)
1576        } else {
1577            anyhow::bail!("not connected");
1578        }
1579    }
1580
1581    pub fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
1582        log::debug!("rpc send. client_id:{}, name:{}", self.id(), T::NAME);
1583        self.peer.send(self.connection_id()?, message)
1584    }
1585
1586    pub fn request<T: RequestMessage>(
1587        &self,
1588        request: T,
1589    ) -> impl Future<Output = Result<T::Response>> + use<T> {
1590        self.request_envelope(request)
1591            .map_ok(|envelope| envelope.payload)
1592    }
1593
1594    pub fn request_stream<T: RequestMessage>(
1595        &self,
1596        request: T,
1597    ) -> impl Future<Output = Result<impl Stream<Item = Result<T::Response>>>> {
1598        let client_id = self.id.load(Ordering::SeqCst);
1599        log::debug!(
1600            "rpc request start. client_id:{}. name:{}",
1601            client_id,
1602            T::NAME
1603        );
1604        let response = self
1605            .connection_id()
1606            .map(|conn_id| self.peer.request_stream(conn_id, request));
1607        async move {
1608            let response = response?.await;
1609            log::debug!(
1610                "rpc request finish. client_id:{}. name:{}",
1611                client_id,
1612                T::NAME
1613            );
1614            response
1615        }
1616    }
1617
1618    pub fn request_envelope<T: RequestMessage>(
1619        &self,
1620        request: T,
1621    ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> + use<T> {
1622        let client_id = self.id();
1623        log::debug!(
1624            "rpc request start. client_id:{}. name:{}",
1625            client_id,
1626            T::NAME
1627        );
1628        let response = self
1629            .connection_id()
1630            .map(|conn_id| self.peer.request_envelope(conn_id, request));
1631        async move {
1632            let response = response?.await;
1633            log::debug!(
1634                "rpc request finish. client_id:{}. name:{}",
1635                client_id,
1636                T::NAME
1637            );
1638            response
1639        }
1640    }
1641
1642    pub fn request_dynamic(
1643        &self,
1644        envelope: proto::Envelope,
1645        request_type: &'static str,
1646    ) -> impl Future<Output = Result<proto::Envelope>> + use<> {
1647        let client_id = self.id();
1648        log::debug!(
1649            "rpc request start. client_id:{}. name:{}",
1650            client_id,
1651            request_type
1652        );
1653        let response = self
1654            .connection_id()
1655            .map(|conn_id| self.peer.request_dynamic(conn_id, envelope, request_type));
1656        async move {
1657            let response = response?.await;
1658            log::debug!(
1659                "rpc request finish. client_id:{}. name:{}",
1660                client_id,
1661                request_type
1662            );
1663            Ok(response?.0)
1664        }
1665    }
1666
1667    fn handle_message(self: &Arc<Client>, message: Box<dyn AnyTypedEnvelope>, cx: &AsyncApp) {
1668        let sender_id = message.sender_id();
1669        let request_id = message.message_id();
1670        let type_name = message.payload_type_name();
1671        let original_sender_id = message.original_sender_id();
1672
1673        if let Some(future) = ProtoMessageHandlerSet::handle_message(
1674            &self.handler_set,
1675            message,
1676            self.clone().into(),
1677            cx.clone(),
1678        ) {
1679            let client_id = self.id();
1680            log::debug!(
1681                "rpc message received. client_id:{}, sender_id:{:?}, type:{}",
1682                client_id,
1683                original_sender_id,
1684                type_name
1685            );
1686            cx.spawn(async move |_| match future.await {
1687                Ok(()) => {
1688                    log::debug!("rpc message handled. client_id:{client_id}, sender_id:{original_sender_id:?}, type:{type_name}");
1689                }
1690                Err(error) => {
1691                    log::error!("error handling message. client_id:{client_id}, sender_id:{original_sender_id:?}, type:{type_name}, error:{error:#}");
1692                }
1693            })
1694            .detach();
1695        } else {
1696            log::info!("unhandled message {}", type_name);
1697            self.peer
1698                .respond_with_unhandled_message(sender_id.into(), request_id, type_name)
1699                .log_err();
1700        }
1701    }
1702
1703    pub fn add_message_to_client_handler(
1704        self: &Arc<Client>,
1705        handler: impl Fn(&MessageToClient, &mut App) + Send + Sync + 'static,
1706    ) {
1707        self.message_to_client_handlers
1708            .lock()
1709            .push(Box::new(handler));
1710    }
1711
1712    fn handle_message_to_client(self: &Arc<Client>, message: MessageToClient, cx: &AsyncApp) {
1713        cx.update(|cx| {
1714            for handler in self.message_to_client_handlers.lock().iter() {
1715                handler(&message, cx);
1716            }
1717        })
1718        .ok();
1719    }
1720
1721    pub fn telemetry(&self) -> &Arc<Telemetry> {
1722        &self.telemetry
1723    }
1724}
1725
1726impl ProtoClient for Client {
1727    fn request(
1728        &self,
1729        envelope: proto::Envelope,
1730        request_type: &'static str,
1731    ) -> BoxFuture<'static, Result<proto::Envelope>> {
1732        self.request_dynamic(envelope, request_type).boxed()
1733    }
1734
1735    fn send(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
1736        log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
1737        let connection_id = self.connection_id()?;
1738        self.peer.send_dynamic(connection_id, envelope)
1739    }
1740
1741    fn send_response(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
1742        log::debug!(
1743            "rpc respond. client_id:{}, name:{}",
1744            self.id(),
1745            message_type
1746        );
1747        let connection_id = self.connection_id()?;
1748        self.peer.send_dynamic(connection_id, envelope)
1749    }
1750
1751    fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet> {
1752        &self.handler_set
1753    }
1754
1755    fn is_via_collab(&self) -> bool {
1756        true
1757    }
1758}
1759
1760/// prefix for the zed:// url scheme
1761pub const ZED_URL_SCHEME: &str = "zed";
1762
1763/// Parses the given link into a Zed link.
1764///
1765/// Returns a [`Some`] containing the unprefixed link if the link is a Zed link.
1766/// Returns [`None`] otherwise.
1767pub fn parse_zed_link<'a>(link: &'a str, cx: &App) -> Option<&'a str> {
1768    let server_url = &ClientSettings::get_global(cx).server_url;
1769    if let Some(stripped) = link
1770        .strip_prefix(server_url)
1771        .and_then(|result| result.strip_prefix('/'))
1772    {
1773        return Some(stripped);
1774    }
1775    if let Some(stripped) = link
1776        .strip_prefix(ZED_URL_SCHEME)
1777        .and_then(|result| result.strip_prefix("://"))
1778    {
1779        return Some(stripped);
1780    }
1781
1782    None
1783}
1784
1785#[cfg(test)]
1786mod tests {
1787    use super::*;
1788    use crate::test::{FakeServer, parse_authorization_header};
1789
1790    use clock::FakeSystemClock;
1791    use gpui::{AppContext as _, BackgroundExecutor, TestAppContext};
1792    use http_client::FakeHttpClient;
1793    use parking_lot::Mutex;
1794    use proto::TypedEnvelope;
1795    use settings::SettingsStore;
1796    use std::future;
1797
1798    #[gpui::test(iterations = 10)]
1799    async fn test_reconnection(cx: &mut TestAppContext) {
1800        init_test(cx);
1801        let user_id = 5;
1802        let client = cx.update(|cx| {
1803            Client::new(
1804                Arc::new(FakeSystemClock::new()),
1805                FakeHttpClient::with_404_response(),
1806                cx,
1807            )
1808        });
1809        let server = FakeServer::for_client(user_id, &client, cx).await;
1810        let mut status = client.status();
1811        assert!(matches!(
1812            status.next().await,
1813            Some(Status::Connected { .. })
1814        ));
1815        assert_eq!(server.auth_count(), 1);
1816
1817        server.forbid_connections();
1818        server.disconnect();
1819        while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1820
1821        server.allow_connections();
1822        cx.executor().advance_clock(Duration::from_secs(10));
1823        while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1824        assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1825
1826        server.forbid_connections();
1827        server.disconnect();
1828        while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1829
1830        // Clear cached credentials after authentication fails
1831        server.roll_access_token();
1832        server.allow_connections();
1833        cx.executor().run_until_parked();
1834        cx.executor().advance_clock(Duration::from_secs(10));
1835        while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1836        assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
1837    }
1838
1839    #[gpui::test(iterations = 10)]
1840    async fn test_auth_failure_during_reconnection(cx: &mut TestAppContext) {
1841        init_test(cx);
1842        let http_client = FakeHttpClient::with_200_response();
1843        let client =
1844            cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client.clone(), cx));
1845        let server = FakeServer::for_client(42, &client, cx).await;
1846        let mut status = client.status();
1847        assert!(matches!(
1848            status.next().await,
1849            Some(Status::Connected { .. })
1850        ));
1851        assert_eq!(server.auth_count(), 1);
1852
1853        // Simulate an auth failure during reconnection.
1854        http_client
1855            .as_fake()
1856            .replace_handler(|_, _request| async move {
1857                Ok(http_client::Response::builder()
1858                    .status(503)
1859                    .body("".into())
1860                    .unwrap())
1861            });
1862        server.disconnect();
1863        while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1864
1865        // Restore the ability to authenticate.
1866        http_client
1867            .as_fake()
1868            .replace_handler(|_, _request| async move {
1869                Ok(http_client::Response::builder()
1870                    .status(200)
1871                    .body("".into())
1872                    .unwrap())
1873            });
1874        cx.executor().advance_clock(Duration::from_secs(10));
1875        while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1876        assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1877    }
1878
1879    #[gpui::test(iterations = 10)]
1880    async fn test_connection_timeout(executor: BackgroundExecutor, cx: &mut TestAppContext) {
1881        init_test(cx);
1882        let user_id = 5;
1883        let client = cx.update(|cx| {
1884            Client::new(
1885                Arc::new(FakeSystemClock::new()),
1886                FakeHttpClient::with_404_response(),
1887                cx,
1888            )
1889        });
1890        let mut status = client.status();
1891
1892        // Time out when client tries to connect.
1893        client.override_authenticate(move |cx| {
1894            cx.background_spawn(async move {
1895                Ok(Credentials {
1896                    user_id,
1897                    access_token: "token".into(),
1898                })
1899            })
1900        });
1901        client.override_establish_connection(|_, cx| {
1902            cx.background_spawn(async move {
1903                future::pending::<()>().await;
1904                unreachable!()
1905            })
1906        });
1907        let auth_and_connect = cx.spawn({
1908            let client = client.clone();
1909            |cx| async move { client.connect(false, &cx).await }
1910        });
1911        executor.run_until_parked();
1912        assert!(matches!(status.next().await, Some(Status::Connecting)));
1913
1914        executor.advance_clock(CONNECTION_TIMEOUT);
1915        assert!(matches!(status.next().await, Some(Status::ConnectionError)));
1916        auth_and_connect.await.into_response().unwrap_err();
1917
1918        // Allow the connection to be established.
1919        let server = FakeServer::for_client(user_id, &client, cx).await;
1920        assert!(matches!(
1921            status.next().await,
1922            Some(Status::Connected { .. })
1923        ));
1924
1925        // Disconnect client.
1926        server.forbid_connections();
1927        server.disconnect();
1928        while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1929
1930        // Time out when re-establishing the connection.
1931        server.allow_connections();
1932        client.override_establish_connection(|_, cx| {
1933            cx.background_spawn(async move {
1934                future::pending::<()>().await;
1935                unreachable!()
1936            })
1937        });
1938        executor.advance_clock(2 * INITIAL_RECONNECTION_DELAY);
1939        assert!(matches!(status.next().await, Some(Status::Reconnecting)));
1940
1941        executor.advance_clock(CONNECTION_TIMEOUT);
1942        assert!(matches!(
1943            status.next().await,
1944            Some(Status::ReconnectionError { .. })
1945        ));
1946    }
1947
1948    #[gpui::test(iterations = 10)]
1949    async fn test_reauthenticate_only_if_unauthorized(cx: &mut TestAppContext) {
1950        init_test(cx);
1951        let auth_count = Arc::new(Mutex::new(0));
1952        let http_client = FakeHttpClient::create(|_request| async move {
1953            Ok(http_client::Response::builder()
1954                .status(200)
1955                .body("".into())
1956                .unwrap())
1957        });
1958        let client =
1959            cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client.clone(), cx));
1960        client.override_authenticate({
1961            let auth_count = auth_count.clone();
1962            move |cx| {
1963                let auth_count = auth_count.clone();
1964                cx.background_spawn(async move {
1965                    *auth_count.lock() += 1;
1966                    Ok(Credentials {
1967                        user_id: 1,
1968                        access_token: auth_count.lock().to_string(),
1969                    })
1970                })
1971            }
1972        });
1973
1974        let credentials = client.sign_in(false, &cx.to_async()).await.unwrap();
1975        assert_eq!(*auth_count.lock(), 1);
1976        assert_eq!(credentials.access_token, "1");
1977
1978        // If credentials are still valid, signing in doesn't trigger authentication.
1979        let credentials = client.sign_in(false, &cx.to_async()).await.unwrap();
1980        assert_eq!(*auth_count.lock(), 1);
1981        assert_eq!(credentials.access_token, "1");
1982
1983        // If the server is unavailable, signing in doesn't trigger authentication.
1984        http_client
1985            .as_fake()
1986            .replace_handler(|_, _request| async move {
1987                Ok(http_client::Response::builder()
1988                    .status(503)
1989                    .body("".into())
1990                    .unwrap())
1991            });
1992        client.sign_in(false, &cx.to_async()).await.unwrap_err();
1993        assert_eq!(*auth_count.lock(), 1);
1994
1995        // If credentials became invalid, signing in triggers authentication.
1996        http_client
1997            .as_fake()
1998            .replace_handler(|_, request| async move {
1999                let credentials = parse_authorization_header(&request).unwrap();
2000                if credentials.access_token == "2" {
2001                    Ok(http_client::Response::builder()
2002                        .status(200)
2003                        .body("".into())
2004                        .unwrap())
2005                } else {
2006                    Ok(http_client::Response::builder()
2007                        .status(401)
2008                        .body("".into())
2009                        .unwrap())
2010                }
2011            });
2012        let credentials = client.sign_in(false, &cx.to_async()).await.unwrap();
2013        assert_eq!(*auth_count.lock(), 2);
2014        assert_eq!(credentials.access_token, "2");
2015    }
2016
2017    #[gpui::test(iterations = 10)]
2018    async fn test_authenticating_more_than_once(
2019        cx: &mut TestAppContext,
2020        executor: BackgroundExecutor,
2021    ) {
2022        init_test(cx);
2023        let auth_count = Arc::new(Mutex::new(0));
2024        let dropped_auth_count = Arc::new(Mutex::new(0));
2025        let client = cx.update(|cx| {
2026            Client::new(
2027                Arc::new(FakeSystemClock::new()),
2028                FakeHttpClient::with_404_response(),
2029                cx,
2030            )
2031        });
2032        client.override_authenticate({
2033            let auth_count = auth_count.clone();
2034            let dropped_auth_count = dropped_auth_count.clone();
2035            move |cx| {
2036                let auth_count = auth_count.clone();
2037                let dropped_auth_count = dropped_auth_count.clone();
2038                cx.background_spawn(async move {
2039                    *auth_count.lock() += 1;
2040                    let _drop = util::defer(move || *dropped_auth_count.lock() += 1);
2041                    future::pending::<()>().await;
2042                    unreachable!()
2043                })
2044            }
2045        });
2046
2047        let _authenticate = cx.spawn({
2048            let client = client.clone();
2049            move |cx| async move { client.connect(false, &cx).await }
2050        });
2051        executor.run_until_parked();
2052        assert_eq!(*auth_count.lock(), 1);
2053        assert_eq!(*dropped_auth_count.lock(), 0);
2054
2055        let _authenticate = cx.spawn(|cx| async move { client.connect(false, &cx).await });
2056        executor.run_until_parked();
2057        assert_eq!(*auth_count.lock(), 2);
2058        assert_eq!(*dropped_auth_count.lock(), 1);
2059    }
2060
2061    #[gpui::test]
2062    async fn test_subscribing_to_entity(cx: &mut TestAppContext) {
2063        init_test(cx);
2064        let user_id = 5;
2065        let client = cx.update(|cx| {
2066            Client::new(
2067                Arc::new(FakeSystemClock::new()),
2068                FakeHttpClient::with_404_response(),
2069                cx,
2070            )
2071        });
2072        let server = FakeServer::for_client(user_id, &client, cx).await;
2073
2074        let (done_tx1, done_rx1) = smol::channel::unbounded();
2075        let (done_tx2, done_rx2) = smol::channel::unbounded();
2076        AnyProtoClient::from(client.clone()).add_entity_message_handler(
2077            move |entity: Entity<TestEntity>, _: TypedEnvelope<proto::JoinProject>, cx| {
2078                match entity.read_with(&cx, |entity, _| entity.id).unwrap() {
2079                    1 => done_tx1.try_send(()).unwrap(),
2080                    2 => done_tx2.try_send(()).unwrap(),
2081                    _ => unreachable!(),
2082                }
2083                async { Ok(()) }
2084            },
2085        );
2086        let entity1 = cx.new(|_| TestEntity {
2087            id: 1,
2088            subscription: None,
2089        });
2090        let entity2 = cx.new(|_| TestEntity {
2091            id: 2,
2092            subscription: None,
2093        });
2094        let entity3 = cx.new(|_| TestEntity {
2095            id: 3,
2096            subscription: None,
2097        });
2098
2099        let _subscription1 = client
2100            .subscribe_to_entity(1)
2101            .unwrap()
2102            .set_entity(&entity1, &cx.to_async());
2103        let _subscription2 = client
2104            .subscribe_to_entity(2)
2105            .unwrap()
2106            .set_entity(&entity2, &cx.to_async());
2107        // Ensure dropping a subscription for the same entity type still allows receiving of
2108        // messages for other entity IDs of the same type.
2109        let subscription3 = client
2110            .subscribe_to_entity(3)
2111            .unwrap()
2112            .set_entity(&entity3, &cx.to_async());
2113        drop(subscription3);
2114
2115        server.send(proto::JoinProject {
2116            project_id: 1,
2117            committer_name: None,
2118            committer_email: None,
2119        });
2120        server.send(proto::JoinProject {
2121            project_id: 2,
2122            committer_name: None,
2123            committer_email: None,
2124        });
2125        done_rx1.recv().await.unwrap();
2126        done_rx2.recv().await.unwrap();
2127    }
2128
2129    #[gpui::test]
2130    async fn test_subscribing_after_dropping_subscription(cx: &mut TestAppContext) {
2131        init_test(cx);
2132        let user_id = 5;
2133        let client = cx.update(|cx| {
2134            Client::new(
2135                Arc::new(FakeSystemClock::new()),
2136                FakeHttpClient::with_404_response(),
2137                cx,
2138            )
2139        });
2140        let server = FakeServer::for_client(user_id, &client, cx).await;
2141
2142        let entity = cx.new(|_| TestEntity::default());
2143        let (done_tx1, _done_rx1) = smol::channel::unbounded();
2144        let (done_tx2, done_rx2) = smol::channel::unbounded();
2145        let subscription1 = client.add_message_handler(
2146            entity.downgrade(),
2147            move |_, _: TypedEnvelope<proto::Ping>, _| {
2148                done_tx1.try_send(()).unwrap();
2149                async { Ok(()) }
2150            },
2151        );
2152        drop(subscription1);
2153        let _subscription2 = client.add_message_handler(
2154            entity.downgrade(),
2155            move |_, _: TypedEnvelope<proto::Ping>, _| {
2156                done_tx2.try_send(()).unwrap();
2157                async { Ok(()) }
2158            },
2159        );
2160        server.send(proto::Ping {});
2161        done_rx2.recv().await.unwrap();
2162    }
2163
2164    #[gpui::test]
2165    async fn test_dropping_subscription_in_handler(cx: &mut TestAppContext) {
2166        init_test(cx);
2167        let user_id = 5;
2168        let client = cx.update(|cx| {
2169            Client::new(
2170                Arc::new(FakeSystemClock::new()),
2171                FakeHttpClient::with_404_response(),
2172                cx,
2173            )
2174        });
2175        let server = FakeServer::for_client(user_id, &client, cx).await;
2176
2177        let entity = cx.new(|_| TestEntity::default());
2178        let (done_tx, done_rx) = smol::channel::unbounded();
2179        let subscription = client.add_message_handler(
2180            entity.clone().downgrade(),
2181            move |entity: Entity<TestEntity>, _: TypedEnvelope<proto::Ping>, mut cx| {
2182                entity
2183                    .update(&mut cx, |entity, _| entity.subscription.take())
2184                    .unwrap();
2185                done_tx.try_send(()).unwrap();
2186                async { Ok(()) }
2187            },
2188        );
2189        entity.update(cx, |entity, _| {
2190            entity.subscription = Some(subscription);
2191        });
2192        server.send(proto::Ping {});
2193        done_rx.recv().await.unwrap();
2194    }
2195
2196    #[derive(Default)]
2197    struct TestEntity {
2198        id: usize,
2199        subscription: Option<Subscription>,
2200    }
2201
2202    fn init_test(cx: &mut TestAppContext) {
2203        cx.update(|cx| {
2204            let settings_store = SettingsStore::test(cx);
2205            cx.set_global(settings_store);
2206            init_settings(cx);
2207        });
2208    }
2209}