From de5f0234773e7dbdedd179d7dfeb44872a3b8a85 Mon Sep 17 00:00:00 2001
From: Agus Zubiaga <agus@zed.dev>
Date: Tue, 14 Jan 2025 16:22:48 -0300
Subject: [PATCH] assistant2: Cancel generation button (#23137)

Turns the "esc to cancel" label into a button so it can be dispatched
via click or keyboard. The keybinding isn't hardcoded anymore.

![CleanShot 2025-01-14 at 13 44
22@2x](https://github.com/user-attachments/assets/a947f58b-7de2-400b-b95a-384b78c79697)


Release Notes:

- N/A
---
 crates/assistant2/src/active_thread.rs   | 31 ++++++++++++++++++++----
 crates/assistant2/src/assistant_panel.rs | 25 +++++++++++--------
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/crates/assistant2/src/active_thread.rs b/crates/assistant2/src/active_thread.rs
index 7d5f7750f6443029dab521b47636eb1d497cff80..6e3c90cee266005a4a81cae274dd2f3226d3d385 100644
--- a/crates/assistant2/src/active_thread.rs
+++ b/crates/assistant2/src/active_thread.rs
@@ -5,16 +5,16 @@ use assistant_tool::ToolWorkingSet;
 use collections::HashMap;
 use gpui::{
     list, percentage, AbsoluteLength, Animation, AnimationExt, AnyElement, AppContext,
-    DefiniteLength, EdgesRefinement, Empty, Length, ListAlignment, ListOffset, ListState, Model,
-    StyleRefinement, Subscription, TextStyleRefinement, Transformation, UnderlineStyle, View,
-    WeakView,
+    DefiniteLength, EdgesRefinement, Empty, FocusHandle, Length, ListAlignment, ListOffset,
+    ListState, Model, StyleRefinement, Subscription, TextStyleRefinement, Transformation,
+    UnderlineStyle, View, WeakView,
 };
 use language::LanguageRegistry;
 use language_model::Role;
 use markdown::{Markdown, MarkdownStyle};
 use settings::Settings as _;
 use theme::ThemeSettings;
-use ui::prelude::*;
+use ui::{prelude::*, ButtonLike, KeyBinding};
 use workspace::Workspace;
 
 use crate::thread::{MessageId, Thread, ThreadError, ThreadEvent};
@@ -29,6 +29,7 @@ pub struct ActiveThread {
     list_state: ListState,
     rendered_messages_by_id: HashMap<MessageId, View<Markdown>>,
     last_error: Option<ThreadError>,
+    focus_handle: FocusHandle,
     _subscriptions: Vec<Subscription>,
 }
 
@@ -38,6 +39,7 @@ impl ActiveThread {
         workspace: WeakView<Workspace>,
         language_registry: Arc<LanguageRegistry>,
         tools: Arc<ToolWorkingSet>,
+        focus_handle: FocusHandle,
         cx: &mut ViewContext<Self>,
     ) -> Self {
         let subscriptions = vec![
@@ -60,6 +62,7 @@ impl ActiveThread {
                 }
             }),
             last_error: None,
+            focus_handle,
             _subscriptions: subscriptions,
         };
 
@@ -345,6 +348,8 @@ impl Render for ActiveThread {
     fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
         let is_streaming_completion = self.thread.read(cx).is_streaming();
 
+        let focus_handle = self.focus_handle.clone();
+
         v_flex()
             .size_full()
             .child(list(self.list_state.clone()).flex_grow())
@@ -363,7 +368,23 @@ impl Render for ActiveThread {
                                 .rounded_md()
                                 .bg(cx.theme().colors().elevated_surface_background)
                                 .child(Label::new("Generating…").size(LabelSize::Small))
-                                .child(Label::new("esc to cancel").size(LabelSize::Small)),
+                                .child(
+                                    ButtonLike::new("cancel-generation")
+                                        .style(ButtonStyle::Filled)
+                                        .child(Label::new("Cancel").size(LabelSize::Small))
+                                        .children(
+                                            KeyBinding::for_action_in(
+                                                &editor::actions::Cancel,
+                                                &self.focus_handle,
+                                                cx,
+                                            )
+                                            .map(|binding| binding.into_any_element()),
+                                        )
+                                        .on_click(move |_event, cx| {
+                                            focus_handle
+                                                .dispatch_action(&editor::actions::Cancel, cx);
+                                        }),
+                                ),
                         )
                     }),
             )
diff --git a/crates/assistant2/src/assistant_panel.rs b/crates/assistant2/src/assistant_panel.rs
index 4c6993312e197317f64137ca582429be965ad339..556b260335a8c85232da4aaf2b857df2c560309a 100644
--- a/crates/assistant2/src/assistant_panel.rs
+++ b/crates/assistant2/src/assistant_panel.rs
@@ -100,6 +100,16 @@ impl AssistantPanel {
         let workspace = workspace.weak_handle();
         let weak_self = cx.view().downgrade();
 
+        let message_editor = cx.new_view(|cx| {
+            MessageEditor::new(
+                fs.clone(),
+                workspace.clone(),
+                thread_store.downgrade(),
+                thread.clone(),
+                cx,
+            )
+        });
+
         Self {
             active_view: ActiveView::Thread,
             workspace: workspace.clone(),
@@ -109,21 +119,14 @@ impl AssistantPanel {
             thread: cx.new_view(|cx| {
                 ActiveThread::new(
                     thread.clone(),
-                    workspace.clone(),
+                    workspace,
                     language_registry,
                     tools.clone(),
+                    message_editor.focus_handle(cx),
                     cx,
                 )
             }),
-            message_editor: cx.new_view(|cx| {
-                MessageEditor::new(
-                    fs.clone(),
-                    workspace,
-                    thread_store.downgrade(),
-                    thread.clone(),
-                    cx,
-                )
-            }),
+            message_editor,
             tools,
             local_timezone: UtcOffset::from_whole_seconds(
                 chrono::Local::now().offset().local_minus_utc(),
@@ -160,6 +163,7 @@ impl AssistantPanel {
                 self.workspace.clone(),
                 self.language_registry.clone(),
                 self.tools.clone(),
+                self.focus_handle(cx),
                 cx,
             )
         });
@@ -196,6 +200,7 @@ impl AssistantPanel {
                 self.workspace.clone(),
                 self.language_registry.clone(),
                 self.tools.clone(),
+                self.focus_handle(cx),
                 cx,
             )
         });