feature_flags.rs

  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 Assistant2FeatureFlag;
 43
 44impl FeatureFlag for Assistant2FeatureFlag {
 45    const NAME: &'static str = "assistant2";
 46
 47    fn enabled_for_staff() -> bool {
 48        false
 49    }
 50}
 51
 52pub struct ToolUseFeatureFlag;
 53
 54impl FeatureFlag for ToolUseFeatureFlag {
 55    const NAME: &'static str = "assistant-tool-use";
 56
 57    fn enabled_for_staff() -> bool {
 58        false
 59    }
 60}
 61
 62pub struct PredictEditsFeatureFlag;
 63impl FeatureFlag for PredictEditsFeatureFlag {
 64    const NAME: &'static str = "predict-edits";
 65}
 66
 67pub struct GitUiFeatureFlag;
 68impl FeatureFlag for GitUiFeatureFlag {
 69    const NAME: &'static str = "git-ui";
 70}
 71
 72pub struct Remoting {}
 73impl FeatureFlag for Remoting {
 74    const NAME: &'static str = "remoting";
 75}
 76
 77pub struct LanguageModels {}
 78impl FeatureFlag for LanguageModels {
 79    const NAME: &'static str = "language-models";
 80}
 81
 82pub struct LlmClosedBeta {}
 83impl FeatureFlag for LlmClosedBeta {
 84    const NAME: &'static str = "llm-closed-beta";
 85}
 86
 87pub struct ZedPro {}
 88impl FeatureFlag for ZedPro {
 89    const NAME: &'static str = "zed-pro";
 90}
 91
 92pub struct NotebookFeatureFlag;
 93
 94impl FeatureFlag for NotebookFeatureFlag {
 95    const NAME: &'static str = "notebooks";
 96}
 97
 98pub struct AutoCommand {}
 99impl FeatureFlag for AutoCommand {
100    const NAME: &'static str = "auto-command";
101
102    fn enabled_for_staff() -> bool {
103        false
104    }
105}
106
107pub trait FeatureFlagViewExt<V: 'static> {
108    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
109    where
110        F: Fn(bool, &mut V, &mut ViewContext<V>) + Send + Sync + 'static;
111}
112
113impl<V> FeatureFlagViewExt<V> for ViewContext<'_, V>
114where
115    V: 'static,
116{
117    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
118    where
119        F: Fn(bool, &mut V, &mut ViewContext<V>) + 'static,
120    {
121        self.observe_global::<FeatureFlags>(move |v, cx| {
122            let feature_flags = cx.global::<FeatureFlags>();
123            callback(feature_flags.has_flag::<T>(), v, cx);
124        })
125    }
126}
127
128pub trait FeatureFlagAppExt {
129    fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag;
130    fn update_flags(&mut self, staff: bool, flags: Vec<String>);
131    fn set_staff(&mut self, staff: bool);
132    fn has_flag<T: FeatureFlag>(&self) -> bool;
133    fn is_staff(&self) -> bool;
134
135    fn observe_flag<T: FeatureFlag, F>(&mut self, callback: F) -> Subscription
136    where
137        F: FnMut(bool, &mut AppContext) + 'static;
138}
139
140impl FeatureFlagAppExt for AppContext {
141    fn update_flags(&mut self, staff: bool, flags: Vec<String>) {
142        let feature_flags = self.default_global::<FeatureFlags>();
143        feature_flags.staff = staff;
144        feature_flags.flags = flags;
145    }
146
147    fn set_staff(&mut self, staff: bool) {
148        let feature_flags = self.default_global::<FeatureFlags>();
149        feature_flags.staff = staff;
150    }
151
152    fn has_flag<T: FeatureFlag>(&self) -> bool {
153        self.try_global::<FeatureFlags>()
154            .map(|flags| flags.has_flag::<T>())
155            .unwrap_or(false)
156    }
157
158    fn is_staff(&self) -> bool {
159        self.try_global::<FeatureFlags>()
160            .map(|flags| flags.staff)
161            .unwrap_or(false)
162    }
163
164    fn observe_flag<T: FeatureFlag, F>(&mut self, mut callback: F) -> Subscription
165    where
166        F: FnMut(bool, &mut AppContext) + 'static,
167    {
168        self.observe_global::<FeatureFlags>(move |cx| {
169            let feature_flags = cx.global::<FeatureFlags>();
170            callback(feature_flags.has_flag::<T>(), cx);
171        })
172    }
173
174    fn wait_for_flag<T: FeatureFlag>(&mut self) -> WaitForFlag {
175        let (tx, rx) = oneshot::channel::<bool>();
176        let mut tx = Some(tx);
177        let subscription: Option<Subscription>;
178
179        match self.try_global::<FeatureFlags>() {
180            Some(feature_flags) => {
181                subscription = None;
182                tx.take().unwrap().send(feature_flags.has_flag::<T>()).ok();
183            }
184            None => {
185                subscription = Some(self.observe_global::<FeatureFlags>(move |cx| {
186                    let feature_flags = cx.global::<FeatureFlags>();
187                    if let Some(tx) = tx.take() {
188                        tx.send(feature_flags.has_flag::<T>()).ok();
189                    }
190                }));
191            }
192        }
193
194        WaitForFlag(rx, subscription)
195    }
196}
197
198pub struct WaitForFlag(oneshot::Receiver<bool>, Option<Subscription>);
199
200impl Future for WaitForFlag {
201    type Output = bool;
202
203    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
204        self.0.poll_unpin(cx).map(|result| {
205            self.1.take();
206            result.unwrap_or(false)
207        })
208    }
209}