From a32505fcc28c2663205a24b9f5c2fd133fd98dfb Mon Sep 17 00:00:00 2001
From: Umesh Yadav <23421535+imumesh18@users.noreply.github.com>
Date: Mon, 23 Jun 2025 17:06:44 +0530
Subject: [PATCH] agent: Fix token limit callout to show burn mode only for zed
 provider (#33096)

The token limit reached callout was shown for all the providers and it
included the burn mode toggle and description. I have made that
conditional to only show for zed provider.

Before this changes
<img width="413" alt="image"
src="https://github.com/user-attachments/assets/b7e177fc-d2e6-4942-8934-39ac8c844649"
/>

After this change:
<img width="404" alt="image"
src="https://github.com/user-attachments/assets/0e091b99-c8e3-4960-8b4d-47170883d224"
/>


Release Notes:

- agent: Fix token limit callout to show burn mode only for zed provider

---------

Co-authored-by: Danilo Leal <daniloleal09@gmail.com>
---
 crates/agent/src/message_editor.rs | 69 ++++++++++++++++--------------
 1 file changed, 37 insertions(+), 32 deletions(-)

diff --git a/crates/agent/src/message_editor.rs b/crates/agent/src/message_editor.rs
index 9e2d0696055f29634f93ea425627af4bca5a3c59..cc20f18e4f9541613323b883ce710bb3fb490d29 100644
--- a/crates/agent/src/message_editor.rs
+++ b/crates/agent/src/message_editor.rs
@@ -1240,16 +1240,17 @@ impl MessageEditor {
             })
     }
 
-    fn render_usage_callout(&self, line_height: Pixels, cx: &mut Context<Self>) -> Option<Div> {
-        let is_using_zed_provider = self
-            .thread
+    fn is_using_zed_provider(&self, cx: &App) -> bool {
+        self.thread
             .read(cx)
             .configured_model()
             .map_or(false, |model| {
                 model.provider.id().0 == ZED_CLOUD_PROVIDER_ID
-            });
+            })
+    }
 
-        if !is_using_zed_provider {
+    fn render_usage_callout(&self, line_height: Pixels, cx: &mut Context<Self>) -> Option<Div> {
+        if !self.is_using_zed_provider(cx) {
             return None;
         }
 
@@ -1303,37 +1304,41 @@ impl MessageEditor {
             "Thread reaching the token limit soon"
         };
 
+        let description = if self.is_using_zed_provider(cx) {
+            "To continue, start a new thread from a summary or turn burn mode on."
+        } else {
+            "To continue, start a new thread from a summary."
+        };
+
+        let mut callout = Callout::new()
+            .line_height(line_height)
+            .icon(icon)
+            .title(title)
+            .description(description)
+            .primary_action(
+                Button::new("start-new-thread", "Start New Thread")
+                    .label_size(LabelSize::Small)
+                    .on_click(cx.listener(|this, _, window, cx| {
+                        let from_thread_id = Some(this.thread.read(cx).id().clone());
+                        window.dispatch_action(Box::new(NewThread { from_thread_id }), cx);
+                    })),
+            );
+
+        if self.is_using_zed_provider(cx) {
+            callout = callout.secondary_action(
+                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);
+                    })),
+            );
+        }
+
         Some(
             div()
                 .border_t_1()
                 .border_color(cx.theme().colors().border)
-                .child(
-                    Callout::new()
-                        .line_height(line_height)
-                        .icon(icon)
-                        .title(title)
-                        .description(
-                            "To continue, start a new thread from a summary or turn burn mode on.",
-                        )
-                        .primary_action(
-                            Button::new("start-new-thread", "Start New Thread")
-                                .label_size(LabelSize::Small)
-                                .on_click(cx.listener(|this, _, window, cx| {
-                                    let from_thread_id = Some(this.thread.read(cx).id().clone());
-                                    window.dispatch_action(
-                                        Box::new(NewThread { from_thread_id }),
-                                        cx,
-                                    );
-                                })),
-                        )
-                        .secondary_action(
-                            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);
-                                })),
-                        ),
-                ),
+                .child(callout),
         )
     }