Acquire LLM token from Cloud instead of Collab for Edit Predictions (#35431)

This PR updates the Zed Edit Prediction provider to acquire the LLM
token from Cloud instead of Collab to allow using Edit Predictions even
when disconnected from or unable to connect to the Collab server.

Release Notes:

- N/A

---------

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

crates/client/src/cloud/user_store.rs

@@ -8,13 +8,14 @@ use cloud_llm_client::Plan;
 use gpui::{Context, Entity, Subscription, Task};
 use util::{ResultExt as _, maybe};
 
-use crate::UserStore;
 use crate::user::Event as RpcUserStoreEvent;
+use crate::{EditPredictionUsage, RequestUsage, UserStore};
 
 pub struct CloudUserStore {
     cloud_client: Arc<CloudApiClient>,
     authenticated_user: Option<Arc<AuthenticatedUser>>,
     plan_info: Option<Arc<PlanInfo>>,
+    edit_prediction_usage: Option<EditPredictionUsage>,
     _maintain_authenticated_user_task: Task<()>,
     _rpc_plan_updated_subscription: Subscription,
 }
@@ -32,6 +33,7 @@ impl CloudUserStore {
             cloud_client: cloud_client.clone(),
             authenticated_user: None,
             plan_info: None,
+            edit_prediction_usage: None,
             _maintain_authenticated_user_task: cx.spawn(async move |this, cx| {
                 maybe!(async move {
                     loop {
@@ -102,8 +104,48 @@ impl CloudUserStore {
             })
     }
 
+    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 edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
+        self.edit_prediction_usage
+    }
+
+    pub fn update_edit_prediction_usage(
+        &mut self,
+        usage: EditPredictionUsage,
+        cx: &mut Context<Self>,
+    ) {
+        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.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));
     }
 

crates/client/src/user.rs

@@ -114,7 +114,6 @@ pub struct UserStore {
     subscription_period: Option<(DateTime<Utc>, DateTime<Utc>)>,
     trial_started_at: Option<DateTime<Utc>>,
     model_request_usage: Option<ModelRequestUsage>,
-    edit_prediction_usage: Option<EditPredictionUsage>,
     is_usage_based_billing_enabled: Option<bool>,
     account_too_young: Option<bool>,
     has_overdue_invoices: Option<bool>,
@@ -193,7 +192,6 @@ impl UserStore {
             subscription_period: None,
             trial_started_at: None,
             model_request_usage: None,
-            edit_prediction_usage: None,
             is_usage_based_billing_enabled: None,
             account_too_young: None,
             has_overdue_invoices: None,
@@ -381,12 +379,6 @@ impl UserStore {
                         RequestUsage::from_proto(usage.model_requests_usage_amount, limit)
                     })
                     .map(ModelRequestUsage);
-                this.edit_prediction_usage = usage
-                    .edit_predictions_usage_limit
-                    .and_then(|limit| {
-                        RequestUsage::from_proto(usage.model_requests_usage_amount, limit)
-                    })
-                    .map(EditPredictionUsage);
             }
 
             cx.emit(Event::PlanUpdated);
@@ -400,15 +392,6 @@ impl UserStore {
         cx.notify();
     }
 
-    pub fn update_edit_prediction_usage(
-        &mut self,
-        usage: EditPredictionUsage,
-        cx: &mut Context<Self>,
-    ) {
-        self.edit_prediction_usage = Some(usage);
-        cx.notify();
-    }
-
     fn update_contacts(&mut self, message: UpdateContacts, cx: &Context<Self>) -> Task<Result<()>> {
         match message {
             UpdateContacts::Wait(barrier) => {
@@ -797,10 +780,6 @@ impl UserStore {
         self.model_request_usage
     }
 
-    pub fn edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
-        self.edit_prediction_usage
-    }
-
     pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
         self.current_user.clone()
     }