diff --git a/Cargo.lock b/Cargo.lock index 6866438deb..250fbba92d 100644 --- a/Cargo.lock +++ b/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", ] diff --git a/crates/inline_completion/Cargo.toml b/crates/inline_completion/Cargo.toml index 84344c2ceb..0094385e16 100644 --- a/crates/inline_completion/Cargo.toml +++ b/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 diff --git a/crates/inline_completion/src/inline_completion.rs b/crates/inline_completion/src/inline_completion.rs index 80077cc169..7733fec1cb 100644 --- a/crates/inline_completion/src/inline_completion.rs +++ b/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) -> Result { + 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::()?; + + 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 { + 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; 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 { + self.read(cx).usage(cx) + } + fn toggle_data_collection(&self, cx: &mut App) { self.update(cx, |this, cx| this.toggle_data_collection(cx)) } diff --git a/crates/inline_completion_button/Cargo.toml b/crates/inline_completion_button/Cargo.toml index 6dc9d5c5b8..c2a619d500 100644 --- a/crates/inline_completion_button/Cargo.toml +++ b/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"] } diff --git a/crates/inline_completion_button/src/inline_completion_button.rs b/crates/inline_completion_button/src/inline_completion_button.rs index fda3e6a080..e4ea9566e7 100644 --- a/crates/inline_completion_button/src/inline_completion_button.rs +++ b/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| { diff --git a/crates/zeta/src/zeta.rs b/crates/zeta/src/zeta.rs index b5367816fd..bec8efbcff 100644 --- a/crates/zeta/src/zeta.rs +++ b/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) -> Result { - 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::()?; - - 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, llm_token: LlmApiToken, _llm_token_subscription: Subscription, + last_usage: Option, /// 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>> where F: FnOnce(PerformPredictEditsParams) -> R + 'static, - R: Future)>> + Send + 'static, + R: Future)>> + + 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 { - UsageLimit::Limited(limit) => limit.to_string(), - UsageLimit::Unlimited => "unlimited".to_string(), - }; - log::info!("edit prediction usage: {} / {}", usage.amount, limit); + this.update(cx, |this, _cx| { + this.last_usage = Some(usage); + }) + .ok(); } Self::process_completion_response( @@ -750,7 +725,7 @@ and then another fn perform_predict_edits( params: PerformPredictEditsParams, - ) -> impl Future)>> { + ) -> impl Future)>> { 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 { + self.zeta.read(cx).last_usage + } + fn is_enabled( &self, _buffer: &Entity,