Show edit predictions usage in status bar menu (#29046)

This PR adds an indicator for edit predictions usage in the edit
predictions menu:

| Free | Zed Pro / Trial |
|
---------------------------------------------------------------------------------------------------------------------------------------------------
|
---------------------------------------------------------------------------------------------------------------------------------------------------
|
| <img width="235" alt="Screenshot 2025-04-18 at 9 53 47 AM"
src="https://github.com/user-attachments/assets/6da001d2-ef9c-49df-86be-03d4c615d45c"
/> | <img width="237" alt="Screenshot 2025-04-18 at 9 54 33 AM"
src="https://github.com/user-attachments/assets/31f5df04-a8e1-43ec-8af7-ebe501516abe"
/> |

Only visible to users on the new billing.

Release Notes:

- N/A

Cargo.lock

@@ -7107,10 +7107,12 @@ dependencies = [
 name = "inline_completion"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "gpui",
  "language",
  "project",
  "workspace-hack",
+ "zed_llm_client",
 ]
 
 [[package]]
@@ -7141,6 +7143,7 @@ dependencies = [
  "workspace",
  "workspace-hack",
  "zed_actions",
+ "zed_llm_client",
  "zeta",
 ]
 

crates/inline_completion/Cargo.toml

@@ -12,7 +12,9 @@ workspace = true
 path = "src/inline_completion.rs"
 
 [dependencies]
+anyhow.workspace = true
 gpui.workspace = true
 language.workspace = true
 project.workspace = true
 workspace-hack.workspace = true
+zed_llm_client.workspace = true

crates/inline_completion/src/inline_completion.rs

@@ -1,7 +1,14 @@
+use std::ops::Range;
+use std::str::FromStr as _;
+
+use anyhow::{Result, anyhow};
+use gpui::http_client::http::{HeaderMap, HeaderValue};
 use gpui::{App, Context, Entity, SharedString};
 use language::Buffer;
 use project::Project;
-use std::ops::Range;
+use zed_llm_client::{
+    EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
+};
 
 // TODO: Find a better home for `Direction`.
 //
@@ -52,6 +59,32 @@ impl DataCollectionState {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct EditPredictionUsage {
+    pub limit: UsageLimit,
+    pub amount: i32,
+}
+
+impl EditPredictionUsage {
+    pub fn from_headers(headers: &HeaderMap<HeaderValue>) -> Result<Self> {
+        let limit = headers
+            .get(EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME)
+            .ok_or_else(|| {
+                anyhow!("missing {EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME:?} header")
+            })?;
+        let limit = UsageLimit::from_str(limit.to_str()?)?;
+
+        let amount = headers
+            .get(EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME)
+            .ok_or_else(|| {
+                anyhow!("missing {EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME:?} header")
+            })?;
+        let amount = amount.to_str()?.parse::<i32>()?;
+
+        Ok(Self { limit, amount })
+    }
+}
+
 pub trait EditPredictionProvider: 'static + Sized {
     fn name() -> &'static str;
     fn display_name() -> &'static str;
@@ -62,6 +95,11 @@ pub trait EditPredictionProvider: 'static + Sized {
     fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
         DataCollectionState::Unsupported
     }
+
+    fn usage(&self, _cx: &App) -> Option<EditPredictionUsage> {
+        None
+    }
+
     fn toggle_data_collection(&mut self, _cx: &mut App) {}
     fn is_enabled(
         &self,
@@ -110,6 +148,7 @@ pub trait InlineCompletionProviderHandle {
     fn show_completions_in_menu(&self) -> bool;
     fn show_tab_accept_marker(&self) -> bool;
     fn data_collection_state(&self, cx: &App) -> DataCollectionState;
+    fn usage(&self, cx: &App) -> Option<EditPredictionUsage>;
     fn toggle_data_collection(&self, cx: &mut App);
     fn needs_terms_acceptance(&self, cx: &App) -> bool;
     fn is_refreshing(&self, cx: &App) -> bool;
@@ -162,6 +201,10 @@ where
         self.read(cx).data_collection_state(cx)
     }
 
+    fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
+        self.read(cx).usage(cx)
+    }
+
     fn toggle_data_collection(&self, cx: &mut App) {
         self.update(cx, |this, cx| this.toggle_data_collection(cx))
     }

crates/inline_completion_button/Cargo.toml

@@ -29,10 +29,11 @@ settings.workspace = true
 supermaven.workspace = true
 telemetry.workspace = true
 ui.workspace = true
+workspace-hack.workspace = true
 workspace.workspace = true
 zed_actions.workspace = true
+zed_llm_client.workspace = true
 zeta.workspace = true
-workspace-hack.workspace = true
 
 [dev-dependencies]
 copilot = { workspace = true, features = ["test-support"] }

crates/inline_completion_button/src/inline_completion_button.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use client::UserStore;
+use client::{UserStore, zed_urls};
 use copilot::{Copilot, Status};
 use editor::{
     Editor,
@@ -27,13 +27,14 @@ use std::{
 use supermaven::{AccountStatus, Supermaven};
 use ui::{
     Clickable, ContextMenu, ContextMenuEntry, IconButton, IconButtonShape, Indicator, PopoverMenu,
-    PopoverMenuHandle, Tooltip, prelude::*,
+    PopoverMenuHandle, ProgressBar, Tooltip, prelude::*,
 };
 use workspace::{
     StatusItemView, Toast, Workspace, create_and_open_local_file, item::ItemHandle,
     notifications::NotificationId,
 };
 use zed_actions::OpenBrowser;
+use zed_llm_client::{Plan, UsageLimit};
 use zeta::RateCompletions;
 
 actions!(edit_prediction, [ToggleMenu]);
@@ -402,6 +403,45 @@ impl InlineCompletionButton {
         let fs = self.fs.clone();
         let line_height = window.line_height();
 
+        if let Some(provider) = self.edit_prediction_provider.as_ref() {
+            if let Some(usage) = provider.usage(cx) {
+                menu = menu.header("Usage");
+                menu = menu.custom_entry(
+                    move |_window, cx| {
+                        let plan = Plan::ZedProTrial;
+                        let edit_predictions_limit = plan.edit_predictions_limit();
+
+                        let used_percentage = match edit_predictions_limit {
+                            UsageLimit::Limited(limit) => {
+                                Some((usage.amount as f32 / limit as f32) * 100.)
+                            }
+                            UsageLimit::Unlimited => None,
+                        };
+
+                        h_flex()
+                            .flex_1()
+                            .gap_1p5()
+                            .children(
+                                used_percentage
+                                    .map(|percent| ProgressBar::new("usage", percent, 100., cx)),
+                            )
+                            .child(
+                                Label::new(match edit_predictions_limit {
+                                    UsageLimit::Limited(limit) => {
+                                        format!("{} / {limit}", usage.amount)
+                                    }
+                                    UsageLimit::Unlimited => format!("{} / ∞", usage.amount),
+                                })
+                                .size(LabelSize::Small)
+                                .color(Color::Muted),
+                            )
+                            .into_any_element()
+                    },
+                    move |_, cx| cx.open_url(&zed_urls::account_url(cx)),
+                );
+            }
+        }
+
         menu = menu.header("Show Edit Predictions For");
 
         let language_state = self.language.as_ref().map(|language| {

crates/zeta/src/zeta.rs

@@ -8,9 +8,8 @@ mod rate_completion_modal;
 
 pub(crate) use completion_diff_element::*;
 use db::kvp::KEY_VALUE_STORE;
-use http_client::http::{HeaderMap, HeaderValue};
 pub use init::*;
-use inline_completion::DataCollectionState;
+use inline_completion::{DataCollectionState, EditPredictionUsage};
 use license_detection::LICENSE_FILES_TO_CHECK;
 pub use license_detection::is_license_eligible_for_data_collection;
 pub use rate_completion_modal::*;
@@ -55,9 +54,8 @@ use workspace::Workspace;
 use workspace::notifications::{ErrorMessagePrompt, NotificationId};
 use worktree::Worktree;
 use zed_llm_client::{
-    EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME,
     EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME, PredictEditsBody,
-    PredictEditsResponse, UsageLimit,
+    PredictEditsResponse,
 };
 
 const CURSOR_MARKER: &'static str = "<|user_cursor_is_here|>";
@@ -76,32 +74,6 @@ const MAX_EVENT_COUNT: usize = 16;
 
 actions!(edit_prediction, [ClearHistory]);
 
-#[derive(Debug, Clone, Copy)]
-pub struct Usage {
-    pub limit: UsageLimit,
-    pub amount: i32,
-}
-
-impl Usage {
-    pub fn from_headers(headers: &HeaderMap<HeaderValue>) -> Result<Self> {
-        let limit = headers
-            .get(EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME)
-            .ok_or_else(|| {
-                anyhow!("missing {EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME:?} header")
-            })?;
-        let limit = UsageLimit::from_str(limit.to_str()?)?;
-
-        let amount = headers
-            .get(EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME)
-            .ok_or_else(|| {
-                anyhow!("missing {EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME:?} header")
-            })?;
-        let amount = amount.to_str()?.parse::<i32>()?;
-
-        Ok(Self { limit, amount })
-    }
-}
-
 #[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 pub struct InlineCompletionId(Uuid);
 
@@ -216,6 +188,7 @@ pub struct Zeta {
     data_collection_choice: Entity<DataCollectionChoice>,
     llm_token: LlmApiToken,
     _llm_token_subscription: Subscription,
+    last_usage: Option<EditPredictionUsage>,
     /// Whether the terms of service have been accepted.
     tos_accepted: bool,
     /// Whether an update to a newer version of Zed is required to continue using Zeta.
@@ -291,6 +264,7 @@ impl Zeta {
                     .detach_and_log_err(cx);
                 },
             ),
+            last_usage: None,
             tos_accepted: user_store
                 .read(cx)
                 .current_user_has_accepted_terms()
@@ -387,7 +361,9 @@ impl Zeta {
     ) -> Task<Result<Option<InlineCompletion>>>
     where
         F: FnOnce(PerformPredictEditsParams) -> R + 'static,
-        R: Future<Output = Result<(PredictEditsResponse, Option<Usage>)>> + Send + 'static,
+        R: Future<Output = Result<(PredictEditsResponse, Option<EditPredictionUsage>)>>
+            + Send
+            + 'static,
     {
         let snapshot = self.report_changes_for_buffer(&buffer, cx);
         let diagnostic_groups = snapshot.diagnostic_groups(None);
@@ -427,7 +403,7 @@ impl Zeta {
             None
         };
 
-        cx.spawn(async move |_, cx| {
+        cx.spawn(async move |this, cx| {
             let request_sent_at = Instant::now();
 
             struct BackgroundValues {
@@ -532,11 +508,10 @@ impl Zeta {
             log::debug!("completion response: {}", &response.output_excerpt);
 
             if let Some(usage) = usage {
-                let limit = match usage.limit {
+                this.update(cx, |this, _cx| {
-                    UsageLimit::Limited(limit) => limit.to_string(),
+                    this.last_usage = Some(usage);
-                    UsageLimit::Unlimited => "unlimited".to_string(),
+                })
-                };
+                .ok();
-                log::info!("edit prediction usage: {} / {}", usage.amount, limit);
             }
 
             Self::process_completion_response(
@@ -750,7 +725,7 @@ and then another
 
     fn perform_predict_edits(
         params: PerformPredictEditsParams,
-    ) -> impl Future<Output = Result<(PredictEditsResponse, Option<Usage>)>> {
+    ) -> impl Future<Output = Result<(PredictEditsResponse, Option<EditPredictionUsage>)>> {
         async move {
             let PerformPredictEditsParams {
                 client,
@@ -796,7 +771,7 @@ and then another
                 }
 
                 if response.status().is_success() {
-                    let usage = Usage::from_headers(response.headers()).ok();
+                    let usage = EditPredictionUsage::from_headers(response.headers()).ok();
 
                     let mut body = String::new();
                     response.body_mut().read_to_string(&mut body).await?;
@@ -1440,6 +1415,10 @@ impl inline_completion::EditPredictionProvider for ZetaInlineCompletionProvider
         self.provider_data_collection.toggle(cx);
     }
 
+    fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
+        self.zeta.read(cx).last_usage
+    }
+
     fn is_enabled(
         &self,
         _buffer: &Entity<Buffer>,