assistant: Polish /workflow and steps UI (#15936)

Fixes #15923
Release Notes:

- Assistant workflow steps can now be applied and reverted directly from
within the assistant panel.

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
Co-authored-by: Antonio <antonio@zed.dev>

assets/icons/undo.svg

@@ -0,0 +1 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-undo"><path d="M3 7v6h6"/><path d="M21 17a9 9 0 0 0-9-9 9 9 0 0 0-6 2.3L3 13"/></svg>

crates/assistant/src/assistant.rs

@@ -53,7 +53,7 @@ actions!(
         DeployPromptLibrary,
         ConfirmCommand,
         ToggleModelSelector,
-        DebugEditSteps
+        DebugWorkflowSteps
     ]
 );
 

crates/assistant/src/context.rs

@@ -284,7 +284,8 @@ pub enum ContextEvent {
     AssistError(String),
     MessagesEdited,
     SummaryChanged,
-    WorkflowStepsChanged,
+    WorkflowStepsRemoved(Vec<Range<language::Anchor>>),
+    WorkflowStepUpdated(Range<language::Anchor>),
     StreamedCompletion,
     PendingSlashCommandsUpdated {
         removed: Vec<Range<language::Anchor>>,
@@ -360,22 +361,17 @@ pub struct ResolvedWorkflowStep {
 pub enum WorkflowStepStatus {
     Pending(Task<Option<()>>),
     Resolved(ResolvedWorkflowStep),
+    Error(Arc<anyhow::Error>),
 }
 
 impl WorkflowStepStatus {
-    pub fn as_resolved(&self) -> Option<&ResolvedWorkflowStep> {
+    pub fn into_resolved(&self) -> Option<Result<ResolvedWorkflowStep, Arc<anyhow::Error>>> {
         match self {
-            WorkflowStepStatus::Resolved(suggestions) => Some(suggestions),
+            WorkflowStepStatus::Resolved(resolved) => Some(Ok(resolved.clone())),
+            WorkflowStepStatus::Error(error) => Some(Err(error.clone())),
             WorkflowStepStatus::Pending(_) => None,
         }
     }
-
-    pub fn is_resolved(&self) -> bool {
-        match self {
-            WorkflowStepStatus::Resolved(_) => true,
-            WorkflowStepStatus::Pending(_) => false,
-        }
-    }
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -583,12 +579,16 @@ impl WorkflowSuggestion {
 impl Debug for WorkflowStepStatus {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            WorkflowStepStatus::Pending(_) => write!(f, "EditStepOperations::Pending"),
+            WorkflowStepStatus::Pending(_) => write!(f, "WorkflowStepStatus::Pending"),
             WorkflowStepStatus::Resolved(ResolvedWorkflowStep { title, suggestions }) => f
-                .debug_struct("EditStepOperations::Parsed")
+                .debug_struct("WorkflowStepStatus::Resolved")
                 .field("title", title)
                 .field("suggestions", suggestions)
                 .finish(),
+            WorkflowStepStatus::Error(error) => f
+                .debug_tuple("WorkflowStepStatus::Error")
+                .field(error)
+                .finish(),
         }
     }
 }
@@ -1058,7 +1058,7 @@ impl Context {
             language::Event::Edited => {
                 self.count_remaining_tokens(cx);
                 self.reparse_slash_commands(cx);
-                self.prune_invalid_edit_steps(cx);
+                self.prune_invalid_workflow_steps(cx);
                 cx.emit(ContextEvent::MessagesEdited);
             }
             _ => {}
@@ -1165,46 +1165,59 @@ impl Context {
         }
     }
 
-    fn prune_invalid_edit_steps(&mut self, cx: &mut ModelContext<Self>) {
+    fn prune_invalid_workflow_steps(&mut self, cx: &mut ModelContext<Self>) {
         let buffer = self.buffer.read(cx);
         let prev_len = self.workflow_steps.len();
+        let mut removed = Vec::new();
         self.workflow_steps.retain(|step| {
-            step.tagged_range.start.is_valid(buffer) && step.tagged_range.end.is_valid(buffer)
+            if step.tagged_range.start.is_valid(buffer) && step.tagged_range.end.is_valid(buffer) {
+                true
+            } else {
+                removed.push(step.tagged_range.clone());
+                false
+            }
         });
         if self.workflow_steps.len() != prev_len {
-            cx.emit(ContextEvent::WorkflowStepsChanged);
+            cx.emit(ContextEvent::WorkflowStepsRemoved(removed));
             cx.notify();
         }
     }
 
-    fn parse_edit_steps_in_range(
+    fn parse_workflow_steps_in_range(
         &mut self,
         range: Range<usize>,
         project: Model<Project>,
         cx: &mut ModelContext<Self>,
     ) {
         let mut new_edit_steps = Vec::new();
+        let mut edits = Vec::new();
 
         let buffer = self.buffer.read(cx).snapshot();
         let mut message_lines = buffer.as_rope().chunks_in_range(range).lines();
         let mut in_step = false;
-        let mut step_start = 0;
+        let mut step_open_tag_start_ix = 0;
         let mut line_start_offset = message_lines.offset();
 
         while let Some(line) = message_lines.next() {
             if let Some(step_start_index) = line.find("<step>") {
                 if !in_step {
                     in_step = true;
-                    step_start = line_start_offset + step_start_index;
+                    step_open_tag_start_ix = line_start_offset + step_start_index;
                 }
             }
 
             if let Some(step_end_index) = line.find("</step>") {
                 if in_step {
-                    let start_anchor = buffer.anchor_after(step_start);
+                    let step_open_tag_end_ix = step_open_tag_start_ix + "<step>".len();
-                    let end_anchor =
+                    let mut step_end_tag_start_ix = line_start_offset + step_end_index;
-                        buffer.anchor_before(line_start_offset + step_end_index + "</step>".len());
+                    let step_end_tag_end_ix = step_end_tag_start_ix + "</step>".len();
-                    let tagged_range = start_anchor..end_anchor;
+                    if buffer.reversed_chars_at(step_end_tag_start_ix).next() == Some('\n') {
+                        step_end_tag_start_ix -= 1;
+                    }
+                    edits.push((step_open_tag_start_ix..step_open_tag_end_ix, ""));
+                    edits.push((step_end_tag_start_ix..step_end_tag_end_ix, ""));
+                    let tagged_range = buffer.anchor_after(step_open_tag_end_ix)
+                        ..buffer.anchor_before(step_end_tag_start_ix);
 
                     // Check if a step with the same range already exists
                     let existing_step_index = self
@@ -1212,14 +1225,11 @@ impl Context {
                         .binary_search_by(|probe| probe.tagged_range.cmp(&tagged_range, &buffer));
 
                     if let Err(ix) = existing_step_index {
-                        // Step doesn't exist, so add it
-                        let task =
-                            self.resolve_workflow_step(tagged_range.clone(), project.clone(), cx);
                         new_edit_steps.push((
                             ix,
                             WorkflowStep {
                                 tagged_range,
-                                status: WorkflowStepStatus::Pending(task),
+                                status: WorkflowStepStatus::Pending(Task::ready(None)),
                             },
                         ));
                     }
@@ -1231,144 +1241,176 @@ impl Context {
             line_start_offset = message_lines.offset();
         }
 
-        // Insert new steps and generate their corresponding tasks
+        let mut updated = Vec::new();
         for (index, step) in new_edit_steps.into_iter().rev() {
+            let step_range = step.tagged_range.clone();
+            updated.push(step_range.clone());
             self.workflow_steps.insert(index, step);
+            self.resolve_workflow_step(step_range, project.clone(), cx);
         }
-
+        self.buffer
-        cx.emit(ContextEvent::WorkflowStepsChanged);
+            .update(cx, |buffer, cx| buffer.edit(edits, None, cx));
-        cx.notify();
     }
 
-    fn resolve_workflow_step(
+    pub fn resolve_workflow_step(
-        &self,
+        &mut self,
         tagged_range: Range<language::Anchor>,
         project: Model<Project>,
         cx: &mut ModelContext<Self>,
-    ) -> Task<Option<()>> {
+    ) {
-        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
+        let Ok(step_index) = self
-            return Task::ready(Err(anyhow!("no active model")).log_err());
+            .workflow_steps
+            .binary_search_by(|step| step.tagged_range.cmp(&tagged_range, self.buffer.read(cx)))
+        else {
+            return;
         };
 
         let mut request = self.to_completion_request(cx);
-        let step_text = self
+        let Some(edit_step) = self.workflow_steps.get_mut(step_index) else {
-            .buffer
+            return;
-            .read(cx)
+        };
-            .text_for_range(tagged_range.clone())
-            .collect::<String>();
 
-        cx.spawn(|this, mut cx| {
+        if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
-            async move {
+            let step_text = self
-                let mut prompt = this.update(&mut cx, |this, _| {
+                .buffer
-                    this.prompt_builder.generate_step_resolution_prompt()
+                .read(cx)
-                })??;
+                .text_for_range(tagged_range.clone())
-                prompt.push_str(&step_text);
+                .collect::<String>();
 
-                request.messages.push(LanguageModelRequestMessage {
+            let tagged_range = tagged_range.clone();
-                    role: Role::User,
+            edit_step.status = WorkflowStepStatus::Pending(cx.spawn(|this, mut cx| {
-                    content: prompt,
+                async move {
-                });
+                    let result = async {
+                        let mut prompt = this.update(&mut cx, |this, _| {
+                            this.prompt_builder.generate_step_resolution_prompt()
+                        })??;
+                        prompt.push_str(&step_text);
 
-                // Invoke the model to get its edit suggestions for this workflow step.
+                        request.messages.push(LanguageModelRequestMessage {
-                let resolution = model
+                            role: Role::User,
-                    .use_tool::<tool::WorkflowStepResolution>(request, &cx)
+                            content: prompt,
-                    .await?;
-
-                // Translate the parsed suggestions to our internal types, which anchor the suggestions to locations in the code.
-                let suggestion_tasks: Vec<_> = resolution
-                    .suggestions
-                    .iter()
-                    .map(|suggestion| suggestion.resolve(project.clone(), cx.clone()))
-                    .collect();
-
-                // Expand the context ranges of each suggestion and group suggestions with overlapping context ranges.
-                let suggestions = future::join_all(suggestion_tasks)
-                    .await
-                    .into_iter()
-                    .filter_map(|task| task.log_err())
-                    .collect::<Vec<_>>();
-
-                let mut suggestions_by_buffer = HashMap::default();
-                for (buffer, suggestion) in suggestions {
-                    suggestions_by_buffer
-                        .entry(buffer)
-                        .or_insert_with(Vec::new)
-                        .push(suggestion);
-                }
-
-                let mut suggestion_groups_by_buffer = HashMap::default();
-                for (buffer, mut suggestions) in suggestions_by_buffer {
-                    let mut suggestion_groups = Vec::<WorkflowSuggestionGroup>::new();
-                    let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
-                    // Sort suggestions by their range so that earlier, larger ranges come first
-                    suggestions.sort_by(|a, b| a.range().cmp(&b.range(), &snapshot));
-
-                    // Merge overlapping suggestions
-                    suggestions.dedup_by(|a, b| b.try_merge(&a, &snapshot));
-
-                    // Create context ranges for each suggestion
-                    for suggestion in suggestions {
-                        let context_range = {
-                            let suggestion_point_range = suggestion.range().to_point(&snapshot);
-                            let start_row = suggestion_point_range.start.row.saturating_sub(5);
-                            let end_row = cmp::min(
-                                suggestion_point_range.end.row + 5,
-                                snapshot.max_point().row,
-                            );
-                            let start = snapshot.anchor_before(Point::new(start_row, 0));
-                            let end = snapshot
-                                .anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
-                            start..end
-                        };
-
-                        if let Some(last_group) = suggestion_groups.last_mut() {
-                            if last_group
-                                .context_range
-                                .end
-                                .cmp(&context_range.start, &snapshot)
-                                .is_ge()
-                            {
-                                // Merge with the previous group if context ranges overlap
-                                last_group.context_range.end = context_range.end;
-                                last_group.suggestions.push(suggestion);
-                            } else {
-                                // Create a new group
-                                suggestion_groups.push(WorkflowSuggestionGroup {
-                                    context_range,
-                                    suggestions: vec![suggestion],
-                                });
-                            }
-                        } else {
-                            // Create the first group
-                            suggestion_groups.push(WorkflowSuggestionGroup {
-                                context_range,
-                                suggestions: vec![suggestion],
-                            });
-                        }
-                    }
-
-                    suggestion_groups_by_buffer.insert(buffer, suggestion_groups);
-                }
-
-                this.update(&mut cx, |this, cx| {
-                    let step_index = this
-                        .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.workflow_steps.get_mut(step_index) {
-                        edit_step.status = WorkflowStepStatus::Resolved(ResolvedWorkflowStep {
-                            title: resolution.step_title,
-                            suggestions: suggestion_groups_by_buffer,
                         });
-                        cx.emit(ContextEvent::WorkflowStepsChanged);
+
-                    }
+                        // Invoke the model to get its edit suggestions for this workflow step.
-                    anyhow::Ok(())
+                        let resolution = model
-                })?
+                            .use_tool::<tool::WorkflowStepResolution>(request, &cx)
-            }
+                            .await?;
-            .log_err()
+
-        })
+                        // Translate the parsed suggestions to our internal types, which anchor the suggestions to locations in the code.
+                        let suggestion_tasks: Vec<_> = resolution
+                            .suggestions
+                            .iter()
+                            .map(|suggestion| suggestion.resolve(project.clone(), cx.clone()))
+                            .collect();
+
+                        // Expand the context ranges of each suggestion and group suggestions with overlapping context ranges.
+                        let suggestions = future::join_all(suggestion_tasks)
+                            .await
+                            .into_iter()
+                            .filter_map(|task| task.log_err())
+                            .collect::<Vec<_>>();
+
+                        let mut suggestions_by_buffer = HashMap::default();
+                        for (buffer, suggestion) in suggestions {
+                            suggestions_by_buffer
+                                .entry(buffer)
+                                .or_insert_with(Vec::new)
+                                .push(suggestion);
+                        }
+
+                        let mut suggestion_groups_by_buffer = HashMap::default();
+                        for (buffer, mut suggestions) in suggestions_by_buffer {
+                            let mut suggestion_groups = Vec::<WorkflowSuggestionGroup>::new();
+                            let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
+                            // Sort suggestions by their range so that earlier, larger ranges come first
+                            suggestions.sort_by(|a, b| a.range().cmp(&b.range(), &snapshot));
+
+                            // Merge overlapping suggestions
+                            suggestions.dedup_by(|a, b| b.try_merge(&a, &snapshot));
+
+                            // Create context ranges for each suggestion
+                            for suggestion in suggestions {
+                                let context_range = {
+                                    let suggestion_point_range =
+                                        suggestion.range().to_point(&snapshot);
+                                    let start_row =
+                                        suggestion_point_range.start.row.saturating_sub(5);
+                                    let end_row = cmp::min(
+                                        suggestion_point_range.end.row + 5,
+                                        snapshot.max_point().row,
+                                    );
+                                    let start = snapshot.anchor_before(Point::new(start_row, 0));
+                                    let end = snapshot.anchor_after(Point::new(
+                                        end_row,
+                                        snapshot.line_len(end_row),
+                                    ));
+                                    start..end
+                                };
+
+                                if let Some(last_group) = suggestion_groups.last_mut() {
+                                    if last_group
+                                        .context_range
+                                        .end
+                                        .cmp(&context_range.start, &snapshot)
+                                        .is_ge()
+                                    {
+                                        // Merge with the previous group if context ranges overlap
+                                        last_group.context_range.end = context_range.end;
+                                        last_group.suggestions.push(suggestion);
+                                    } else {
+                                        // Create a new group
+                                        suggestion_groups.push(WorkflowSuggestionGroup {
+                                            context_range,
+                                            suggestions: vec![suggestion],
+                                        });
+                                    }
+                                } else {
+                                    // Create the first group
+                                    suggestion_groups.push(WorkflowSuggestionGroup {
+                                        context_range,
+                                        suggestions: vec![suggestion],
+                                    });
+                                }
+                            }
+
+                            suggestion_groups_by_buffer.insert(buffer, suggestion_groups);
+                        }
+
+                        Ok((resolution.step_title, suggestion_groups_by_buffer))
+                    };
+
+                    let result = result.await;
+                    this.update(&mut cx, |this, cx| {
+                        let step_index = this
+                            .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.workflow_steps.get_mut(step_index) {
+                            edit_step.status = match result {
+                                Ok((title, suggestions)) => {
+                                    WorkflowStepStatus::Resolved(ResolvedWorkflowStep {
+                                        title,
+                                        suggestions,
+                                    })
+                                }
+                                Err(error) => WorkflowStepStatus::Error(Arc::new(error)),
+                            };
+                            cx.emit(ContextEvent::WorkflowStepUpdated(tagged_range));
+                            cx.notify();
+                        }
+                        anyhow::Ok(())
+                    })?
+                }
+                .log_err()
+            }));
+        } else {
+            edit_step.status = WorkflowStepStatus::Error(Arc::new(anyhow!("no active model")));
+        }
+
+        cx.emit(ContextEvent::WorkflowStepUpdated(tagged_range));
+        cx.notify();
     }
 
     pub fn pending_command_for_position(
@@ -1587,7 +1629,7 @@ impl Context {
                                 message_start_offset..message_new_end_offset
                             });
                             if let Some(project) = this.project.clone() {
-                                this.parse_edit_steps_in_range(message_range, project, cx);
+                                this.parse_workflow_steps_in_range(message_range, project, cx);
                             }
                             cx.emit(ContextEvent::StreamedCompletion);
 
@@ -3011,13 +3053,13 @@ mod tests {
                 vec![
                     (
                         Point::new(response_start_row + 2, 0)
-                            ..Point::new(response_start_row + 14, 7),
+                            ..Point::new(response_start_row + 13, 3),
-                        WorkflowStepEditSuggestionStatus::Pending
+                        WorkflowStepTestStatus::Pending
                     ),
                     (
-                        Point::new(response_start_row + 16, 0)
+                        Point::new(response_start_row + 15, 0)
-                            ..Point::new(response_start_row + 28, 7),
+                            ..Point::new(response_start_row + 26, 3),
-                        WorkflowStepEditSuggestionStatus::Pending
+                        WorkflowStepTestStatus::Pending
                     ),
                 ]
             );
@@ -3041,45 +3083,45 @@ mod tests {
         // Wait for tool use to be processed.
         cx.run_until_parked();
 
-        // Verify that the last edit step is not pending anymore.
+        // Verify that the first edit step is not pending anymore.
         context.read_with(cx, |context, cx| {
             assert_eq!(
                 workflow_steps(context, cx),
                 vec![
                     (
                         Point::new(response_start_row + 2, 0)
-                            ..Point::new(response_start_row + 14, 7),
+                            ..Point::new(response_start_row + 13, 3),
-                        WorkflowStepEditSuggestionStatus::Pending
+                        WorkflowStepTestStatus::Resolved
                     ),
                     (
-                        Point::new(response_start_row + 16, 0)
+                        Point::new(response_start_row + 15, 0)
-                            ..Point::new(response_start_row + 28, 7),
+                            ..Point::new(response_start_row + 26, 3),
-                        WorkflowStepEditSuggestionStatus::Resolved
+                        WorkflowStepTestStatus::Pending
                     ),
                 ]
             );
         });
 
         #[derive(Copy, Clone, Debug, Eq, PartialEq)]
-        enum WorkflowStepEditSuggestionStatus {
+        enum WorkflowStepTestStatus {
             Pending,
             Resolved,
+            Error,
         }
 
         fn workflow_steps(
             context: &Context,
             cx: &AppContext,
-        ) -> Vec<(Range<Point>, WorkflowStepEditSuggestionStatus)> {
+        ) -> Vec<(Range<Point>, WorkflowStepTestStatus)> {
             context
                 .workflow_steps
                 .iter()
                 .map(|step| {
                     let buffer = context.buffer.read(cx);
                     let status = match &step.status {
-                        WorkflowStepStatus::Pending(_) => WorkflowStepEditSuggestionStatus::Pending,
+                        WorkflowStepStatus::Pending(_) => WorkflowStepTestStatus::Pending,
-                        WorkflowStepStatus::Resolved { .. } => {
+                        WorkflowStepStatus::Resolved { .. } => WorkflowStepTestStatus::Resolved,
-                            WorkflowStepEditSuggestionStatus::Resolved
+                        WorkflowStepStatus::Error(_) => WorkflowStepTestStatus::Error,
-                        }
                     };
                     (step.tagged_range.to_point(buffer), status)
                 })

crates/assistant/src/inline_assistant.rs

@@ -68,6 +68,9 @@ pub struct InlineAssistant {
     assists: HashMap<InlineAssistId, InlineAssist>,
     assists_by_editor: HashMap<WeakView<Editor>, EditorInlineAssists>,
     assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
+    assist_observations:
+        HashMap<InlineAssistId, (async_watch::Sender<()>, async_watch::Receiver<()>)>,
+    confirmed_assists: HashMap<InlineAssistId, Model<Codegen>>,
     prompt_history: VecDeque<String>,
     prompt_builder: Arc<PromptBuilder>,
     telemetry: Option<Arc<Telemetry>>,
@@ -88,6 +91,8 @@ impl InlineAssistant {
             assists: HashMap::default(),
             assists_by_editor: HashMap::default(),
             assist_groups: HashMap::default(),
+            assist_observations: HashMap::default(),
+            confirmed_assists: HashMap::default(),
             prompt_history: VecDeque::default(),
             prompt_builder,
             telemetry: Some(telemetry),
@@ -343,6 +348,7 @@ impl InlineAssistant {
                 height: prompt_editor_height,
                 render: build_assist_editor_renderer(prompt_editor),
                 disposition: BlockDisposition::Above,
+                priority: 0,
             },
             BlockProperties {
                 style: BlockStyle::Sticky,
@@ -357,6 +363,7 @@ impl InlineAssistant {
                         .into_any_element()
                 }),
                 disposition: BlockDisposition::Below,
+                priority: 0,
             },
         ];
 
@@ -654,8 +661,21 @@ impl InlineAssistant {
 
             if undo {
                 assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
+            } else {
+                self.confirmed_assists.insert(assist_id, assist.codegen);
             }
         }
+
+        // Remove the assist from the status updates map
+        self.assist_observations.remove(&assist_id);
+    }
+
+    pub fn undo_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
+        let Some(codegen) = self.confirmed_assists.remove(&assist_id) else {
+            return false;
+        };
+        codegen.update(cx, |this, cx| this.undo(cx));
+        true
     }
 
     fn dismiss_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
@@ -854,6 +874,10 @@ impl InlineAssistant {
                 )
             })
             .log_err();
+
+        if let Some((tx, _)) = self.assist_observations.get(&assist_id) {
+            tx.send(()).ok();
+        }
     }
 
     pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
@@ -864,19 +888,24 @@ impl InlineAssistant {
         };
 
         assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
+
+        if let Some((tx, _)) = self.assist_observations.get(&assist_id) {
+            tx.send(()).ok();
+        }
     }
 
-    pub fn status_for_assist(
+    pub fn assist_status(&self, assist_id: InlineAssistId, cx: &AppContext) -> InlineAssistStatus {
-        &self,
+        if let Some(assist) = self.assists.get(&assist_id) {
-        assist_id: InlineAssistId,
+            match &assist.codegen.read(cx).status {
-        cx: &WindowContext,
+                CodegenStatus::Idle => InlineAssistStatus::Idle,
-    ) -> Option<CodegenStatus> {
+                CodegenStatus::Pending => InlineAssistStatus::Pending,
-        let assist = self.assists.get(&assist_id)?;
+                CodegenStatus::Done => InlineAssistStatus::Done,
-        match &assist.codegen.read(cx).status {
+                CodegenStatus::Error(_) => InlineAssistStatus::Error,
-            CodegenStatus::Idle => Some(CodegenStatus::Idle),
+            }
-            CodegenStatus::Pending => Some(CodegenStatus::Pending),
+        } else if self.confirmed_assists.contains_key(&assist_id) {
-            CodegenStatus::Done => Some(CodegenStatus::Done),
+            InlineAssistStatus::Confirmed
-            CodegenStatus::Error(error) => Some(CodegenStatus::Error(anyhow!("{:?}", error))),
+        } else {
+            InlineAssistStatus::Canceled
         }
     }
 
@@ -1051,6 +1080,7 @@ impl InlineAssistant {
                             .into_any_element()
                     }),
                     disposition: BlockDisposition::Above,
+                    priority: 0,
                 });
             }
 
@@ -1060,6 +1090,37 @@ impl InlineAssistant {
                 .collect();
         })
     }
+
+    pub fn observe_assist(&mut self, assist_id: InlineAssistId) -> async_watch::Receiver<()> {
+        if let Some((_, rx)) = self.assist_observations.get(&assist_id) {
+            rx.clone()
+        } else {
+            let (tx, rx) = async_watch::channel(());
+            self.assist_observations.insert(assist_id, (tx, rx.clone()));
+            rx
+        }
+    }
+}
+
+pub enum InlineAssistStatus {
+    Idle,
+    Pending,
+    Done,
+    Error,
+    Confirmed,
+    Canceled,
+}
+
+impl InlineAssistStatus {
+    pub(crate) fn is_pending(&self) -> bool {
+        matches!(self, Self::Pending)
+    }
+    pub(crate) fn is_confirmed(&self) -> bool {
+        matches!(self, Self::Confirmed)
+    }
+    pub(crate) fn is_done(&self) -> bool {
+        matches!(self, Self::Done)
+    }
 }
 
 struct EditorInlineAssists {
@@ -1964,6 +2025,8 @@ impl InlineAssist {
 
                             if assist.decorations.is_none() {
                                 this.finish_assist(assist_id, false, cx);
+                            } else if let Some(tx) = this.assist_observations.get(&assist_id) {
+                                tx.0.send(()).ok();
                             }
                         }
                     })
@@ -2037,7 +2100,7 @@ pub struct Codegen {
     builder: Arc<PromptBuilder>,
 }
 
-pub enum CodegenStatus {
+enum CodegenStatus {
     Idle,
     Pending,
     Done,

crates/diagnostics/src/diagnostics.rs

@@ -449,6 +449,7 @@ impl ProjectDiagnosticsEditor {
                                     style: BlockStyle::Sticky,
                                     render: diagnostic_header_renderer(primary),
                                     disposition: BlockDisposition::Above,
+                                    priority: 0,
                                 });
                             }
 
@@ -470,6 +471,7 @@ impl ProjectDiagnosticsEditor {
                                             diagnostic, None, true, true,
                                         ),
                                         disposition: BlockDisposition::Below,
+                                        priority: 0,
                                     });
                                 }
                             }
@@ -508,6 +510,7 @@ impl ProjectDiagnosticsEditor {
                         style: block.style,
                         render: block.render,
                         disposition: block.disposition,
+                        priority: 0,
                     })
                 }),
                 Some(Autoscroll::fit()),

crates/editor/src/display_map.rs

@@ -1281,12 +1281,14 @@ pub mod tests {
                                         position.to_point(&buffer),
                                         height
                                     );
+                                    let priority = rng.gen_range(1..100);
                                     BlockProperties {
                                         style: BlockStyle::Fixed,
                                         position,
                                         height,
                                         disposition,
                                         render: Box::new(|_| div().into_any()),
+                                        priority: priority,
                                     }
                                 })
                                 .collect::<Vec<_>>();

crates/editor/src/display_map/block_map.rs

@@ -84,6 +84,7 @@ pub struct CustomBlock {
     style: BlockStyle,
     render: Arc<Mutex<RenderBlock>>,
     disposition: BlockDisposition,
+    priority: usize,
 }
 
 pub struct BlockProperties<P> {
@@ -92,6 +93,7 @@ pub struct BlockProperties<P> {
     pub style: BlockStyle,
     pub render: RenderBlock,
     pub disposition: BlockDisposition,
+    pub priority: usize,
 }
 
 impl<P: Debug> Debug for BlockProperties<P> {
@@ -182,6 +184,7 @@ pub(crate) enum BlockType {
 pub(crate) trait BlockLike {
     fn block_type(&self) -> BlockType;
     fn disposition(&self) -> BlockDisposition;
+    fn priority(&self) -> usize;
 }
 
 #[allow(clippy::large_enum_variant)]
@@ -215,6 +218,14 @@ impl BlockLike for Block {
     fn disposition(&self) -> BlockDisposition {
         self.disposition()
     }
+
+    fn priority(&self) -> usize {
+        match self {
+            Block::Custom(block) => block.priority,
+            Block::ExcerptHeader { .. } => usize::MAX,
+            Block::ExcerptFooter { .. } => 0,
+        }
+    }
 }
 
 impl Block {
@@ -660,7 +671,10 @@ impl BlockMap {
                         (BlockType::Header, BlockType::Header) => Ordering::Equal,
                         (BlockType::Header, _) => Ordering::Less,
                         (_, BlockType::Header) => Ordering::Greater,
-                        (BlockType::Custom(a_id), BlockType::Custom(b_id)) => a_id.cmp(&b_id),
+                        (BlockType::Custom(a_id), BlockType::Custom(b_id)) => block_b
+                            .priority()
+                            .cmp(&block_a.priority())
+                            .then_with(|| a_id.cmp(&b_id)),
                     })
             })
         });
@@ -802,6 +816,7 @@ impl<'a> BlockMapWriter<'a> {
                 render: Arc::new(Mutex::new(block.render)),
                 disposition: block.disposition,
                 style: block.style,
+                priority: block.priority,
             });
             self.0.custom_blocks.insert(block_ix, new_block.clone());
             self.0.custom_blocks_by_id.insert(id, new_block);
@@ -832,6 +847,7 @@ impl<'a> BlockMapWriter<'a> {
                         style: block.style,
                         render: block.render.clone(),
                         disposition: block.disposition,
+                        priority: block.priority,
                     };
                     let new_block = Arc::new(new_block);
                     *block = new_block.clone();
@@ -1463,6 +1479,7 @@ mod tests {
                 height: 1,
                 disposition: BlockDisposition::Above,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             },
             BlockProperties {
                 style: BlockStyle::Fixed,
@@ -1470,6 +1487,7 @@ mod tests {
                 height: 2,
                 disposition: BlockDisposition::Above,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             },
             BlockProperties {
                 style: BlockStyle::Fixed,
@@ -1477,6 +1495,7 @@ mod tests {
                 height: 3,
                 disposition: BlockDisposition::Below,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             },
         ]);
 
@@ -1716,6 +1735,7 @@ mod tests {
                 height: 1,
                 disposition: BlockDisposition::Above,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             },
             BlockProperties {
                 style: BlockStyle::Fixed,
@@ -1723,6 +1743,7 @@ mod tests {
                 height: 2,
                 disposition: BlockDisposition::Above,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             },
             BlockProperties {
                 style: BlockStyle::Fixed,
@@ -1730,6 +1751,7 @@ mod tests {
                 height: 3,
                 disposition: BlockDisposition::Below,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             },
         ]);
 
@@ -1819,6 +1841,7 @@ mod tests {
                 disposition: BlockDisposition::Above,
                 render: Box::new(|_| div().into_any()),
                 height: 1,
+                priority: 0,
             },
             BlockProperties {
                 style: BlockStyle::Fixed,
@@ -1826,6 +1849,7 @@ mod tests {
                 disposition: BlockDisposition::Below,
                 render: Box::new(|_| div().into_any()),
                 height: 1,
+                priority: 0,
             },
         ]);
 
@@ -1924,6 +1948,7 @@ mod tests {
                                 height,
                                 disposition,
                                 render: Box::new(|_| div().into_any()),
+                                priority: 0,
                             }
                         })
                         .collect::<Vec<_>>();
@@ -1944,6 +1969,7 @@ mod tests {
                             style: props.style,
                             render: Box::new(|_| div().into_any()),
                             disposition: props.disposition,
+                            priority: 0,
                         }));
                     for (block_id, props) in block_ids.into_iter().zip(block_properties) {
                         custom_blocks.push((block_id, props));
@@ -2014,6 +2040,7 @@ mod tests {
                         disposition: block.disposition,
                         id: *id,
                         height: block.height,
+                        priority: block.priority,
                     },
                 )
             }));
@@ -2235,6 +2262,7 @@ mod tests {
                 disposition: BlockDisposition,
                 id: CustomBlockId,
                 height: u32,
+                priority: usize,
             },
         }
 
@@ -2250,6 +2278,14 @@ mod tests {
             fn disposition(&self) -> BlockDisposition {
                 self.disposition()
             }
+
+            fn priority(&self) -> usize {
+                match self {
+                    ExpectedBlock::Custom { priority, .. } => *priority,
+                    ExpectedBlock::ExcerptHeader { .. } => usize::MAX,
+                    ExpectedBlock::ExcerptFooter { .. } => 0,
+                }
+            }
         }
 
         impl ExpectedBlock {
@@ -2277,6 +2313,7 @@ mod tests {
                         id: block.id,
                         disposition: block.disposition,
                         height: block.height,
+                        priority: block.priority,
                     },
                     Block::ExcerptHeader {
                         height,

crates/editor/src/editor.rs

@@ -9614,6 +9614,7 @@ impl Editor {
                                 }
                             }),
                             disposition: BlockDisposition::Below,
+                            priority: 0,
                         }],
                         Some(Autoscroll::fit()),
                         cx,
@@ -9877,6 +9878,7 @@ impl Editor {
                             height: message_height,
                             render: diagnostic_block_renderer(diagnostic, None, true, true),
                             disposition: BlockDisposition::Below,
+                            priority: 0,
                         }
                     }),
                     cx,
@@ -10182,6 +10184,7 @@ impl Editor {
         if let Some(autoscroll) = autoscroll {
             self.request_autoscroll(autoscroll, cx);
         }
+        cx.notify();
         blocks
     }
 
@@ -10196,6 +10199,7 @@ impl Editor {
         if let Some(autoscroll) = autoscroll {
             self.request_autoscroll(autoscroll, cx);
         }
+        cx.notify();
     }
 
     pub fn replace_blocks(
@@ -10208,9 +10212,8 @@ impl Editor {
             .update(cx, |display_map, _cx| display_map.replace_blocks(renderers));
         if let Some(autoscroll) = autoscroll {
             self.request_autoscroll(autoscroll, cx);
-        } else {
-            cx.notify();
         }
+        cx.notify();
     }
 
     pub fn remove_blocks(
@@ -10225,6 +10228,7 @@ impl Editor {
         if let Some(autoscroll) = autoscroll {
             self.request_autoscroll(autoscroll, cx);
         }
+        cx.notify();
     }
 
     pub fn row_for_block(

crates/editor/src/editor_tests.rs

@@ -3785,6 +3785,7 @@ fn test_move_line_up_down_with_blocks(cx: &mut TestAppContext) {
                 disposition: BlockDisposition::Below,
                 height: 1,
                 render: Box::new(|_| div().into_any()),
+                priority: 0,
             }],
             Some(Autoscroll::fit()),
             cx,

crates/editor/src/element.rs

@@ -6478,6 +6478,7 @@ mod tests {
                         height: 3,
                         position: Anchor::min(),
                         render: Box::new(|cx| div().h(3. * cx.line_height()).into_any()),
+                        priority: 0,
                     }],
                     None,
                     cx,

crates/editor/src/hunk_diff.rs

@@ -525,6 +525,7 @@ impl Editor {
                         .child(editor_with_deleted_text.clone())
                         .into_any_element()
                 }),
+                priority: 0,
             }),
             None,
             cx,

crates/repl/src/session.rs

@@ -87,6 +87,7 @@ impl EditorBlock {
                 style: BlockStyle::Sticky,
                 render: Self::create_output_area_renderer(execution_view.clone(), on_close.clone()),
                 disposition: BlockDisposition::Below,
+                priority: 0,
             };
 
             let block_id = editor.insert_blocks([block], None, cx)[0];

crates/ui/src/components/button/button_like.rs

@@ -50,6 +50,7 @@ pub enum TintColor {
     Accent,
     Negative,
     Warning,
+    Positive,
 }
 
 impl TintColor {
@@ -73,6 +74,12 @@ impl TintColor {
                 label_color: cx.theme().colors().text,
                 icon_color: cx.theme().colors().text,
             },
+            TintColor::Positive => ButtonLikeStyles {
+                background: cx.theme().status().success_background,
+                border_color: cx.theme().status().success_border,
+                label_color: cx.theme().colors().text,
+                icon_color: cx.theme().colors().text,
+            },
         }
     }
 }
@@ -83,6 +90,7 @@ impl From<TintColor> for Color {
             TintColor::Accent => Color::Accent,
             TintColor::Negative => Color::Error,
             TintColor::Warning => Color::Warning,
+            TintColor::Positive => Color::Success,
         }
     }
 }

crates/ui/src/components/icon.rs

@@ -256,6 +256,7 @@ pub enum IconName {
     TextSearch,
     Trash,
     TriangleRight,
+    Undo,
     Update,
     WholeWord,
     XCircle,
@@ -419,6 +420,7 @@ impl IconName {
             IconName::Trash => "icons/trash.svg",
             IconName::TriangleRight => "icons/triangle_right.svg",
             IconName::Update => "icons/update.svg",
+            IconName::Undo => "icons/undo.svg",
             IconName::WholeWord => "icons/word_search.svg",
             IconName::XCircle => "icons/error.svg",
             IconName::ZedAssistant => "icons/zed_assistant.svg",