Merge pull request #1231 from zed-industries/prompt-to-save-on-window-close

Prompt user to save changed buffers when closing window via the mouse

crates/gpui/src/app.rs

@@ -521,16 +521,8 @@ impl TestAppContext {
     pub fn simulate_prompt_answer(&self, window_id: usize, answer: usize) {
         use postage::prelude::Sink as _;
 
-        let mut state = self.cx.borrow_mut();
-        let (_, window) = state
-            .presenters_and_platform_windows
-            .get_mut(&window_id)
-            .unwrap();
-        let test_window = window
-            .as_any_mut()
-            .downcast_mut::<platform::test::Window>()
-            .unwrap();
-        let mut done_tx = test_window
+        let mut done_tx = self
+            .window_mut(window_id)
             .pending_prompts
             .borrow_mut()
             .pop_front()
@@ -539,30 +531,28 @@ impl TestAppContext {
     }
 
     pub fn has_pending_prompt(&self, window_id: usize) -> bool {
-        let mut state = self.cx.borrow_mut();
-        let (_, window) = state
-            .presenters_and_platform_windows
-            .get_mut(&window_id)
-            .unwrap();
-        let test_window = window
-            .as_any_mut()
-            .downcast_mut::<platform::test::Window>()
-            .unwrap();
-        let prompts = test_window.pending_prompts.borrow_mut();
+        let window = self.window_mut(window_id);
+        let prompts = window.pending_prompts.borrow_mut();
         !prompts.is_empty()
     }
 
     pub fn current_window_title(&self, window_id: usize) -> Option<String> {
-        let mut state = self.cx.borrow_mut();
-        let (_, window) = state
-            .presenters_and_platform_windows
-            .get_mut(&window_id)
-            .unwrap();
-        let test_window = window
-            .as_any_mut()
-            .downcast_mut::<platform::test::Window>()
-            .unwrap();
-        test_window.title.clone()
+        self.window_mut(window_id).title.clone()
+    }
+
+    pub fn simulate_window_close(&self, window_id: usize) -> bool {
+        let handler = self.window_mut(window_id).should_close_handler.take();
+        if let Some(mut handler) = handler {
+            let should_close = handler();
+            self.window_mut(window_id).should_close_handler = Some(handler);
+            should_close
+        } else {
+            false
+        }
+    }
+
+    pub fn is_window_edited(&self, window_id: usize) -> bool {
+        self.window_mut(window_id).edited
     }
 
     pub fn leak_detector(&self) -> Arc<Mutex<LeakDetector>> {
@@ -576,6 +566,20 @@ impl TestAppContext {
             .lock()
             .assert_dropped(handle.id())
     }
+
+    fn window_mut(&self, window_id: usize) -> std::cell::RefMut<platform::test::Window> {
+        std::cell::RefMut::map(self.cx.borrow_mut(), |state| {
+            let (_, window) = state
+                .presenters_and_platform_windows
+                .get_mut(&window_id)
+                .unwrap();
+            let test_window = window
+                .as_any_mut()
+                .downcast_mut::<platform::test::Window>()
+                .unwrap();
+            test_window
+        })
+    }
 }
 
 impl AsyncAppContext {
@@ -780,6 +784,7 @@ type GlobalObservationCallback = Box<dyn FnMut(&mut MutableAppContext)>;
 type ReleaseObservationCallback = Box<dyn FnOnce(&dyn Any, &mut MutableAppContext)>;
 type ActionObservationCallback = Box<dyn FnMut(TypeId, &mut MutableAppContext)>;
 type DeserializeActionCallback = fn(json: &str) -> anyhow::Result<Box<dyn Action>>;
+type WindowShouldCloseSubscriptionCallback = Box<dyn FnMut(&mut MutableAppContext) -> bool>;
 
 pub struct MutableAppContext {
     weak_self: Option<rc::Weak<RefCell<Self>>>,
@@ -2004,6 +2009,12 @@ impl MutableAppContext {
                         Effect::ActionDispatchNotification { action_id } => {
                             self.handle_action_dispatch_notification_effect(action_id)
                         }
+                        Effect::WindowShouldCloseSubscription {
+                            window_id,
+                            callback,
+                        } => {
+                            self.handle_window_should_close_subscription_effect(window_id, callback)
+                        }
                     }
                     self.pending_notifications.clear();
                     self.remove_dropped_entities();
@@ -2451,6 +2462,17 @@ impl MutableAppContext {
         self.action_dispatch_observations.lock().extend(callbacks);
     }
 
+    fn handle_window_should_close_subscription_effect(
+        &mut self,
+        window_id: usize,
+        mut callback: WindowShouldCloseSubscriptionCallback,
+    ) {
+        let mut app = self.upgrade();
+        if let Some((_, window)) = self.presenters_and_platform_windows.get_mut(&window_id) {
+            window.on_should_close(Box::new(move || app.update(|cx| callback(cx))))
+        }
+    }
+
     pub fn focus(&mut self, window_id: usize, view_id: Option<usize>) {
         if let Some(pending_focus_index) = self.pending_focus_index {
             self.pending_effects.remove(pending_focus_index);
@@ -2828,6 +2850,10 @@ pub enum Effect {
     ActionDispatchNotification {
         action_id: TypeId,
     },
+    WindowShouldCloseSubscription {
+        window_id: usize,
+        callback: WindowShouldCloseSubscriptionCallback,
+    },
 }
 
 impl Debug for Effect {
@@ -2921,6 +2947,10 @@ impl Debug for Effect {
                 .field("is_active", is_active)
                 .finish(),
             Effect::RefreshWindows => f.debug_struct("Effect::FullViewRefresh").finish(),
+            Effect::WindowShouldCloseSubscription { window_id, .. } => f
+                .debug_struct("Effect::WindowShouldCloseSubscription")
+                .field("window_id", window_id)
+                .finish(),
         }
     }
 }
@@ -3339,6 +3369,32 @@ impl<'a, T: View> ViewContext<'a, T> {
         }
     }
 
+    pub fn set_window_edited(&mut self, edited: bool) {
+        let window_id = self.window_id();
+        if let Some((_, window)) = self.presenters_and_platform_windows.get_mut(&window_id) {
+            window.set_edited(edited);
+        }
+    }
+
+    pub fn on_window_should_close<F>(&mut self, mut callback: F)
+    where
+        F: 'static + FnMut(&mut T, &mut ViewContext<T>) -> bool,
+    {
+        let window_id = self.window_id();
+        let view = self.weak_handle();
+        self.pending_effects
+            .push_back(Effect::WindowShouldCloseSubscription {
+                window_id,
+                callback: Box::new(move |cx| {
+                    if let Some(view) = view.upgrade(cx) {
+                        view.update(cx, |view, cx| callback(view, cx))
+                    } else {
+                        true
+                    }
+                }),
+            });
+    }
+
     pub fn add_model<S, F>(&mut self, build_model: F) -> ModelHandle<S>
     where
         S: Entity,

crates/gpui/src/platform.rs

@@ -93,10 +93,12 @@ pub trait Window: WindowContext {
     fn on_event(&mut self, callback: Box<dyn FnMut(Event) -> bool>);
     fn on_active_status_change(&mut self, callback: Box<dyn FnMut(bool)>);
     fn on_resize(&mut self, callback: Box<dyn FnMut()>);
+    fn on_should_close(&mut self, callback: Box<dyn FnMut() -> bool>);
     fn on_close(&mut self, callback: Box<dyn FnOnce()>);
     fn prompt(&self, level: PromptLevel, msg: &str, answers: &[&str]) -> oneshot::Receiver<usize>;
     fn activate(&self);
     fn set_title(&mut self, title: &str);
+    fn set_edited(&mut self, edited: bool);
 }
 
 pub trait WindowContext {

crates/gpui/src/platform/mac/window.rs

@@ -81,6 +81,10 @@ unsafe fn build_classes() {
             sel!(windowDidResignKey:),
             window_did_change_key_status as extern "C" fn(&Object, Sel, id),
         );
+        decl.add_method(
+            sel!(windowShouldClose:),
+            window_should_close as extern "C" fn(&Object, Sel, id) -> BOOL,
+        );
         decl.add_method(sel!(close), close_window as extern "C" fn(&Object, Sel));
         decl.register()
     };
@@ -167,6 +171,7 @@ struct WindowState {
     event_callback: Option<Box<dyn FnMut(Event) -> bool>>,
     activate_callback: Option<Box<dyn FnMut(bool)>>,
     resize_callback: Option<Box<dyn FnMut()>>,
+    should_close_callback: Option<Box<dyn FnMut() -> bool>>,
     close_callback: Option<Box<dyn FnOnce()>>,
     synthetic_drag_counter: usize,
     executor: Rc<executor::Foreground>,
@@ -242,6 +247,7 @@ impl Window {
                 native_window,
                 event_callback: None,
                 resize_callback: None,
+                should_close_callback: None,
                 close_callback: None,
                 activate_callback: None,
                 synthetic_drag_counter: 0,
@@ -339,6 +345,10 @@ impl platform::Window for Window {
         self.0.as_ref().borrow_mut().resize_callback = Some(callback);
     }
 
+    fn on_should_close(&mut self, callback: Box<dyn FnMut() -> bool>) {
+        self.0.as_ref().borrow_mut().should_close_callback = Some(callback);
+    }
+
     fn on_close(&mut self, callback: Box<dyn FnOnce()>) {
         self.0.as_ref().borrow_mut().close_callback = Some(callback);
     }
@@ -397,6 +407,17 @@ impl platform::Window for Window {
             msg_send![app, changeWindowsItem:window title:title filename:false]
         }
     }
+
+    fn set_edited(&mut self, edited: bool) {
+        unsafe {
+            let window = self.0.borrow().native_window;
+            msg_send![window, setDocumentEdited: edited as BOOL]
+        }
+
+        // Changing the document edited state resets the traffic light position,
+        // so we have to move it again.
+        self.0.borrow().move_traffic_light();
+    }
 }
 
 impl platform::WindowContext for Window {
@@ -663,6 +684,19 @@ extern "C" fn window_did_change_key_status(this: &Object, selector: Sel, _: id)
         .detach();
 }
 
+extern "C" fn window_should_close(this: &Object, _: Sel, _: id) -> BOOL {
+    let window_state = unsafe { get_window_state(this) };
+    let mut window_state_borrow = window_state.as_ref().borrow_mut();
+    if let Some(mut callback) = window_state_borrow.should_close_callback.take() {
+        drop(window_state_borrow);
+        let should_close = callback();
+        window_state.borrow_mut().should_close_callback = Some(callback);
+        should_close as BOOL
+    } else {
+        YES
+    }
+}
+
 extern "C" fn close_window(this: &Object, _: Sel) {
     unsafe {
         let close_callback = {

crates/gpui/src/platform/test.rs

@@ -37,7 +37,9 @@ pub struct Window {
     event_handlers: Vec<Box<dyn FnMut(super::Event) -> bool>>,
     resize_handlers: Vec<Box<dyn FnMut()>>,
     close_handlers: Vec<Box<dyn FnOnce()>>,
+    pub(crate) should_close_handler: Option<Box<dyn FnMut() -> bool>>,
     pub(crate) title: Option<String>,
+    pub(crate) edited: bool,
     pub(crate) pending_prompts: RefCell<VecDeque<oneshot::Sender<usize>>>,
 }
 
@@ -185,12 +187,14 @@ impl Window {
     fn new(size: Vector2F) -> Self {
         Self {
             size,
-            event_handlers: Vec::new(),
-            resize_handlers: Vec::new(),
-            close_handlers: Vec::new(),
+            event_handlers: Default::default(),
+            resize_handlers: Default::default(),
+            close_handlers: Default::default(),
+            should_close_handler: Default::default(),
             scale_factor: 1.0,
             current_scene: None,
             title: None,
+            edited: false,
             pending_prompts: Default::default(),
         }
     }
@@ -258,6 +262,14 @@ impl super::Window for Window {
     fn set_title(&mut self, title: &str) {
         self.title = Some(title.to_string())
     }
+
+    fn set_edited(&mut self, edited: bool) {
+        self.edited = edited;
+    }
+
+    fn on_should_close(&mut self, callback: Box<dyn FnMut() -> bool>) {
+        self.should_close_handler = Some(callback);
+    }
 }
 
 pub fn platform() -> Platform {

crates/workspace/src/pane.rs

@@ -110,6 +110,7 @@ pub enum Event {
     Activate,
     ActivateItem { local: bool },
     Remove,
+    RemoveItem,
     Split(SplitDirection),
     ChangeItemTitle,
 }
@@ -575,6 +576,7 @@ impl Pane {
                         }
 
                         let item = pane.items.remove(item_ix);
+                        cx.emit(Event::RemoveItem);
                         if pane.items.is_empty() {
                             item.deactivated(cx);
                             pane.update_toolbar(cx);

crates/workspace/src/workspace.rs

@@ -731,6 +731,7 @@ pub struct Workspace {
     leader_state: LeaderState,
     follower_states_by_leader: FollowerStatesByLeader,
     last_leaders_by_pane: HashMap<WeakViewHandle<Pane>, PeerId>,
+    window_edited: bool,
     _observe_current_user: Task<()>,
 }
 
@@ -846,6 +847,7 @@ impl Workspace {
             leader_state: Default::default(),
             follower_states_by_leader: Default::default(),
             last_leaders_by_pane: Default::default(),
+            window_edited: false,
             _observe_current_user,
         };
         this.project_remote_id_changed(this.project.read(cx).remote_id(), cx);
@@ -905,7 +907,11 @@ impl Workspace {
         }
     }
 
-    fn close(&mut self, _: &CloseWindow, cx: &mut ViewContext<Self>) -> Option<Task<Result<()>>> {
+    pub fn close(
+        &mut self,
+        _: &CloseWindow,
+        cx: &mut ViewContext<Self>,
+    ) -> Option<Task<Result<()>>> {
         let prepare = self.prepare_to_close(cx);
         Some(cx.spawn(|this, mut cx| async move {
             if prepare.await? {
@@ -1474,6 +1480,10 @@ impl Workspace {
                     if pane == self.active_pane {
                         self.active_item_path_changed(cx);
                     }
+                    self.update_window_edited(cx);
+                }
+                pane::Event::RemoveItem => {
+                    self.update_window_edited(cx);
                 }
             }
         } else {
@@ -1778,6 +1788,16 @@ impl Workspace {
         cx.set_window_title(&title);
     }
 
+    fn update_window_edited(&mut self, cx: &mut ViewContext<Self>) {
+        let is_edited = self
+            .items(cx)
+            .any(|item| item.has_conflict(cx) || item.is_dirty(cx));
+        if is_edited != self.window_edited {
+            self.window_edited = is_edited;
+            cx.set_window_edited(self.window_edited)
+        }
+    }
+
     fn render_collaborators(&self, theme: &Theme, cx: &mut RenderContext<Self>) -> Vec<ElementBox> {
         let mut collaborators = self
             .project

crates/zed/src/zed.rs

@@ -222,6 +222,13 @@ pub fn initialize_workspace(
     });
 
     auto_update::notify_of_any_new_update(cx.weak_handle(), cx);
+
+    cx.on_window_should_close(|workspace, cx| {
+        if let Some(task) = workspace.close(&Default::default(), cx) {
+            task.detach_and_log_err(cx);
+        }
+        false
+    });
 }
 
 pub fn build_window_options() -> WindowOptions<'static> {
@@ -364,7 +371,9 @@ mod tests {
     use super::*;
     use assets::Assets;
     use editor::{Autoscroll, DisplayPoint, Editor};
-    use gpui::{AssetSource, MutableAppContext, TestAppContext, ViewHandle};
+    use gpui::{
+        executor::Deterministic, AssetSource, MutableAppContext, TestAppContext, ViewHandle,
+    };
     use project::ProjectPath;
     use serde_json::json;
     use std::{
@@ -432,6 +441,80 @@ mod tests {
         assert_eq!(cx.window_ids().len(), 2);
     }
 
+    #[gpui::test]
+    async fn test_window_edit_state(executor: Arc<Deterministic>, cx: &mut TestAppContext) {
+        let app_state = init(cx);
+        app_state
+            .fs
+            .as_fake()
+            .insert_tree("/root", json!({"a": "hey"}))
+            .await;
+
+        cx.update(|cx| open_paths(&[PathBuf::from("/root/a")], &app_state, cx))
+            .await;
+        assert_eq!(cx.window_ids().len(), 1);
+
+        // When opening the workspace, the window is not in a edited state.
+        let workspace = cx.root_view::<Workspace>(cx.window_ids()[0]).unwrap();
+        let editor = workspace.read_with(cx, |workspace, cx| {
+            workspace
+                .active_item(cx)
+                .unwrap()
+                .downcast::<Editor>()
+                .unwrap()
+        });
+        assert!(!cx.is_window_edited(workspace.window_id()));
+
+        // Editing a buffer marks the window as edited.
+        editor.update(cx, |editor, cx| editor.insert("EDIT", cx));
+        assert!(cx.is_window_edited(workspace.window_id()));
+
+        // Undoing the edit restores the window's edited state.
+        editor.update(cx, |editor, cx| editor.undo(&Default::default(), cx));
+        assert!(!cx.is_window_edited(workspace.window_id()));
+
+        // Redoing the edit marks the window as edited again.
+        editor.update(cx, |editor, cx| editor.redo(&Default::default(), cx));
+        assert!(cx.is_window_edited(workspace.window_id()));
+
+        // Closing the item restores the window's edited state.
+        let close = workspace.update(cx, |workspace, cx| {
+            drop(editor);
+            Pane::close_active_item(workspace, &Default::default(), cx).unwrap()
+        });
+        executor.run_until_parked();
+        cx.simulate_prompt_answer(workspace.window_id(), 1);
+        close.await.unwrap();
+        assert!(!cx.is_window_edited(workspace.window_id()));
+
+        // Opening the buffer again doesn't impact the window's edited state.
+        cx.update(|cx| open_paths(&[PathBuf::from("/root/a")], &app_state, cx))
+            .await;
+        let editor = workspace.read_with(cx, |workspace, cx| {
+            workspace
+                .active_item(cx)
+                .unwrap()
+                .downcast::<Editor>()
+                .unwrap()
+        });
+        assert!(!cx.is_window_edited(workspace.window_id()));
+
+        // Editing the buffer marks the window as edited.
+        editor.update(cx, |editor, cx| editor.insert("EDIT", cx));
+        assert!(cx.is_window_edited(workspace.window_id()));
+
+        // Ensure closing the window via the mouse gets preempted due to the
+        // buffer having unsaved changes.
+        assert!(!cx.simulate_window_close(workspace.window_id()));
+        executor.run_until_parked();
+        assert_eq!(cx.window_ids().len(), 1);
+
+        // The window is successfully closed after the user dismisses the prompt.
+        cx.simulate_prompt_answer(workspace.window_id(), 1);
+        executor.run_until_parked();
+        assert_eq!(cx.window_ids().len(), 0);
+    }
+
     #[gpui::test]
     async fn test_new_empty_workspace(cx: &mut TestAppContext) {
         let app_state = init(cx);