Sketch in assistant edit button (#19705)

Add an edit button to the assistant. This is totally hacked in for now,
just to see how this would feel rendered simply in the UI.

![CleanShot 2024-10-24 at 16 26
14@2x](https://github.com/user-attachments/assets/e630d078-78b7-42d7-93f1-cf61c00bd20e)

cc @as-cii @danilo-leal 

Release Notes:

- N/A

---------

Co-authored-by: Danilo Leal <67129314+danilo-leal@users.noreply.github.com>
Co-authored-by: Richard Feldman <oss@rtfeldman.com>

assets/keymaps/default-linux.json

@@ -532,6 +532,7 @@
     "context": "ContextEditor > Editor",
     "bindings": {
       "ctrl-enter": "assistant::Assist",
+      "ctrl-shift-enter": "assistant::Edit",
       "ctrl-s": "workspace::Save",
       "ctrl->": "assistant::QuoteSelection",
       "ctrl-<": "assistant::InsertIntoEditor",

assets/keymaps/default-macos.json

@@ -201,6 +201,7 @@
     "context": "ContextEditor > Editor",
     "bindings": {
       "cmd-enter": "assistant::Assist",
+      "cmd-shift-enter": "assistant::Edit",
       "cmd-s": "workspace::Save",
       "cmd->": "assistant::QuoteSelection",
       "cmd-<": "assistant::InsertIntoEditor",

crates/assistant/src/assistant.rs

@@ -59,6 +59,7 @@ actions!(
     assistant,
     [
         Assist,
+        Edit,
         Split,
         CopyCode,
         CycleMessageRole,

crates/assistant/src/assistant_panel.rs

@@ -13,10 +13,11 @@ use crate::{
     terminal_inline_assistant::TerminalInlineAssistant,
     Assist, AssistantPatch, AssistantPatchStatus, CacheStatus, ConfirmCommand, Content, Context,
     ContextEvent, ContextId, ContextStore, ContextStoreEvent, CopyCode, CycleMessageRole,
-    DeployHistory, DeployPromptLibrary, InlineAssistant, InsertDraggedFiles, InsertIntoEditor,
+    DeployHistory, DeployPromptLibrary, Edit, InlineAssistant, InsertDraggedFiles,
-    Message, MessageId, MessageMetadata, MessageStatus, ModelPickerDelegate, ModelSelector,
+    InsertIntoEditor, Message, MessageId, MessageMetadata, MessageStatus, ModelPickerDelegate,
-    NewContext, PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection,
+    ModelSelector, NewContext, PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection,
-    RemoteContextMetadata, SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector,
+    RemoteContextMetadata, RequestType, SavedContextMetadata, Split, ToggleFocus,
+    ToggleModelSelector,
 };
 use anyhow::Result;
 use assistant_slash_command::{SlashCommand, SlashCommandOutputSection};
@@ -1588,23 +1589,11 @@ impl ContextEditor {
     }
 
     fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
-        let provider = LanguageModelRegistry::read_global(cx).active_provider();
+        self.send_to_model(RequestType::Chat, cx);
-        if provider
-            .as_ref()
-            .map_or(false, |provider| provider.must_accept_terms(cx))
-        {
-            self.show_accept_terms = true;
-            cx.notify();
-            return;
     }
 
-        if self.focus_active_patch(cx) {
+    fn edit(&mut self, _: &Edit, cx: &mut ViewContext<Self>) {
-            return;
+        self.send_to_model(RequestType::SuggestEdits, cx);
-        }
-
-        self.last_error = None;
-        self.send_to_model(cx);
-        cx.notify();
     }
 
     fn focus_active_patch(&mut self, cx: &mut ViewContext<Self>) -> bool {
@@ -1622,8 +1611,27 @@ impl ContextEditor {
         false
     }
 
-    fn send_to_model(&mut self, cx: &mut ViewContext<Self>) {
+    fn send_to_model(&mut self, request_type: RequestType, cx: &mut ViewContext<Self>) {
-        if let Some(user_message) = self.context.update(cx, |context, cx| context.assist(cx)) {
+        let provider = LanguageModelRegistry::read_global(cx).active_provider();
+        if provider
+            .as_ref()
+            .map_or(false, |provider| provider.must_accept_terms(cx))
+        {
+            self.show_accept_terms = true;
+            cx.notify();
+            return;
+        }
+
+        if self.focus_active_patch(cx) {
+            return;
+        }
+
+        self.last_error = None;
+
+        if let Some(user_message) = self
+            .context
+            .update(cx, |context, cx| context.assist(request_type, cx))
+        {
             let new_selection = {
                 let cursor = user_message
                     .start
@@ -1640,6 +1648,8 @@ impl ContextEditor {
             // Avoid scrolling to the new cursor position so the assistant's output is stable.
             cx.defer(|this, _| this.scroll_position = None);
         }
+
+        cx.notify();
     }
 
     fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
@@ -3644,7 +3654,13 @@ impl ContextEditor {
                 button.tooltip(move |_| tooltip.clone())
             })
             .layer(ElevationIndex::ModalSurface)
-            .child(Label::new("Send"))
+            .child(Label::new(
+                if AssistantSettings::get_global(cx).are_live_diffs_enabled(cx) {
+                    "Chat"
+                } else {
+                    "Send"
+                },
+            ))
             .children(
                 KeyBinding::for_action_in(&Assist, &focus_handle, cx)
                     .map(|binding| binding.into_any_element()),
@@ -3654,6 +3670,57 @@ impl ContextEditor {
             })
     }
 
+    fn render_edit_button(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let focus_handle = self.focus_handle(cx).clone();
+
+        let (style, tooltip) = match token_state(&self.context, cx) {
+            Some(TokenState::NoTokensLeft { .. }) => (
+                ButtonStyle::Tinted(TintColor::Negative),
+                Some(Tooltip::text("Token limit reached", cx)),
+            ),
+            Some(TokenState::HasMoreTokens {
+                over_warn_threshold,
+                ..
+            }) => {
+                let (style, tooltip) = if over_warn_threshold {
+                    (
+                        ButtonStyle::Tinted(TintColor::Warning),
+                        Some(Tooltip::text("Token limit is close to exhaustion", cx)),
+                    )
+                } else {
+                    (ButtonStyle::Filled, None)
+                };
+                (style, tooltip)
+            }
+            None => (ButtonStyle::Filled, None),
+        };
+
+        let provider = LanguageModelRegistry::read_global(cx).active_provider();
+
+        let has_configuration_error = configuration_error(cx).is_some();
+        let needs_to_accept_terms = self.show_accept_terms
+            && provider
+                .as_ref()
+                .map_or(false, |provider| provider.must_accept_terms(cx));
+        let disabled = has_configuration_error || needs_to_accept_terms;
+
+        ButtonLike::new("edit_button")
+            .disabled(disabled)
+            .style(style)
+            .when_some(tooltip, |button, tooltip| {
+                button.tooltip(move |_| tooltip.clone())
+            })
+            .layer(ElevationIndex::ModalSurface)
+            .child(Label::new("Suggest Edits"))
+            .children(
+                KeyBinding::for_action_in(&Edit, &focus_handle, cx)
+                    .map(|binding| binding.into_any_element()),
+            )
+            .on_click(move |_event, cx| {
+                focus_handle.dispatch_action(&Edit, cx);
+            })
+    }
+
     fn render_last_error(&self, cx: &mut ViewContext<Self>) -> Option<AnyElement> {
         let last_error = self.last_error.as_ref()?;
 
@@ -3910,6 +3977,7 @@ impl Render for ContextEditor {
             .capture_action(cx.listener(ContextEditor::paste))
             .capture_action(cx.listener(ContextEditor::cycle_message_role))
             .capture_action(cx.listener(ContextEditor::confirm_command))
+            .on_action(cx.listener(ContextEditor::edit))
             .on_action(cx.listener(ContextEditor::assist))
             .on_action(cx.listener(ContextEditor::split))
             .size_full()
@@ -3974,7 +4042,21 @@ impl Render for ContextEditor {
                             h_flex()
                                 .w_full()
                                 .justify_end()
-                                .child(div().child(self.render_send_button(cx))),
+                                .when(
+                                    AssistantSettings::get_global(cx).are_live_diffs_enabled(cx),
+                                    |buttons| {
+                                        buttons
+                                            .items_center()
+                                            .gap_1p5()
+                                            .child(self.render_edit_button(cx))
+                                            .child(
+                                                Label::new("or")
+                                                    .size(LabelSize::Small)
+                                                    .color(Color::Muted),
+                                            )
+                                    },
+                                )
+                                .child(self.render_send_button(cx)),
                         ),
                 ),
             )

crates/assistant/src/context.rs

@@ -66,6 +66,14 @@ impl ContextId {
     }
 }
 
+#[derive(Clone, Copy, Debug)]
+pub enum RequestType {
+    /// Request a normal chat response from the model.
+    Chat,
+    /// Add a preamble to the message, which tells the model to return a structured response that suggests edits.
+    SuggestEdits,
+}
+
 #[derive(Clone, Debug)]
 pub enum ContextOperation {
     InsertMessage {
@@ -1028,7 +1036,7 @@ impl Context {
     }
 
     pub(crate) fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
-        let request = self.to_completion_request(cx);
+        let request = self.to_completion_request(RequestType::SuggestEdits, cx); // Conservatively assume SuggestEdits, since it takes more tokens.
         let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
             return;
         };
@@ -1171,7 +1179,7 @@ impl Context {
         }
 
         let request = {
-            let mut req = self.to_completion_request(cx);
+            let mut req = self.to_completion_request(RequestType::Chat, cx);
             // Skip the last message because it's likely to change and
             // therefore would be a waste to cache.
             req.messages.pop();
@@ -1859,7 +1867,11 @@ impl Context {
         })
     }
 
-    pub fn assist(&mut self, cx: &mut ModelContext<Self>) -> Option<MessageAnchor> {
+    pub fn assist(
+        &mut self,
+        request_type: RequestType,
+        cx: &mut ModelContext<Self>,
+    ) -> Option<MessageAnchor> {
         let model_registry = LanguageModelRegistry::read_global(cx);
         let provider = model_registry.active_provider()?;
         let model = model_registry.active_model()?;
@@ -1872,7 +1884,7 @@ impl Context {
         // Compute which messages to cache, including the last one.
         self.mark_cache_anchors(&model.cache_configuration(), false, cx);
 
-        let mut request = self.to_completion_request(cx);
+        let mut request = self.to_completion_request(request_type, cx);
 
         if cx.has_flag::<ToolUseFeatureFlag>() {
             let tool_registry = ToolRegistry::global(cx);
@@ -2074,7 +2086,11 @@ impl Context {
         Some(user_message)
     }
 
-    pub fn to_completion_request(&self, cx: &AppContext) -> LanguageModelRequest {
+    pub fn to_completion_request(
+        &self,
+        request_type: RequestType,
+        cx: &AppContext,
+    ) -> LanguageModelRequest {
         let buffer = self.buffer.read(cx);
 
         let mut contents = self.contents(cx).peekable();
@@ -2163,6 +2179,25 @@ impl Context {
             completion_request.messages.push(request_message);
         }
 
+        if let RequestType::SuggestEdits = request_type {
+            if let Ok(preamble) = self.prompt_builder.generate_workflow_prompt() {
+                let last_elem_index = completion_request.messages.len();
+
+                completion_request
+                    .messages
+                    .push(LanguageModelRequestMessage {
+                        role: Role::User,
+                        content: vec![MessageContent::Text(preamble)],
+                        cache: false,
+                    });
+
+                // The preamble message should be sent right before the last actual user message.
+                completion_request
+                    .messages
+                    .swap(last_elem_index, last_elem_index.saturating_sub(1));
+            }
+        }
+
         completion_request
     }
 
@@ -2477,7 +2512,7 @@ impl Context {
                 return;
             }
 
-            let mut request = self.to_completion_request(cx);
+            let mut request = self.to_completion_request(RequestType::Chat, cx);
             request.messages.push(LanguageModelRequestMessage {
                 role: Role::User,
                 content: vec![

crates/assistant/src/inline_assistant.rs

@@ -1,7 +1,7 @@
 use crate::{
     assistant_settings::AssistantSettings, humanize_token_count, prompts::PromptBuilder,
     AssistantPanel, AssistantPanelEvent, CharOperation, CycleNextInlineAssist,
-    CyclePreviousInlineAssist, LineDiff, LineOperation, ModelSelector, StreamingDiff,
+    CyclePreviousInlineAssist, LineDiff, LineOperation, ModelSelector, RequestType, StreamingDiff,
 };
 use anyhow::{anyhow, Context as _, Result};
 use client::{telemetry::Telemetry, ErrorExt};
@@ -2234,7 +2234,7 @@ impl InlineAssist {
                     .read(cx)
                     .active_context(cx)?
                     .read(cx)
-                    .to_completion_request(cx),
+                    .to_completion_request(RequestType::Chat, cx),
             )
         } else {
             None

crates/assistant/src/terminal_inline_assistant.rs

@@ -1,6 +1,6 @@
 use crate::{
     humanize_token_count, prompts::PromptBuilder, AssistantPanel, AssistantPanelEvent,
-    ModelSelector, DEFAULT_CONTEXT_LINES,
+    ModelSelector, RequestType, DEFAULT_CONTEXT_LINES,
 };
 use anyhow::{Context as _, Result};
 use client::telemetry::Telemetry;
@@ -251,7 +251,7 @@ impl TerminalInlineAssistant {
                         .read(cx)
                         .active_context(cx)?
                         .read(cx)
-                        .to_completion_request(cx),
+                        .to_completion_request(RequestType::Chat, cx),
                 )
             })
         } else {

crates/ui/src/components/keybinding.rs

@@ -184,7 +184,7 @@ pub struct KeyIcon {
 impl RenderOnce for KeyIcon {
     fn render(self, _cx: &mut WindowContext) -> impl IntoElement {
         Icon::new(self.icon)
-            .size(IconSize::Small)
+            .size(IconSize::XSmall)
             .color(Color::Muted)
     }
 }