agent: Load usage eagerly (#29937)

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <git@maxdeviant.com>

crates/agent/src/assistant_panel.rs

@@ -25,7 +25,7 @@ use gpui::{
     Pixels, Subscription, Task, UpdateGlobal, WeakEntity, prelude::*, pulsating_between,
 };
 use language::LanguageRegistry;
-use language_model::{LanguageModelProviderTosView, LanguageModelRegistry};
+use language_model::{LanguageModelProviderTosView, LanguageModelRegistry, RequestUsage};
 use language_model_selector::ToggleModelSelector;
 use project::Project;
 use prompt_store::{PromptBuilder, PromptStore, UserPromptId};
@@ -38,7 +38,7 @@ use ui::{
     Banner, ContextMenu, KeyBinding, PopoverMenu, PopoverMenuHandle, ProgressBar, Tab, Tooltip,
     prelude::*,
 };
-use util::ResultExt as _;
+use util::{ResultExt as _, maybe};
 use workspace::dock::{DockPosition, Panel, PanelEvent};
 use workspace::{CollaboratorId, ToolbarItemView, Workspace};
 use zed_actions::agent::OpenConfiguration;
@@ -1388,10 +1388,29 @@ impl AssistantPanel {
 
     fn render_toolbar(&self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         let active_thread = self.thread.read(cx);
+        let user_store = self.user_store.read(cx);
         let thread = active_thread.thread().read(cx);
         let thread_id = thread.id().clone();
         let is_empty = active_thread.is_empty();
-        let last_usage = active_thread.thread().read(cx).last_usage();
+        let last_usage = active_thread.thread().read(cx).last_usage().or_else(|| {
+            maybe!({
+                let amount = user_store.model_request_usage_amount()?;
+                let limit = user_store.model_request_usage_limit()?.variant?;
+
+                Some(RequestUsage {
+                    amount: amount as i32,
+                    limit: match limit {
+                        proto::usage_limit::Variant::Limited(limited) => {
+                            zed_llm_client::UsageLimit::Limited(limited.limit as i32)
+                        }
+                        proto::usage_limit::Variant::Unlimited(_) => {
+                            zed_llm_client::UsageLimit::Unlimited
+                        }
+                    },
+                })
+            })
+        });
+
         let account_url = zed_urls::account_url(cx);
 
         let show_token_count = match &self.active_view {

crates/agent/src/message_editor.rs

@@ -26,7 +26,7 @@ use gpui::{
     Task, TextStyle, WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
 };
 use language::{Buffer, Language};
-use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent};
+use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent, RequestUsage};
 use language_model_selector::ToggleModelSelector;
 use multi_buffer;
 use project::Project;
@@ -36,7 +36,7 @@ use settings::Settings;
 use std::time::Duration;
 use theme::ThemeSettings;
 use ui::{Disclosure, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*};
-use util::ResultExt as _;
+use util::{ResultExt as _, maybe};
 use workspace::{CollaboratorId, Workspace};
 use zed_llm_client::CompletionMode;
 
@@ -1043,9 +1043,17 @@ impl MessageEditor {
             return None;
         }
 
-        let plan = self
+        let user_store = self.user_store.read(cx);
-            .user_store
+
-            .read(cx)
+        let ubb_enable = user_store
+            .usage_based_billing_enabled()
+            .map_or(false, |enabled| enabled);
+
+        if ubb_enable {
+            return None;
+        }
+
+        let plan = user_store
             .current_plan()
             .map(|plan| match plan {
                 Plan::Free => zed_llm_client::Plan::Free,
@@ -1053,7 +1061,24 @@ impl MessageEditor {
                 Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
             })
             .unwrap_or(zed_llm_client::Plan::Free);
-        let usage = self.thread.read(cx).last_usage()?;
+        let usage = self.thread.read(cx).last_usage().or_else(|| {
+            maybe!({
+                let amount = user_store.model_request_usage_amount()?;
+                let limit = user_store.model_request_usage_limit()?.variant?;
+
+                Some(RequestUsage {
+                    amount: amount as i32,
+                    limit: match limit {
+                        proto::usage_limit::Variant::Limited(limited) => {
+                            zed_llm_client::UsageLimit::Limited(limited.limit as i32)
+                        }
+                        proto::usage_limit::Variant::Unlimited(_) => {
+                            zed_llm_client::UsageLimit::Unlimited
+                        }
+                    },
+                })
+            })
+        })?;
 
         Some(
             div()

crates/client/src/user.rs

@@ -102,6 +102,10 @@ pub struct UserStore {
     update_contacts_tx: mpsc::UnboundedSender<UpdateContacts>,
     current_plan: Option<proto::Plan>,
     trial_started_at: Option<DateTime<Utc>>,
+    model_request_usage_amount: Option<u32>,
+    model_request_usage_limit: Option<proto::UsageLimit>,
+    edit_predictions_usage_amount: Option<u32>,
+    edit_predictions_usage_limit: Option<proto::UsageLimit>,
     is_usage_based_billing_enabled: Option<bool>,
     current_user: watch::Receiver<Option<Arc<User>>>,
     accepted_tos_at: Option<Option<DateTime<Utc>>>,
@@ -163,6 +167,10 @@ impl UserStore {
             current_user: current_user_rx,
             current_plan: None,
             trial_started_at: None,
+            model_request_usage_amount: None,
+            model_request_usage_limit: None,
+            edit_predictions_usage_amount: None,
+            edit_predictions_usage_limit: None,
             is_usage_based_billing_enabled: None,
             accepted_tos_at: None,
             contacts: Default::default(),
@@ -330,6 +338,14 @@ impl UserStore {
                 .trial_started_at
                 .and_then(|trial_started_at| DateTime::from_timestamp(trial_started_at as i64, 0));
             this.is_usage_based_billing_enabled = message.payload.is_usage_based_billing_enabled;
+
+            if let Some(usage) = message.payload.usage {
+                this.model_request_usage_amount = Some(usage.model_requests_usage_amount);
+                this.model_request_usage_limit = usage.model_requests_usage_limit;
+                this.edit_predictions_usage_amount = Some(usage.edit_predictions_usage_amount);
+                this.edit_predictions_usage_limit = usage.edit_predictions_usage_limit;
+            }
+
             cx.notify();
         })?;
         Ok(())
@@ -697,6 +713,26 @@ impl UserStore {
         self.current_plan
     }
 
+    pub fn usage_based_billing_enabled(&self) -> Option<bool> {
+        self.is_usage_based_billing_enabled
+    }
+
+    pub fn model_request_usage_amount(&self) -> Option<u32> {
+        self.model_request_usage_amount
+    }
+
+    pub fn model_request_usage_limit(&self) -> Option<proto::UsageLimit> {
+        self.model_request_usage_limit.clone()
+    }
+
+    pub fn edit_predictions_usage_amount(&self) -> Option<u32> {
+        self.edit_predictions_usage_amount
+    }
+
+    pub fn edit_predictions_usage_limit(&self) -> Option<proto::UsageLimit> {
+        self.edit_predictions_usage_limit.clone()
+    }
+
     pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
         self.current_user.clone()
     }