Update Agent panel to work with CloudUserStore (#35436)

This PR updates the Agent panel to work with the `CloudUserStore`
instead of the `UserStore`, reducing its reliance on being connected to
Collab to function.

Release Notes:

- N/A

---------

Co-authored-by: Richard Feldman <oss@rtfeldman.com>

crates/client/src/cloud/user_store.rs

@@ -9,12 +9,13 @@ use gpui::{Context, Entity, Subscription, Task};
 use util::{ResultExt as _, maybe};
 
 use crate::user::Event as RpcUserStoreEvent;
-use crate::{EditPredictionUsage, RequestUsage, UserStore};
+use crate::{EditPredictionUsage, ModelRequestUsage, RequestUsage, UserStore};
 
 pub struct CloudUserStore {
     cloud_client: Arc<CloudApiClient>,
     authenticated_user: Option<Arc<AuthenticatedUser>>,
     plan_info: Option<Arc<PlanInfo>>,
+    model_request_usage: Option<ModelRequestUsage>,
     edit_prediction_usage: Option<EditPredictionUsage>,
     _maintain_authenticated_user_task: Task<()>,
     _rpc_plan_updated_subscription: Subscription,
@@ -33,6 +34,7 @@ impl CloudUserStore {
             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 {
@@ -104,6 +106,13 @@ impl CloudUserStore {
             })
     }
 
+    pub fn trial_started_at(&self) -> Option<DateTime<Utc>> {
+        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()
@@ -127,6 +136,22 @@ impl CloudUserStore {
             .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<ModelRequestUsage> {
+        self.model_request_usage
+    }
+
+    pub fn update_model_request_usage(&mut self, usage: ModelRequestUsage, cx: &mut Context<Self>) {
+        self.model_request_usage = Some(usage);
+        cx.notify();
+    }
+
     pub fn edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
         self.edit_prediction_usage
     }
@@ -142,6 +167,10 @@ impl CloudUserStore {
 
     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,

crates/client/src/user.rs

@@ -113,7 +113,6 @@ pub struct UserStore {
     current_plan: Option<proto::Plan>,
     subscription_period: Option<(DateTime<Utc>, DateTime<Utc>)>,
     trial_started_at: Option<DateTime<Utc>>,
-    model_request_usage: Option<ModelRequestUsage>,
     is_usage_based_billing_enabled: Option<bool>,
     account_too_young: Option<bool>,
     has_overdue_invoices: Option<bool>,
@@ -191,7 +190,6 @@ impl UserStore {
             current_plan: None,
             subscription_period: None,
             trial_started_at: None,
-            model_request_usage: None,
             is_usage_based_billing_enabled: None,
             account_too_young: None,
             has_overdue_invoices: None,
@@ -371,27 +369,12 @@ impl UserStore {
             this.account_too_young = message.payload.account_too_young;
             this.has_overdue_invoices = message.payload.has_overdue_invoices;
 
-            if let Some(usage) = message.payload.usage {
-                // limits are always present even though they are wrapped in Option
-                this.model_request_usage = usage
-                    .model_requests_usage_limit
-                    .and_then(|limit| {
-                        RequestUsage::from_proto(usage.model_requests_usage_amount, limit)
-                    })
-                    .map(ModelRequestUsage);
-            }
-
             cx.emit(Event::PlanUpdated);
             cx.notify();
         })?;
         Ok(())
     }
 
-    pub fn update_model_request_usage(&mut self, usage: ModelRequestUsage, cx: &mut Context<Self>) {
-        self.model_request_usage = Some(usage);
-        cx.notify();
-    }
-
     fn update_contacts(&mut self, message: UpdateContacts, cx: &Context<Self>) -> Task<Result<()>> {
         match message {
             UpdateContacts::Wait(barrier) => {
@@ -776,10 +759,6 @@ impl UserStore {
         self.is_usage_based_billing_enabled
     }
 
-    pub fn model_request_usage(&self) -> Option<ModelRequestUsage> {
-        self.model_request_usage
-    }
-
     pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
         self.current_user.clone()
     }