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