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