feature_flags: Add FeatureFlagAppExt::wait_for_flag_or_timeout (#24055)

This PR adds a new `wait_for_flag_or_timeout` method to the
`FeatureFlagAppExt` trait.

This encapsulates the somewhat gnarly code for using `wait_for_flag`
with a timeout.

A side benefit of this is that the tasks waiting on the feature flags
run in parallel, so in the case where the feature flags do not resolve
we don't end up having to wait on consecutive timeouts. This should help
a bit with https://github.com/zed-industries/zed/issues/23922.

Release Notes:

- N/A

Cargo.lock

@@ -4547,6 +4547,7 @@ version = "0.1.0"
 dependencies = [
  "futures 0.3.31",
  "gpui",
+ "smol",
 ]
 
 [[package]]

crates/feature_flags/Cargo.toml

@@ -12,5 +12,6 @@ workspace = true
 path = "src/feature_flags.rs"
 
 [dependencies]
-gpui.workspace = true
 futures.workspace = true
+gpui.workspace = true
+smol.workspace = true

crates/feature_flags/src/feature_flags.rs

@@ -1,5 +1,7 @@
-use futures::{channel::oneshot, FutureExt as _};
-use gpui::{App, Context, Global, Subscription, Window};
+use futures::channel::oneshot;
+use futures::{select_biased, FutureExt};
+use gpui::{App, Context, Global, Subscription, Task, Window};
+use std::time::Duration;
 use std::{future::Future, pin::Pin, task::Poll};
 
 #[derive(Default)]
@@ -114,6 +116,10 @@ where
 
 pub trait FeatureFlagAppExt {
     fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag;
+
+    /// Waits for the specified feature flag to resolve, up to the given timeout.
+    fn wait_for_flag_or_timeout<T: FeatureFlag>(&mut self, timeout: Duration) -> Task<bool>;
+
     fn update_flags(&mut self, staff: bool, flags: Vec<String>);
     fn set_staff(&mut self, staff: bool);
     fn has_flag<T: FeatureFlag>(&self) -> bool;
@@ -180,6 +186,20 @@ impl FeatureFlagAppExt for App {
 
         WaitForFlag(rx, subscription)
     }
+
+    fn wait_for_flag_or_timeout<T: FeatureFlag>(&mut self, timeout: Duration) -> Task<bool> {
+        let wait_for_flag = self.wait_for_flag::<T>();
+
+        self.spawn(|_cx| async move {
+            let mut wait_for_flag = wait_for_flag.fuse();
+            let mut timeout = FutureExt::fuse(smol::Timer::after(timeout));
+
+            select_biased! {
+                is_enabled = wait_for_flag => is_enabled,
+                _ = timeout => false,
+            }
+        })
+    }
 }
 
 pub struct WaitForFlag(oneshot::Receiver<bool>, Option<Subscription>);

crates/zed/src/zed.rs

@@ -20,7 +20,6 @@ use command_palette_hooks::CommandPaletteFilter;
 use editor::ProposedChangesEditorToolbar;
 use editor::{scroll::Autoscroll, Editor, MultiBuffer};
 use feature_flags::FeatureFlagAppExt;
-use futures::FutureExt;
 use futures::{channel::mpsc, select_biased, StreamExt};
 use gpui::{
     actions, point, px, Action, App, AppContext as _, AsyncApp, Context, DismissEvent, Element,
@@ -45,6 +44,7 @@ use settings::{
 };
 use std::any::TypeId;
 use std::path::PathBuf;
+use std::time::Duration;
 use std::{borrow::Cow, ops::Deref, path::Path, sync::Arc};
 use terminal_view::terminal_panel::{self, TerminalPanel};
 use theme::{ActiveTheme, ThemeSettings};
@@ -362,8 +362,10 @@ fn initialize_panels(
     window: &mut Window,
     cx: &mut Context<Workspace>,
 ) {
-    let assistant2_feature_flag = cx.wait_for_flag::<feature_flags::Assistant2FeatureFlag>();
-    let git_ui_feature_flag = cx.wait_for_flag::<feature_flags::GitUiFeatureFlag>();
+    let assistant2_feature_flag =
+        cx.wait_for_flag_or_timeout::<feature_flags::Assistant2FeatureFlag>(Duration::from_secs(5));
+    let git_ui_feature_flag =
+        cx.wait_for_flag_or_timeout::<feature_flags::GitUiFeatureFlag>(Duration::from_secs(5));
 
     let prompt_builder = prompt_builder.clone();
 
@@ -405,16 +407,7 @@ fn initialize_panels(
             workspace.add_panel(notification_panel, window, cx);
         })?;
 
-        let git_ui_enabled = {
-            let mut git_ui_feature_flag = git_ui_feature_flag.fuse();
-            let mut timeout =
-                FutureExt::fuse(smol::Timer::after(std::time::Duration::from_secs(5)));
-
-            select_biased! {
-                is_git_ui_enabled = git_ui_feature_flag => is_git_ui_enabled,
-                _ = timeout => false,
-            }
-        };
+        let git_ui_enabled = git_ui_feature_flag.await;
 
         let git_panel = if git_ui_enabled {
             Some(git_ui::git_panel::GitPanel::load(workspace_handle.clone(), cx.clone()).await?)
@@ -430,14 +423,7 @@ fn initialize_panels(
         let is_assistant2_enabled = if cfg!(test) {
             false
         } else {
-            let mut assistant2_feature_flag = assistant2_feature_flag.fuse();
-            let mut timeout =
-                FutureExt::fuse(smol::Timer::after(std::time::Duration::from_secs(5)));
-
-            select_biased! {
-                is_assistant2_enabled = assistant2_feature_flag => is_assistant2_enabled,
-                _ = timeout => false,
-            }
+            assistant2_feature_flag.await
         };
 
         let (assistant_panel, assistant2_panel) = if is_assistant2_enabled {