Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent: Fix issues with usage display sometimes showing initially fetched usage (#33125)

Browse source 
Having `Thread::last_usage` as an override of the initially fetched
usage could cause the initial usage to be displayed when the current
thread is empty or in text threads. Fix is to just store last usage info
in `UserStore` and not have these overrides

Release Notes:

- Agent: Fixed request usage display to always include the most recently
known usage - there were some cases where it would show the initially
requested usage.
This commit is contained in:
Michael Sloan 2025-06-20 15:28:48 -06:00 • committed by GitHub
parent e0c0b6f95d
commit 7e801dccb0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 188 additions and 211 deletions
Download patch file Download diff file Expand all files Collapse all files

39
crates/zeta/src/zeta.rs
View file

@ -9,14 +9,14 @@ mod rate_completion_modal;
pub(crate) use completion_diff_element::*;
use db::kvp::KEY_VALUE_STORE;
pub use init::*;
use inline_completion::{DataCollectionState, EditPredictionUsage};
use inline_completion::DataCollectionState;
use license_detection::LICENSE_FILES_TO_CHECK;
pub use license_detection::is_license_eligible_for_data_collection;
pub use rate_completion_modal::*;
use anyhow::{Context as _, Result, anyhow};
use arrayvec::ArrayVec;
use client::{Client, UserStore};
use client::{Client, EditPredictionUsage, UserStore};
use collections::{HashMap, HashSet, VecDeque};
use futures::AsyncReadExt;
use gpui::{
@ -48,7 +48,7 @@ use std::{
};
use telemetry_events::InlineCompletionRating;
use thiserror::Error;
use util::{ResultExt, maybe};
use util::ResultExt;
use uuid::Uuid;
use workspace::Workspace;
use workspace::notifications::{ErrorMessagePrompt, NotificationId};
@ -188,7 +188,6 @@ 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.
@ -234,25 +233,7 @@ impl Zeta {
}
pub fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
self.last_usage.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
}
},
})
})
})
self.user_store.read(cx).edit_prediction_usage()
}
fn new(
@ -287,7 +268,6 @@ impl Zeta {
.detach_and_log_err(cx);
},
),
last_usage: None,
tos_accepted: user_store
.read(cx)
.current_user_has_accepted_terms()
@ -533,8 +513,10 @@ impl Zeta {
log::debug!("completion response: {}", &response.output_excerpt);
if let Some(usage) = usage {
this.update(cx, |this, _cx| {
this.last_usage = Some(usage);
this.update(cx, |this, cx| {
this.user_store.update(cx, |user_store, cx| {
user_store.update_edit_prediction_usage(usage, cx);
});
})
.ok();
}
@ -874,8 +856,9 @@ and then another
if response.status().is_success() {
if let Some(usage) = EditPredictionUsage::from_headers(response.headers()).ok() {
this.update(cx, |this, cx| {
this.last_usage = Some(usage);
cx.notify();
this.user_store.update(cx, |user_store, cx| {
user_store.update_edit_prediction_usage(usage, cx);
});
})?;
}