agent: Show the usage limits returned from the API (#29236)

Marshall Bowers created 8 months ago

This PR updates the usage banners in the Agent panel to use the limits
returned from the API instead of basing it off the plan.

This will allow limits to be updated from the server rather than being
embedded in the client.

Release Notes:

- N/A

Change summary

crates/agent/src/assistant_panel.rs |   2 
crates/agent/src/ui/usage_banner.rs | 100 ++++++++++++++++++++++++------
2 files changed, 79 insertions(+), 23 deletions(-)

Detailed changes

crates/agent/src/assistant_panel.rs

@@ -1550,7 +1550,7 @@ 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())
+        Some(UsageBanner::new(zed_llm_client::Plan::ZedProTrial, usage).into_any_element())
     }
 
     fn render_last_error(&self, cx: &mut Context<Self>) -> Option<AnyElement> {

crates/agent/src/ui/usage_banner.rs

@@ -1,31 +1,30 @@
 use client::zed_urls;
+use language_model::RequestUsage;
 use ui::{Banner, ProgressBar, Severity, prelude::*};
 use zed_llm_client::{Plan, UsageLimit};
 
 #[derive(IntoElement, RegisterComponent)]
 pub struct UsageBanner {
     plan: Plan,
-    requests: i32,
+    usage: RequestUsage,
 }
 
 impl UsageBanner {
-    pub fn new(plan: Plan, requests: i32) -> Self {
-        Self { plan, requests }
+    pub fn new(plan: Plan, usage: RequestUsage) -> Self {
+        Self { plan, usage }
     }
 }
 
 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.),
+        let used_percentage = match self.usage.limit {
+            UsageLimit::Limited(limit) => Some((self.usage.amount as f32 / limit as f32) * 100.),
             UsageLimit::Unlimited => None,
         };
 
-        let (severity, message) = match request_limit {
+        let (severity, message) = match self.usage.limit {
             UsageLimit::Limited(limit) => {
-                if self.requests >= limit {
+                if self.usage.amount >= limit {
                     let message = match self.plan {
                         Plan::ZedPro => "Monthly request limit reached",
                         Plan::ZedProTrial => "Trial request limit reached",
@@ -33,7 +32,7 @@ impl RenderOnce for UsageBanner {
                     };
 
                     (Severity::Error, message)
-                } else if (self.requests as f32 / limit as f32) >= 0.9 {
+                } else if (self.usage.amount as f32 / limit as f32) >= 0.9 {
                     (Severity::Warning, "Approaching request limit")
                 } else {
                     let message = match self.plan {
@@ -81,11 +80,11 @@ impl RenderOnce for UsageBanner {
                             .child(ProgressBar::new("usage", percent, 100., cx))
                     }))
                     .child(
-                        Label::new(match request_limit {
+                        Label::new(match self.usage.limit {
                             UsageLimit::Limited(limit) => {
-                                format!("{} / {limit}", self.requests)
+                                format!("{} / {limit}", self.usage.amount)
                             }
-                            UsageLimit::Unlimited => format!("{} / ∞", self.requests),
+                            UsageLimit::Unlimited => format!("{} / ∞", self.usage.amount),
                         })
                         .size(LabelSize::Small)
                         .color(Color::Muted),
@@ -104,74 +103,131 @@ impl Component for UsageBanner {
     }
 
     fn preview(_window: &mut Window, _cx: &mut App) -> Option<AnyElement> {
+        let trial_limit = Plan::ZedProTrial.model_requests_limit();
         let trial_examples = vec![
             single_example(
                 "Zed Pro Trial - New User",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedProTrial, 10))
+                    .child(UsageBanner::new(
+                        Plan::ZedProTrial,
+                        RequestUsage {
+                            limit: trial_limit,
+                            amount: 10,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro Trial - Approaching Limit",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedProTrial, 135))
+                    .child(UsageBanner::new(
+                        Plan::ZedProTrial,
+                        RequestUsage {
+                            limit: trial_limit,
+                            amount: 135,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro Trial - Request Limit Reached",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedProTrial, 150))
+                    .child(UsageBanner::new(
+                        Plan::ZedProTrial,
+                        RequestUsage {
+                            limit: trial_limit,
+                            amount: 150,
+                        },
+                    ))
                     .into_any_element(),
             ),
         ];
 
+        let free_limit = Plan::Free.model_requests_limit();
         let free_examples = vec![
             single_example(
                 "Free - Normal Usage",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::Free, 25))
+                    .child(UsageBanner::new(
+                        Plan::Free,
+                        RequestUsage {
+                            limit: free_limit,
+                            amount: 25,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Free - Approaching Limit",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::Free, 45))
+                    .child(UsageBanner::new(
+                        Plan::Free,
+                        RequestUsage {
+                            limit: free_limit,
+                            amount: 45,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Free - Request Limit Reached",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::Free, 50))
+                    .child(UsageBanner::new(
+                        Plan::Free,
+                        RequestUsage {
+                            limit: free_limit,
+                            amount: 50,
+                        },
+                    ))
                     .into_any_element(),
             ),
         ];
 
+        let zed_pro_limit = Plan::ZedPro.model_requests_limit();
         let zed_pro_examples = vec![
             single_example(
                 "Zed Pro - Normal Usage",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedPro, 250))
+                    .child(UsageBanner::new(
+                        Plan::ZedPro,
+                        RequestUsage {
+                            limit: zed_pro_limit,
+                            amount: 250,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro - Approaching Limit",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedPro, 450))
+                    .child(UsageBanner::new(
+                        Plan::ZedPro,
+                        RequestUsage {
+                            limit: zed_pro_limit,
+                            amount: 450,
+                        },
+                    ))
                     .into_any_element(),
             ),
             single_example(
                 "Zed Pro - Request Limit Reached",
                 div()
                     .size_full()
-                    .child(UsageBanner::new(Plan::ZedPro, 500))
+                    .child(UsageBanner::new(
+                        Plan::ZedPro,
+                        RequestUsage {
+                            limit: zed_pro_limit,
+                            amount: 500,
+                        },
+                    ))
                     .into_any_element(),
             ),
         ];