Make zed --wait work with directories (#44936)

Fixes #23347

Release Notes:

- Implemented the `zed --wait` flag so that it works when opening a
directory. The command will block until the window is closed.

crates/cli/src/main.rs

@@ -61,6 +61,8 @@ Examples:
 )]
 struct Args {
     /// Wait for all of the given paths to be opened/closed before exiting.
+    ///
+    /// When opening a directory, waits until the created window is closed.
     #[arg(short, long)]
     wait: bool,
     /// Add files to the currently open workspace

crates/zed/src/zed/open_listener.rs

@@ -10,6 +10,7 @@ use editor::Editor;
 use fs::Fs;
 use futures::channel::mpsc::{UnboundedReceiver, UnboundedSender};
 use futures::channel::{mpsc, oneshot};
+use futures::future;
 use futures::future::join_all;
 use futures::{FutureExt, SinkExt, StreamExt};
 use git_ui::file_diff_view::FileDiffView;
@@ -514,33 +515,27 @@ async fn open_local_workspace(
     app_state: &Arc<AppState>,
     cx: &mut AsyncApp,
 ) -> bool {
-    let mut errored = false;
-
     let paths_with_position =
         derive_paths_with_position(app_state.fs.as_ref(), workspace_paths).await;
 
-    // Handle reuse flag by finding existing window to replace
-    let replace_window = if reuse {
-        cx.update(|cx| workspace::local_workspace_windows(cx).into_iter().next())
-            .ok()
-            .flatten()
+    // If reuse flag is passed, open a new workspace in an existing window.
+    let (open_new_workspace, replace_window) = if reuse {
+        (
+            Some(true),
+            cx.update(|cx| workspace::local_workspace_windows(cx).into_iter().next())
+                .ok()
+                .flatten(),
+        )
     } else {
-        None
+        (open_new_workspace, None)
     };
 
-    // For reuse, force new workspace creation but with replace_window set
-    let effective_open_new_workspace = if reuse {
-        Some(true)
-    } else {
-        open_new_workspace
-    };
-
-    match open_paths_with_positions(
+    let (workspace, items) = match open_paths_with_positions(
         &paths_with_position,
         &diff_paths,
         app_state.clone(),
         workspace::OpenOptions {
-            open_new_workspace: effective_open_new_workspace,
+            open_new_workspace,
             replace_window,
             prefer_focused_window: wait,
             env: env.cloned(),
@@ -550,80 +545,95 @@ async fn open_local_workspace(
     )
     .await
     {
-        Ok((workspace, items)) => {
-            let mut item_release_futures = Vec::new();
-
-            for item in items {
-                match item {
-                    Some(Ok(item)) => {
-                        cx.update(|cx| {
-                            let released = oneshot::channel();
-                            item.on_release(
-                                cx,
-                                Box::new(move |_| {
-                                    let _ = released.0.send(());
-                                }),
-                            )
-                            .detach();
-                            item_release_futures.push(released.1);
-                        })
-                        .log_err();
-                    }
-                    Some(Err(err)) => {
-                        responses
-                            .send(CliResponse::Stderr {
-                                message: err.to_string(),
-                            })
-                            .log_err();
-                        errored = true;
-                    }
-                    None => {}
-                }
-            }
-
-            if wait {
-                let background = cx.background_executor().clone();
-                let wait = async move {
-                    if paths_with_position.is_empty() && diff_paths.is_empty() {
-                        let (done_tx, done_rx) = oneshot::channel();
-                        let _subscription = workspace.update(cx, |_, _, cx| {
-                            cx.on_release(move |_, _| {
-                                let _ = done_tx.send(());
-                            })
-                        });
-                        let _ = done_rx.await;
-                    } else {
-                        let _ = futures::future::try_join_all(item_release_futures).await;
-                    };
-                }
-                .fuse();
-
-                futures::pin_mut!(wait);
-
-                loop {
-                    // Repeatedly check if CLI is still open to avoid wasting resources
-                    // waiting for files or workspaces to close.
-                    let mut timer = background.timer(Duration::from_secs(1)).fuse();
-                    futures::select_biased! {
-                        _ = wait => break,
-                        _ = timer => {
-                            if responses.send(CliResponse::Ping).is_err() {
-                                break;
-                            }
-                        }
-                    }
-                }
-            }
-        }
+        Ok(result) => result,
         Err(error) => {
-            errored = true;
             responses
                 .send(CliResponse::Stderr {
                     message: format!("error opening {paths_with_position:?}: {error}"),
                 })
                 .log_err();
+            return true;
+        }
+    };
+
+    let mut errored = false;
+    let mut item_release_futures = Vec::new();
+    let mut subscriptions = Vec::new();
+
+    // If --wait flag is used with no paths, or a directory, then wait until
+    // the entire workspace is closed.
+    if wait {
+        let mut wait_for_window_close = paths_with_position.is_empty() && diff_paths.is_empty();
+        for path_with_position in &paths_with_position {
+            if app_state.fs.is_dir(&path_with_position.path).await {
+                wait_for_window_close = true;
+                break;
+            }
+        }
+
+        if wait_for_window_close {
+            let (release_tx, release_rx) = oneshot::channel();
+            item_release_futures.push(release_rx);
+            subscriptions.push(workspace.update(cx, |_, _, cx| {
+                cx.on_release(move |_, _| {
+                    let _ = release_tx.send(());
+                })
+            }));
         }
     }
+
+    for item in items {
+        match item {
+            Some(Ok(item)) => {
+                if wait {
+                    let (release_tx, release_rx) = oneshot::channel();
+                    item_release_futures.push(release_rx);
+                    subscriptions.push(cx.update(|cx| {
+                        item.on_release(
+                            cx,
+                            Box::new(move |_| {
+                                release_tx.send(()).ok();
+                            }),
+                        )
+                    }));
+                }
+            }
+            Some(Err(err)) => {
+                responses
+                    .send(CliResponse::Stderr {
+                        message: err.to_string(),
+                    })
+                    .log_err();
+                errored = true;
+            }
+            None => {}
+        }
+    }
+
+    if wait {
+        let wait = async move {
+            let _subscriptions = subscriptions;
+            let _ = future::try_join_all(item_release_futures).await;
+        }
+        .fuse();
+        futures::pin_mut!(wait);
+
+        let background = cx.background_executor().clone();
+        loop {
+            // Repeatedly check if CLI is still open to avoid wasting resources
+            // waiting for files or workspaces to close.
+            let mut timer = background.timer(Duration::from_secs(1)).fuse();
+            futures::select_biased! {
+                _ = wait => break,
+                _ = timer => {
+                    if responses.send(CliResponse::Ping).is_err() {
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
     errored
 }
 
@@ -653,12 +663,13 @@ mod tests {
         ipc::{self},
     };
     use editor::Editor;
-    use gpui::TestAppContext;
+    use futures::poll;
+    use gpui::{AppContext as _, TestAppContext};
     use language::LineEnding;
     use remote::SshConnectionOptions;
     use rope::Rope;
     use serde_json::json;
-    use std::sync::Arc;
+    use std::{sync::Arc, task::Poll};
     use util::path;
     use workspace::{AppState, Workspace};
 
@@ -754,6 +765,60 @@ mod tests {
             .unwrap();
     }
 
+    #[gpui::test]
+    async fn test_wait_with_directory_waits_for_window_close(cx: &mut TestAppContext) {
+        let app_state = init_test(cx);
+
+        app_state
+            .fs
+            .as_fake()
+            .insert_tree(
+                path!("/root"),
+                json!({
+                    "dir1": {
+                        "file1.txt": "content1",
+                    },
+                }),
+            )
+            .await;
+
+        let (response_tx, _) = ipc::channel::<CliResponse>().unwrap();
+        let workspace_paths = vec![path!("/root/dir1").to_owned()];
+
+        let (done_tx, mut done_rx) = futures::channel::oneshot::channel();
+        cx.spawn({
+            let app_state = app_state.clone();
+            move |mut cx| async move {
+                let errored = open_local_workspace(
+                    workspace_paths,
+                    vec![],
+                    None,
+                    false,
+                    true,
+                    &response_tx,
+                    None,
+                    &app_state,
+                    &mut cx,
+                )
+                .await;
+                let _ = done_tx.send(errored);
+            }
+        })
+        .detach();
+
+        cx.background_executor.run_until_parked();
+        assert_eq!(cx.windows().len(), 1);
+        assert!(matches!(poll!(&mut done_rx), Poll::Pending));
+
+        let window = cx.windows()[0];
+        cx.update_window(window, |_, window, _| window.remove_window())
+            .unwrap();
+        cx.background_executor.run_until_parked();
+
+        let errored = done_rx.await.unwrap();
+        assert!(!errored);
+    }
+
     #[gpui::test]
     async fn test_open_workspace_with_nonexistent_files(cx: &mut TestAppContext) {
         let app_state = init_test(cx);