inline_completion_button: Show the initial usage data from the server (#29952)

Marshall Bowers created 7 months ago

This PR updates the usage indicator for edit predictions to show the
initial usage data returned from the server.

Release Notes:

- N/A

Change summary

Cargo.lock                                                      |  2 
crates/inline_completion_button/Cargo.toml                      |  2 
crates/inline_completion_button/src/inline_completion_button.rs | 25 ++
3 files changed, 28 insertions(+), 1 deletion(-)

Detailed changes

Cargo.lock

@@ -7296,6 +7296,7 @@ dependencies = [
  "lsp",
  "paths",
  "project",
+ "proto",
  "regex",
  "serde_json",
  "settings",
@@ -7303,6 +7304,7 @@ dependencies = [
  "telemetry",
  "theme",
  "ui",
+ "util",
  "workspace",
  "workspace-hack",
  "zed_actions",

crates/inline_completion_button/Cargo.toml

@@ -24,11 +24,13 @@ indoc.workspace = true
 inline_completion.workspace = true
 language.workspace = true
 paths.workspace = true
+proto.workspace = true
 regex.workspace = true
 settings.workspace = true
 supermaven.workspace = true
 telemetry.workspace = true
 ui.workspace = true
+util.workspace = true
 workspace-hack.workspace = true
 workspace.workspace = true
 zed_actions.workspace = true

crates/inline_completion_button/src/inline_completion_button.rs

@@ -14,6 +14,7 @@ use gpui::{
     pulsating_between,
 };
 use indoc::indoc;
+use inline_completion::EditPredictionUsage;
 use language::{
     EditPredictionsMode, File, Language,
     language_settings::{self, AllLanguageSettings, EditPredictionProvider, all_language_settings},
@@ -29,6 +30,7 @@ use ui::{
     Clickable, ContextMenu, ContextMenuEntry, IconButton, IconButtonShape, Indicator, PopoverMenu,
     PopoverMenuHandle, ProgressBar, Tooltip, prelude::*,
 };
+use util::maybe;
 use workspace::{
     StatusItemView, Toast, Workspace, create_and_open_local_file, item::ItemHandle,
     notifications::NotificationId,
@@ -404,7 +406,28 @@ impl InlineCompletionButton {
         let line_height = window.line_height();
 
         if let Some(provider) = self.edit_prediction_provider.as_ref() {
-            if let Some(usage) = provider.usage(cx) {
+            let usage = provider.usage(cx).or_else(|| {
+                let user_store = self.user_store.read(cx);
+
+                maybe!({
+                    let amount = user_store.edit_predictions_usage_amount()?;
+                    let limit = user_store.edit_predictions_usage_limit()?.variant?;
+
+                    Some(EditPredictionUsage {
+                        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
+                            }
+                        },
+                    })
+                })
+            });
+
+            if let Some(usage) = usage {
                 menu = menu.header("Usage");
                 menu = menu.custom_entry(
                     move |_window, cx| {