From e64f5ff358e5a6ba406b2a9c740eb0dcee5c62e9 Mon Sep 17 00:00:00 2001 From: Nate Butler Date: Mon, 5 May 2025 14:42:21 -0400 Subject: [PATCH] agent: Load usage eagerly (#29937) Release Notes: - N/A --------- Co-authored-by: Marshall Bowers --- crates/agent/src/assistant_panel.rs | 25 ++++++++++++++++--- crates/agent/src/message_editor.rs | 37 ++++++++++++++++++++++++----- crates/client/src/user.rs | 36 ++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 9 deletions(-) diff --git a/crates/agent/src/assistant_panel.rs b/crates/agent/src/assistant_panel.rs index 64484b14344031cb07d590f263e3ef828926f359..74d7ea3197bc52ce6b348558b95e82a9e51d8e07 100644 --- a/crates/agent/src/assistant_panel.rs +++ b/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) -> 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 { diff --git a/crates/agent/src/message_editor.rs b/crates/agent/src/message_editor.rs index 5c88769f9922de82c986fd7262e3fb420b25015a..52326e3cd1328cd028f1a449aed3bf43df00fce6 100644 --- a/crates/agent/src/message_editor.rs +++ b/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 - .user_store - .read(cx) + let user_store = self.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() diff --git a/crates/client/src/user.rs b/crates/client/src/user.rs index 55d4e9b214e0a852fe97bdee7e74588da5a0a74b..75a15a3e193d5027c71d97f438de302ab8c8cd9c 100644 --- a/crates/client/src/user.rs +++ b/crates/client/src/user.rs @@ -102,6 +102,10 @@ pub struct UserStore { update_contacts_tx: mpsc::UnboundedSender, current_plan: Option, trial_started_at: Option>, + model_request_usage_amount: Option, + model_request_usage_limit: Option, + edit_predictions_usage_amount: Option, + edit_predictions_usage_limit: Option, is_usage_based_billing_enabled: Option, current_user: watch::Receiver>>, accepted_tos_at: Option>>, @@ -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 { + self.is_usage_based_billing_enabled + } + + pub fn model_request_usage_amount(&self) -> Option { + self.model_request_usage_amount + } + + pub fn model_request_usage_limit(&self) -> Option { + self.model_request_usage_limit.clone() + } + + pub fn edit_predictions_usage_amount(&self) -> Option { + self.edit_predictions_usage_amount + } + + pub fn edit_predictions_usage_limit(&self) -> Option { + self.edit_predictions_usage_limit.clone() + } + pub fn watch_current_user(&self) -> watch::Receiver>> { self.current_user.clone() }