1use std::sync::Arc;
2use std::time::Duration;
3
4use anyhow::Context as _;
5use chrono::{DateTime, Utc};
6use cloud_api_client::{AuthenticatedUser, CloudApiClient, GetAuthenticatedUserResponse, PlanInfo};
7use cloud_llm_client::Plan;
8use gpui::{Context, Entity, Subscription, Task};
9use util::{ResultExt as _, maybe};
10
11use crate::user::Event as RpcUserStoreEvent;
12use crate::{EditPredictionUsage, ModelRequestUsage, RequestUsage, UserStore};
13
14pub struct CloudUserStore {
15 cloud_client: Arc<CloudApiClient>,
16 authenticated_user: Option<Arc<AuthenticatedUser>>,
17 plan_info: Option<Arc<PlanInfo>>,
18 model_request_usage: Option<ModelRequestUsage>,
19 edit_prediction_usage: Option<EditPredictionUsage>,
20 _maintain_authenticated_user_task: Task<()>,
21 _rpc_plan_updated_subscription: Subscription,
22}
23
24impl CloudUserStore {
25 pub fn new(
26 cloud_client: Arc<CloudApiClient>,
27 rpc_user_store: Entity<UserStore>,
28 cx: &mut Context<Self>,
29 ) -> Self {
30 let rpc_plan_updated_subscription =
31 cx.subscribe(&rpc_user_store, Self::handle_rpc_user_store_event);
32
33 Self {
34 cloud_client: cloud_client.clone(),
35 authenticated_user: None,
36 plan_info: None,
37 model_request_usage: None,
38 edit_prediction_usage: None,
39 _maintain_authenticated_user_task: cx.spawn(async move |this, cx| {
40 maybe!(async move {
41 loop {
42 let Some(this) = this.upgrade() else {
43 return anyhow::Ok(());
44 };
45
46 if cloud_client.has_credentials() {
47 let already_fetched_authenticated_user = this
48 .read_with(cx, |this, _cx| this.authenticated_user().is_some())
49 .unwrap_or(false);
50
51 if already_fetched_authenticated_user {
52 // We already fetched the authenticated user; nothing to do.
53 } else {
54 let authenticated_user_result = cloud_client
55 .get_authenticated_user()
56 .await
57 .context("failed to fetch authenticated user");
58 if let Some(response) = authenticated_user_result.log_err() {
59 this.update(cx, |this, _cx| {
60 this.update_authenticated_user(response);
61 })
62 .ok();
63 }
64 }
65 } else {
66 this.update(cx, |this, _cx| {
67 this.authenticated_user.take();
68 this.plan_info.take();
69 })
70 .ok();
71 }
72
73 cx.background_executor()
74 .timer(Duration::from_millis(100))
75 .await;
76 }
77 })
78 .await
79 .log_err();
80 }),
81 _rpc_plan_updated_subscription: rpc_plan_updated_subscription,
82 }
83 }
84
85 pub fn is_authenticated(&self) -> bool {
86 self.authenticated_user.is_some()
87 }
88
89 pub fn authenticated_user(&self) -> Option<Arc<AuthenticatedUser>> {
90 self.authenticated_user.clone()
91 }
92
93 pub fn plan(&self) -> Option<Plan> {
94 self.plan_info.as_ref().map(|plan| plan.plan)
95 }
96
97 pub fn subscription_period(&self) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
98 self.plan_info
99 .as_ref()
100 .and_then(|plan| plan.subscription_period)
101 .map(|subscription_period| {
102 (
103 subscription_period.started_at.0,
104 subscription_period.ended_at.0,
105 )
106 })
107 }
108
109 pub fn trial_started_at(&self) -> Option<DateTime<Utc>> {
110 self.plan_info
111 .as_ref()
112 .and_then(|plan| plan.trial_started_at)
113 .map(|trial_started_at| trial_started_at.0)
114 }
115
116 pub fn has_accepted_tos(&self) -> bool {
117 self.authenticated_user
118 .as_ref()
119 .map(|user| user.accepted_tos_at.is_some())
120 .unwrap_or_default()
121 }
122
123 /// Returns whether the user's account is too new to use the service.
124 pub fn account_too_young(&self) -> bool {
125 self.plan_info
126 .as_ref()
127 .map(|plan| plan.is_account_too_young)
128 .unwrap_or_default()
129 }
130
131 /// Returns whether the current user has overdue invoices and usage should be blocked.
132 pub fn has_overdue_invoices(&self) -> bool {
133 self.plan_info
134 .as_ref()
135 .map(|plan| plan.has_overdue_invoices)
136 .unwrap_or_default()
137 }
138
139 pub fn is_usage_based_billing_enabled(&self) -> bool {
140 self.plan_info
141 .as_ref()
142 .map(|plan| plan.is_usage_based_billing_enabled)
143 .unwrap_or_default()
144 }
145
146 pub fn model_request_usage(&self) -> Option<ModelRequestUsage> {
147 self.model_request_usage
148 }
149
150 pub fn update_model_request_usage(&mut self, usage: ModelRequestUsage, cx: &mut Context<Self>) {
151 self.model_request_usage = Some(usage);
152 cx.notify();
153 }
154
155 pub fn edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
156 self.edit_prediction_usage
157 }
158
159 pub fn update_edit_prediction_usage(
160 &mut self,
161 usage: EditPredictionUsage,
162 cx: &mut Context<Self>,
163 ) {
164 self.edit_prediction_usage = Some(usage);
165 cx.notify();
166 }
167
168 fn update_authenticated_user(&mut self, response: GetAuthenticatedUserResponse) {
169 self.authenticated_user = Some(Arc::new(response.user));
170 self.model_request_usage = Some(ModelRequestUsage(RequestUsage {
171 limit: response.plan.usage.model_requests.limit,
172 amount: response.plan.usage.model_requests.used as i32,
173 }));
174 self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage {
175 limit: response.plan.usage.edit_predictions.limit,
176 amount: response.plan.usage.edit_predictions.used as i32,
177 }));
178 self.plan_info = Some(Arc::new(response.plan));
179 }
180
181 fn handle_rpc_user_store_event(
182 &mut self,
183 _: Entity<UserStore>,
184 event: &RpcUserStoreEvent,
185 cx: &mut Context<Self>,
186 ) {
187 match event {
188 RpcUserStoreEvent::PlanUpdated => {
189 cx.spawn(async move |this, cx| {
190 let cloud_client =
191 cx.update(|cx| this.read_with(cx, |this, _cx| this.cloud_client.clone()))??;
192
193 let response = cloud_client
194 .get_authenticated_user()
195 .await
196 .context("failed to fetch authenticated user")?;
197
198 cx.update(|cx| {
199 this.update(cx, |this, _cx| {
200 this.update_authenticated_user(response);
201 })
202 })??;
203
204 anyhow::Ok(())
205 })
206 .detach_and_log_err(cx);
207 }
208 _ => {}
209 }
210 }
211}