Cancel workflow step automatically when moving outside <step> tag (#15842)

Antonio Scandurra created

Also, open edit suggestions automatically as soon as the edit step is
resolved.

Release Notes:

- N/A

Change summary

crates/assistant/src/assistant_panel.rs  | 226 +++++++++++++++++--------
crates/assistant/src/context.rs          |  82 +++++---
crates/assistant/src/inline_assistant.rs |  32 +++
3 files changed, 228 insertions(+), 112 deletions(-)

Detailed changes

crates/assistant/src/assistant_panel.rs 🔗

@@ -8,12 +8,12 @@ use crate::{
         SlashCommandCompletionProvider, SlashCommandRegistry,
     },
     terminal_inline_assistant::TerminalInlineAssistant,
-    Assist, ConfirmCommand, Context, ContextEvent, ContextId, ContextStore, CycleMessageRole,
-    DebugEditSteps, DeployHistory, DeployPromptLibrary, EditSuggestionGroup, InlineAssist,
-    InlineAssistId, InlineAssistant, InsertIntoEditor, MessageStatus, ModelSelector,
+    Assist, CodegenStatus, ConfirmCommand, Context, ContextEvent, ContextId, ContextStore,
+    CycleMessageRole, DebugEditSteps, DeployHistory, DeployPromptLibrary, EditSuggestionGroup,
+    InlineAssist, InlineAssistId, InlineAssistant, InsertIntoEditor, MessageStatus, ModelSelector,
     PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection, RemoteContextMetadata,
-    SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector, WorkflowStep,
-    WorkflowStepEditSuggestions,
+    ResolvedWorkflowStepEditSuggestions, SavedContextMetadata, Split, ToggleFocus,
+    ToggleModelSelector, WorkflowStepEditSuggestions,
 };
 use crate::{ContextStoreEvent, ShowConfiguration};
 use anyhow::{anyhow, Result};
@@ -35,8 +35,8 @@ use gpui::{
     div, percentage, point, Action, Animation, AnimationExt, AnyElement, AnyView, AppContext,
     AsyncWindowContext, ClipboardItem, Context as _, DismissEvent, Empty, Entity, EventEmitter,
     FocusHandle, FocusableView, FontWeight, InteractiveElement, IntoElement, Model, ParentElement,
-    Pixels, Render, SharedString, StatefulInteractiveElement, Styled, Subscription, Task,
-    Transformation, UpdateGlobal, View, ViewContext, VisualContext, WeakView, WindowContext,
+    Pixels, ReadGlobal, Render, SharedString, StatefulInteractiveElement, Styled, Subscription,
+    Task, Transformation, UpdateGlobal, View, ViewContext, VisualContext, WeakView, WindowContext,
 };
 use indexed_docs::IndexedDocsStore;
 use language::{
@@ -1295,10 +1295,15 @@ struct ScrollPosition {
     cursor: Anchor,
 }
 
-struct ActiveEditStep {
-    start: language::Anchor,
+struct StepAssists {
     assist_ids: Vec<InlineAssistId>,
-    editor: Option<WeakView<Editor>>,
+    editor: WeakView<Editor>,
+}
+
+#[derive(Debug, Eq, PartialEq)]
+struct ActiveWorkflowStep {
+    range: Range<language::Anchor>,
+    suggestions: Option<ResolvedWorkflowStepEditSuggestions>,
 }
 
 pub struct ContextEditor {
@@ -1314,7 +1319,8 @@ pub struct ContextEditor {
     pending_slash_command_creases: HashMap<Range<language::Anchor>, CreaseId>,
     pending_slash_command_blocks: HashMap<Range<language::Anchor>, CustomBlockId>,
     _subscriptions: Vec<Subscription>,
-    active_edit_step: Option<ActiveEditStep>,
+    assists_by_step: HashMap<Range<language::Anchor>, StepAssists>,
+    active_workflow_step: Option<ActiveWorkflowStep>,
     assistant_panel: WeakView<AssistantPanel>,
     error_message: Option<SharedString>,
 }
@@ -1372,7 +1378,8 @@ impl ContextEditor {
             pending_slash_command_creases: HashMap::default(),
             pending_slash_command_blocks: HashMap::default(),
             _subscriptions,
-            active_edit_step: None,
+            assists_by_step: HashMap::default(),
+            active_workflow_step: None,
             assistant_panel,
             error_message: None,
         };
@@ -1415,17 +1422,21 @@ impl ContextEditor {
     }
 
     fn apply_edit_step(&mut self, cx: &mut ViewContext<Self>) -> bool {
-        if let Some(step) = self.active_edit_step.as_ref() {
-            let assist_ids = step.assist_ids.clone();
-            cx.window_context().defer(|cx| {
-                InlineAssistant::update_global(cx, |assistant, cx| {
-                    for assist_id in assist_ids {
-                        assistant.start_assist(assist_id, cx);
-                    }
-                })
-            });
+        if let Some(step) = self.active_workflow_step.as_ref() {
+            if let Some(assists) = self.assists_by_step.get(&step.range) {
+                let assist_ids = assists.assist_ids.clone();
+                cx.window_context().defer(|cx| {
+                    InlineAssistant::update_global(cx, |assistant, cx| {
+                        for assist_id in assist_ids {
+                            assistant.start_assist(assist_id, cx);
+                        }
+                    })
+                });
 
-            !step.assist_ids.is_empty()
+                !assists.assist_ids.is_empty()
+            } else {
+                false
+            }
         } else {
             false
         }
@@ -1462,7 +1473,7 @@ impl ContextEditor {
 
     fn debug_edit_steps(&mut self, _: &DebugEditSteps, cx: &mut ViewContext<Self>) {
         let mut output = String::new();
-        for (i, step) in self.context.read(cx).edit_steps().iter().enumerate() {
+        for (i, step) in self.context.read(cx).workflow_steps().iter().enumerate() {
             output.push_str(&format!("Step {}:\n", i + 1));
             output.push_str(&format!(
                 "Content: {}\n",
@@ -1474,10 +1485,10 @@ impl ContextEditor {
                     .collect::<String>()
             ));
             match &step.edit_suggestions {
-                WorkflowStepEditSuggestions::Resolved {
+                WorkflowStepEditSuggestions::Resolved(ResolvedWorkflowStepEditSuggestions {
                     title,
                     edit_suggestions,
-                } => {
+                }) => {
                     output.push_str("Resolution:\n");
                     output.push_str(&format!("  {:?}\n", title));
                     output.push_str(&format!("  {:?}\n", edit_suggestions));
@@ -1629,7 +1640,8 @@ impl ContextEditor {
                     context.save(Some(Duration::from_millis(500)), self.fs.clone(), cx);
                 });
             }
-            ContextEvent::EditStepsChanged => {
+            ContextEvent::WorkflowStepsChanged => {
+                self.update_active_workflow_step(cx);
                 cx.notify();
             }
             ContextEvent::SummaryChanged => {
@@ -1902,57 +1914,114 @@ impl ContextEditor {
     }
 
     fn update_active_workflow_step(&mut self, cx: &mut ViewContext<Self>) {
-        if self
-            .workflow_step_for_cursor(cx)
-            .map(|step| step.tagged_range.start)
-            != self.active_edit_step.as_ref().map(|step| step.start)
-        {
-            if let Some(old_active_edit_step) = self.active_edit_step.take() {
-                if let Some(editor) = old_active_edit_step
-                    .editor
-                    .and_then(|editor| editor.upgrade())
-                {
-                    self.workspace
-                        .update(cx, |workspace, cx| {
-                            if let Some(pane) = workspace.pane_for(&editor) {
-                                pane.update(cx, |pane, cx| {
-                                    let item_id = editor.entity_id();
-                                    if pane.is_active_preview_item(item_id) {
-                                        pane.close_item_by_id(item_id, SaveIntent::Skip, cx)
-                                            .detach_and_log_err(cx);
-                                    }
-                                });
-                            }
-                        })
-                        .ok();
-                }
+        let new_step = self
+            .workflow_step_range_for_cursor(cx)
+            .as_ref()
+            .and_then(|step_range| {
+                let workflow_step = self
+                    .context
+                    .read(cx)
+                    .workflow_step_for_range(step_range.clone())?;
+                Some(ActiveWorkflowStep {
+                    range: workflow_step.tagged_range.clone(),
+                    suggestions: workflow_step.edit_suggestions.as_resolved().cloned(),
+                })
+            });
+        if new_step.as_ref() != self.active_workflow_step.as_ref() {
+            if let Some(old_step) = self.active_workflow_step.take() {
+                self.cancel_workflow_step_if_idle(old_step.range, cx);
             }
 
-            if let Some(new_active_step) = self.workflow_step_for_cursor(cx) {
-                let start = new_active_step.tagged_range.start;
+            if let Some(new_step) = new_step {
+                self.activate_workflow_step(new_step, cx);
+            }
+        }
+    }
 
-                let mut editor = None;
-                let mut assist_ids = Vec::new();
-                if let WorkflowStepEditSuggestions::Resolved {
-                    title,
-                    edit_suggestions,
-                } = &new_active_step.edit_suggestions
-                {
-                    if let Some((opened_editor, inline_assist_ids)) =
-                        self.suggest_edits(title.clone(), edit_suggestions.clone(), cx)
-                    {
-                        editor = Some(opened_editor.downgrade());
-                        assist_ids = inline_assist_ids;
+    fn cancel_workflow_step_if_idle(
+        &mut self,
+        step_range: Range<language::Anchor>,
+        cx: &mut ViewContext<Self>,
+    ) {
+        let Some(step_assists) = self.assists_by_step.get_mut(&step_range) else {
+            return;
+        };
+        let Some(editor) = step_assists.editor.upgrade() else {
+            self.assists_by_step.remove(&step_range);
+            return;
+        };
+
+        InlineAssistant::update_global(cx, |assistant, cx| {
+            step_assists.assist_ids.retain(|assist_id| {
+                match assistant.status_for_assist(*assist_id, cx) {
+                    Some(CodegenStatus::Idle) | None => {
+                        assistant.finish_assist(*assist_id, true, cx);
+                        false
                     }
+                    _ => true,
                 }
+            });
+        });
 
-                self.active_edit_step = Some(ActiveEditStep {
-                    start,
-                    assist_ids,
-                    editor,
-                });
+        if step_assists.assist_ids.is_empty() {
+            self.assists_by_step.remove(&step_range);
+            self.workspace
+                .update(cx, |workspace, cx| {
+                    if let Some(pane) = workspace.pane_for(&editor) {
+                        pane.update(cx, |pane, cx| {
+                            let item_id = editor.entity_id();
+                            if pane.is_active_preview_item(item_id) {
+                                pane.close_item_by_id(item_id, SaveIntent::Skip, cx)
+                                    .detach_and_log_err(cx);
+                            }
+                        });
+                    }
+                })
+                .ok();
+        }
+    }
+
+    fn activate_workflow_step(&mut self, step: ActiveWorkflowStep, cx: &mut ViewContext<Self>) {
+        if let Some(step_assists) = self.assists_by_step.get(&step.range) {
+            if let Some(editor) = step_assists.editor.upgrade() {
+                for assist_id in &step_assists.assist_ids {
+                    match InlineAssistant::global(cx).status_for_assist(*assist_id, cx) {
+                        Some(CodegenStatus::Idle) | None => {}
+                        _ => {
+                            self.workspace
+                                .update(cx, |workspace, cx| {
+                                    workspace.activate_item(&editor, false, false, cx);
+                                })
+                                .ok();
+                            InlineAssistant::update_global(cx, |assistant, cx| {
+                                assistant.scroll_to_assist(*assist_id, cx)
+                            });
+                            return;
+                        }
+                    }
+                }
+            }
+        }
+
+        if let Some(ResolvedWorkflowStepEditSuggestions {
+            title,
+            edit_suggestions,
+        }) = step.suggestions.as_ref()
+        {
+            if let Some((editor, assist_ids)) =
+                self.suggest_edits(title.clone(), edit_suggestions.clone(), cx)
+            {
+                self.assists_by_step.insert(
+                    step.range.clone(),
+                    StepAssists {
+                        assist_ids,
+                        editor: editor.downgrade(),
+                    },
+                );
             }
         }
+
+        self.active_workflow_step = Some(step);
     }
 
     fn suggest_edits(
@@ -2436,11 +2505,14 @@ impl ContextEditor {
 
     fn render_send_button(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
         let focus_handle = self.focus_handle(cx).clone();
-        let button_text = match self.workflow_step_for_cursor(cx) {
-            Some(edit_step) => match &edit_step.edit_suggestions {
-                WorkflowStepEditSuggestions::Pending(_) => "Computing Changes...",
-                WorkflowStepEditSuggestions::Resolved { .. } => "Apply Changes",
-            },
+        let button_text = match self.active_workflow_step.as_ref() {
+            Some(step) => {
+                if step.suggestions.is_none() {
+                    "Computing Changes..."
+                } else {
+                    "Apply Changes"
+                }
+            }
             None => "Send",
         };
 
@@ -2482,7 +2554,7 @@ impl ContextEditor {
             })
     }
 
-    fn workflow_step_for_cursor<'a>(&'a self, cx: &'a AppContext) -> Option<&'a WorkflowStep> {
+    fn workflow_step_range_for_cursor(&self, cx: &AppContext) -> Option<Range<language::Anchor>> {
         let newest_cursor = self
             .editor
             .read(cx)
@@ -2493,7 +2565,7 @@ impl ContextEditor {
         let context = self.context.read(cx);
         let buffer = context.buffer().read(cx);
 
-        let edit_steps = context.edit_steps();
+        let edit_steps = context.workflow_steps();
         edit_steps
             .binary_search_by(|step| {
                 let step_range = step.tagged_range.clone();
@@ -2506,7 +2578,7 @@ impl ContextEditor {
                 }
             })
             .ok()
-            .map(|index| &edit_steps[index])
+            .map(|index| edit_steps[index].tagged_range.clone())
     }
 }
 

crates/assistant/src/context.rs 🔗

@@ -284,7 +284,7 @@ pub enum ContextEvent {
     AssistError(String),
     MessagesEdited,
     SummaryChanged,
-    EditStepsChanged,
+    WorkflowStepsChanged,
     StreamedCompletion,
     PendingSlashCommandsUpdated {
         removed: Vec<Range<language::Anchor>>,
@@ -351,21 +351,33 @@ pub struct WorkflowStep {
     pub edit_suggestions: WorkflowStepEditSuggestions,
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct ResolvedWorkflowStepEditSuggestions {
+    pub title: String,
+    pub edit_suggestions: HashMap<Model<Buffer>, Vec<EditSuggestionGroup>>,
+}
+
 pub enum WorkflowStepEditSuggestions {
     Pending(Task<Option<()>>),
-    Resolved {
-        title: String,
-        edit_suggestions: HashMap<Model<Buffer>, Vec<EditSuggestionGroup>>,
-    },
+    Resolved(ResolvedWorkflowStepEditSuggestions),
 }
 
-#[derive(Clone, Debug)]
+impl WorkflowStepEditSuggestions {
+    pub fn as_resolved(&self) -> Option<&ResolvedWorkflowStepEditSuggestions> {
+        match self {
+            WorkflowStepEditSuggestions::Resolved(suggestions) => Some(suggestions),
+            WorkflowStepEditSuggestions::Pending(_) => None,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub struct EditSuggestionGroup {
     pub context_range: Range<language::Anchor>,
     pub suggestions: Vec<EditSuggestion>,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub enum EditSuggestion {
     Update {
         range: Range<language::Anchor>,
@@ -561,10 +573,10 @@ impl Debug for WorkflowStepEditSuggestions {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             WorkflowStepEditSuggestions::Pending(_) => write!(f, "EditStepOperations::Pending"),
-            WorkflowStepEditSuggestions::Resolved {
+            WorkflowStepEditSuggestions::Resolved(ResolvedWorkflowStepEditSuggestions {
                 title,
                 edit_suggestions,
-            } => f
+            }) => f
                 .debug_struct("EditStepOperations::Parsed")
                 .field("title", title)
                 .field("edit_suggestions", edit_suggestions)
@@ -597,7 +609,7 @@ pub struct Context {
     _subscriptions: Vec<Subscription>,
     telemetry: Option<Arc<Telemetry>>,
     language_registry: Arc<LanguageRegistry>,
-    edit_steps: Vec<WorkflowStep>,
+    workflow_steps: Vec<WorkflowStep>,
     project: Option<Model<Project>>,
 }
 
@@ -667,7 +679,7 @@ impl Context {
             telemetry,
             project,
             language_registry,
-            edit_steps: Vec::new(),
+            workflow_steps: Vec::new(),
         };
 
         let first_message_id = MessageId(clock::Lamport {
@@ -987,8 +999,14 @@ impl Context {
         self.summary.as_ref()
     }
 
-    pub fn edit_steps(&self) -> &[WorkflowStep] {
-        &self.edit_steps
+    pub fn workflow_steps(&self) -> &[WorkflowStep] {
+        &self.workflow_steps
+    }
+
+    pub fn workflow_step_for_range(&self, range: Range<language::Anchor>) -> Option<&WorkflowStep> {
+        self.workflow_steps
+            .iter()
+            .find(|step| step.tagged_range == range)
     }
 
     pub fn pending_slash_commands(&self) -> &[PendingSlashCommand] {
@@ -1133,12 +1151,12 @@ impl Context {
 
     fn prune_invalid_edit_steps(&mut self, cx: &mut ModelContext<Self>) {
         let buffer = self.buffer.read(cx);
-        let prev_len = self.edit_steps.len();
-        self.edit_steps.retain(|step| {
+        let prev_len = self.workflow_steps.len();
+        self.workflow_steps.retain(|step| {
             step.tagged_range.start.is_valid(buffer) && step.tagged_range.end.is_valid(buffer)
         });
-        if self.edit_steps.len() != prev_len {
-            cx.emit(ContextEvent::EditStepsChanged);
+        if self.workflow_steps.len() != prev_len {
+            cx.emit(ContextEvent::WorkflowStepsChanged);
             cx.notify();
         }
     }
@@ -1174,7 +1192,7 @@ impl Context {
 
                     // Check if a step with the same range already exists
                     let existing_step_index = self
-                        .edit_steps
+                        .workflow_steps
                         .binary_search_by(|probe| probe.tagged_range.cmp(&tagged_range, &buffer));
 
                     if let Err(ix) = existing_step_index {
@@ -1202,10 +1220,10 @@ impl Context {
 
         // Insert new steps and generate their corresponding tasks
         for (index, step) in new_edit_steps.into_iter().rev() {
-            self.edit_steps.insert(index, step);
+            self.workflow_steps.insert(index, step);
         }
 
-        cx.emit(ContextEvent::EditStepsChanged);
+        cx.emit(ContextEvent::WorkflowStepsChanged);
         cx.notify();
     }
 
@@ -1321,17 +1339,19 @@ impl Context {
 
                 this.update(&mut cx, |this, cx| {
                     let step_index = this
-                        .edit_steps
+                        .workflow_steps
                         .binary_search_by(|step| {
                             step.tagged_range.cmp(&tagged_range, this.buffer.read(cx))
                         })
                         .map_err(|_| anyhow!("edit step not found"))?;
-                    if let Some(edit_step) = this.edit_steps.get_mut(step_index) {
-                        edit_step.edit_suggestions = WorkflowStepEditSuggestions::Resolved {
-                            title: step_suggestions.step_title,
-                            edit_suggestions: suggestion_groups_by_buffer,
-                        };
-                        cx.emit(ContextEvent::EditStepsChanged);
+                    if let Some(edit_step) = this.workflow_steps.get_mut(step_index) {
+                        edit_step.edit_suggestions = WorkflowStepEditSuggestions::Resolved(
+                            ResolvedWorkflowStepEditSuggestions {
+                                title: step_suggestions.step_title,
+                                edit_suggestions: suggestion_groups_by_buffer,
+                            },
+                        );
+                        cx.emit(ContextEvent::WorkflowStepsChanged);
                     }
                     anyhow::Ok(())
                 })?
@@ -2959,7 +2979,7 @@ mod tests {
         // Verify that the edit steps were parsed correctly
         context.read_with(cx, |context, cx| {
             assert_eq!(
-                edit_steps(context, cx),
+                workflow_steps(context, cx),
                 vec![
                     (
                         Point::new(response_start_row + 2, 0)
@@ -2998,7 +3018,7 @@ mod tests {
         // Verify that the last edit step is not pending anymore.
         context.read_with(cx, |context, cx| {
             assert_eq!(
-                edit_steps(context, cx),
+                workflow_steps(context, cx),
                 vec![
                     (
                         Point::new(response_start_row + 2, 0)
@@ -3020,12 +3040,12 @@ mod tests {
             Resolved,
         }
 
-        fn edit_steps(
+        fn workflow_steps(
             context: &Context,
             cx: &AppContext,
         ) -> Vec<(Range<Point>, WorkflowStepEditSuggestionStatus)> {
             context
-                .edit_steps
+                .workflow_steps
                 .iter()
                 .map(|step| {
                     let buffer = context.buffer.read(cx);

crates/assistant/src/inline_assistant.rs 🔗

@@ -604,7 +604,7 @@ impl InlineAssistant {
         }
     }
 
-    fn finish_assist(&mut self, assist_id: InlineAssistId, undo: bool, cx: &mut WindowContext) {
+    pub fn finish_assist(&mut self, assist_id: InlineAssistId, undo: bool, cx: &mut WindowContext) {
         if let Some(assist) = self.assists.get(&assist_id) {
             let assist_group_id = assist.group_id;
             if self.assist_groups[&assist_group_id].linked {
@@ -715,8 +715,7 @@ impl InlineAssistant {
     }
 
     fn focus_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
-        let assist = &self.assists[&assist_id];
-        let Some(editor) = assist.editor.upgrade() else {
+        let Some(assist) = self.assists.get(&assist_id) else {
             return;
         };
 
@@ -729,6 +728,17 @@ impl InlineAssistant {
             });
         }
 
+        self.scroll_to_assist(assist_id, cx);
+    }
+
+    pub fn scroll_to_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
+        let Some(assist) = self.assists.get(&assist_id) else {
+            return;
+        };
+        let Some(editor) = assist.editor.upgrade() else {
+            return;
+        };
+
         let position = assist.range.start;
         editor.update(cx, |editor, cx| {
             editor.change_selections(None, cx, |selections| {
@@ -844,6 +854,20 @@ impl InlineAssistant {
         assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
     }
 
+    pub fn status_for_assist(
+        &self,
+        assist_id: InlineAssistId,
+        cx: &WindowContext,
+    ) -> Option<CodegenStatus> {
+        let assist = self.assists.get(&assist_id)?;
+        match &assist.codegen.read(cx).status {
+            CodegenStatus::Idle => Some(CodegenStatus::Idle),
+            CodegenStatus::Pending => Some(CodegenStatus::Pending),
+            CodegenStatus::Done => Some(CodegenStatus::Done),
+            CodegenStatus::Error(error) => Some(CodegenStatus::Error(anyhow!("{:?}", error))),
+        }
+    }
+
     fn update_editor_highlights(&self, editor: &View<Editor>, cx: &mut WindowContext) {
         let mut gutter_pending_ranges = Vec::new();
         let mut gutter_transformed_ranges = Vec::new();
@@ -2000,7 +2024,7 @@ pub struct Codegen {
     _subscription: gpui::Subscription,
 }
 
-enum CodegenStatus {
+pub enum CodegenStatus {
     Idle,
     Pending,
     Done,