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