1use futures::{channel::oneshot, FutureExt as _};
  2use gpui::{AppContext, Global, Subscription, ViewContext};
  3use std::{
  4    future::Future,
  5    pin::Pin,
  6    task::{Context, Poll},
  7};
  8
  9#[derive(Default)]
 10struct FeatureFlags {
 11    flags: Vec<String>,
 12    staff: bool,
 13}
 14
 15impl FeatureFlags {
 16    fn has_flag<T: FeatureFlag>(&self) -> bool {
 17        if self.staff && T::enabled_for_staff() {
 18            return true;
 19        }
 20
 21        self.flags.iter().any(|f| f.as_str() == T::NAME)
 22    }
 23}
 24
 25impl Global for FeatureFlags {}
 26
 27/// To create a feature flag, implement this trait on a trivial type and use it as
 28/// a generic parameter when called [`FeatureFlagAppExt::has_flag`].
 29///
 30/// Feature flags are enabled for members of Zed staff by default. To disable this behavior
 31/// so you can test flags being disabled, set ZED_DISABLE_STAFF=1 in your environment,
 32/// which will force Zed to treat the current user as non-staff.
 33pub trait FeatureFlag {
 34    const NAME: &'static str;
 35
 36    /// Returns whether this feature flag is enabled for Zed staff.
 37    fn enabled_for_staff() -> bool {
 38        true
 39    }
 40}
 41
 42pub struct Remoting {}
 43impl FeatureFlag for Remoting {
 44    const NAME: &'static str = "remoting";
 45}
 46
 47pub struct LanguageModels {}
 48impl FeatureFlag for LanguageModels {
 49    const NAME: &'static str = "language-models";
 50}
 51
 52pub struct LlmClosedBeta {}
 53impl FeatureFlag for LlmClosedBeta {
 54    const NAME: &'static str = "llm-closed-beta";
 55}
 56
 57pub struct ZedPro {}
 58impl FeatureFlag for ZedPro {
 59    const NAME: &'static str = "zed-pro";
 60}
 61
 62pub struct AutoCommand {}
 63impl FeatureFlag for AutoCommand {
 64    const NAME: &'static str = "auto-command";
 65
 66    fn enabled_for_staff() -> bool {
 67        false
 68    }
 69}
 70
 71pub trait FeatureFlagViewExt<V: 'static> {
 72    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
 73    where
 74        F: Fn(bool, &mut V, &mut ViewContext<V>) + Send + Sync + 'static;
 75}
 76
 77impl<V> FeatureFlagViewExt<V> for ViewContext<'_, V>
 78where
 79    V: 'static,
 80{
 81    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
 82    where
 83        F: Fn(bool, &mut V, &mut ViewContext<V>) + 'static,
 84    {
 85        self.observe_global::<FeatureFlags>(move |v, cx| {
 86            let feature_flags = cx.global::<FeatureFlags>();
 87            callback(feature_flags.has_flag::<T>(), v, cx);
 88        })
 89    }
 90}
 91
 92pub trait FeatureFlagAppExt {
 93    fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag;
 94    fn update_flags(&mut self, staff: bool, flags: Vec<String>);
 95    fn set_staff(&mut self, staff: bool);
 96    fn has_flag<T: FeatureFlag>(&self) -> bool;
 97    fn is_staff(&self) -> bool;
 98
 99    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
100    where
101        F: FnMut(bool, &mut AppContext) + 'static;
102}
103
104impl FeatureFlagAppExt for AppContext {
105    fn update_flags(&mut self, staff: bool, flags: Vec<String>) {
106        let feature_flags = self.default_global::<FeatureFlags>();
107        feature_flags.staff = staff;
108        feature_flags.flags = flags;
109    }
110
111    fn set_staff(&mut self, staff: bool) {
112        let feature_flags = self.default_global::<FeatureFlags>();
113        feature_flags.staff = staff;
114    }
115
116    fn has_flag<T: FeatureFlag>(&self) -> bool {
117        self.try_global::<FeatureFlags>()
118            .map(|flags| flags.has_flag::<T>())
119            .unwrap_or(false)
120    }
121
122    fn is_staff(&self) -> bool {
123        self.try_global::<FeatureFlags>()
124            .map(|flags| flags.staff)
125            .unwrap_or(false)
126    }
127
128    fn observe_flag<T: FeatureFlag, F>(&mut self, mut callback: F) -> Subscription
129    where
130        F: FnMut(bool, &mut AppContext) + 'static,
131    {
132        self.observe_global::<FeatureFlags>(move |cx| {
133            let feature_flags = cx.global::<FeatureFlags>();
134            callback(feature_flags.has_flag::<T>(), cx);
135        })
136    }
137
138    fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag {
139        let (tx, rx) = oneshot::channel::<bool>();
140        let mut tx = Some(tx);
141        let subscription: Option<Subscription>;
142
143        match self.try_global::<FeatureFlags>() {
144            Some(feature_flags) => {
145                subscription = None;
146                tx.take().unwrap().send(feature_flags.has_flag::<T>()).ok();
147            }
148            None => {
149                subscription = Some(self.observe_global::<FeatureFlags>(move |cx| {
150                    let feature_flags = cx.global::<FeatureFlags>();
151                    if let Some(tx) = tx.take() {
152                        tx.send(feature_flags.has_flag::<T>()).ok();
153                    }
154                }));
155            }
156        }
157
158        WaitForFlag(rx, subscription)
159    }
160}
161
162pub struct WaitForFlag(oneshot::Receiver<bool>, Option<Subscription>);
163
164impl Future for WaitForFlag {
165    type Output = bool;
166
167    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
168        self.0.poll_unpin(cx).map(|result| {
169            self.1.take();
170            result.unwrap_or(false)
171        })
172    }
173}