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