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
This commit is contained in:
Marshall Bowers 2025-04-18 10:15:19 -04:00 committed by GitHub
parent 62b8ef980b
commit 0dc0701967
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 111 additions and 43 deletions

3
Cargo.lock generated
View File

@ -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",
]

View File

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

View File

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

View File

@ -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"] }

View File

@ -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| {

View File

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