Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent: Show request usage in the panel (#29006)

Browse source 
This PR adds a banner showing request usage in the Agent panel:

<img width="640" alt="Screenshot 2025-04-17 at 5 51 46 PM"
src="https://github.com/user-attachments/assets/e0eb036c-57c1-441c-bbab-7dab1c6e56d9"
/>

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 <nate@zed.dev>
This commit is contained in:
Marshall Bowers 2025-04-17 18:16:57 -04:00 • committed by GitHub
parent 4095011af5
commit c2cd4fd7a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 248 additions and 223 deletions
Download patch file Download diff file Expand all files Collapse all files

12
crates/agent/src/active_thread.rs
View file

@ -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<ThreadError>,
last_usage: Option<RequestUsage>,
notifications: Vec<WindowHandle<AgentNotification>>,
copied_code_block_ids: HashSet<(MessageId, usize)>,
_subscriptions: Vec<Subscription>,
@ -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<RequestUsage> {
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 => {

8
crates/agent/src/assistant_panel.rs
View file

@ -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<Self>) -> Option<AnyElement> {
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<Self>) -> Option<AnyElement> {
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()),

15
crates/agent/src/thread.rs
View file

@ -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),

4
crates/agent/src/ui.rs
View file

@ -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::*;

202
crates/agent/src/ui/usage_banner.rs Normal file
View file

@ -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<AnyElement> {
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(),
)
}
}

186
crates/agent/src/ui/user_spending.rs
View file

@ -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<ProgressBar>,
over_tier_progress: Entity<ProgressBar>,
}
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<Self>) -> 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<AnyElement> {
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(),
)
}
}

3
crates/eval/src/example.rs
View file

@ -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);
}

41
crates/ui/src/components/progress/progress_bar.rs
View file

@ -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<ElementId>,
value: f32,
max_value: f32,
cx: &mut Context<Self>,
) -> Self {
pub fn new(id: impl Into<ElementId>, 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<Self>) -> 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<AnyElement> {
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(),
)