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 NotebookFeatureFlag;
 63
 64impl FeatureFlag for NotebookFeatureFlag {
 65    const NAME: &'static str = "notebooks";
 66}
 67
 68pub struct AutoCommand {}
 69impl FeatureFlag for AutoCommand {
 70    const NAME: &'static str = "auto-command";
 71
 72    fn enabled_for_staff() -> bool {
 73        false
 74    }
 75}
 76
 77pub trait FeatureFlagViewExt<V: 'static> {
 78    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
 79    where
 80        F: Fn(bool, &mut V, &mut ViewContext<V>) + Send + Sync + 'static;
 81}
 82
 83impl<V> FeatureFlagViewExt<V> for ViewContext<'_, V>
 84where
 85    V: 'static,
 86{
 87    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
 88    where
 89        F: Fn(bool, &mut V, &mut ViewContext<V>) + 'static,
 90    {
 91        self.observe_global::<FeatureFlags>(move |v, cx| {
 92            let feature_flags = cx.global::<FeatureFlags>();
 93            callback(feature_flags.has_flag::<T>(), v, cx);
 94        })
 95    }
 96}
 97
 98pub trait FeatureFlagAppExt {
 99    fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag;
100    fn update_flags(&mut self, staff: bool, flags: Vec<String>);
101    fn set_staff(&mut self, staff: bool);
102    fn has_flag<T: FeatureFlag>(&self) -> bool;
103    fn is_staff(&self) -> bool;
104
105    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
106    where
107        F: FnMut(bool, &mut AppContext) + 'static;
108}
109
110impl FeatureFlagAppExt for AppContext {
111    fn update_flags(&mut self, staff: bool, flags: Vec<String>) {
112        let feature_flags = self.default_global::<FeatureFlags>();
113        feature_flags.staff = staff;
114        feature_flags.flags = flags;
115    }
116
117    fn set_staff(&mut self, staff: bool) {
118        let feature_flags = self.default_global::<FeatureFlags>();
119        feature_flags.staff = staff;
120    }
121
122    fn has_flag<T: FeatureFlag>(&self) -> bool {
123        self.try_global::<FeatureFlags>()
124            .map(|flags| flags.has_flag::<T>())
125            .unwrap_or(false)
126    }
127
128    fn is_staff(&self) -> bool {
129        self.try_global::<FeatureFlags>()
130            .map(|flags| flags.staff)
131            .unwrap_or(false)
132    }
133
134    fn observe_flag<T: FeatureFlag, F>(&mut self, mut callback: F) -> Subscription
135    where
136        F: FnMut(bool, &mut AppContext) + 'static,
137    {
138        self.observe_global::<FeatureFlags>(move |cx| {
139            let feature_flags = cx.global::<FeatureFlags>();
140            callback(feature_flags.has_flag::<T>(), cx);
141        })
142    }
143
144    fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag {
145        let (tx, rx) = oneshot::channel::<bool>();
146        let mut tx = Some(tx);
147        let subscription: Option<Subscription>;
148
149        match self.try_global::<FeatureFlags>() {
150            Some(feature_flags) => {
151                subscription = None;
152                tx.take().unwrap().send(feature_flags.has_flag::<T>()).ok();
153            }
154            None => {
155                subscription = Some(self.observe_global::<FeatureFlags>(move |cx| {
156                    let feature_flags = cx.global::<FeatureFlags>();
157                    if let Some(tx) = tx.take() {
158                        tx.send(feature_flags.has_flag::<T>()).ok();
159                    }
160                }));
161            }
162        }
163
164        WaitForFlag(rx, subscription)
165    }
166}
167
168pub struct WaitForFlag(oneshot::Receiver<bool>, Option<Subscription>);
169
170impl Future for WaitForFlag {
171    type Output = bool;
172
173    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
174        self.0.poll_unpin(cx).map(|result| {
175            self.1.take();
176            result.unwrap_or(false)
177        })
178    }
179}