use std::sync::Arc; use std::time::Duration; use anyhow::Context as _; use chrono::{DateTime, Utc}; use cloud_api_client::{AuthenticatedUser, CloudApiClient, GetAuthenticatedUserResponse, PlanInfo}; use cloud_llm_client::Plan; use gpui::{Context, Entity, Subscription, Task}; use util::{ResultExt as _, maybe}; use crate::user::Event as RpcUserStoreEvent; use crate::{EditPredictionUsage, ModelRequestUsage, RequestUsage, UserStore}; pub struct CloudUserStore { cloud_client: Arc, authenticated_user: Option>, plan_info: Option>, model_request_usage: Option, edit_prediction_usage: Option, _maintain_authenticated_user_task: Task<()>, _rpc_plan_updated_subscription: Subscription, } impl CloudUserStore { pub fn new( cloud_client: Arc, rpc_user_store: Entity, cx: &mut Context, ) -> Self { let rpc_plan_updated_subscription = cx.subscribe(&rpc_user_store, Self::handle_rpc_user_store_event); Self { cloud_client: cloud_client.clone(), authenticated_user: None, plan_info: None, model_request_usage: None, edit_prediction_usage: None, _maintain_authenticated_user_task: cx.spawn(async move |this, cx| { maybe!(async move { loop { let Some(this) = this.upgrade() else { return anyhow::Ok(()); }; if cloud_client.has_credentials() { let already_fetched_authenticated_user = this .read_with(cx, |this, _cx| this.authenticated_user().is_some()) .unwrap_or(false); if already_fetched_authenticated_user { // We already fetched the authenticated user; nothing to do. } else { let authenticated_user_result = cloud_client .get_authenticated_user() .await .context("failed to fetch authenticated user"); if let Some(response) = authenticated_user_result.log_err() { this.update(cx, |this, _cx| { this.update_authenticated_user(response); }) .ok(); } } } else { this.update(cx, |this, _cx| { this.authenticated_user.take(); this.plan_info.take(); }) .ok(); } cx.background_executor() .timer(Duration::from_millis(100)) .await; } }) .await .log_err(); }), _rpc_plan_updated_subscription: rpc_plan_updated_subscription, } } pub fn is_authenticated(&self) -> bool { self.authenticated_user.is_some() } pub fn authenticated_user(&self) -> Option> { self.authenticated_user.clone() } pub fn plan(&self) -> Option { self.plan_info.as_ref().map(|plan| plan.plan) } pub fn subscription_period(&self) -> Option<(DateTime, DateTime)> { self.plan_info .as_ref() .and_then(|plan| plan.subscription_period) .map(|subscription_period| { ( subscription_period.started_at.0, subscription_period.ended_at.0, ) }) } pub fn trial_started_at(&self) -> Option> { self.plan_info .as_ref() .and_then(|plan| plan.trial_started_at) .map(|trial_started_at| trial_started_at.0) } pub fn has_accepted_tos(&self) -> bool { self.authenticated_user .as_ref() .map(|user| user.accepted_tos_at.is_some()) .unwrap_or_default() } /// Returns whether the user's account is too new to use the service. pub fn account_too_young(&self) -> bool { self.plan_info .as_ref() .map(|plan| plan.is_account_too_young) .unwrap_or_default() } /// Returns whether the current user has overdue invoices and usage should be blocked. pub fn has_overdue_invoices(&self) -> bool { self.plan_info .as_ref() .map(|plan| plan.has_overdue_invoices) .unwrap_or_default() } pub fn is_usage_based_billing_enabled(&self) -> bool { self.plan_info .as_ref() .map(|plan| plan.is_usage_based_billing_enabled) .unwrap_or_default() } pub fn model_request_usage(&self) -> Option { self.model_request_usage } pub fn update_model_request_usage(&mut self, usage: ModelRequestUsage, cx: &mut Context) { self.model_request_usage = Some(usage); cx.notify(); } pub fn edit_prediction_usage(&self) -> Option { self.edit_prediction_usage } pub fn update_edit_prediction_usage( &mut self, usage: EditPredictionUsage, cx: &mut Context, ) { self.edit_prediction_usage = Some(usage); cx.notify(); } fn update_authenticated_user(&mut self, response: GetAuthenticatedUserResponse) { self.authenticated_user = Some(Arc::new(response.user)); self.model_request_usage = Some(ModelRequestUsage(RequestUsage { limit: response.plan.usage.model_requests.limit, amount: response.plan.usage.model_requests.used as i32, })); self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage { limit: response.plan.usage.edit_predictions.limit, amount: response.plan.usage.edit_predictions.used as i32, })); self.plan_info = Some(Arc::new(response.plan)); } fn handle_rpc_user_store_event( &mut self, _: Entity, event: &RpcUserStoreEvent, cx: &mut Context, ) { match event { RpcUserStoreEvent::PlanUpdated => { cx.spawn(async move |this, cx| { let cloud_client = cx.update(|cx| this.read_with(cx, |this, _cx| this.cloud_client.clone()))??; let response = cloud_client .get_authenticated_user() .await .context("failed to fetch authenticated user")?; cx.update(|cx| { this.update(cx, |this, _cx| { this.update_authenticated_user(response); }) })??; anyhow::Ok(()) }) .detach_and_log_err(cx); } _ => {} } } }