Make method only available in tests

Co-authored-by: Antonio <antonio@zed.dev>

crates/editor/src/editor_tests.rs

@@ -7308,6 +7308,22 @@ async fn test_document_format_manual_trigger(cx: &mut gpui::TestAppContext) {
     );
 }
 
+#[gpui::test]
+async fn test_lsp_shutdown(cx: &mut gpui::TestAppContext) {
+    init_test(cx, |_| {});
+
+    let cx = EditorLspTestContext::new_rust(
+        lsp::ServerCapabilities {
+            document_formatting_provider: Some(lsp::OneOf::Left(true)),
+            ..Default::default()
+        },
+        cx,
+    )
+    .await;
+
+    drop(cx);
+}
+
 #[gpui::test]
 async fn test_concurrent_format_requests(cx: &mut gpui::TestAppContext) {
     init_test(cx, |_| {});

crates/gpui/src/app.rs

@@ -344,7 +344,7 @@ impl AppContext {
         platform.on_quit(Box::new({
             let cx = app.clone();
             move || {
-                cx.borrow_mut().shutdown();
+                Self::shutdown(&cx);
             }
         }));
 
@@ -353,23 +353,27 @@ impl AppContext {
 
     /// Quit the application gracefully. Handlers registered with [`ModelContext::on_app_quit`]
     /// will be given 100ms to complete before exiting.
-    pub fn shutdown(&mut self) {
+    pub fn shutdown(this: &Rc<AppCell>) {
         let mut futures = Vec::new();
+        let background_executor = {
+            let mut this = this.borrow_mut();
 
-        for observer in self.quit_observers.remove(&()) {
-            futures.push(observer(self));
-        }
+            for observer in this.quit_observers.remove(&()) {
+                futures.push(observer(&mut this));
+            }
 
-        self.windows.clear();
-        self.window_handles.clear();
-        self.flush_effects();
+            this.windows.clear();
+            this.window_handles.clear();
+            this.background_executor.clone()
+        };
 
         let futures = futures::future::join_all(futures);
-        if self
-            .background_executor
-            .block_with_timeout(SHUTDOWN_TIMEOUT, futures)
-            .is_err()
-        {
+        #[cfg(any(test, feature = "test-support"))]
+        let result = background_executor.block_to_shutdown(SHUTDOWN_TIMEOUT, futures);
+        #[cfg(not(any(test, feature = "test-support")))]
+        let result = background_executor.block_with_timeout(SHUTDOWN_TIMEOUT, futures);
+
+        if result.is_err() {
             log::error!("timed out waiting on app_will_quit");
         }
     }

crates/gpui/src/app/async_context.rs

@@ -106,6 +106,17 @@ impl Context for AsyncAppContext {
 }
 
 impl AsyncAppContext {
+    /// Quit the application gracefully. Handlers registered with [`ModelContext::on_app_quit`]
+    /// will be given 100ms to complete before exiting.
+    pub fn shutdown(&self) -> Result<()> {
+        let app = self
+            .app
+            .upgrade()
+            .ok_or_else(|| anyhow!("app was released"))?;
+        AppContext::shutdown(&app);
+        Ok(())
+    }
+
     /// Schedules all windows in the application to be redrawn.
     pub fn refresh(&self) -> Result<()> {
         let app = self

crates/gpui/src/app/test_context.rs

@@ -138,7 +138,7 @@ impl TestAppContext {
     /// public so the macro can call it.
     pub fn quit(&self) {
         self.on_quit.borrow_mut().drain(..).for_each(|f| f());
-        self.app.borrow_mut().shutdown();
+        AppContext::shutdown(&self.app);
     }
 
     /// Register cleanup to run when the test ends.

crates/gpui/src/executor.rs

@@ -299,6 +299,17 @@ impl BackgroundExecutor {
         self.block_internal(true, future, Some(duration))
     }
 
+    /// Block the current thread until the given future resolves
+    /// or `duration` has elapsed.
+    #[cfg(any(test, feature = "test-support"))]
+    pub(crate) fn block_to_shutdown<R>(
+        &self,
+        duration: Duration,
+        future: impl Future<Output = R>,
+    ) -> Result<R, impl Future<Output = R>> {
+        self.block_internal(false, future, Some(duration))
+    }
+
     /// Scoped lets you start a number of tasks and waits
     /// for all of them to complete before returning.
     pub async fn scoped<'scope, F>(&self, scheduler: F)

crates/live_kit_client/examples/test_app.rs

@@ -161,7 +161,7 @@ fn main() {
                 panic!("unexpected message");
             }
 
-            cx.update(|cx| cx.shutdown()).ok();
+            cx.shutdown().unwrap();
         })
         .detach();
     });

crates/lsp/src/input_handler.rs

@@ -23,7 +23,7 @@ pub struct LspStdoutHandler {
     pub(super) notifications_channel: UnboundedReceiver<AnyNotification>,
 }
 
-async fn read_headers<Stdout>(reader: &mut BufReader<Stdout>, buffer: &mut Vec<u8>) -> Result<()>
+async fn read_headers<Stdout>(reader: &mut BufReader<Stdout>, buffer: &mut Vec<u8>) -> Result<bool>
 where
     Stdout: AsyncRead + Unpin + Send + 'static,
 {
@@ -31,11 +31,12 @@ where
         if buffer.len() >= HEADER_DELIMITER.len()
             && buffer[(buffer.len() - HEADER_DELIMITER.len())..] == HEADER_DELIMITER[..]
         {
-            return Ok(());
+            return Ok(true);
         }
 
+        // If read_until returns 0, we reached EOF and exit.
         if reader.read_until(b'\n', buffer).await? == 0 {
-            return Err(anyhow!("cannot read LSP message headers"));
+            return Ok(false);
         }
     }
 }
@@ -74,7 +75,9 @@ impl LspStdoutHandler {
         loop {
             buffer.clear();
 
-            read_headers(&mut stdout, &mut buffer).await?;
+            if !read_headers(&mut stdout, &mut buffer).await? {
+                return Ok(());
+            }
 
             let headers = std::str::from_utf8(&buffer)?;
 

crates/lsp/src/lsp.rs

@@ -749,8 +749,6 @@ impl LanguageServer {
                 &executor,
                 (),
             );
-            let exit = Self::notify_internal::<notification::Exit>(&outbound_tx, ());
-            outbound_tx.close();
 
             let server = self.server.clone();
             let name = self.name.clone();
@@ -762,15 +760,16 @@ impl LanguageServer {
                     select! {
                         request_result = shutdown_request.fuse() => {
                             request_result?;
+                            let exit = Self::notify_internal::<notification::Exit>(&outbound_tx, ());
+                            outbound_tx.close();
+                            exit?;
                         }
 
                         _ = timer => {
                             log::info!("timeout waiting for language server {name} to shutdown");
                         },
                     }
-
                     response_handlers.lock().take();
-                    exit?;
                     output_done.recv().await;
                     server.lock().take().map(|mut child| child.kill());
                     log::debug!("language server shutdown finished");
@@ -1238,6 +1237,8 @@ impl FakeLanguageServer {
             }
         });
 
+        fake.handle_request::<request::Shutdown, _, _>(|_, _| async move { Ok(()) });
+
         (server, fake)
     }
 

crates/remote_server/src/headless_project.rs

@@ -24,7 +24,7 @@ use std::{
     path::{Path, PathBuf},
     sync::{atomic::AtomicUsize, Arc},
 };
-use util::ResultExt;
+use util::{ResultExt, TryFutureExt};
 use worktree::Worktree;
 
 pub struct HeadlessProject {
@@ -554,13 +554,14 @@ impl HeadlessProject {
         _envelope: TypedEnvelope<proto::ShutdownRemoteServer>,
         cx: AsyncAppContext,
     ) -> Result<proto::Ack> {
-        cx.spawn(|cx| async move {
-            cx.update(|cx| {
+        cx.spawn(|cx| {
+            async move {
                 // TODO: This is a hack, because in a headless project, shutdown isn't executed
                 // when calling quit, but it should be.
-                cx.shutdown();
-                cx.quit();
-            })
+                AsyncAppContext::shutdown(&cx)?;
+                cx.update(|cx| cx.quit())
+            }
+            .log_err()
         })
         .detach();
 

crates/remote_server/src/unix.rs

@@ -284,12 +284,10 @@ fn start_server(
                 }
                 _ = futures::FutureExt::fuse(smol::Timer::after(IDLE_TIMEOUT)) => {
                     log::warn!("timed out waiting for new connections after {:?}. exiting.", IDLE_TIMEOUT);
-                    cx.update(|cx| {
-                        // TODO: This is a hack, because in a headless project, shutdown isn't executed
-                        // when calling quit, but it should be.
-                        cx.shutdown();
-                        cx.quit();
-                    })?;
+                    // TODO: This is a hack, because in a headless project, shutdown isn't executed
+                    // when calling quit, but it should be.
+                    cx.shutdown()?;
+                    cx.update(|cx| cx.quit())?;
                     break;
                 }
                 _ = app_quit_rx.next().fuse() => {