assistant: Undo workflow step when buffer is discarded (#16465)

This fixes a weird bug:

1. Use `/workflow` in assistant
2. Have it generate a step that modifies a file
3. Either (a) select the step in the assistant and have it auto-insert
newlines (b) select "Transform" to have the step applied
4. Close the modified file in the editor ("Discard")
5. Re-open the file
6. BUG: the changes made by assistant are still there!

The reason for the bug is that the assistant keeps references to buffers
and they're not closed/reloaded when closed/reopened.

To fix the bug we now rollback the applied workflow steps when
discarding a buffer.

(This does *not* yet fix the issue where a workflow step inserts a new
buffer into the project/worktree that does not show up on the file
system yet but in `/file` and hangs around until Zed is closed.)


Release Notes:

- N/A

Co-authored-by: Bennet <bennet@zed.dev>

crates/assistant/src/assistant_panel.rs

@@ -2479,6 +2479,18 @@ impl ContextEditor {
         };
 
         let resolved_step = step.read(cx).resolution.clone();
+
+        if let Some(Ok(resolution)) = resolved_step.as_ref() {
+            for (buffer, _) in resolution.suggestion_groups.iter() {
+                let step_range = step_range.clone();
+                cx.subscribe(buffer, move |this, _, event, cx| match event {
+                    language::Event::Discarded => this.undo_workflow_step(step_range.clone(), cx),
+                    _ => {}
+                })
+                .detach();
+            }
+        }
+
         if let Some(existing_step) = self.workflow_steps.get_mut(&step_range) {
             existing_step.resolved_step = resolved_step;
         } else {

crates/editor/src/items.rs

@@ -680,6 +680,12 @@ impl Item for Editor {
         self.nav_history = Some(history);
     }
 
+    fn discarded(&self, _project: Model<Project>, cx: &mut ViewContext<Self>) {
+        for buffer in self.buffer().clone().read(cx).all_buffers() {
+            buffer.update(cx, |buffer, cx| buffer.discarded(cx))
+        }
+    }
+
     fn deactivated(&mut self, cx: &mut ViewContext<Self>) {
         let selection = self.selections.newest_anchor();
         self.push_to_nav_history(selection.head(), None, cx);

crates/language/src/buffer.rs

@@ -332,6 +332,8 @@ pub enum Event {
     CapabilityChanged,
     /// The buffer was explicitly requested to close.
     Closed,
+    /// The buffer was discarded when closing.
+    Discarded,
 }
 
 /// The file associated with a buffer.
@@ -827,6 +829,12 @@ impl Buffer {
         cx.notify();
     }
 
+    /// This method is called to signal that the buffer has been discarded.
+    pub fn discarded(&mut self, cx: &mut ModelContext<Self>) {
+        cx.emit(Event::Discarded);
+        cx.notify();
+    }
+
     /// Reloads the contents of the buffer from disk.
     pub fn reload(
         &mut self,

crates/multi_buffer/src/multi_buffer.rs

@@ -106,6 +106,7 @@ pub enum Event {
     Saved,
     FileHandleChanged,
     Closed,
+    Discarded,
     DirtyChanged,
     DiagnosticsUpdated,
 }
@@ -1691,6 +1692,7 @@ impl MultiBuffer {
             language::Event::Reparsed => Event::Reparsed(buffer.read(cx).remote_id()),
             language::Event::DiagnosticsUpdated => Event::DiagnosticsUpdated,
             language::Event::Closed => Event::Closed,
+            language::Event::Discarded => Event::Discarded,
             language::Event::CapabilityChanged => {
                 self.capability = buffer.read(cx).capability();
                 Event::CapabilityChanged

crates/workspace/src/item.rs

@@ -184,6 +184,7 @@ pub trait Item: FocusableView + EventEmitter<Self::Event> {
     fn to_item_events(_event: &Self::Event, _f: impl FnMut(ItemEvent)) {}
 
     fn deactivated(&mut self, _: &mut ViewContext<Self>) {}
+    fn discarded(&self, _project: Model<Project>, _cx: &mut ViewContext<Self>) {}
     fn workspace_deactivated(&mut self, _: &mut ViewContext<Self>) {}
     fn navigate(&mut self, _: Box<dyn Any>, _: &mut ViewContext<Self>) -> bool {
         false
@@ -394,6 +395,7 @@ pub trait ItemHandle: 'static + Send {
         cx: &mut ViewContext<Workspace>,
     );
     fn deactivated(&self, cx: &mut WindowContext);
+    fn discarded(&self, project: Model<Project>, cx: &mut WindowContext);
     fn workspace_deactivated(&self, cx: &mut WindowContext);
     fn navigate(&self, data: Box<dyn Any>, cx: &mut WindowContext) -> bool;
     fn item_id(&self) -> EntityId;
@@ -735,6 +737,10 @@ impl<T: Item> ItemHandle for View<T> {
         });
     }
 
+    fn discarded(&self, project: Model<Project>, cx: &mut WindowContext) {
+        self.update(cx, |this, cx| this.discarded(project, cx));
+    }
+
     fn deactivated(&self, cx: &mut WindowContext) {
         self.update(cx, |this, cx| this.deactivated(cx));
     }

crates/workspace/src/pane.rs

@@ -1500,8 +1500,13 @@ impl Pane {
                         })?;
                         match answer {
                             Ok(0) => {}
-                            Ok(1) => return Ok(true), // Don't save this file
-                            _ => return Ok(false),    // Cancel
+                            Ok(1) => {
+                                // Don't save this file
+                                pane.update(cx, |_, cx| item.discarded(project, cx))
+                                    .log_err();
+                                return Ok(true);
+                            }
+                            _ => return Ok(false), // Cancel
                         }
                     } else {
                         return Ok(false);