1#[cfg(any(test, feature = "test-support"))]
2pub mod test;
3
4pub mod channel;
5pub mod http;
6pub mod telemetry;
7pub mod user;
8
9use anyhow::{anyhow, Context, Result};
10use async_recursion::async_recursion;
11use async_tungstenite::tungstenite::{
12 error::Error as WebsocketError,
13 http::{Request, StatusCode},
14};
15use db::Db;
16use futures::{future::LocalBoxFuture, FutureExt, SinkExt, StreamExt, TryStreamExt};
17use gpui::{
18 actions, serde_json::Value, AnyModelHandle, AnyViewHandle, AnyWeakModelHandle,
19 AnyWeakViewHandle, AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle,
20 MutableAppContext, Task, View, ViewContext, ViewHandle,
21};
22use http::HttpClient;
23use lazy_static::lazy_static;
24use parking_lot::RwLock;
25use postage::watch;
26use rand::prelude::*;
27use rpc::proto::{AnyTypedEnvelope, EntityMessage, EnvelopedMessage, RequestMessage};
28use std::{
29 any::TypeId,
30 collections::HashMap,
31 convert::TryFrom,
32 fmt::Write as _,
33 future::Future,
34 path::PathBuf,
35 sync::{Arc, Weak},
36 time::{Duration, Instant},
37};
38use telemetry::Telemetry;
39use thiserror::Error;
40use url::Url;
41use util::{ResultExt, TryFutureExt};
42
43pub use channel::*;
44pub use rpc::*;
45pub use user::*;
46
47lazy_static! {
48 pub static ref ZED_SERVER_URL: String =
49 std::env::var("ZED_SERVER_URL").unwrap_or_else(|_| "https://zed.dev".to_string());
50 pub static ref IMPERSONATE_LOGIN: Option<String> = std::env::var("ZED_IMPERSONATE")
51 .ok()
52 .and_then(|s| if s.is_empty() { None } else { Some(s) });
53}
54
55pub const ZED_SECRET_CLIENT_TOKEN: &str = "618033988749894";
56pub const INITIAL_RECONNECTION_DELAY: Duration = Duration::from_millis(100);
57pub const CONNECTION_TIMEOUT: Duration = Duration::from_secs(5);
58
59actions!(client, [Authenticate]);
60
61pub fn init(client: Arc<Client>, cx: &mut MutableAppContext) {
62 cx.add_global_action({
63 let client = client.clone();
64 move |_: &Authenticate, cx| {
65 let client = client.clone();
66 cx.spawn(
67 |cx| async move { client.authenticate_and_connect(true, &cx).log_err().await },
68 )
69 .detach();
70 }
71 });
72}
73
74pub struct Client {
75 id: usize,
76 peer: Arc<Peer>,
77 http: Arc<dyn HttpClient>,
78 telemetry: Arc<Telemetry>,
79 state: RwLock<ClientState>,
80
81 #[allow(clippy::type_complexity)]
82 #[cfg(any(test, feature = "test-support"))]
83 authenticate: RwLock<
84 Option<Box<dyn 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>>>,
85 >,
86
87 #[allow(clippy::type_complexity)]
88 #[cfg(any(test, feature = "test-support"))]
89 establish_connection: RwLock<
90 Option<
91 Box<
92 dyn 'static
93 + Send
94 + Sync
95 + Fn(
96 &Credentials,
97 &AsyncAppContext,
98 ) -> Task<Result<Connection, EstablishConnectionError>>,
99 >,
100 >,
101 >,
102}
103
104#[derive(Error, Debug)]
105pub enum EstablishConnectionError {
106 #[error("upgrade required")]
107 UpgradeRequired,
108 #[error("unauthorized")]
109 Unauthorized,
110 #[error("{0}")]
111 Other(#[from] anyhow::Error),
112 #[error("{0}")]
113 Http(#[from] http::Error),
114 #[error("{0}")]
115 Io(#[from] std::io::Error),
116 #[error("{0}")]
117 Websocket(#[from] async_tungstenite::tungstenite::http::Error),
118}
119
120impl From<WebsocketError> for EstablishConnectionError {
121 fn from(error: WebsocketError) -> Self {
122 if let WebsocketError::Http(response) = &error {
123 match response.status() {
124 StatusCode::UNAUTHORIZED => return EstablishConnectionError::Unauthorized,
125 StatusCode::UPGRADE_REQUIRED => return EstablishConnectionError::UpgradeRequired,
126 _ => {}
127 }
128 }
129 EstablishConnectionError::Other(error.into())
130 }
131}
132
133impl EstablishConnectionError {
134 pub fn other(error: impl Into<anyhow::Error> + Send + Sync) -> Self {
135 Self::Other(error.into())
136 }
137}
138
139#[derive(Copy, Clone, Debug, Eq, PartialEq)]
140pub enum Status {
141 SignedOut,
142 UpgradeRequired,
143 Authenticating,
144 Connecting,
145 ConnectionError,
146 Connected {
147 peer_id: PeerId,
148 connection_id: ConnectionId,
149 },
150 ConnectionLost,
151 Reauthenticating,
152 Reconnecting,
153 ReconnectionError {
154 next_reconnection: Instant,
155 },
156}
157
158impl Status {
159 pub fn is_connected(&self) -> bool {
160 matches!(self, Self::Connected { .. })
161 }
162}
163
164struct ClientState {
165 credentials: Option<Credentials>,
166 status: (watch::Sender<Status>, watch::Receiver<Status>),
167 entity_id_extractors: HashMap<TypeId, fn(&dyn AnyTypedEnvelope) -> u64>,
168 _reconnect_task: Option<Task<()>>,
169 reconnect_interval: Duration,
170 entities_by_type_and_remote_id: HashMap<(TypeId, u64), AnyWeakEntityHandle>,
171 models_by_message_type: HashMap<TypeId, AnyWeakModelHandle>,
172 entity_types_by_message_type: HashMap<TypeId, TypeId>,
173 #[allow(clippy::type_complexity)]
174 message_handlers: HashMap<
175 TypeId,
176 Arc<
177 dyn Send
178 + Sync
179 + Fn(
180 AnyEntityHandle,
181 Box<dyn AnyTypedEnvelope>,
182 &Arc<Client>,
183 AsyncAppContext,
184 ) -> LocalBoxFuture<'static, Result<()>>,
185 >,
186 >,
187}
188
189enum AnyWeakEntityHandle {
190 Model(AnyWeakModelHandle),
191 View(AnyWeakViewHandle),
192}
193
194enum AnyEntityHandle {
195 Model(AnyModelHandle),
196 View(AnyViewHandle),
197}
198
199#[derive(Clone, Debug)]
200pub struct Credentials {
201 pub user_id: u64,
202 pub access_token: String,
203}
204
205impl Default for ClientState {
206 fn default() -> Self {
207 Self {
208 credentials: None,
209 status: watch::channel_with(Status::SignedOut),
210 entity_id_extractors: Default::default(),
211 _reconnect_task: None,
212 reconnect_interval: Duration::from_secs(5),
213 models_by_message_type: Default::default(),
214 entities_by_type_and_remote_id: Default::default(),
215 entity_types_by_message_type: Default::default(),
216 message_handlers: Default::default(),
217 }
218 }
219}
220
221pub enum Subscription {
222 Entity {
223 client: Weak<Client>,
224 id: (TypeId, u64),
225 },
226 Message {
227 client: Weak<Client>,
228 id: TypeId,
229 },
230}
231
232impl Drop for Subscription {
233 fn drop(&mut self) {
234 match self {
235 Subscription::Entity { client, id } => {
236 if let Some(client) = client.upgrade() {
237 let mut state = client.state.write();
238 let _ = state.entities_by_type_and_remote_id.remove(id);
239 }
240 }
241 Subscription::Message { client, id } => {
242 if let Some(client) = client.upgrade() {
243 let mut state = client.state.write();
244 let _ = state.entity_types_by_message_type.remove(id);
245 let _ = state.message_handlers.remove(id);
246 }
247 }
248 }
249 }
250}
251
252impl Client {
253 pub fn new(http: Arc<dyn HttpClient>, cx: &AppContext) -> Arc<Self> {
254 Arc::new(Self {
255 id: 0,
256 peer: Peer::new(),
257 telemetry: Telemetry::new(http.clone(), cx),
258 http,
259 state: Default::default(),
260
261 #[cfg(any(test, feature = "test-support"))]
262 authenticate: Default::default(),
263 #[cfg(any(test, feature = "test-support"))]
264 establish_connection: Default::default(),
265 })
266 }
267
268 pub fn id(&self) -> usize {
269 self.id
270 }
271
272 pub fn http_client(&self) -> Arc<dyn HttpClient> {
273 self.http.clone()
274 }
275
276 #[cfg(any(test, feature = "test-support"))]
277 pub fn set_id(&mut self, id: usize) -> &Self {
278 self.id = id;
279 self
280 }
281
282 #[cfg(any(test, feature = "test-support"))]
283 pub fn tear_down(&self) {
284 let mut state = self.state.write();
285 state._reconnect_task.take();
286 state.message_handlers.clear();
287 state.models_by_message_type.clear();
288 state.entities_by_type_and_remote_id.clear();
289 state.entity_id_extractors.clear();
290 self.peer.reset();
291 }
292
293 #[cfg(any(test, feature = "test-support"))]
294 pub fn override_authenticate<F>(&self, authenticate: F) -> &Self
295 where
296 F: 'static + Send + Sync + Fn(&AsyncAppContext) -> Task<Result<Credentials>>,
297 {
298 *self.authenticate.write() = Some(Box::new(authenticate));
299 self
300 }
301
302 #[cfg(any(test, feature = "test-support"))]
303 pub fn override_establish_connection<F>(&self, connect: F) -> &Self
304 where
305 F: 'static
306 + Send
307 + Sync
308 + Fn(&Credentials, &AsyncAppContext) -> Task<Result<Connection, EstablishConnectionError>>,
309 {
310 *self.establish_connection.write() = Some(Box::new(connect));
311 self
312 }
313
314 pub fn user_id(&self) -> Option<u64> {
315 self.state
316 .read()
317 .credentials
318 .as_ref()
319 .map(|credentials| credentials.user_id)
320 }
321
322 pub fn peer_id(&self) -> Option<PeerId> {
323 if let Status::Connected { peer_id, .. } = &*self.status().borrow() {
324 Some(*peer_id)
325 } else {
326 None
327 }
328 }
329
330 pub fn status(&self) -> watch::Receiver<Status> {
331 self.state.read().status.1.clone()
332 }
333
334 fn set_status(self: &Arc<Self>, status: Status, cx: &AsyncAppContext) {
335 log::info!("set status on client {}: {:?}", self.id, status);
336 let mut state = self.state.write();
337 *state.status.0.borrow_mut() = status;
338
339 match status {
340 Status::Connected { .. } => {
341 state._reconnect_task = None;
342 }
343 Status::ConnectionLost => {
344 let this = self.clone();
345 let reconnect_interval = state.reconnect_interval;
346 state._reconnect_task = Some(cx.spawn(|cx| async move {
347 let mut rng = StdRng::from_entropy();
348 let mut delay = INITIAL_RECONNECTION_DELAY;
349 while let Err(error) = this.authenticate_and_connect(true, &cx).await {
350 log::error!("failed to connect {}", error);
351 if matches!(*this.status().borrow(), Status::ConnectionError) {
352 this.set_status(
353 Status::ReconnectionError {
354 next_reconnection: Instant::now() + delay,
355 },
356 &cx,
357 );
358 cx.background().timer(delay).await;
359 delay = delay
360 .mul_f32(rng.gen_range(1.0..=2.0))
361 .min(reconnect_interval);
362 } else {
363 break;
364 }
365 }
366 }));
367 }
368 Status::SignedOut | Status::UpgradeRequired => {
369 self.telemetry.set_authenticated_user_info(None, false);
370 state._reconnect_task.take();
371 }
372 _ => {}
373 }
374 }
375
376 pub fn add_view_for_remote_entity<T: View>(
377 self: &Arc<Self>,
378 remote_id: u64,
379 cx: &mut ViewContext<T>,
380 ) -> Subscription {
381 let id = (TypeId::of::<T>(), remote_id);
382 self.state
383 .write()
384 .entities_by_type_and_remote_id
385 .insert(id, AnyWeakEntityHandle::View(cx.weak_handle().into()));
386 Subscription::Entity {
387 client: Arc::downgrade(self),
388 id,
389 }
390 }
391
392 pub fn add_model_for_remote_entity<T: Entity>(
393 self: &Arc<Self>,
394 remote_id: u64,
395 cx: &mut ModelContext<T>,
396 ) -> Subscription {
397 let id = (TypeId::of::<T>(), remote_id);
398 self.state
399 .write()
400 .entities_by_type_and_remote_id
401 .insert(id, AnyWeakEntityHandle::Model(cx.weak_handle().into()));
402 Subscription::Entity {
403 client: Arc::downgrade(self),
404 id,
405 }
406 }
407
408 pub fn add_message_handler<M, E, H, F>(
409 self: &Arc<Self>,
410 model: ModelHandle<E>,
411 handler: H,
412 ) -> Subscription
413 where
414 M: EnvelopedMessage,
415 E: Entity,
416 H: 'static
417 + Send
418 + Sync
419 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
420 F: 'static + Future<Output = Result<()>>,
421 {
422 let message_type_id = TypeId::of::<M>();
423
424 let mut state = self.state.write();
425 state
426 .models_by_message_type
427 .insert(message_type_id, model.downgrade().into());
428
429 let prev_handler = state.message_handlers.insert(
430 message_type_id,
431 Arc::new(move |handle, envelope, client, cx| {
432 let handle = if let AnyEntityHandle::Model(handle) = handle {
433 handle
434 } else {
435 unreachable!();
436 };
437 let model = handle.downcast::<E>().unwrap();
438 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
439 handler(model, *envelope, client.clone(), cx).boxed_local()
440 }),
441 );
442 if prev_handler.is_some() {
443 panic!("registered handler for the same message twice");
444 }
445
446 Subscription::Message {
447 client: Arc::downgrade(self),
448 id: message_type_id,
449 }
450 }
451
452 pub fn add_request_handler<M, E, H, F>(
453 self: &Arc<Self>,
454 model: ModelHandle<E>,
455 handler: H,
456 ) -> Subscription
457 where
458 M: RequestMessage,
459 E: Entity,
460 H: 'static
461 + Send
462 + Sync
463 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
464 F: 'static + Future<Output = Result<M::Response>>,
465 {
466 self.add_message_handler(model, move |handle, envelope, this, cx| {
467 Self::respond_to_request(
468 envelope.receipt(),
469 handler(handle, envelope, this.clone(), cx),
470 this,
471 )
472 })
473 }
474
475 pub fn add_view_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
476 where
477 M: EntityMessage,
478 E: View,
479 H: 'static
480 + Send
481 + Sync
482 + Fn(ViewHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
483 F: 'static + Future<Output = Result<()>>,
484 {
485 self.add_entity_message_handler::<M, E, _, _>(move |handle, message, client, cx| {
486 if let AnyEntityHandle::View(handle) = handle {
487 handler(handle.downcast::<E>().unwrap(), message, client, cx)
488 } else {
489 unreachable!();
490 }
491 })
492 }
493
494 pub fn add_model_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
495 where
496 M: EntityMessage,
497 E: Entity,
498 H: 'static
499 + Send
500 + Sync
501 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
502 F: 'static + Future<Output = Result<()>>,
503 {
504 self.add_entity_message_handler::<M, E, _, _>(move |handle, message, client, cx| {
505 if let AnyEntityHandle::Model(handle) = handle {
506 handler(handle.downcast::<E>().unwrap(), message, client, cx)
507 } else {
508 unreachable!();
509 }
510 })
511 }
512
513 fn add_entity_message_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
514 where
515 M: EntityMessage,
516 E: Entity,
517 H: 'static
518 + Send
519 + Sync
520 + Fn(AnyEntityHandle, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
521 F: 'static + Future<Output = Result<()>>,
522 {
523 let model_type_id = TypeId::of::<E>();
524 let message_type_id = TypeId::of::<M>();
525
526 let mut state = self.state.write();
527 state
528 .entity_types_by_message_type
529 .insert(message_type_id, model_type_id);
530 state
531 .entity_id_extractors
532 .entry(message_type_id)
533 .or_insert_with(|| {
534 |envelope| {
535 envelope
536 .as_any()
537 .downcast_ref::<TypedEnvelope<M>>()
538 .unwrap()
539 .payload
540 .remote_entity_id()
541 }
542 });
543 let prev_handler = state.message_handlers.insert(
544 message_type_id,
545 Arc::new(move |handle, envelope, client, cx| {
546 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
547 handler(handle, *envelope, client.clone(), cx).boxed_local()
548 }),
549 );
550 if prev_handler.is_some() {
551 panic!("registered handler for the same message twice");
552 }
553 }
554
555 pub fn add_model_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
556 where
557 M: EntityMessage + RequestMessage,
558 E: Entity,
559 H: 'static
560 + Send
561 + Sync
562 + Fn(ModelHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
563 F: 'static + Future<Output = Result<M::Response>>,
564 {
565 self.add_model_message_handler(move |entity, envelope, client, cx| {
566 Self::respond_to_request::<M, _>(
567 envelope.receipt(),
568 handler(entity, envelope, client.clone(), cx),
569 client,
570 )
571 })
572 }
573
574 pub fn add_view_request_handler<M, E, H, F>(self: &Arc<Self>, handler: H)
575 where
576 M: EntityMessage + RequestMessage,
577 E: View,
578 H: 'static
579 + Send
580 + Sync
581 + Fn(ViewHandle<E>, TypedEnvelope<M>, Arc<Self>, AsyncAppContext) -> F,
582 F: 'static + Future<Output = Result<M::Response>>,
583 {
584 self.add_view_message_handler(move |entity, envelope, client, cx| {
585 Self::respond_to_request::<M, _>(
586 envelope.receipt(),
587 handler(entity, envelope, client.clone(), cx),
588 client,
589 )
590 })
591 }
592
593 async fn respond_to_request<T: RequestMessage, F: Future<Output = Result<T::Response>>>(
594 receipt: Receipt<T>,
595 response: F,
596 client: Arc<Self>,
597 ) -> Result<()> {
598 match response.await {
599 Ok(response) => {
600 client.respond(receipt, response)?;
601 Ok(())
602 }
603 Err(error) => {
604 client.respond_with_error(
605 receipt,
606 proto::Error {
607 message: format!("{:?}", error),
608 },
609 )?;
610 Err(error)
611 }
612 }
613 }
614
615 pub fn has_keychain_credentials(&self, cx: &AsyncAppContext) -> bool {
616 read_credentials_from_keychain(cx).is_some()
617 }
618
619 #[async_recursion(?Send)]
620 pub async fn authenticate_and_connect(
621 self: &Arc<Self>,
622 try_keychain: bool,
623 cx: &AsyncAppContext,
624 ) -> anyhow::Result<()> {
625 let was_disconnected = match *self.status().borrow() {
626 Status::SignedOut => true,
627 Status::ConnectionError
628 | Status::ConnectionLost
629 | Status::Authenticating { .. }
630 | Status::Reauthenticating { .. }
631 | Status::ReconnectionError { .. } => false,
632 Status::Connected { .. } | Status::Connecting { .. } | Status::Reconnecting { .. } => {
633 return Ok(())
634 }
635 Status::UpgradeRequired => return Err(EstablishConnectionError::UpgradeRequired)?,
636 };
637
638 if was_disconnected {
639 self.set_status(Status::Authenticating, cx);
640 } else {
641 self.set_status(Status::Reauthenticating, cx)
642 }
643
644 let mut read_from_keychain = false;
645 let mut credentials = self.state.read().credentials.clone();
646 if credentials.is_none() && try_keychain {
647 credentials = read_credentials_from_keychain(cx);
648 read_from_keychain = credentials.is_some();
649 if read_from_keychain {
650 self.report_event("read credentials from keychain", Default::default());
651 }
652 }
653 if credentials.is_none() {
654 let mut status_rx = self.status();
655 let _ = status_rx.next().await;
656 futures::select_biased! {
657 authenticate = self.authenticate(cx).fuse() => {
658 match authenticate {
659 Ok(creds) => credentials = Some(creds),
660 Err(err) => {
661 self.set_status(Status::ConnectionError, cx);
662 return Err(err);
663 }
664 }
665 }
666 _ = status_rx.next().fuse() => {
667 return Err(anyhow!("authentication canceled"));
668 }
669 }
670 }
671 let credentials = credentials.unwrap();
672
673 if was_disconnected {
674 self.set_status(Status::Connecting, cx);
675 } else {
676 self.set_status(Status::Reconnecting, cx);
677 }
678
679 let mut timeout = cx.background().timer(CONNECTION_TIMEOUT).fuse();
680 futures::select_biased! {
681 connection = self.establish_connection(&credentials, cx).fuse() => {
682 match connection {
683 Ok(conn) => {
684 self.state.write().credentials = Some(credentials.clone());
685 if !read_from_keychain && IMPERSONATE_LOGIN.is_none() {
686 write_credentials_to_keychain(&credentials, cx).log_err();
687 }
688
689 futures::select_biased! {
690 result = self.set_connection(conn, cx).fuse() => result,
691 _ = timeout => {
692 self.set_status(Status::ConnectionError, cx);
693 Err(anyhow!("timed out waiting on hello message from server"))
694 }
695 }
696 }
697 Err(EstablishConnectionError::Unauthorized) => {
698 self.state.write().credentials.take();
699 if read_from_keychain {
700 cx.platform().delete_credentials(&ZED_SERVER_URL).log_err();
701 self.set_status(Status::SignedOut, cx);
702 self.authenticate_and_connect(false, cx).await
703 } else {
704 self.set_status(Status::ConnectionError, cx);
705 Err(EstablishConnectionError::Unauthorized)?
706 }
707 }
708 Err(EstablishConnectionError::UpgradeRequired) => {
709 self.set_status(Status::UpgradeRequired, cx);
710 Err(EstablishConnectionError::UpgradeRequired)?
711 }
712 Err(error) => {
713 self.set_status(Status::ConnectionError, cx);
714 Err(error)?
715 }
716 }
717 }
718 _ = &mut timeout => {
719 self.set_status(Status::ConnectionError, cx);
720 Err(anyhow!("timed out trying to establish connection"))
721 }
722 }
723 }
724
725 async fn set_connection(
726 self: &Arc<Self>,
727 conn: Connection,
728 cx: &AsyncAppContext,
729 ) -> Result<()> {
730 let executor = cx.background();
731 log::info!("add connection to peer");
732 let (connection_id, handle_io, mut incoming) = self
733 .peer
734 .add_connection(conn, move |duration| executor.timer(duration));
735 let handle_io = cx.background().spawn(handle_io);
736
737 let peer_id = async {
738 log::info!("waiting for server hello");
739 let message = incoming
740 .next()
741 .await
742 .ok_or_else(|| anyhow!("no hello message received"))?;
743 log::info!("got server hello");
744 let hello_message_type_name = message.payload_type_name().to_string();
745 let hello = message
746 .into_any()
747 .downcast::<TypedEnvelope<proto::Hello>>()
748 .map_err(|_| {
749 anyhow!(
750 "invalid hello message received: {:?}",
751 hello_message_type_name
752 )
753 })?;
754 Ok(PeerId(hello.payload.peer_id))
755 };
756
757 let peer_id = match peer_id.await {
758 Ok(peer_id) => peer_id,
759 Err(error) => {
760 self.peer.disconnect(connection_id);
761 return Err(error);
762 }
763 };
764
765 log::info!(
766 "set status to connected (connection id: {}, peer id: {})",
767 connection_id,
768 peer_id
769 );
770 self.set_status(
771 Status::Connected {
772 peer_id,
773 connection_id,
774 },
775 cx,
776 );
777 cx.foreground()
778 .spawn({
779 let cx = cx.clone();
780 let this = self.clone();
781 async move {
782 let mut message_id = 0_usize;
783 while let Some(message) = incoming.next().await {
784 let mut state = this.state.write();
785 message_id += 1;
786 let type_name = message.payload_type_name();
787 let payload_type_id = message.payload_type_id();
788 let sender_id = message.original_sender_id().map(|id| id.0);
789
790 let model = state
791 .models_by_message_type
792 .get(&payload_type_id)
793 .and_then(|model| model.upgrade(&cx))
794 .map(AnyEntityHandle::Model)
795 .or_else(|| {
796 let entity_type_id =
797 *state.entity_types_by_message_type.get(&payload_type_id)?;
798 let entity_id = state
799 .entity_id_extractors
800 .get(&message.payload_type_id())
801 .map(|extract_entity_id| {
802 (extract_entity_id)(message.as_ref())
803 })?;
804
805 let entity = state
806 .entities_by_type_and_remote_id
807 .get(&(entity_type_id, entity_id))?;
808 if let Some(entity) = entity.upgrade(&cx) {
809 Some(entity)
810 } else {
811 state
812 .entities_by_type_and_remote_id
813 .remove(&(entity_type_id, entity_id));
814 None
815 }
816 });
817
818 let model = if let Some(model) = model {
819 model
820 } else {
821 log::info!("unhandled message {}", type_name);
822 continue;
823 };
824
825 if let Some(handler) = state.message_handlers.get(&payload_type_id).cloned()
826 {
827 drop(state); // Avoid deadlocks if the handler interacts with rpc::Client
828 let future = handler(model, message, &this, cx.clone());
829
830 let client_id = this.id;
831 log::debug!(
832 "rpc message received. client_id:{}, message_id:{}, sender_id:{:?}, type:{}",
833 client_id,
834 message_id,
835 sender_id,
836 type_name
837 );
838 cx.foreground()
839 .spawn(async move {
840 match future.await {
841 Ok(()) => {
842 log::debug!(
843 "rpc message handled. client_id:{}, message_id:{}, sender_id:{:?}, type:{}",
844 client_id,
845 message_id,
846 sender_id,
847 type_name
848 );
849 }
850 Err(error) => {
851 log::error!(
852 "error handling message. client_id:{}, message_id:{}, sender_id:{:?}, type:{}, error:{:?}",
853 client_id,
854 message_id,
855 sender_id,
856 type_name,
857 error
858 );
859 }
860 }
861 })
862 .detach();
863 } else {
864 log::info!("unhandled message {}", type_name);
865 }
866
867 // Don't starve the main thread when receiving lots of messages at once.
868 smol::future::yield_now().await;
869 }
870 }
871 })
872 .detach();
873
874 let this = self.clone();
875 let cx = cx.clone();
876 cx.foreground()
877 .spawn(async move {
878 match handle_io.await {
879 Ok(()) => {
880 if *this.status().borrow()
881 == (Status::Connected {
882 connection_id,
883 peer_id,
884 })
885 {
886 this.set_status(Status::SignedOut, &cx);
887 }
888 }
889 Err(err) => {
890 log::error!("connection error: {:?}", err);
891 this.set_status(Status::ConnectionLost, &cx);
892 }
893 }
894 })
895 .detach();
896
897 Ok(())
898 }
899
900 fn authenticate(self: &Arc<Self>, cx: &AsyncAppContext) -> Task<Result<Credentials>> {
901 #[cfg(any(test, feature = "test-support"))]
902 if let Some(callback) = self.authenticate.read().as_ref() {
903 return callback(cx);
904 }
905
906 self.authenticate_with_browser(cx)
907 }
908
909 fn establish_connection(
910 self: &Arc<Self>,
911 credentials: &Credentials,
912 cx: &AsyncAppContext,
913 ) -> Task<Result<Connection, EstablishConnectionError>> {
914 #[cfg(any(test, feature = "test-support"))]
915 if let Some(callback) = self.establish_connection.read().as_ref() {
916 return callback(credentials, cx);
917 }
918
919 self.establish_websocket_connection(credentials, cx)
920 }
921
922 fn establish_websocket_connection(
923 self: &Arc<Self>,
924 credentials: &Credentials,
925 cx: &AsyncAppContext,
926 ) -> Task<Result<Connection, EstablishConnectionError>> {
927 let request = Request::builder()
928 .header(
929 "Authorization",
930 format!("{} {}", credentials.user_id, credentials.access_token),
931 )
932 .header("x-zed-protocol-version", rpc::PROTOCOL_VERSION);
933
934 let http = self.http.clone();
935 cx.background().spawn(async move {
936 let mut rpc_url = format!("{}/rpc", *ZED_SERVER_URL);
937 let rpc_response = http.get(&rpc_url, Default::default(), false).await?;
938 if rpc_response.status().is_redirection() {
939 rpc_url = rpc_response
940 .headers()
941 .get("Location")
942 .ok_or_else(|| anyhow!("missing location header in /rpc response"))?
943 .to_str()
944 .map_err(EstablishConnectionError::other)?
945 .to_string();
946 }
947 // Until we switch the zed.dev domain to point to the new Next.js app, there
948 // will be no redirect required, and the app will connect directly to
949 // wss://zed.dev/rpc.
950 else if rpc_response.status() != StatusCode::UPGRADE_REQUIRED {
951 Err(anyhow!(
952 "unexpected /rpc response status {}",
953 rpc_response.status()
954 ))?
955 }
956
957 let mut rpc_url = Url::parse(&rpc_url).context("invalid rpc url")?;
958 let rpc_host = rpc_url
959 .host_str()
960 .zip(rpc_url.port_or_known_default())
961 .ok_or_else(|| anyhow!("missing host in rpc url"))?;
962 let stream = smol::net::TcpStream::connect(rpc_host).await?;
963
964 log::info!("connected to rpc endpoint {}", rpc_url);
965
966 match rpc_url.scheme() {
967 "https" => {
968 rpc_url.set_scheme("wss").unwrap();
969 let request = request.uri(rpc_url.as_str()).body(())?;
970 let (stream, _) =
971 async_tungstenite::async_tls::client_async_tls(request, stream).await?;
972 Ok(Connection::new(
973 stream
974 .map_err(|error| anyhow!(error))
975 .sink_map_err(|error| anyhow!(error)),
976 ))
977 }
978 "http" => {
979 rpc_url.set_scheme("ws").unwrap();
980 let request = request.uri(rpc_url.as_str()).body(())?;
981 let (stream, _) = async_tungstenite::client_async(request, stream).await?;
982 Ok(Connection::new(
983 stream
984 .map_err(|error| anyhow!(error))
985 .sink_map_err(|error| anyhow!(error)),
986 ))
987 }
988 _ => Err(anyhow!("invalid rpc url: {}", rpc_url))?,
989 }
990 })
991 }
992
993 pub fn authenticate_with_browser(
994 self: &Arc<Self>,
995 cx: &AsyncAppContext,
996 ) -> Task<Result<Credentials>> {
997 let platform = cx.platform();
998 let executor = cx.background();
999 let telemetry = self.telemetry.clone();
1000 executor.clone().spawn(async move {
1001 // Generate a pair of asymmetric encryption keys. The public key will be used by the
1002 // zed server to encrypt the user's access token, so that it can'be intercepted by
1003 // any other app running on the user's device.
1004 let (public_key, private_key) =
1005 rpc::auth::keypair().expect("failed to generate keypair for auth");
1006 let public_key_string =
1007 String::try_from(public_key).expect("failed to serialize public key for auth");
1008
1009 // Start an HTTP server to receive the redirect from Zed's sign-in page.
1010 let server = tiny_http::Server::http("127.0.0.1:0").expect("failed to find open port");
1011 let port = server.server_addr().port();
1012
1013 // Open the Zed sign-in page in the user's browser, with query parameters that indicate
1014 // that the user is signing in from a Zed app running on the same device.
1015 let mut url = format!(
1016 "{}/native_app_signin?native_app_port={}&native_app_public_key={}",
1017 *ZED_SERVER_URL, port, public_key_string
1018 );
1019
1020 if let Some(impersonate_login) = IMPERSONATE_LOGIN.as_ref() {
1021 log::info!("impersonating user @{}", impersonate_login);
1022 write!(&mut url, "&impersonate={}", impersonate_login).unwrap();
1023 }
1024
1025 platform.open_url(&url);
1026
1027 // Receive the HTTP request from the user's browser. Retrieve the user id and encrypted
1028 // access token from the query params.
1029 //
1030 // TODO - Avoid ever starting more than one HTTP server. Maybe switch to using a
1031 // custom URL scheme instead of this local HTTP server.
1032 let (user_id, access_token) = executor
1033 .spawn(async move {
1034 for _ in 0..100 {
1035 if let Some(req) = server.recv_timeout(Duration::from_secs(1))? {
1036 let path = req.url();
1037 let mut user_id = None;
1038 let mut access_token = None;
1039 let url = Url::parse(&format!("http://example.com{}", path))
1040 .context("failed to parse login notification url")?;
1041 for (key, value) in url.query_pairs() {
1042 if key == "access_token" {
1043 access_token = Some(value.to_string());
1044 } else if key == "user_id" {
1045 user_id = Some(value.to_string());
1046 }
1047 }
1048
1049 let post_auth_url =
1050 format!("{}/native_app_signin_succeeded", *ZED_SERVER_URL);
1051 req.respond(
1052 tiny_http::Response::empty(302).with_header(
1053 tiny_http::Header::from_bytes(
1054 &b"Location"[..],
1055 post_auth_url.as_bytes(),
1056 )
1057 .unwrap(),
1058 ),
1059 )
1060 .context("failed to respond to login http request")?;
1061 return Ok((
1062 user_id.ok_or_else(|| anyhow!("missing user_id parameter"))?,
1063 access_token
1064 .ok_or_else(|| anyhow!("missing access_token parameter"))?,
1065 ));
1066 }
1067 }
1068
1069 Err(anyhow!("didn't receive login redirect"))
1070 })
1071 .await?;
1072
1073 let access_token = private_key
1074 .decrypt_string(&access_token)
1075 .context("failed to decrypt access token")?;
1076 platform.activate(true);
1077
1078 telemetry.report_event("authenticate with browser", Default::default());
1079
1080 Ok(Credentials {
1081 user_id: user_id.parse()?,
1082 access_token,
1083 })
1084 })
1085 }
1086
1087 pub fn disconnect(self: &Arc<Self>, cx: &AsyncAppContext) -> Result<()> {
1088 let conn_id = self.connection_id()?;
1089 self.peer.disconnect(conn_id);
1090 self.set_status(Status::SignedOut, cx);
1091 Ok(())
1092 }
1093
1094 fn connection_id(&self) -> Result<ConnectionId> {
1095 if let Status::Connected { connection_id, .. } = *self.status().borrow() {
1096 Ok(connection_id)
1097 } else {
1098 Err(anyhow!("not connected"))
1099 }
1100 }
1101
1102 pub fn send<T: EnvelopedMessage>(&self, message: T) -> Result<()> {
1103 log::debug!("rpc send. client_id:{}, name:{}", self.id, T::NAME);
1104 self.peer.send(self.connection_id()?, message)
1105 }
1106
1107 pub fn request<T: RequestMessage>(
1108 &self,
1109 request: T,
1110 ) -> impl Future<Output = Result<T::Response>> {
1111 let client_id = self.id;
1112 log::debug!(
1113 "rpc request start. client_id:{}. name:{}",
1114 client_id,
1115 T::NAME
1116 );
1117 let response = self
1118 .connection_id()
1119 .map(|conn_id| self.peer.request(conn_id, request));
1120 async move {
1121 let response = response?.await;
1122 log::debug!(
1123 "rpc request finish. client_id:{}. name:{}",
1124 client_id,
1125 T::NAME
1126 );
1127 response
1128 }
1129 }
1130
1131 fn respond<T: RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) -> Result<()> {
1132 log::debug!("rpc respond. client_id:{}. name:{}", self.id, T::NAME);
1133 self.peer.respond(receipt, response)
1134 }
1135
1136 fn respond_with_error<T: RequestMessage>(
1137 &self,
1138 receipt: Receipt<T>,
1139 error: proto::Error,
1140 ) -> Result<()> {
1141 log::debug!("rpc respond. client_id:{}. name:{}", self.id, T::NAME);
1142 self.peer.respond_with_error(receipt, error)
1143 }
1144
1145 pub fn start_telemetry(&self, db: Db) {
1146 self.telemetry.start(db);
1147 }
1148
1149 pub fn report_event(&self, kind: &str, properties: Value) {
1150 self.telemetry.report_event(kind, properties)
1151 }
1152
1153 pub fn telemetry_log_file_path(&self) -> Option<PathBuf> {
1154 self.telemetry.log_file_path()
1155 }
1156}
1157
1158impl AnyWeakEntityHandle {
1159 fn upgrade(&self, cx: &AsyncAppContext) -> Option<AnyEntityHandle> {
1160 match self {
1161 AnyWeakEntityHandle::Model(handle) => handle.upgrade(cx).map(AnyEntityHandle::Model),
1162 AnyWeakEntityHandle::View(handle) => handle.upgrade(cx).map(AnyEntityHandle::View),
1163 }
1164 }
1165}
1166
1167fn read_credentials_from_keychain(cx: &AsyncAppContext) -> Option<Credentials> {
1168 if IMPERSONATE_LOGIN.is_some() {
1169 return None;
1170 }
1171
1172 let (user_id, access_token) = cx
1173 .platform()
1174 .read_credentials(&ZED_SERVER_URL)
1175 .log_err()
1176 .flatten()?;
1177 Some(Credentials {
1178 user_id: user_id.parse().ok()?,
1179 access_token: String::from_utf8(access_token).ok()?,
1180 })
1181}
1182
1183fn write_credentials_to_keychain(credentials: &Credentials, cx: &AsyncAppContext) -> Result<()> {
1184 cx.platform().write_credentials(
1185 &ZED_SERVER_URL,
1186 &credentials.user_id.to_string(),
1187 credentials.access_token.as_bytes(),
1188 )
1189}
1190
1191const WORKTREE_URL_PREFIX: &str = "zed://worktrees/";
1192
1193pub fn encode_worktree_url(id: u64, access_token: &str) -> String {
1194 format!("{}{}/{}", WORKTREE_URL_PREFIX, id, access_token)
1195}
1196
1197pub fn decode_worktree_url(url: &str) -> Option<(u64, String)> {
1198 let path = url.trim().strip_prefix(WORKTREE_URL_PREFIX)?;
1199 let mut parts = path.split('/');
1200 let id = parts.next()?.parse::<u64>().ok()?;
1201 let access_token = parts.next()?;
1202 if access_token.is_empty() {
1203 return None;
1204 }
1205 Some((id, access_token.to_string()))
1206}
1207
1208#[cfg(test)]
1209mod tests {
1210 use super::*;
1211 use crate::test::{FakeHttpClient, FakeServer};
1212 use gpui::{executor::Deterministic, TestAppContext};
1213 use parking_lot::Mutex;
1214 use std::future;
1215
1216 #[gpui::test(iterations = 10)]
1217 async fn test_reconnection(cx: &mut TestAppContext) {
1218 cx.foreground().forbid_parking();
1219
1220 let user_id = 5;
1221 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1222 let server = FakeServer::for_client(user_id, &client, cx).await;
1223 let mut status = client.status();
1224 assert!(matches!(
1225 status.next().await,
1226 Some(Status::Connected { .. })
1227 ));
1228 assert_eq!(server.auth_count(), 1);
1229
1230 server.forbid_connections();
1231 server.disconnect();
1232 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1233
1234 server.allow_connections();
1235 cx.foreground().advance_clock(Duration::from_secs(10));
1236 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1237 assert_eq!(server.auth_count(), 1); // Client reused the cached credentials when reconnecting
1238
1239 server.forbid_connections();
1240 server.disconnect();
1241 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1242
1243 // Clear cached credentials after authentication fails
1244 server.roll_access_token();
1245 server.allow_connections();
1246 cx.foreground().advance_clock(Duration::from_secs(10));
1247 while !matches!(status.next().await, Some(Status::Connected { .. })) {}
1248 assert_eq!(server.auth_count(), 2); // Client re-authenticated due to an invalid token
1249 }
1250
1251 #[gpui::test(iterations = 10)]
1252 async fn test_connection_timeout(deterministic: Arc<Deterministic>, cx: &mut TestAppContext) {
1253 deterministic.forbid_parking();
1254
1255 let user_id = 5;
1256 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1257 let mut status = client.status();
1258
1259 // Time out when client tries to connect.
1260 client.override_authenticate(move |cx| {
1261 cx.foreground().spawn(async move {
1262 Ok(Credentials {
1263 user_id,
1264 access_token: "token".into(),
1265 })
1266 })
1267 });
1268 client.override_establish_connection(|_, cx| {
1269 cx.foreground().spawn(async move {
1270 future::pending::<()>().await;
1271 unreachable!()
1272 })
1273 });
1274 let auth_and_connect = cx.spawn({
1275 let client = client.clone();
1276 |cx| async move { client.authenticate_and_connect(false, &cx).await }
1277 });
1278 deterministic.run_until_parked();
1279 assert!(matches!(status.next().await, Some(Status::Connecting)));
1280
1281 deterministic.advance_clock(CONNECTION_TIMEOUT);
1282 assert!(matches!(
1283 status.next().await,
1284 Some(Status::ConnectionError { .. })
1285 ));
1286 auth_and_connect.await.unwrap_err();
1287
1288 // Allow the connection to be established.
1289 let server = FakeServer::for_client(user_id, &client, cx).await;
1290 assert!(matches!(
1291 status.next().await,
1292 Some(Status::Connected { .. })
1293 ));
1294
1295 // Disconnect client.
1296 server.forbid_connections();
1297 server.disconnect();
1298 while !matches!(status.next().await, Some(Status::ReconnectionError { .. })) {}
1299
1300 // Time out when re-establishing the connection.
1301 server.allow_connections();
1302 client.override_establish_connection(|_, cx| {
1303 cx.foreground().spawn(async move {
1304 future::pending::<()>().await;
1305 unreachable!()
1306 })
1307 });
1308 deterministic.advance_clock(2 * INITIAL_RECONNECTION_DELAY);
1309 assert!(matches!(
1310 status.next().await,
1311 Some(Status::Reconnecting { .. })
1312 ));
1313
1314 deterministic.advance_clock(CONNECTION_TIMEOUT);
1315 assert!(matches!(
1316 status.next().await,
1317 Some(Status::ReconnectionError { .. })
1318 ));
1319 }
1320
1321 #[gpui::test(iterations = 10)]
1322 async fn test_authenticating_more_than_once(
1323 cx: &mut TestAppContext,
1324 deterministic: Arc<Deterministic>,
1325 ) {
1326 cx.foreground().forbid_parking();
1327
1328 let auth_count = Arc::new(Mutex::new(0));
1329 let dropped_auth_count = Arc::new(Mutex::new(0));
1330 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1331 client.override_authenticate({
1332 let auth_count = auth_count.clone();
1333 let dropped_auth_count = dropped_auth_count.clone();
1334 move |cx| {
1335 let auth_count = auth_count.clone();
1336 let dropped_auth_count = dropped_auth_count.clone();
1337 cx.foreground().spawn(async move {
1338 *auth_count.lock() += 1;
1339 let _drop = util::defer(move || *dropped_auth_count.lock() += 1);
1340 future::pending::<()>().await;
1341 unreachable!()
1342 })
1343 }
1344 });
1345
1346 let _authenticate = cx.spawn(|cx| {
1347 let client = client.clone();
1348 async move { client.authenticate_and_connect(false, &cx).await }
1349 });
1350 deterministic.run_until_parked();
1351 assert_eq!(*auth_count.lock(), 1);
1352 assert_eq!(*dropped_auth_count.lock(), 0);
1353
1354 let _authenticate = cx.spawn(|cx| {
1355 let client = client.clone();
1356 async move { client.authenticate_and_connect(false, &cx).await }
1357 });
1358 deterministic.run_until_parked();
1359 assert_eq!(*auth_count.lock(), 2);
1360 assert_eq!(*dropped_auth_count.lock(), 1);
1361 }
1362
1363 #[test]
1364 fn test_encode_and_decode_worktree_url() {
1365 let url = encode_worktree_url(5, "deadbeef");
1366 assert_eq!(decode_worktree_url(&url), Some((5, "deadbeef".to_string())));
1367 assert_eq!(
1368 decode_worktree_url(&format!("\n {}\t", url)),
1369 Some((5, "deadbeef".to_string()))
1370 );
1371 assert_eq!(decode_worktree_url("not://the-right-format"), None);
1372 }
1373
1374 #[gpui::test]
1375 async fn test_subscribing_to_entity(cx: &mut TestAppContext) {
1376 cx.foreground().forbid_parking();
1377
1378 let user_id = 5;
1379 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1380 let server = FakeServer::for_client(user_id, &client, cx).await;
1381
1382 let (done_tx1, mut done_rx1) = smol::channel::unbounded();
1383 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1384 client.add_model_message_handler(
1385 move |model: ModelHandle<Model>, _: TypedEnvelope<proto::JoinProject>, _, cx| {
1386 match model.read_with(&cx, |model, _| model.id) {
1387 1 => done_tx1.try_send(()).unwrap(),
1388 2 => done_tx2.try_send(()).unwrap(),
1389 _ => unreachable!(),
1390 }
1391 async { Ok(()) }
1392 },
1393 );
1394 let model1 = cx.add_model(|_| Model {
1395 id: 1,
1396 subscription: None,
1397 });
1398 let model2 = cx.add_model(|_| Model {
1399 id: 2,
1400 subscription: None,
1401 });
1402 let model3 = cx.add_model(|_| Model {
1403 id: 3,
1404 subscription: None,
1405 });
1406
1407 let _subscription1 = model1.update(cx, |_, cx| client.add_model_for_remote_entity(1, cx));
1408 let _subscription2 = model2.update(cx, |_, cx| client.add_model_for_remote_entity(2, cx));
1409 // Ensure dropping a subscription for the same entity type still allows receiving of
1410 // messages for other entity IDs of the same type.
1411 let subscription3 = model3.update(cx, |_, cx| client.add_model_for_remote_entity(3, cx));
1412 drop(subscription3);
1413
1414 server.send(proto::JoinProject { project_id: 1 });
1415 server.send(proto::JoinProject { project_id: 2 });
1416 done_rx1.next().await.unwrap();
1417 done_rx2.next().await.unwrap();
1418 }
1419
1420 #[gpui::test]
1421 async fn test_subscribing_after_dropping_subscription(cx: &mut TestAppContext) {
1422 cx.foreground().forbid_parking();
1423
1424 let user_id = 5;
1425 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1426 let server = FakeServer::for_client(user_id, &client, cx).await;
1427
1428 let model = cx.add_model(|_| Model::default());
1429 let (done_tx1, _done_rx1) = smol::channel::unbounded();
1430 let (done_tx2, mut done_rx2) = smol::channel::unbounded();
1431 let subscription1 = client.add_message_handler(
1432 model.clone(),
1433 move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1434 done_tx1.try_send(()).unwrap();
1435 async { Ok(()) }
1436 },
1437 );
1438 drop(subscription1);
1439 let _subscription2 =
1440 client.add_message_handler(model, move |_, _: TypedEnvelope<proto::Ping>, _, _| {
1441 done_tx2.try_send(()).unwrap();
1442 async { Ok(()) }
1443 });
1444 server.send(proto::Ping {});
1445 done_rx2.next().await.unwrap();
1446 }
1447
1448 #[gpui::test]
1449 async fn test_dropping_subscription_in_handler(cx: &mut TestAppContext) {
1450 cx.foreground().forbid_parking();
1451
1452 let user_id = 5;
1453 let client = cx.update(|cx| Client::new(FakeHttpClient::with_404_response(), cx));
1454 let server = FakeServer::for_client(user_id, &client, cx).await;
1455
1456 let model = cx.add_model(|_| Model::default());
1457 let (done_tx, mut done_rx) = smol::channel::unbounded();
1458 let subscription = client.add_message_handler(
1459 model.clone(),
1460 move |model, _: TypedEnvelope<proto::Ping>, _, mut cx| {
1461 model.update(&mut cx, |model, _| model.subscription.take());
1462 done_tx.try_send(()).unwrap();
1463 async { Ok(()) }
1464 },
1465 );
1466 model.update(cx, |model, _| {
1467 model.subscription = Some(subscription);
1468 });
1469 server.send(proto::Ping {});
1470 done_rx.next().await.unwrap();
1471 }
1472
1473 #[derive(Default)]
1474 struct Model {
1475 id: usize,
1476 subscription: Option<Subscription>,
1477 }
1478
1479 impl Entity for Model {
1480 type Event = ();
1481 }
1482}