assistant2: Show token count in the Prompt Editor view (#25091)

Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <hi@aguz.me>

crates/assistant2/src/assistant_panel.rs

@@ -3,8 +3,8 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, Result};
 use assistant_context_editor::{
-    make_lsp_adapter_delegate, AssistantPanelDelegate, ConfigurationError, ContextEditor,
+    make_lsp_adapter_delegate, render_remaining_tokens, AssistantPanelDelegate, ConfigurationError,
-    ContextHistory, SlashCommandCompletionProvider,
+    ContextEditor, ContextHistory, SlashCommandCompletionProvider,
 };
 use assistant_settings::{AssistantDockPosition, AssistantSettings};
 use assistant_slash_command::SlashCommandWorkingSet;
@@ -635,20 +635,33 @@ impl AssistantPanel {
             .border_color(cx.theme().colors().border)
             .child(
                 h_flex()
-                    .child(Label::new(title))
+                    .w_full()
-                    .when(sub_title.is_some(), |this| {
+                    .gap_1()
-                        this.child(
+                    .justify_between()
-                            h_flex()
+                    .child(
-                                .pl_1p5()
+                        h_flex()
-                                .gap_1p5()
+                            .child(Label::new(title))
-                                .child(
+                            .when(sub_title.is_some(), |this| {
-                                    Label::new("/")
+                                this.child(
-                                        .size(LabelSize::Small)
+                                    h_flex()
-                                        .color(Color::Disabled)
+                                        .pl_1p5()
-                                        .alpha(0.5),
+                                        .gap_1p5()
+                                        .child(
+                                            Label::new("/")
+                                                .size(LabelSize::Small)
+                                                .color(Color::Disabled)
+                                                .alpha(0.5),
+                                        )
+                                        .child(Label::new(sub_title.unwrap())),
                                 )
-                                .child(Label::new(sub_title.unwrap())),
+                            }),
-                        )
+                    )
+                    .children(if matches!(self.active_view, ActiveView::PromptEditor) {
+                        self.context_editor
+                            .as_ref()
+                            .and_then(|editor| render_remaining_tokens(editor, cx))
+                    } else {
+                        None
                     }),
             )
             .child(

crates/assistant_context_editor/src/context.rs

@@ -1190,11 +1190,14 @@ impl AssistantContext {
         let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
             return;
         };
+        let debounce = self.token_count.is_some();
         self.pending_token_count = cx.spawn(|this, mut cx| {
             async move {
-                cx.background_executor()
+                if debounce {
-                    .timer(Duration::from_millis(200))
+                    cx.background_executor()
-                    .await;
+                        .timer(Duration::from_millis(200))
+                        .await;
+                }
 
                 let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
                 this.update(&mut cx, |this, cx| {

crates/assistant_context_editor/src/context_editor.rs

@@ -2867,7 +2867,6 @@ impl EventEmitter<SearchEvent> for ContextEditor {}
 impl Render for ContextEditor {
     fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         let provider = LanguageModelRegistry::read_global(cx).active_provider();
-
         let accept_terms = if self.show_accept_terms {
             provider.as_ref().and_then(|provider| {
                 provider.render_accept_terms(LanguageModelProviderTosView::PromptEditorPopup, cx)
@@ -3250,48 +3249,58 @@ impl ContextEditorToolbarItem {
             model_summary_editor,
         }
     }
+}
 
-    fn render_remaining_tokens(&self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
+pub fn render_remaining_tokens(
-        let context = &self
+    context_editor: &Entity<ContextEditor>,
-            .active_context_editor
+    cx: &App,
-            .as_ref()?
+) -> Option<impl IntoElement> {
-            .upgrade()?
+    let context = &context_editor.read(cx).context;
-            .read(cx)
+
-            .context;
+    let (token_count_color, token_count, max_token_count, tooltip) = match token_state(context, cx)?
-        let (token_count_color, token_count, max_token_count) = match token_state(context, cx)? {
+    {
-            TokenState::NoTokensLeft {
+        TokenState::NoTokensLeft {
-                max_token_count,
+            max_token_count,
-                token_count,
+            token_count,
-            } => (Color::Error, token_count, max_token_count),
+        } => (
-            TokenState::HasMoreTokens {
+            Color::Error,
-                max_token_count,
+            token_count,
-                token_count,
+            max_token_count,
-                over_warn_threshold,
+            Some("Token Limit Reached"),
-            } => {
+        ),
-                let color = if over_warn_threshold {
+        TokenState::HasMoreTokens {
-                    Color::Warning
+            max_token_count,
-                } else {
+            token_count,
-                    Color::Muted
+            over_warn_threshold,
-                };
+        } => {
-                (color, token_count, max_token_count)
+            let (color, tooltip) = if over_warn_threshold {
-            }
+                (Color::Warning, Some("Token Limit is Close to Exhaustion"))
-        };
+            } else {
-        Some(
+                (Color::Muted, None)
-            h_flex()
+            };
-                .gap_0p5()
+            (color, token_count, max_token_count, tooltip)
-                .child(
+        }
-                    Label::new(humanize_token_count(token_count))
+    };
-                        .size(LabelSize::Small)
+
-                        .color(token_count_color),
+    Some(
-                )
+        h_flex()
-                .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
+            .id("token-count")
-                .child(
+            .gap_0p5()
-                    Label::new(humanize_token_count(max_token_count))
+            .child(
-                        .size(LabelSize::Small)
+                Label::new(humanize_token_count(token_count))
-                        .color(Color::Muted),
+                    .size(LabelSize::Small)
-                ),
+                    .color(token_count_color),
-        )
+            )
-    }
+            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
+            .child(
+                Label::new(humanize_token_count(max_token_count))
+                    .size(LabelSize::Small)
+                    .color(Color::Muted),
+            )
+            .when_some(tooltip, |element, tooltip| {
+                element.tooltip(Tooltip::text(tooltip))
+            }),
+    )
 }
 
 impl Render for ContextEditorToolbarItem {
@@ -3334,7 +3343,12 @@ impl Render for ContextEditorToolbarItem {
             //     scan_items_remaining
             //         .map(|remaining_items| format!("Files to scan: {}", remaining_items))
             // })
-            .children(self.render_remaining_tokens(cx));
+            .children(
+                self.active_context_editor
+                    .as_ref()
+                    .and_then(|editor| editor.upgrade())
+                    .and_then(|editor| render_remaining_tokens(&editor, cx)),
+            );
 
         h_flex()
             .px_0p5()