diff --git a/crates/acp_thread/src/acp_thread.rs b/crates/acp_thread/src/acp_thread.rs index 5d3b35d018..e58f0a291f 100644 --- a/crates/acp_thread/src/acp_thread.rs +++ b/crates/acp_thread/src/acp_thread.rs @@ -674,6 +674,37 @@ pub struct TokenUsage { pub used_tokens: u64, } +impl TokenUsage { + pub fn ratio(&self) -> TokenUsageRatio { + #[cfg(debug_assertions)] + let warning_threshold: f32 = std::env::var("ZED_THREAD_WARNING_THRESHOLD") + .unwrap_or("0.8".to_string()) + .parse() + .unwrap(); + #[cfg(not(debug_assertions))] + let warning_threshold: f32 = 0.8; + + // When the maximum is unknown because there is no selected model, + // avoid showing the token limit warning. + if self.max_tokens == 0 { + TokenUsageRatio::Normal + } else if self.used_tokens >= self.max_tokens { + TokenUsageRatio::Exceeded + } else if self.used_tokens as f32 / self.max_tokens as f32 >= warning_threshold { + TokenUsageRatio::Warning + } else { + TokenUsageRatio::Normal + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TokenUsageRatio { + Normal, + Warning, + Exceeded, +} + #[derive(Debug, Clone)] pub struct RetryStatus { pub last_error: SharedString, diff --git a/crates/agent_ui/src/acp/thread_view.rs b/crates/agent_ui/src/acp/thread_view.rs index 3be88ee3c3..b93df3a5db 100644 --- a/crates/agent_ui/src/acp/thread_view.rs +++ b/crates/agent_ui/src/acp/thread_view.rs @@ -54,6 +54,7 @@ use crate::acp::entry_view_state::{EntryViewEvent, ViewEvent}; use crate::acp::message_editor::{MessageEditor, MessageEditorEvent}; use crate::agent_diff::AgentDiff; use crate::profile_selector::{ProfileProvider, ProfileSelector}; +use crate::ui::preview::UsageCallout; use crate::ui::{AgentNotification, AgentNotificationEvent, BurnModeTooltip}; use crate::{ AgentDiffPane, AgentPanel, ContinueThread, ContinueWithBurnMode, ExpandMessageEditor, Follow, @@ -2940,6 +2941,12 @@ impl AcpThreadView { .thread(acp_thread.session_id(), cx) } + fn is_using_zed_ai_models(&self, cx: &App) -> bool { + self.as_native_thread(cx) + .and_then(|thread| thread.read(cx).model()) + .is_some_and(|model| model.provider_id() == language_model::ZED_CLOUD_PROVIDER_ID) + } + fn render_token_usage(&self, cx: &mut Context) -> Option
{ let thread = self.thread()?.read(cx); let usage = thread.token_usage()?; @@ -3587,6 +3594,88 @@ impl AcpThreadView { .children(Scrollbar::vertical(self.scrollbar_state.clone()).map(|s| s.auto_hide(cx))) } + fn render_token_limit_callout( + &self, + line_height: Pixels, + cx: &mut Context, + ) -> Option { + let token_usage = self.thread()?.read(cx).token_usage()?; + let ratio = token_usage.ratio(); + + let (severity, title) = match ratio { + acp_thread::TokenUsageRatio::Normal => return None, + acp_thread::TokenUsageRatio::Warning => { + (Severity::Warning, "Thread reaching the token limit soon") + } + acp_thread::TokenUsageRatio::Exceeded => { + (Severity::Error, "Thread reached the token limit") + } + }; + + let burn_mode_available = self.as_native_thread(cx).is_some_and(|thread| { + thread.read(cx).completion_mode() == CompletionMode::Normal + && thread + .read(cx) + .model() + .is_some_and(|model| model.supports_burn_mode()) + }); + + let description = if burn_mode_available { + "To continue, start a new thread from a summary or turn Burn Mode on." + } else { + "To continue, start a new thread from a summary." + }; + + Some( + Callout::new() + .severity(severity) + .line_height(line_height) + .title(title) + .description(description) + .actions_slot( + h_flex() + .gap_0p5() + .child( + Button::new("start-new-thread", "Start New Thread") + .label_size(LabelSize::Small) + .on_click(cx.listener(|_this, _, _window, _cx| { + // todo: Once thread summarization is implemented, start a new thread from a summary. + })), + ) + .when(burn_mode_available, |this| { + this.child( + IconButton::new("burn-mode-callout", IconName::ZedBurnMode) + .icon_size(IconSize::XSmall) + .on_click(cx.listener(|this, _event, window, cx| { + this.toggle_burn_mode(&ToggleBurnMode, window, cx); + })), + ) + }), + ), + ) + } + + fn render_usage_callout(&self, line_height: Pixels, cx: &mut Context) -> Option
{ + if !self.is_using_zed_ai_models(cx) { + return None; + } + + let user_store = self.project.read(cx).user_store().read(cx); + if user_store.is_usage_based_billing_enabled() { + return None; + } + + let plan = user_store.plan().unwrap_or(cloud_llm_client::Plan::ZedFree); + + let usage = user_store.model_request_usage()?; + + Some( + div() + .child(UsageCallout::new(plan, usage)) + .line_height(line_height), + ) + } + fn settings_changed(&mut self, _window: &mut Window, cx: &mut Context) { self.entry_view_state.update(cx, |entry_view_state, cx| { entry_view_state.settings_changed(cx); @@ -3843,6 +3932,7 @@ impl Focusable for AcpThreadView { impl Render for AcpThreadView { fn render(&mut self, window: &mut Window, cx: &mut Context) -> impl IntoElement { let has_messages = self.list_state.item_count() > 0; + let line_height = TextSize::Small.rems(cx).to_pixels(window.rem_size()) * 1.5; v_flex() .size_full() @@ -3921,6 +4011,17 @@ impl Render for AcpThreadView { }) .children(self.render_thread_retry_status_callout(window, cx)) .children(self.render_thread_error(window, cx)) + .children( + if let Some(usage_callout) = self.render_usage_callout(line_height, cx) { + Some(usage_callout.into_any_element()) + } else if let Some(token_limit_callout) = + self.render_token_limit_callout(line_height, cx) + { + Some(token_limit_callout.into_any_element()) + } else { + None + }, + ) .child(self.render_message_editor(window, cx)) } } diff --git a/crates/ui/src/components/callout.rs b/crates/ui/src/components/callout.rs index 22ba0468cd..7ffeda881c 100644 --- a/crates/ui/src/components/callout.rs +++ b/crates/ui/src/components/callout.rs @@ -81,7 +81,8 @@ impl Callout { self } - /// Sets an optional tertiary call-to-action button. + /// Sets an optional dismiss button, which is usually an icon button with a close icon. + /// This button is always rendered as the last one to the far right. pub fn dismiss_action(mut self, action: impl IntoElement) -> Self { self.dismiss_action = Some(action.into_any_element()); self