tenant.openagents/omega

1use super::{Client, Status, TypedEnvelope, proto};
2use anyhow::{Context as _, Result};
3use chrono::{DateTime, Utc};
4use cloud_api_client::websocket_protocol::MessageToClient;
5use cloud_api_client::{
6    GetAuthenticatedUserResponse, KnownOrUnknown, Organization, OrganizationId, Plan, PlanInfo,
7    UpdateSystemSettingsBody,
8};
9use cloud_api_types::OrganizationConfiguration;
10use cloud_llm_client::{
11    EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
12};
13use collections::{HashMap, HashSet, hash_map::Entry};
14use derive_more::Deref;
15use feature_flags::FeatureFlagAppExt;
16use futures::{Future, StreamExt, channel::mpsc};
17use gpui::{
18    App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, SharedString, SharedUri, Task,
19    TaskExt, WeakEntity,
20};
21use http_client::http::{HeaderMap, HeaderValue};
22use postage::{sink::Sink, watch};
23use rpc::proto::{RequestMessage, UsersResponse};
24use std::{
25    str::FromStr as _,
26    sync::{Arc, Weak},
27};
28use text::ReplicaId;
29use util::{ResultExt, TryFutureExt as _};
30
31pub type LegacyUserId = u64;
32
33#[derive(
34    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
35)]
36pub struct ChannelId(pub u64);
37
38impl std::fmt::Display for ChannelId {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        self.0.fmt(f)
41    }
42}
43
44#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
45pub struct ProjectId(pub u64);
46
47impl ProjectId {
48    pub fn to_proto(self) -> u64 {
49        self.0
50    }
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct ParticipantIndex(pub u32);
55
56#[derive(Default, Debug)]
57pub struct User {
58    pub legacy_id: LegacyUserId,
59    pub username: SharedString,
60    pub avatar_uri: SharedUri,
61    pub name: Option<String>,
62}
63
64#[derive(Clone, Debug, PartialEq, Eq)]
65pub struct Collaborator {
66    pub peer_id: proto::PeerId,
67    pub replica_id: ReplicaId,
68    pub user_id: LegacyUserId,
69    pub is_host: bool,
70    pub committer_name: Option<String>,
71    pub committer_email: Option<String>,
72}
73
74impl PartialOrd for User {
75    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
76        Some(self.cmp(other))
77    }
78}
79
80impl Ord for User {
81    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
82        self.username.cmp(&other.username)
83    }
84}
85
86impl PartialEq for User {
87    fn eq(&self, other: &Self) -> bool {
88        self.legacy_id == other.legacy_id && self.username == other.username
89    }
90}
91
92impl Eq for User {}
93
94#[derive(Debug, PartialEq)]
95pub struct Contact {
96    pub user: Arc<User>,
97    pub online: bool,
98    pub busy: bool,
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub enum ContactRequestStatus {
103    None,
104    RequestSent,
105    RequestReceived,
106    RequestAccepted,
107}
108
109pub struct UserStore {
110    users: HashMap<u64, Arc<User>>,
111    participant_indices: HashMap<u64, ParticipantIndex>,
112    update_contacts_tx: mpsc::UnboundedSender<UpdateContacts>,
113    edit_prediction_usage: Option<EditPredictionUsage>,
114    plan_info: Option<PlanInfo>,
115    current_user: watch::Receiver<Option<Arc<User>>>,
116    current_organization: Option<Arc<Organization>>,
117    organizations: Vec<Arc<Organization>>,
118    plans_by_organization: HashMap<OrganizationId, Plan>,
119    configuration_by_organization: HashMap<OrganizationId, OrganizationConfiguration>,
120    contacts: Vec<Arc<Contact>>,
121    incoming_contact_requests: Vec<Arc<User>>,
122    outgoing_contact_requests: Vec<Arc<User>>,
123    pending_contact_requests: HashMap<u64, usize>,
124    client: Weak<Client>,
125    _maintain_contacts: Task<()>,
126    _maintain_current_user: Task<Result<()>>,
127    _handle_sign_out: Task<()>,
128    weak_self: WeakEntity<Self>,
129}
130
131#[derive(Clone)]
132pub struct InviteInfo {
133    pub count: u32,
134    pub url: Arc<str>,
135}
136
137pub enum Event {
138    Contact {
139        user: Arc<User>,
140        kind: ContactEventKind,
141    },
142    ShowContacts,
143    ParticipantIndicesChanged,
144    PrivateUserInfoUpdated,
145    PlanUpdated,
146    OrganizationChanged,
147}
148
149#[derive(Clone, Copy)]
150pub enum ContactEventKind {
151    Requested,
152    Accepted,
153    Cancelled,
154}
155
156impl EventEmitter<Event> for UserStore {}
157
158enum UpdateContacts {
159    Update(proto::UpdateContacts),
160    Wait(postage::barrier::Sender),
161    Clear(postage::barrier::Sender),
162}
163
164#[derive(Debug, Clone, Copy, Deref)]
165pub struct EditPredictionUsage(pub RequestUsage);
166
167#[derive(Debug, Clone, Copy)]
168pub struct RequestUsage {
169    pub limit: UsageLimit,
170    pub amount: i32,
171}
172
173impl UserStore {
174    pub fn new(client: Arc<Client>, cx: &Context<Self>) -> Self {
175        let (mut current_user_tx, current_user_rx) = watch::channel();
176        let (sign_out_tx, mut sign_out_rx) = mpsc::unbounded();
177        let (update_contacts_tx, mut update_contacts_rx) = mpsc::unbounded();
178        let rpc_subscriptions = vec![
179            client.add_message_handler(cx.weak_entity(), Self::handle_update_contacts),
180            client.add_message_handler(cx.weak_entity(), Self::handle_show_contacts),
181        ];
182
183        client.sign_out_tx.lock().replace(sign_out_tx);
184        client.add_message_to_client_handler({
185            let this = cx.weak_entity();
186            move |message, cx| Self::handle_message_to_client(this.clone(), message, cx)
187        });
188
189        Self {
190            users: Default::default(),
191            current_user: current_user_rx,
192            current_organization: None,
193            organizations: Vec::new(),
194            plans_by_organization: HashMap::default(),
195            configuration_by_organization: HashMap::default(),
196            plan_info: None,
197            edit_prediction_usage: None,
198            contacts: Default::default(),
199            incoming_contact_requests: Default::default(),
200            participant_indices: Default::default(),
201            outgoing_contact_requests: Default::default(),
202            client: Arc::downgrade(&client),
203            update_contacts_tx,
204            _maintain_contacts: cx.spawn(async move |this, cx| {
205                let _subscriptions = rpc_subscriptions;
206                while let Some(message) = update_contacts_rx.next().await {
207                    if let Ok(task) = this.update(cx, |this, cx| this.update_contacts(message, cx))
208                    {
209                        task.log_err().await;
210                    } else {
211                        break;
212                    }
213                }
214            }),
215            _maintain_current_user: cx.spawn(async move |this, cx| {
216                let mut status = client.status();
217                let weak = Arc::downgrade(&client);
218                drop(client);
219                while let Some(status) = status.next().await {
220                    // If the client is dropped, the app is shutting down.
221                    let Some(client) = weak.upgrade() else {
222                        return Ok(());
223                    };
224                    match status {
225                        Status::Authenticated
226                        | Status::Reauthenticated
227                        | Status::Connected { .. } => {
228                            if let Some(user_id) = client.user_id() {
229                                let system_id =
230                                    client.telemetry().system_id().map(|id| id.to_string());
231                                let response = client
232                                    .cloud_client()
233                                    .get_authenticated_user(system_id)
234                                    .await
235                                    .log_err();
236
237                                let current_user_and_response = if let Some(response) = response {
238                                    let user = Arc::new(User {
239                                        legacy_id: user_id,
240                                        username: response.user.username.clone().into(),
241                                        avatar_uri: response.user.avatar_url.clone().into(),
242                                        name: response.user.name.clone(),
243                                    });
244
245                                    Some((user, response))
246                                } else {
247                                    None
248                                };
249                                current_user_tx
250                                    .send(
251                                        current_user_and_response
252                                            .as_ref()
253                                            .map(|(user, _)| user.clone()),
254                                    )
255                                    .await
256                                    .ok();
257
258                                cx.update(|cx| {
259                                    if let Some((user, response)) = current_user_and_response {
260                                        this.update(cx, |this, cx| {
261                                            this.users.insert(user_id, user);
262                                            this.update_authenticated_user(response, cx)
263                                        })
264                                    } else {
265                                        anyhow::Ok(())
266                                    }
267                                })?;
268
269                                this.update(cx, |_, cx| cx.notify())?;
270                            }
271                        }
272                        Status::SignedOut => {
273                            current_user_tx.send(None).await.ok();
274                            this.update(cx, |this, cx| {
275                                this.clear_organizations();
276                                this.clear_plan_and_usage();
277                                cx.emit(Event::PrivateUserInfoUpdated);
278                                cx.notify();
279                                this.clear_contacts()
280                            })?
281                            .await;
282                        }
283                        Status::ConnectionLost => {
284                            this.update(cx, |this, cx| {
285                                cx.notify();
286                                this.clear_contacts()
287                            })?
288                            .await;
289                        }
290                        _ => {}
291                    }
292                }
293                Ok(())
294            }),
295            _handle_sign_out: cx.spawn(async move |this, cx| {
296                while let Some(()) = sign_out_rx.next().await {
297                    let Some(client) = this
298                        .read_with(cx, |this, _cx| this.client.upgrade())
299                        .ok()
300                        .flatten()
301                    else {
302                        break;
303                    };
304
305                    client.sign_out(cx).await;
306                }
307            }),
308            pending_contact_requests: Default::default(),
309            weak_self: cx.weak_entity(),
310        }
311    }
312
313    #[cfg(feature = "test-support")]
314    pub fn clear_cache(&mut self) {
315        self.users.clear();
316    }
317
318    async fn handle_show_contacts(
319        this: Entity<Self>,
320        _: TypedEnvelope<proto::ShowContacts>,
321        mut cx: AsyncApp,
322    ) -> Result<()> {
323        this.update(&mut cx, |_, cx| cx.emit(Event::ShowContacts));
324        Ok(())
325    }
326
327    async fn handle_update_contacts(
328        this: Entity<Self>,
329        message: TypedEnvelope<proto::UpdateContacts>,
330        cx: AsyncApp,
331    ) -> Result<()> {
332        this.read_with(&cx, |this, _| {
333            this.update_contacts_tx
334                .unbounded_send(UpdateContacts::Update(message.payload))
335                .unwrap();
336        });
337        Ok(())
338    }
339
340    fn update_contacts(&mut self, message: UpdateContacts, cx: &Context<Self>) -> Task<Result<()>> {
341        match message {
342            UpdateContacts::Wait(barrier) => {
343                drop(barrier);
344                Task::ready(Ok(()))
345            }
346            UpdateContacts::Clear(barrier) => {
347                self.contacts.clear();
348                self.incoming_contact_requests.clear();
349                self.outgoing_contact_requests.clear();
350                drop(barrier);
351                Task::ready(Ok(()))
352            }
353            UpdateContacts::Update(message) => {
354                let mut user_ids = HashSet::default();
355                for contact in &message.contacts {
356                    user_ids.insert(contact.user_id);
357                }
358                user_ids.extend(message.incoming_requests.iter().map(|req| req.requester_id));
359                user_ids.extend(message.outgoing_requests.iter());
360
361                let load_users = self.get_users(user_ids.into_iter().collect(), cx);
362                cx.spawn(async move |this, cx| {
363                    load_users.await?;
364
365                    // Users are fetched in parallel above and cached in call to get_users
366                    // No need to parallelize here
367                    let mut updated_contacts = Vec::new();
368                    let this = this.upgrade().context("can't upgrade user store handle")?;
369                    for contact in message.contacts {
370                        updated_contacts
371                            .push(Arc::new(Contact::from_proto(contact, &this, cx).await?));
372                    }
373
374                    let mut incoming_requests = Vec::new();
375                    for request in message.incoming_requests {
376                        incoming_requests.push({
377                            this.update(cx, |this, cx| this.get_user(request.requester_id, cx))
378                                .await?
379                        });
380                    }
381
382                    let mut outgoing_requests = Vec::new();
383                    for requested_user_id in message.outgoing_requests {
384                        outgoing_requests.push(
385                            this.update(cx, |this, cx| this.get_user(requested_user_id, cx))
386                                .await?,
387                        );
388                    }
389
390                    let removed_contacts =
391                        HashSet::<u64>::from_iter(message.remove_contacts.iter().copied());
392                    let removed_incoming_requests =
393                        HashSet::<u64>::from_iter(message.remove_incoming_requests.iter().copied());
394                    let removed_outgoing_requests =
395                        HashSet::<u64>::from_iter(message.remove_outgoing_requests.iter().copied());
396
397                    this.update(cx, |this, cx| {
398                        // Remove contacts
399                        this.contacts
400                            .retain(|contact| !removed_contacts.contains(&contact.user.legacy_id));
401                        // Update existing contacts and insert new ones
402                        for updated_contact in updated_contacts {
403                            match this
404                                .contacts
405                                .binary_search_by_key(&&updated_contact.user.username, |contact| {
406                                    &contact.user.username
407                                }) {
408                                Ok(ix) => this.contacts[ix] = updated_contact,
409                                Err(ix) => this.contacts.insert(ix, updated_contact),
410                            }
411                        }
412
413                        // Remove incoming contact requests
414                        this.incoming_contact_requests.retain(|user| {
415                            if removed_incoming_requests.contains(&user.legacy_id) {
416                                cx.emit(Event::Contact {
417                                    user: user.clone(),
418                                    kind: ContactEventKind::Cancelled,
419                                });
420                                false
421                            } else {
422                                true
423                            }
424                        });
425                        // Update existing incoming requests and insert new ones
426                        for user in incoming_requests {
427                            match this
428                                .incoming_contact_requests
429                                .binary_search_by_key(&&user.username, |contact| &contact.username)
430                            {
431                                Ok(ix) => this.incoming_contact_requests[ix] = user,
432                                Err(ix) => this.incoming_contact_requests.insert(ix, user),
433                            }
434                        }
435
436                        // Remove outgoing contact requests
437                        this.outgoing_contact_requests
438                            .retain(|user| !removed_outgoing_requests.contains(&user.legacy_id));
439                        // Update existing incoming requests and insert new ones
440                        for request in outgoing_requests {
441                            match this
442                                .outgoing_contact_requests
443                                .binary_search_by_key(&&request.username, |contact| {
444                                    &contact.username
445                                }) {
446                                Ok(ix) => this.outgoing_contact_requests[ix] = request,
447                                Err(ix) => this.outgoing_contact_requests.insert(ix, request),
448                            }
449                        }
450
451                        cx.notify();
452                    });
453
454                    Ok(())
455                })
456            }
457        }
458    }
459
460    pub fn contacts(&self) -> &[Arc<Contact>] {
461        &self.contacts
462    }
463
464    pub fn has_contact(&self, user: &Arc<User>) -> bool {
465        self.contacts
466            .binary_search_by_key(&&user.username, |contact| &contact.user.username)
467            .is_ok()
468    }
469
470    pub fn incoming_contact_requests(&self) -> &[Arc<User>] {
471        &self.incoming_contact_requests
472    }
473
474    pub fn outgoing_contact_requests(&self) -> &[Arc<User>] {
475        &self.outgoing_contact_requests
476    }
477
478    pub fn is_contact_request_pending(&self, user: &User) -> bool {
479        self.pending_contact_requests.contains_key(&user.legacy_id)
480    }
481
482    pub fn contact_request_status(&self, user: &User) -> ContactRequestStatus {
483        if self
484            .contacts
485            .binary_search_by_key(&&user.username, |contact| &contact.user.username)
486            .is_ok()
487        {
488            ContactRequestStatus::RequestAccepted
489        } else if self
490            .outgoing_contact_requests
491            .binary_search_by_key(&&user.username, |user| &user.username)
492            .is_ok()
493        {
494            ContactRequestStatus::RequestSent
495        } else if self
496            .incoming_contact_requests
497            .binary_search_by_key(&&user.username, |user| &user.username)
498            .is_ok()
499        {
500            ContactRequestStatus::RequestReceived
501        } else {
502            ContactRequestStatus::None
503        }
504    }
505
506    pub fn request_contact(
507        &mut self,
508        responder_id: u64,
509        cx: &mut Context<Self>,
510    ) -> Task<Result<()>> {
511        self.perform_contact_request(responder_id, proto::RequestContact { responder_id }, cx)
512    }
513
514    pub fn remove_contact(&mut self, user_id: u64, cx: &mut Context<Self>) -> Task<Result<()>> {
515        self.perform_contact_request(user_id, proto::RemoveContact { user_id }, cx)
516    }
517
518    pub fn has_incoming_contact_request(&self, user_id: u64) -> bool {
519        self.incoming_contact_requests
520            .iter()
521            .any(|user| user.legacy_id == user_id)
522    }
523
524    pub fn respond_to_contact_request(
525        &mut self,
526        requester_id: u64,
527        accept: bool,
528        cx: &mut Context<Self>,
529    ) -> Task<Result<()>> {
530        self.perform_contact_request(
531            requester_id,
532            proto::RespondToContactRequest {
533                requester_id,
534                response: if accept {
535                    proto::ContactRequestResponse::Accept
536                } else {
537                    proto::ContactRequestResponse::Decline
538                } as i32,
539            },
540            cx,
541        )
542    }
543
544    pub fn dismiss_contact_request(
545        &self,
546        requester_id: u64,
547        cx: &Context<Self>,
548    ) -> Task<Result<()>> {
549        let client = self.client.upgrade();
550        cx.spawn(async move |_, _| {
551            client
552                .context("can't upgrade client reference")?
553                .request(proto::RespondToContactRequest {
554                    requester_id,
555                    response: proto::ContactRequestResponse::Dismiss as i32,
556                })
557                .await?;
558            Ok(())
559        })
560    }
561
562    fn perform_contact_request<T: RequestMessage>(
563        &mut self,
564        user_id: u64,
565        request: T,
566        cx: &mut Context<Self>,
567    ) -> Task<Result<()>> {
568        let client = self.client.upgrade();
569        *self.pending_contact_requests.entry(user_id).or_insert(0) += 1;
570        cx.notify();
571
572        cx.spawn(async move |this, cx| {
573            let response = client
574                .context("can't upgrade client reference")?
575                .request(request)
576                .await;
577            this.update(cx, |this, cx| {
578                if let Entry::Occupied(mut request_count) =
579                    this.pending_contact_requests.entry(user_id)
580                {
581                    *request_count.get_mut() -= 1;
582                    if *request_count.get() == 0 {
583                        request_count.remove();
584                    }
585                }
586                cx.notify();
587            })?;
588            response?;
589            Ok(())
590        })
591    }
592
593    pub fn clear_contacts(&self) -> impl Future<Output = ()> + use<> {
594        let (tx, mut rx) = postage::barrier::channel();
595        self.update_contacts_tx
596            .unbounded_send(UpdateContacts::Clear(tx))
597            .unwrap();
598        async move {
599            rx.next().await;
600        }
601    }
602
603    pub fn contact_updates_done(&self) -> impl Future<Output = ()> {
604        let (tx, mut rx) = postage::barrier::channel();
605        self.update_contacts_tx
606            .unbounded_send(UpdateContacts::Wait(tx))
607            .unwrap();
608        async move {
609            rx.next().await;
610        }
611    }
612
613    pub fn get_users(
614        &self,
615        user_ids: Vec<u64>,
616        cx: &Context<Self>,
617    ) -> Task<Result<Vec<Arc<User>>>> {
618        let mut user_ids_to_fetch = user_ids.clone();
619        user_ids_to_fetch.retain(|id| !self.users.contains_key(id));
620
621        cx.spawn(async move |this, cx| {
622            if !user_ids_to_fetch.is_empty() {
623                this.update(cx, |this, cx| {
624                    this.load_users(
625                        proto::GetUsers {
626                            user_ids: user_ids_to_fetch,
627                        },
628                        cx,
629                    )
630                })?
631                .await?;
632            }
633
634            this.read_with(cx, |this, _| {
635                user_ids
636                    .iter()
637                    .map(|user_id| {
638                        this.users
639                            .get(user_id)
640                            .cloned()
641                            .with_context(|| format!("user {user_id} not found"))
642                    })
643                    .collect()
644            })?
645        })
646    }
647
648    pub fn fuzzy_search_users(
649        &self,
650        query: String,
651        cx: &Context<Self>,
652    ) -> Task<Result<Vec<Arc<User>>>> {
653        self.load_users(proto::FuzzySearchUsers { query }, cx)
654    }
655
656    pub fn get_cached_user(&self, user_id: u64) -> Option<Arc<User>> {
657        self.users.get(&user_id).cloned()
658    }
659
660    pub fn get_user_optimistic(&self, user_id: u64, cx: &Context<Self>) -> Option<Arc<User>> {
661        if let Some(user) = self.users.get(&user_id).cloned() {
662            return Some(user);
663        }
664
665        self.get_user(user_id, cx).detach_and_log_err(cx);
666        None
667    }
668
669    pub fn get_user(&self, user_id: u64, cx: &Context<Self>) -> Task<Result<Arc<User>>> {
670        if let Some(user) = self.users.get(&user_id).cloned() {
671            return Task::ready(Ok(user));
672        }
673
674        let load_users = self.get_users(vec![user_id], cx);
675        cx.spawn(async move |this, cx| {
676            load_users.await?;
677            this.read_with(cx, |this, _| {
678                this.users
679                    .get(&user_id)
680                    .cloned()
681                    .context("server responded with no users")
682            })?
683        })
684    }
685
686    pub fn current_user(&self) -> Option<Arc<User>> {
687        self.current_user.borrow().clone()
688    }
689
690    pub fn current_organization(&self) -> Option<Arc<Organization>> {
691        self.current_organization.clone()
692    }
693
694    pub fn set_current_organization(
695        &mut self,
696        organization: Arc<Organization>,
697        cx: &mut Context<Self>,
698    ) -> Task<Result<()>> {
699        let is_same_organization = self
700            .current_organization
701            .as_ref()
702            .is_some_and(|current| current.id == organization.id);
703
704        if is_same_organization {
705            return Task::ready(Ok(()));
706        }
707
708        let organization_id = organization.id.clone();
709        self.current_organization.replace(organization);
710        cx.emit(Event::OrganizationChanged);
711        cx.notify();
712
713        let Some(client) = self.client.upgrade() else {
714            return Task::ready(Ok(()));
715        };
716        let Some(system_id) = client.telemetry().system_id().map(|id| id.to_string()) else {
717            // Without a system ID we have no addressable target row on the
718            // server, so the selection stays purely session-local.
719            return Task::ready(Ok(()));
720        };
721        let cloud_client = client.cloud_client();
722
723        cx.background_spawn(async move {
724            let body = UpdateSystemSettingsBody {
725                selected_organization_id: Some(organization_id),
726            };
727            cloud_client
728                .update_system_settings(system_id, body)
729                .await
730                .context("failed to persist selected organization")?;
731            Ok(())
732        })
733    }
734
735    pub fn organizations(&self) -> &Vec<Arc<Organization>> {
736        &self.organizations
737    }
738
739    pub fn plan_for_organization(&self, organization_id: &OrganizationId) -> Option<Plan> {
740        self.plans_by_organization.get(organization_id).copied()
741    }
742
743    pub fn current_organization_configuration(&self) -> Option<&OrganizationConfiguration> {
744        let current_organization = self.current_organization.as_ref()?;
745
746        self.configuration_by_organization
747            .get(&current_organization.id)
748    }
749
750    #[cfg(any(test, feature = "test-support"))]
751    pub fn set_current_organization_configuration_for_test(
752        &mut self,
753        organization: Arc<Organization>,
754        configuration: OrganizationConfiguration,
755        cx: &mut Context<Self>,
756    ) {
757        self.current_organization = Some(organization.clone());
758        self.organizations = vec![organization.clone()];
759        self.configuration_by_organization
760            .insert(organization.id.clone(), configuration);
761        cx.emit(Event::OrganizationChanged);
762        cx.notify();
763    }
764
765    pub fn plan(&self) -> Option<Plan> {
766        #[cfg(debug_assertions)]
767        if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {
768            use cloud_api_client::Plan;
769
770            return match plan.as_str() {
771                "free" => Some(Plan::ZedFree),
772                "trial" => Some(Plan::ZedProTrial),
773                "pro" => Some(Plan::ZedPro),
774                _ => {
775                    panic!("ZED_SIMULATE_PLAN must be one of 'free', 'trial', or 'pro'");
776                }
777            };
778        }
779
780        if let Some(organization) = &self.current_organization {
781            return self.plan_for_organization(&organization.id);
782        }
783
784        self.plan_info.as_ref().map(|info| info.plan())
785    }
786
787    pub fn subscription_period(&self) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
788        self.plan_info
789            .as_ref()
790            .and_then(|plan| plan.subscription_period)
791            .map(|subscription_period| {
792                (
793                    subscription_period.started_at.0,
794                    subscription_period.ended_at.0,
795                )
796            })
797    }
798
799    pub fn trial_started_at(&self) -> Option<DateTime<Utc>> {
800        self.plan_info
801            .as_ref()
802            .and_then(|plan| plan.trial_started_at)
803            .map(|trial_started_at| trial_started_at.0)
804    }
805
806    /// Returns whether the user's account is too new to use the service.
807    ///
808    /// This only applies when operating under the user's personal organization,
809    /// not a business organization.
810    pub fn account_too_young(&self) -> bool {
811        if let Some(org) = &self.current_organization {
812            if !org.is_personal {
813                return false;
814            }
815        }
816
817        self.plan_info
818            .as_ref()
819            .map(|plan| plan.is_account_too_young)
820            .unwrap_or_default()
821    }
822
823    /// Returns whether the current user has overdue invoices and usage should be blocked.
824    pub fn has_overdue_invoices(&self) -> bool {
825        self.plan_info
826            .as_ref()
827            .map(|plan| plan.has_overdue_invoices)
828            .unwrap_or_default()
829    }
830
831    pub fn edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
832        self.edit_prediction_usage
833    }
834
835    pub fn update_edit_prediction_usage(
836        &mut self,
837        usage: EditPredictionUsage,
838        cx: &mut Context<Self>,
839    ) {
840        self.edit_prediction_usage = Some(usage);
841        cx.notify();
842    }
843
844    pub fn clear_organizations(&mut self) {
845        self.organizations.clear();
846        self.current_organization = None;
847    }
848
849    pub fn clear_plan_and_usage(&mut self) {
850        self.plan_info = None;
851        self.edit_prediction_usage = None;
852    }
853
854    fn update_authenticated_user(
855        &mut self,
856        response: GetAuthenticatedUserResponse,
857        cx: &mut Context<Self>,
858    ) {
859        let staff = response.user.is_staff && !*feature_flags::ZED_DISABLE_STAFF;
860        cx.update_flags(staff, response.feature_flags);
861        if let Some(client) = self.client.upgrade() {
862            client
863                .telemetry
864                .set_authenticated_user_info(Some(response.user.metrics_id.clone()), staff);
865        }
866
867        self.organizations = response.organizations.into_iter().map(Arc::new).collect();
868
869        self.current_organization = response
870            .default_organization_id
871            .and_then(|default_organization_id| {
872                self.organizations
873                    .iter()
874                    .find(|organization| organization.id == default_organization_id)
875                    .cloned()
876            })
877            .or_else(|| self.organizations.first().cloned());
878        self.plans_by_organization = response
879            .plans_by_organization
880            .into_iter()
881            .map(|(organization_id, plan)| {
882                let plan = match plan {
883                    KnownOrUnknown::Known(plan) => plan,
884                    KnownOrUnknown::Unknown(_) => {
885                        // If we get a plan that we don't recognize, fall back to the Free plan.
886                        Plan::ZedFree
887                    }
888                };
889
890                (organization_id, plan)
891            })
892            .collect();
893        self.configuration_by_organization =
894            response.configuration_by_organization.into_iter().collect();
895
896        self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage {
897            limit: response.plan.usage.edit_predictions.limit,
898            amount: response.plan.usage.edit_predictions.used as i32,
899        }));
900        self.plan_info = Some(response.plan);
901        cx.emit(Event::PrivateUserInfoUpdated);
902    }
903
904    fn handle_message_to_client(this: WeakEntity<Self>, message: &MessageToClient, cx: &App) {
905        cx.spawn(async move |cx| {
906            match message {
907                MessageToClient::UserUpdated => {
908                    let (cloud_client, system_id) = cx
909                        .update(|cx| {
910                            this.read_with(cx, |this, _cx| {
911                                this.client.upgrade().map(|client| {
912                                    let system_id =
913                                        client.telemetry().system_id().map(|id| id.to_string());
914                                    (client.cloud_client(), system_id)
915                                })
916                            })
917                        })?
918                        .ok_or(anyhow::anyhow!("Failed to get Cloud client"))?;
919
920                    let response = cloud_client.get_authenticated_user(system_id).await?;
921                    cx.update(|cx| {
922                        this.update(cx, |this, cx| {
923                            this.update_authenticated_user(response, cx);
924                        })
925                    })?;
926                }
927            }
928
929            anyhow::Ok(())
930        })
931        .detach_and_log_err(cx);
932    }
933
934    pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
935        self.current_user.clone()
936    }
937
938    fn load_users(
939        &self,
940        request: impl RequestMessage<Response = UsersResponse>,
941        cx: &Context<Self>,
942    ) -> Task<Result<Vec<Arc<User>>>> {
943        let client = self.client.clone();
944        cx.spawn(async move |this, cx| {
945            if let Some(rpc) = client.upgrade() {
946                let response = rpc.request(request).await.context("error loading users")?;
947                let users = response.users;
948
949                this.update(cx, |this, _| this.insert(users))
950            } else {
951                Ok(Vec::new())
952            }
953        })
954    }
955
956    pub fn insert(&mut self, users: Vec<proto::User>) -> Vec<Arc<User>> {
957        let mut ret = Vec::with_capacity(users.len());
958        for user in users {
959            let user = User::new(user);
960            self.users.insert(user.legacy_id, user.clone());
961            ret.push(user)
962        }
963        ret
964    }
965
966    pub fn set_participant_indices(
967        &mut self,
968        participant_indices: HashMap<u64, ParticipantIndex>,
969        cx: &mut Context<Self>,
970    ) {
971        if participant_indices != self.participant_indices {
972            self.participant_indices = participant_indices;
973            cx.emit(Event::ParticipantIndicesChanged);
974        }
975    }
976
977    pub fn participant_indices(&self) -> &HashMap<u64, ParticipantIndex> {
978        &self.participant_indices
979    }
980
981    pub fn participant_names(
982        &self,
983        user_ids: impl Iterator<Item = u64>,
984        cx: &App,
985    ) -> HashMap<u64, SharedString> {
986        let mut ret = HashMap::default();
987        let mut missing_user_ids = Vec::new();
988        for id in user_ids {
989            if let Some(username) = self.get_cached_user(id).map(|u| u.username.clone()) {
990                ret.insert(id, username);
991            } else {
992                missing_user_ids.push(id)
993            }
994        }
995        if !missing_user_ids.is_empty() {
996            let this = self.weak_self.clone();
997            cx.spawn(async move |cx| {
998                this.update(cx, |this, cx| this.get_users(missing_user_ids, cx))?
999                    .await
1000            })
1001            .detach_and_log_err(cx);
1002        }
1003        ret
1004    }
1005}
1006
1007impl User {
1008    fn new(message: proto::User) -> Arc<Self> {
1009        Arc::new(User {
1010            legacy_id: message.id,
1011            username: message.username.into(),
1012            avatar_uri: message.avatar_url.into(),
1013            name: message.name,
1014        })
1015    }
1016}
1017
1018impl Contact {
1019    async fn from_proto(
1020        contact: proto::Contact,
1021        user_store: &Entity<UserStore>,
1022        cx: &mut AsyncApp,
1023    ) -> Result<Self> {
1024        let user = user_store
1025            .update(cx, |user_store, cx| {
1026                user_store.get_user(contact.user_id, cx)
1027            })
1028            .await?;
1029        Ok(Self {
1030            user,
1031            online: contact.online,
1032            busy: contact.busy,
1033        })
1034    }
1035}
1036
1037impl Collaborator {
1038    pub fn from_proto(message: proto::Collaborator) -> Result<Self> {
1039        Ok(Self {
1040            peer_id: message.peer_id.context("invalid peer id")?,
1041            replica_id: ReplicaId::new(message.replica_id as u16),
1042            user_id: message.user_id as LegacyUserId,
1043            is_host: message.is_host,
1044            committer_name: message.committer_name,
1045            committer_email: message.committer_email,
1046        })
1047    }
1048}
1049
1050impl RequestUsage {
1051    pub fn over_limit(&self) -> bool {
1052        match self.limit {
1053            UsageLimit::Limited(limit) => self.amount >= limit,
1054            UsageLimit::Unlimited => false,
1055        }
1056    }
1057
1058    fn from_headers(
1059        limit_name: &str,
1060        amount_name: &str,
1061        headers: &HeaderMap<HeaderValue>,
1062    ) -> Result<Self> {
1063        let limit = headers
1064            .get(limit_name)
1065            .with_context(|| format!("missing {limit_name:?} header"))?;
1066        let limit = UsageLimit::from_str(limit.to_str()?)?;
1067
1068        let amount = headers
1069            .get(amount_name)
1070            .with_context(|| format!("missing {amount_name:?} header"))?;
1071        let amount = amount.to_str()?.parse::<i32>()?;
1072
1073        Ok(Self { limit, amount })
1074    }
1075}
1076
1077impl EditPredictionUsage {
1078    pub fn from_headers(headers: &HeaderMap<HeaderValue>) -> Result<Self> {
1079        Ok(Self(RequestUsage::from_headers(
1080            EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME,
1081            EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME,
1082            headers,
1083        )?))
1084    }
1085}
1086