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