1085 lines · 36.8 KB · Rust
RawDownload
use super::{Client, Status, TypedEnvelope, proto};use anyhow::{Context as _, Result};use chrono::{DateTime, Utc};use cloud_api_client::websocket_protocol::MessageToClient;use cloud_api_client::{    GetAuthenticatedUserResponse, KnownOrUnknown, Organization, OrganizationId, Plan, PlanInfo,    UpdateSystemSettingsBody,};use cloud_api_types::OrganizationConfiguration;use cloud_llm_client::{    EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,};use collections::{HashMap, HashSet, hash_map::Entry};use derive_more::Deref;use feature_flags::FeatureFlagAppExt;use futures::{Future, StreamExt, channel::mpsc};use gpui::{    App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, SharedString, SharedUri, Task,    TaskExt, WeakEntity,};use http_client::http::{HeaderMap, HeaderValue};use postage::{sink::Sink, watch};use rpc::proto::{RequestMessage, UsersResponse};use std::{    str::FromStr as _,    sync::{Arc, Weak},};use text::ReplicaId;use util::{ResultExt, TryFutureExt as _}; pub type LegacyUserId = u64; #[derive(    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,)]pub struct ChannelId(pub u64); impl std::fmt::Display for ChannelId {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        self.0.fmt(f)    }} #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]pub struct ProjectId(pub u64); impl ProjectId {    pub fn to_proto(self) -> u64 {        self.0    }} #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub struct ParticipantIndex(pub u32); #[derive(Default, Debug)]pub struct User {    pub legacy_id: LegacyUserId,    pub username: SharedString,    pub avatar_uri: SharedUri,    pub name: Option<String>,} #[derive(Clone, Debug, PartialEq, Eq)]pub struct Collaborator {    pub peer_id: proto::PeerId,    pub replica_id: ReplicaId,    pub user_id: LegacyUserId,    pub is_host: bool,    pub committer_name: Option<String>,    pub committer_email: Option<String>,} impl PartialOrd for User {    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {        Some(self.cmp(other))    }} impl Ord for User {    fn cmp(&self, other: &Self) -> std::cmp::Ordering {        self.username.cmp(&other.username)    }} impl PartialEq for User {    fn eq(&self, other: &Self) -> bool {        self.legacy_id == other.legacy_id && self.username == other.username    }} impl Eq for User {} #[derive(Debug, PartialEq)]pub struct Contact {    pub user: Arc<User>,    pub online: bool,    pub busy: bool,} #[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum ContactRequestStatus {    None,    RequestSent,    RequestReceived,    RequestAccepted,} pub struct UserStore {    users: HashMap<u64, Arc<User>>,    participant_indices: HashMap<u64, ParticipantIndex>,    update_contacts_tx: mpsc::UnboundedSender<UpdateContacts>,    edit_prediction_usage: Option<EditPredictionUsage>,    plan_info: Option<PlanInfo>,    current_user: watch::Receiver<Option<Arc<User>>>,    current_organization: Option<Arc<Organization>>,    organizations: Vec<Arc<Organization>>,    plans_by_organization: HashMap<OrganizationId, Plan>,    configuration_by_organization: HashMap<OrganizationId, OrganizationConfiguration>,    contacts: Vec<Arc<Contact>>,    incoming_contact_requests: Vec<Arc<User>>,    outgoing_contact_requests: Vec<Arc<User>>,    pending_contact_requests: HashMap<u64, usize>,    client: Weak<Client>,    _maintain_contacts: Task<()>,    _maintain_current_user: Task<Result<()>>,    _handle_sign_out: Task<()>,    weak_self: WeakEntity<Self>,} #[derive(Clone)]pub struct InviteInfo {    pub count: u32,    pub url: Arc<str>,} pub enum Event {    Contact {        user: Arc<User>,        kind: ContactEventKind,    },    ShowContacts,    ParticipantIndicesChanged,    PrivateUserInfoUpdated,    PlanUpdated,    OrganizationChanged,} #[derive(Clone, Copy)]pub enum ContactEventKind {    Requested,    Accepted,    Cancelled,} impl EventEmitter<Event> for UserStore {} enum UpdateContacts {    Update(proto::UpdateContacts),    Wait(postage::barrier::Sender),    Clear(postage::barrier::Sender),} #[derive(Debug, Clone, Copy, Deref)]pub struct EditPredictionUsage(pub RequestUsage); #[derive(Debug, Clone, Copy)]pub struct RequestUsage {    pub limit: UsageLimit,    pub amount: i32,} impl UserStore {    pub fn new(client: Arc<Client>, cx: &Context<Self>) -> Self {        let (mut current_user_tx, current_user_rx) = watch::channel();        let (sign_out_tx, mut sign_out_rx) = mpsc::unbounded();        let (update_contacts_tx, mut update_contacts_rx) = mpsc::unbounded();        let rpc_subscriptions = vec![            client.add_message_handler(cx.weak_entity(), Self::handle_update_contacts),            client.add_message_handler(cx.weak_entity(), Self::handle_show_contacts),        ];         client.sign_out_tx.lock().replace(sign_out_tx);        client.add_message_to_client_handler({            let this = cx.weak_entity();            move |message, cx| Self::handle_message_to_client(this.clone(), message, cx)        });         Self {            users: Default::default(),            current_user: current_user_rx,            current_organization: None,            organizations: Vec::new(),            plans_by_organization: HashMap::default(),            configuration_by_organization: HashMap::default(),            plan_info: None,            edit_prediction_usage: None,            contacts: Default::default(),            incoming_contact_requests: Default::default(),            participant_indices: Default::default(),            outgoing_contact_requests: Default::default(),            client: Arc::downgrade(&client),            update_contacts_tx,            _maintain_contacts: cx.spawn(async move |this, cx| {                let _subscriptions = rpc_subscriptions;                while let Some(message) = update_contacts_rx.next().await {                    if let Ok(task) = this.update(cx, |this, cx| this.update_contacts(message, cx))                    {                        task.log_err().await;                    } else {                        break;                    }                }            }),            _maintain_current_user: cx.spawn(async move |this, cx| {                let mut status = client.status();                let weak = Arc::downgrade(&client);                drop(client);                while let Some(status) = status.next().await {                    // If the client is dropped, the app is shutting down.                    let Some(client) = weak.upgrade() else {                        return Ok(());                    };                    match status {                        Status::Authenticated                        | Status::Reauthenticated                        | Status::Connected { .. } => {                            if let Some(user_id) = client.user_id() {                                let system_id =                                    client.telemetry().system_id().map(|id| id.to_string());                                let response = client                                    .cloud_client()                                    .get_authenticated_user(system_id)                                    .await                                    .log_err();                                 let current_user_and_response = if let Some(response) = response {                                    let user = Arc::new(User {                                        legacy_id: user_id,                                        username: response.user.username.clone().into(),                                        avatar_uri: response.user.avatar_url.clone().into(),                                        name: response.user.name.clone(),                                    });                                     Some((user, response))                                } else {                                    None                                };                                current_user_tx                                    .send(                                        current_user_and_response                                            .as_ref()                                            .map(|(user, _)| user.clone()),                                    )                                    .await                                    .ok();                                 cx.update(|cx| {                                    if let Some((user, response)) = current_user_and_response {                                        this.update(cx, |this, cx| {                                            this.users.insert(user_id, user);                                            this.update_authenticated_user(response, cx)                                        })                                    } else {                                        anyhow::Ok(())                                    }                                })?;                                 this.update(cx, |_, cx| cx.notify())?;                            }                        }                        Status::SignedOut => {                            current_user_tx.send(None).await.ok();                            this.update(cx, |this, cx| {                                this.clear_organizations();                                this.clear_plan_and_usage();                                cx.emit(Event::PrivateUserInfoUpdated);                                cx.notify();                                this.clear_contacts()                            })?                            .await;                        }                        Status::ConnectionLost => {                            this.update(cx, |this, cx| {                                cx.notify();                                this.clear_contacts()                            })?                            .await;                        }                        _ => {}                    }                }                Ok(())            }),            _handle_sign_out: cx.spawn(async move |this, cx| {                while let Some(()) = sign_out_rx.next().await {                    let Some(client) = this                        .read_with(cx, |this, _cx| this.client.upgrade())                        .ok()                        .flatten()                    else {                        break;                    };                     client.sign_out(cx).await;                }            }),            pending_contact_requests: Default::default(),            weak_self: cx.weak_entity(),        }    }     #[cfg(feature = "test-support")]    pub fn clear_cache(&mut self) {        self.users.clear();    }     async fn handle_show_contacts(        this: Entity<Self>,        _: TypedEnvelope<proto::ShowContacts>,        mut cx: AsyncApp,    ) -> Result<()> {        this.update(&mut cx, |_, cx| cx.emit(Event::ShowContacts));        Ok(())    }     async fn handle_update_contacts(        this: Entity<Self>,        message: TypedEnvelope<proto::UpdateContacts>,        cx: AsyncApp,    ) -> Result<()> {        this.read_with(&cx, |this, _| {            this.update_contacts_tx                .unbounded_send(UpdateContacts::Update(message.payload))                .unwrap();        });        Ok(())    }     fn update_contacts(&mut self, message: UpdateContacts, cx: &Context<Self>) -> Task<Result<()>> {        match message {            UpdateContacts::Wait(barrier) => {                drop(barrier);                Task::ready(Ok(()))            }            UpdateContacts::Clear(barrier) => {                self.contacts.clear();                self.incoming_contact_requests.clear();                self.outgoing_contact_requests.clear();                drop(barrier);                Task::ready(Ok(()))            }            UpdateContacts::Update(message) => {                let mut user_ids = HashSet::default();                for contact in &message.contacts {                    user_ids.insert(contact.user_id);                }                user_ids.extend(message.incoming_requests.iter().map(|req| req.requester_id));                user_ids.extend(message.outgoing_requests.iter());                 let load_users = self.get_users(user_ids.into_iter().collect(), cx);                cx.spawn(async move |this, cx| {                    load_users.await?;                     // Users are fetched in parallel above and cached in call to get_users                    // No need to parallelize here                    let mut updated_contacts = Vec::new();                    let this = this.upgrade().context("can't upgrade user store handle")?;                    for contact in message.contacts {                        updated_contacts                            .push(Arc::new(Contact::from_proto(contact, &this, cx).await?));                    }                     let mut incoming_requests = Vec::new();                    for request in message.incoming_requests {                        incoming_requests.push({                            this.update(cx, |this, cx| this.get_user(request.requester_id, cx))                                .await?                        });                    }                     let mut outgoing_requests = Vec::new();                    for requested_user_id in message.outgoing_requests {                        outgoing_requests.push(                            this.update(cx, |this, cx| this.get_user(requested_user_id, cx))                                .await?,                        );                    }                     let removed_contacts =                        HashSet::<u64>::from_iter(message.remove_contacts.iter().copied());                    let removed_incoming_requests =                        HashSet::<u64>::from_iter(message.remove_incoming_requests.iter().copied());                    let removed_outgoing_requests =                        HashSet::<u64>::from_iter(message.remove_outgoing_requests.iter().copied());                     this.update(cx, |this, cx| {                        // Remove contacts                        this.contacts                            .retain(|contact| !removed_contacts.contains(&contact.user.legacy_id));                        // Update existing contacts and insert new ones                        for updated_contact in updated_contacts {                            match this                                .contacts                                .binary_search_by_key(&&updated_contact.user.username, |contact| {                                    &contact.user.username                                }) {                                Ok(ix) => this.contacts[ix] = updated_contact,                                Err(ix) => this.contacts.insert(ix, updated_contact),                            }                        }                         // Remove incoming contact requests                        this.incoming_contact_requests.retain(|user| {                            if removed_incoming_requests.contains(&user.legacy_id) {                                cx.emit(Event::Contact {                                    user: user.clone(),                                    kind: ContactEventKind::Cancelled,                                });                                false                            } else {                                true                            }                        });                        // Update existing incoming requests and insert new ones                        for user in incoming_requests {                            match this                                .incoming_contact_requests                                .binary_search_by_key(&&user.username, |contact| &contact.username)                            {                                Ok(ix) => this.incoming_contact_requests[ix] = user,                                Err(ix) => this.incoming_contact_requests.insert(ix, user),                            }                        }                         // Remove outgoing contact requests                        this.outgoing_contact_requests                            .retain(|user| !removed_outgoing_requests.contains(&user.legacy_id));                        // Update existing incoming requests and insert new ones                        for request in outgoing_requests {                            match this                                .outgoing_contact_requests                                .binary_search_by_key(&&request.username, |contact| {                                    &contact.username                                }) {                                Ok(ix) => this.outgoing_contact_requests[ix] = request,                                Err(ix) => this.outgoing_contact_requests.insert(ix, request),                            }                        }                         cx.notify();                    });                     Ok(())                })            }        }    }     pub fn contacts(&self) -> &[Arc<Contact>] {        &self.contacts    }     pub fn has_contact(&self, user: &Arc<User>) -> bool {        self.contacts            .binary_search_by_key(&&user.username, |contact| &contact.user.username)            .is_ok()    }     pub fn incoming_contact_requests(&self) -> &[Arc<User>] {        &self.incoming_contact_requests    }     pub fn outgoing_contact_requests(&self) -> &[Arc<User>] {        &self.outgoing_contact_requests    }     pub fn is_contact_request_pending(&self, user: &User) -> bool {        self.pending_contact_requests.contains_key(&user.legacy_id)    }     pub fn contact_request_status(&self, user: &User) -> ContactRequestStatus {        if self            .contacts            .binary_search_by_key(&&user.username, |contact| &contact.user.username)            .is_ok()        {            ContactRequestStatus::RequestAccepted        } else if self            .outgoing_contact_requests            .binary_search_by_key(&&user.username, |user| &user.username)            .is_ok()        {            ContactRequestStatus::RequestSent        } else if self            .incoming_contact_requests            .binary_search_by_key(&&user.username, |user| &user.username)            .is_ok()        {            ContactRequestStatus::RequestReceived

Lines 1–500 of 1085.