tenant.openagents/omega

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