Detailed changes
@@ -25,7 +25,7 @@ use gpui::{
Pixels, Subscription, Task, UpdateGlobal, WeakEntity, prelude::*, pulsating_between,
};
use language::LanguageRegistry;
-use language_model::{LanguageModelProviderTosView, LanguageModelRegistry};
+use language_model::{LanguageModelProviderTosView, LanguageModelRegistry, RequestUsage};
use language_model_selector::ToggleModelSelector;
use project::Project;
use prompt_store::{PromptBuilder, PromptStore, UserPromptId};
@@ -38,7 +38,7 @@ use ui::{
Banner, ContextMenu, KeyBinding, PopoverMenu, PopoverMenuHandle, ProgressBar, Tab, Tooltip,
prelude::*,
};
-use util::ResultExt as _;
+use util::{ResultExt as _, maybe};
use workspace::dock::{DockPosition, Panel, PanelEvent};
use workspace::{CollaboratorId, ToolbarItemView, Workspace};
use zed_actions::agent::OpenConfiguration;
@@ -1388,10 +1388,29 @@ impl AssistantPanel {
fn render_toolbar(&self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
let active_thread = self.thread.read(cx);
+ let user_store = self.user_store.read(cx);
let thread = active_thread.thread().read(cx);
let thread_id = thread.id().clone();
let is_empty = active_thread.is_empty();
- let last_usage = active_thread.thread().read(cx).last_usage();
+ let last_usage = active_thread.thread().read(cx).last_usage().or_else(|| {
+ maybe!({
+ let amount = user_store.model_request_usage_amount()?;
+ let limit = user_store.model_request_usage_limit()?.variant?;
+
+ Some(RequestUsage {
+ 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
+ }
+ },
+ })
+ })
+ });
+
let account_url = zed_urls::account_url(cx);
let show_token_count = match &self.active_view {
@@ -26,7 +26,7 @@ use gpui::{
Task, TextStyle, WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
};
use language::{Buffer, Language};
-use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent};
+use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent, RequestUsage};
use language_model_selector::ToggleModelSelector;
use multi_buffer;
use project::Project;
@@ -36,7 +36,7 @@ use settings::Settings;
use std::time::Duration;
use theme::ThemeSettings;
use ui::{Disclosure, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*};
-use util::ResultExt as _;
+use util::{ResultExt as _, maybe};
use workspace::{CollaboratorId, Workspace};
use zed_llm_client::CompletionMode;
@@ -1043,9 +1043,17 @@ impl MessageEditor {
return None;
}
- let plan = self
- .user_store
- .read(cx)
+ let user_store = self.user_store.read(cx);
+
+ let ubb_enable = user_store
+ .usage_based_billing_enabled()
+ .map_or(false, |enabled| enabled);
+
+ if ubb_enable {
+ return None;
+ }
+
+ let plan = user_store
.current_plan()
.map(|plan| match plan {
Plan::Free => zed_llm_client::Plan::Free,
@@ -1053,7 +1061,24 @@ impl MessageEditor {
Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
})
.unwrap_or(zed_llm_client::Plan::Free);
- let usage = self.thread.read(cx).last_usage()?;
+ let usage = self.thread.read(cx).last_usage().or_else(|| {
+ maybe!({
+ let amount = user_store.model_request_usage_amount()?;
+ let limit = user_store.model_request_usage_limit()?.variant?;
+
+ Some(RequestUsage {
+ 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
+ }
+ },
+ })
+ })
+ })?;
Some(
div()
@@ -102,6 +102,10 @@ pub struct UserStore {
update_contacts_tx: mpsc::UnboundedSender<UpdateContacts>,
current_plan: Option<proto::Plan>,
trial_started_at: Option<DateTime<Utc>>,
+ model_request_usage_amount: Option<u32>,
+ model_request_usage_limit: Option<proto::UsageLimit>,
+ edit_predictions_usage_amount: Option<u32>,
+ edit_predictions_usage_limit: Option<proto::UsageLimit>,
is_usage_based_billing_enabled: Option<bool>,
current_user: watch::Receiver<Option<Arc<User>>>,
accepted_tos_at: Option<Option<DateTime<Utc>>>,
@@ -163,6 +167,10 @@ impl UserStore {
current_user: current_user_rx,
current_plan: None,
trial_started_at: None,
+ model_request_usage_amount: None,
+ model_request_usage_limit: None,
+ edit_predictions_usage_amount: None,
+ edit_predictions_usage_limit: None,
is_usage_based_billing_enabled: None,
accepted_tos_at: None,
contacts: Default::default(),
@@ -330,6 +338,14 @@ impl UserStore {
.trial_started_at
.and_then(|trial_started_at| DateTime::from_timestamp(trial_started_at as i64, 0));
this.is_usage_based_billing_enabled = message.payload.is_usage_based_billing_enabled;
+
+ if let Some(usage) = message.payload.usage {
+ this.model_request_usage_amount = Some(usage.model_requests_usage_amount);
+ this.model_request_usage_limit = usage.model_requests_usage_limit;
+ this.edit_predictions_usage_amount = Some(usage.edit_predictions_usage_amount);
+ this.edit_predictions_usage_limit = usage.edit_predictions_usage_limit;
+ }
+
cx.notify();
})?;
Ok(())
@@ -697,6 +713,26 @@ impl UserStore {
self.current_plan
}
+ pub fn usage_based_billing_enabled(&self) -> Option<bool> {
+ self.is_usage_based_billing_enabled
+ }
+
+ pub fn model_request_usage_amount(&self) -> Option<u32> {
+ self.model_request_usage_amount
+ }
+
+ pub fn model_request_usage_limit(&self) -> Option<proto::UsageLimit> {
+ self.model_request_usage_limit.clone()
+ }
+
+ pub fn edit_predictions_usage_amount(&self) -> Option<u32> {
+ self.edit_predictions_usage_amount
+ }
+
+ pub fn edit_predictions_usage_limit(&self) -> Option<proto::UsageLimit> {
+ self.edit_predictions_usage_limit.clone()
+ }
+
pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
self.current_user.clone()
}