From c2cd4fd7a1dcc6331bc1537c6de633e65204b1f9 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Thu, 17 Apr 2025 18:16:57 -0400 Subject: [PATCH] agent: Show request usage in the panel (#29006) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds a banner showing request usage in the Agent panel: Screenshot 2025-04-17 at 5 51 46 PM Only visible to users on the new billing. Note to Joseph: Doesn't need to be cherry-picked to Preview. Release Notes: - N/A --------- Co-authored-by: Nate --- crates/agent/src/active_thread.rs | 12 +- crates/agent/src/assistant_panel.rs | 8 + crates/agent/src/thread.rs | 15 +- crates/agent/src/ui.rs | 4 +- crates/agent/src/ui/usage_banner.rs | 202 ++++++++++++++++++ crates/agent/src/ui/user_spending.rs | 186 ---------------- crates/eval/src/example.rs | 3 +- .../src/components/progress/progress_bar.rs | 41 ++-- 8 files changed, 248 insertions(+), 223 deletions(-) create mode 100644 crates/agent/src/ui/usage_banner.rs delete mode 100644 crates/agent/src/ui/user_spending.rs diff --git a/crates/agent/src/active_thread.rs b/crates/agent/src/active_thread.rs index 79fb331f30..991906a236 100644 --- a/crates/agent/src/active_thread.rs +++ b/crates/agent/src/active_thread.rs @@ -23,7 +23,8 @@ use gpui::{ }; use language::{Buffer, LanguageRegistry}; use language_model::{ - LanguageModelRegistry, LanguageModelRequestMessage, LanguageModelToolUseId, Role, StopReason, + LanguageModelRegistry, LanguageModelRequestMessage, LanguageModelToolUseId, RequestUsage, Role, + StopReason, }; use markdown::parser::{CodeBlockKind, CodeBlockMetadata}; use markdown::{HeadingLevelStyles, Markdown, MarkdownElement, MarkdownStyle, ParsedMarkdown}; @@ -63,6 +64,7 @@ pub struct ActiveThread { expanded_thinking_segments: HashMap<(MessageId, usize), bool>, expanded_code_blocks: HashMap<(MessageId, usize), bool>, last_error: Option, + last_usage: Option, notifications: Vec>, copied_code_block_ids: HashSet<(MessageId, usize)>, _subscriptions: Vec, @@ -734,6 +736,7 @@ impl ActiveThread { hide_scrollbar_task: None, editing_message: None, last_error: None, + last_usage: None, copied_code_block_ids: HashSet::default(), notifications: Vec::new(), _subscriptions: subscriptions, @@ -792,6 +795,10 @@ impl ActiveThread { self.last_error.take(); } + pub fn last_usage(&self) -> Option { + self.last_usage + } + /// Returns the editing message id and the estimated token count in the content pub fn editing_message_id(&self) -> Option<(MessageId, usize)> { self.editing_message @@ -876,6 +883,9 @@ impl ActiveThread { ThreadEvent::ShowError(error) => { self.last_error = Some(error.clone()); } + ThreadEvent::UsageUpdated(usage) => { + self.last_usage = Some(*usage); + } ThreadEvent::StreamedCompletion | ThreadEvent::SummaryGenerated | ThreadEvent::SummaryChanged => { diff --git a/crates/agent/src/assistant_panel.rs b/crates/agent/src/assistant_panel.rs index 002f3ebcbe..6a72677eaa 100644 --- a/crates/agent/src/assistant_panel.rs +++ b/crates/agent/src/assistant_panel.rs @@ -45,6 +45,7 @@ use crate::message_editor::{MessageEditor, MessageEditorEvent}; use crate::thread::{Thread, ThreadError, ThreadId, TokenUsageRatio}; use crate::thread_history::{PastContext, PastThread, ThreadHistory}; use crate::thread_store::ThreadStore; +use crate::ui::UsageBanner; use crate::{ AgentDiff, ExpandMessageEditor, InlineAssistant, NewTextThread, NewThread, OpenActiveThreadAsMarkdown, OpenAgentDiff, OpenHistory, ThreadEvent, ToggleContextPicker, @@ -1541,6 +1542,12 @@ impl AssistantPanel { }) } + fn render_usage_banner(&self, cx: &mut Context) -> Option { + let usage = self.thread.read(cx).last_usage()?; + + Some(UsageBanner::new(zed_llm_client::Plan::ZedProTrial, usage.amount).into_any_element()) + } + fn render_last_error(&self, cx: &mut Context) -> Option { let last_error = self.thread.read(cx).last_error()?; @@ -1802,6 +1809,7 @@ impl Render for AssistantPanel { .map(|parent| match &self.active_view { ActiveView::Thread { .. } => parent .child(self.render_active_thread_or_empty_state(window, cx)) + .children(self.render_usage_banner(cx)) .child(h_flex().child(self.message_editor.clone())) .children(self.render_last_error(cx)), ActiveView::History => parent.child(self.history.clone()), diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index 94882f8cbd..af704fdfe3 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -19,7 +19,8 @@ use language_model::{ LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult, LanguageModelToolUseId, MaxMonthlySpendReachedError, MessageContent, - ModelRequestLimitReachedError, PaymentRequiredError, Role, StopReason, TokenUsage, + ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, StopReason, + TokenUsage, }; use project::Project; use project::git_store::{GitStore, GitStoreCheckpoint, RepositoryState}; @@ -31,7 +32,6 @@ use settings::Settings; use thiserror::Error; use util::{ResultExt as _, TryFutureExt as _, post_inc}; use uuid::Uuid; -use zed_llm_client::UsageLimit; use crate::context::{AssistantContext, ContextId, format_context_as_string}; use crate::thread_store::{ @@ -1080,11 +1080,11 @@ impl Thread { let mut current_token_usage = TokenUsage::default(); if let Some(usage) = usage { - let limit = match usage.limit { - UsageLimit::Limited(limit) => limit.to_string(), - UsageLimit::Unlimited => "unlimited".to_string(), - }; - log::info!("model request usage: {} / {}", usage.amount, limit); + thread + .update(cx, |_thread, cx| { + cx.emit(ThreadEvent::UsageUpdated(usage)); + }) + .ok(); } while let Some(event) = events.next().await { @@ -2050,6 +2050,7 @@ pub enum ThreadError { #[derive(Debug, Clone)] pub enum ThreadEvent { ShowError(ThreadError), + UsageUpdated(RequestUsage), StreamedCompletion, StreamedAssistantText(MessageId, String), StreamedAssistantThinking(MessageId, String), diff --git a/crates/agent/src/ui.rs b/crates/agent/src/ui.rs index d5b374208e..e31733ecca 100644 --- a/crates/agent/src/ui.rs +++ b/crates/agent/src/ui.rs @@ -1,7 +1,7 @@ mod agent_notification; mod context_pill; -mod user_spending; +mod usage_banner; pub use agent_notification::*; pub use context_pill::*; -// pub use user_spending::*; +pub use usage_banner::*; diff --git a/crates/agent/src/ui/usage_banner.rs b/crates/agent/src/ui/usage_banner.rs new file mode 100644 index 0000000000..a29258f76d --- /dev/null +++ b/crates/agent/src/ui/usage_banner.rs @@ -0,0 +1,202 @@ +use client::zed_urls; +use ui::{Banner, ProgressBar, Severity, prelude::*}; +use zed_llm_client::{Plan, UsageLimit}; + +#[derive(IntoElement, RegisterComponent)] +pub struct UsageBanner { + plan: Plan, + requests: i32, +} + +impl UsageBanner { + pub fn new(plan: Plan, requests: i32) -> Self { + Self { plan, requests } + } +} + +impl RenderOnce for UsageBanner { + fn render(self, _window: &mut Window, cx: &mut App) -> impl IntoElement { + let request_limit = self.plan.model_requests_limit(); + + let used_percentage = match request_limit { + UsageLimit::Limited(limit) => Some((self.requests as f32 / limit as f32) * 100.), + UsageLimit::Unlimited => None, + }; + + let (severity, message) = match request_limit { + UsageLimit::Limited(limit) => { + if self.requests >= limit { + let message = match self.plan { + Plan::ZedPro => "Monthly request limit reached", + Plan::ZedProTrial => "Trial request limit reached", + Plan::Free => "Free tier request limit reached", + }; + + (Severity::Error, message) + } else if (self.requests as f32 / limit as f32) >= 0.9 { + (Severity::Warning, "Approaching request limit") + } else { + let message = match self.plan { + Plan::ZedPro => "Zed Pro", + Plan::ZedProTrial => "Zed Pro (Trial)", + Plan::Free => "Zed Free", + }; + + (Severity::Info, message) + } + } + UsageLimit::Unlimited => { + let message = match self.plan { + Plan::ZedPro => "Zed Pro", + Plan::ZedProTrial => "Zed Pro (Trial)", + Plan::Free => "Zed Free", + }; + + (Severity::Info, message) + } + }; + + let action = match self.plan { + Plan::ZedProTrial | Plan::Free => { + Button::new("upgrade", "Upgrade").on_click(|_, _window, cx| { + cx.open_url(&zed_urls::account_url(cx)); + }) + } + Plan::ZedPro => Button::new("manage", "Manage").on_click(|_, _window, cx| { + cx.open_url(&zed_urls::account_url(cx)); + }), + }; + + Banner::new().severity(severity).children( + h_flex().flex_1().gap_1().child(Label::new(message)).child( + h_flex() + .flex_1() + .justify_end() + .gap_1p5() + .children(used_percentage.map(|percent| { + h_flex() + .items_center() + .w_full() + .max_w(px(180.)) + .child(ProgressBar::new("usage", percent, 100., cx)) + })) + .child( + Label::new(match request_limit { + UsageLimit::Limited(limit) => { + format!("{} / {limit}", self.requests) + } + UsageLimit::Unlimited => format!("{} / ∞", self.requests), + }) + .size(LabelSize::Small) + .color(Color::Muted), + ) + // Note: This should go in the banner's `action_slot`, but doing that messes with the size of the + // progress bar. + .child(action), + ), + ) + } +} + +impl Component for UsageBanner { + fn sort_name() -> &'static str { + "AgentUsageBanner" + } + + fn preview(_window: &mut Window, _cx: &mut App) -> Option { + let trial_examples = vec![ + single_example( + "Zed Pro Trial - New User", + div() + .size_full() + .child(UsageBanner::new(Plan::ZedProTrial, 10)) + .into_any_element(), + ), + single_example( + "Zed Pro Trial - Approaching Limit", + div() + .size_full() + .child(UsageBanner::new(Plan::ZedProTrial, 135)) + .into_any_element(), + ), + single_example( + "Zed Pro Trial - Request Limit Reached", + div() + .size_full() + .child(UsageBanner::new(Plan::ZedProTrial, 150)) + .into_any_element(), + ), + ]; + + let free_examples = vec![ + single_example( + "Free - Normal Usage", + div() + .size_full() + .child(UsageBanner::new(Plan::Free, 25)) + .into_any_element(), + ), + single_example( + "Free - Approaching Limit", + div() + .size_full() + .child(UsageBanner::new(Plan::Free, 45)) + .into_any_element(), + ), + single_example( + "Free - Request Limit Reached", + div() + .size_full() + .child(UsageBanner::new(Plan::Free, 50)) + .into_any_element(), + ), + ]; + + let zed_pro_examples = vec![ + single_example( + "Zed Pro - Normal Usage", + div() + .size_full() + .child(UsageBanner::new(Plan::ZedPro, 250)) + .into_any_element(), + ), + single_example( + "Zed Pro - Approaching Limit", + div() + .size_full() + .child(UsageBanner::new(Plan::ZedPro, 450)) + .into_any_element(), + ), + single_example( + "Zed Pro - Request Limit Reached", + div() + .size_full() + .child(UsageBanner::new(Plan::ZedPro, 500)) + .into_any_element(), + ), + ]; + + Some( + v_flex() + .gap_6() + .p_4() + .children(vec![ + Label::new("Trial Plan") + .size(LabelSize::Large) + .into_any_element(), + example_group(trial_examples).vertical().into_any_element(), + Label::new("Free Plan") + .size(LabelSize::Large) + .into_any_element(), + example_group(free_examples).vertical().into_any_element(), + Label::new("Pro Plan") + .size(LabelSize::Large) + .into_any_element(), + example_group(zed_pro_examples) + .vertical() + .into_any_element(), + ]) + .into_any_element(), + ) + } +} diff --git a/crates/agent/src/ui/user_spending.rs b/crates/agent/src/ui/user_spending.rs deleted file mode 100644 index 59bcbc9671..0000000000 --- a/crates/agent/src/ui/user_spending.rs +++ /dev/null @@ -1,186 +0,0 @@ -use gpui::{Entity, Render}; -use ui::{ProgressBar, prelude::*}; - -#[derive(RegisterComponent)] -pub struct UserSpending { - free_tier_current: u32, - free_tier_cap: u32, - over_tier_current: u32, - over_tier_cap: u32, - free_tier_progress: Entity, - over_tier_progress: Entity, -} - -impl UserSpending { - pub fn new( - free_tier_current: u32, - free_tier_cap: u32, - over_tier_current: u32, - over_tier_cap: u32, - cx: &mut App, - ) -> Self { - let free_tier_capped = free_tier_current == free_tier_cap; - let free_tier_near_capped = - free_tier_current as f32 / 100.0 >= free_tier_cap as f32 / 100.0 * 0.9; - let over_tier_capped = over_tier_current == over_tier_cap; - let over_tier_near_capped = - over_tier_current as f32 / 100.0 >= over_tier_cap as f32 / 100.0 * 0.9; - - let free_tier_progress = cx.new(|cx| { - ProgressBar::new( - "free_tier", - free_tier_current as f32, - free_tier_cap as f32, - cx, - ) - }); - let over_tier_progress = cx.new(|cx| { - ProgressBar::new( - "over_tier", - over_tier_current as f32, - over_tier_cap as f32, - cx, - ) - }); - - if free_tier_capped { - free_tier_progress.update(cx, |progress_bar, cx| { - progress_bar.fg_color(cx.theme().status().error); - }); - } else if free_tier_near_capped { - free_tier_progress.update(cx, |progress_bar, cx| { - progress_bar.fg_color(cx.theme().status().warning); - }); - } - - if over_tier_capped { - over_tier_progress.update(cx, |progress_bar, cx| { - progress_bar.fg_color(cx.theme().status().error); - }); - } else if over_tier_near_capped { - over_tier_progress.update(cx, |progress_bar, cx| { - progress_bar.fg_color(cx.theme().status().warning); - }); - } - - Self { - free_tier_current, - free_tier_cap, - over_tier_current, - over_tier_cap, - free_tier_progress, - over_tier_progress, - } - } -} - -impl Render for UserSpending { - fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { - let formatted_free_tier = format!( - "${} / ${}", - self.free_tier_current as f32 / 100.0, - self.free_tier_cap as f32 / 100.0 - ); - let formatted_over_tier = format!( - "${} / ${}", - self.over_tier_current as f32 / 100.0, - self.over_tier_cap as f32 / 100.0 - ); - - v_group() - .elevation_2(cx) - .py_1p5() - .px_2p5() - .w(px(360.)) - .child( - v_flex() - .child( - v_flex() - .p_1p5() - .gap_0p5() - .child( - h_flex() - .justify_between() - .child(Label::new("Free Tier Usage").size(LabelSize::Small)) - .child( - Label::new(formatted_free_tier) - .size(LabelSize::Small) - .color(Color::Muted), - ), - ) - .child(self.free_tier_progress.clone()), - ) - .child( - v_flex() - .p_1p5() - .gap_0p5() - .child( - h_flex() - .justify_between() - .child(Label::new("Current Spending").size(LabelSize::Small)) - .child( - Label::new(formatted_over_tier) - .size(LabelSize::Small) - .color(Color::Muted), - ), - ) - .child(self.over_tier_progress.clone()), - ), - ) - } -} - -impl Component for UserSpending { - fn scope() -> ComponentScope { - ComponentScope::None - } - - fn preview(_window: &mut Window, cx: &mut App) -> Option { - let new_user = cx.new(|cx| UserSpending::new(0, 2000, 0, 2000, cx)); - let free_capped = cx.new(|cx| UserSpending::new(2000, 2000, 0, 2000, cx)); - let free_near_capped = cx.new(|cx| UserSpending::new(1800, 2000, 0, 2000, cx)); - let over_near_capped = cx.new(|cx| UserSpending::new(2000, 2000, 1800, 2000, cx)); - let over_capped = cx.new(|cx| UserSpending::new(1000, 2000, 2000, 2000, cx)); - - Some( - v_flex() - .gap_6() - .p_4() - .children(vec![example_group(vec![ - single_example( - "New User", - div().size_full().child(new_user.clone()).into_any_element(), - ), - single_example( - "Free Tier Capped", - div() - .size_full() - .child(free_capped.clone()) - .into_any_element(), - ), - single_example( - "Free Tier Near Capped", - div() - .size_full() - .child(free_near_capped.clone()) - .into_any_element(), - ), - single_example( - "Over Tier Near Capped", - div() - .size_full() - .child(over_near_capped.clone()) - .into_any_element(), - ), - single_example( - "Over Tier Capped", - div() - .size_full() - .child(over_capped.clone()) - .into_any_element(), - ), - ])]) - .into_any_element(), - ) - } -} diff --git a/crates/eval/src/example.rs b/crates/eval/src/example.rs index 0968b62423..cbd8ee7a8d 100644 --- a/crates/eval/src/example.rs +++ b/crates/eval/src/example.rs @@ -421,7 +421,8 @@ impl Example { ThreadEvent::MessageDeleted(_) | ThreadEvent::SummaryChanged | ThreadEvent::SummaryGenerated | - ThreadEvent::CheckpointChanged => { + ThreadEvent::CheckpointChanged | + ThreadEvent::UsageUpdated(_) => { if std::env::var("ZED_EVAL_DEBUG").is_ok() { println!("{}Event: {:#?}", log_prefix, event); } diff --git a/crates/ui/src/components/progress/progress_bar.rs b/crates/ui/src/components/progress/progress_bar.rs index ec9ee2d08f..a151222277 100644 --- a/crates/ui/src/components/progress/progress_bar.rs +++ b/crates/ui/src/components/progress/progress_bar.rs @@ -7,7 +7,7 @@ use crate::prelude::*; /// A progress bar is a horizontal bar that communicates the status of a process. /// /// A progress bar should not be used to represent indeterminate progress. -#[derive(RegisterComponent, Documented)] +#[derive(IntoElement, RegisterComponent, Documented)] pub struct ProgressBar { id: ElementId, value: f32, @@ -17,13 +17,7 @@ pub struct ProgressBar { } impl ProgressBar { - /// Create a new progress bar with the given value and maximum value. - pub fn new( - id: impl Into, - value: f32, - max_value: f32, - cx: &mut Context, - ) -> Self { + pub fn new(id: impl Into, value: f32, max_value: f32, cx: &App) -> Self { Self { id: id.into(), value, @@ -33,33 +27,33 @@ impl ProgressBar { } } - /// Set the current value of the progress bar. - pub fn value(&mut self, value: f32) -> &mut Self { + /// Sets the current value of the progress bar. + pub fn value(mut self, value: f32) -> Self { self.value = value; self } - /// Set the maximum value of the progress bar. - pub fn max_value(&mut self, max_value: f32) -> &mut Self { + /// Sets the maximum value of the progress bar. + pub fn max_value(mut self, max_value: f32) -> Self { self.max_value = max_value; self } - /// Set the background color of the progress bar. - pub fn bg_color(&mut self, color: Hsla) -> &mut Self { + /// Sets the background color of the progress bar. + pub fn bg_color(mut self, color: Hsla) -> Self { self.bg_color = color; self } - /// Set the foreground color of the progress bar. - pub fn fg_color(&mut self, color: Hsla) -> &mut Self { + /// Sets the foreground color of the progress bar. + pub fn fg_color(mut self, color: Hsla) -> Self { self.fg_color = color; self } } -impl Render for ProgressBar { - fn render(&mut self, _window: &mut Window, _cx: &mut Context) -> impl IntoElement { +impl RenderOnce for ProgressBar { + fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement { let fill_width = (self.value / self.max_value).clamp(0.02, 1.0); div() @@ -98,11 +92,6 @@ impl Component for ProgressBar { fn preview(_window: &mut Window, cx: &mut App) -> Option { let max_value = 180.0; - let empty_progress_bar = cx.new(|cx| ProgressBar::new("empty", 0.0, max_value, cx)); - let partial_progress_bar = - cx.new(|cx| ProgressBar::new("partial", max_value * 0.35, max_value, cx)); - let filled_progress_bar = cx.new(|cx| ProgressBar::new("filled", max_value, max_value, cx)); - Some( div() .flex() @@ -123,7 +112,7 @@ impl Component for ProgressBar { .child(Label::new("0%")) .child(Label::new("Empty")), ) - .child(empty_progress_bar.clone()), + .child(ProgressBar::new("empty", 0.0, max_value, cx)), ) .child( div() @@ -137,7 +126,7 @@ impl Component for ProgressBar { .child(Label::new("38%")) .child(Label::new("Partial")), ) - .child(partial_progress_bar.clone()), + .child(ProgressBar::new("partial", max_value * 0.35, max_value, cx)), ) .child( div() @@ -151,7 +140,7 @@ impl Component for ProgressBar { .child(Label::new("100%")) .child(Label::new("Complete")), ) - .child(filled_progress_bar.clone()), + .child(ProgressBar::new("filled", max_value, max_value, cx)), ) .into_any_element(), )