Check for `predict-edits` feature flag, remove `is_staff` check (#23165)

Agus Zubiaga and Thorsten Ball created 1 year ago

Release Notes:

- N/A

---------

Co-authored-by: Thorsten Ball <[e-mail]>

Change summary

crates/collab/src/llm.rs                                        | 7 +
crates/collab/src/llm/token.rs                                  | 4 +
crates/collab/src/rpc.rs                                        | 2 
crates/feature_flags/src/feature_flags.rs                       | 6 +-
crates/inline_completion_button/src/inline_completion_button.rs | 4 
crates/zed/src/zed/inline_completion_registry.rs                | 9 +-
6 files changed, 21 insertions(+), 11 deletions(-)

Detailed changes

crates/collab/src/llm.rs

@@ -440,8 +440,11 @@ async fn predict_edits(
     _country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
     Json(params): Json<PredictEditsParams>,
 ) -> Result<impl IntoResponse> {
-    if !claims.is_staff {
-        return Err(anyhow!("not found"))?;
+    if !claims.is_staff && !claims.has_predict_edits_feature_flag {
+        return Err(Error::http(
+            StatusCode::FORBIDDEN,
+            "no access to Zed's edit prediction feature".to_string(),
+        ));
     }
 
     let api_url = state

crates/collab/src/llm/token.rs

@@ -22,6 +22,8 @@ pub struct LlmTokenClaims {
     pub github_user_login: String,
     pub is_staff: bool,
     pub has_llm_closed_beta_feature_flag: bool,
+    #[serde(default)]
+    pub has_predict_edits_feature_flag: bool,
     pub has_llm_subscription: bool,
     pub max_monthly_spend_in_cents: u32,
     pub custom_llm_monthly_allowance_in_cents: Option<u32>,
@@ -37,6 +39,7 @@ impl LlmTokenClaims {
         is_staff: bool,
         billing_preferences: Option<billing_preference::Model>,
         has_llm_closed_beta_feature_flag: bool,
+        has_predict_edits_feature_flag: bool,
         has_llm_subscription: bool,
         plan: rpc::proto::Plan,
         system_id: Option<String>,
@@ -58,6 +61,7 @@ impl LlmTokenClaims {
             github_user_login: user.github_login.clone(),
             is_staff,
             has_llm_closed_beta_feature_flag,
+            has_predict_edits_feature_flag,
             has_llm_subscription,
             max_monthly_spend_in_cents: billing_preferences
                 .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {

crates/collab/src/rpc.rs

@@ -4025,6 +4025,7 @@ async fn get_llm_api_token(
     let flags = db.get_user_flags(session.user_id()).await?;
     let has_language_models_feature_flag = flags.iter().any(|flag| flag == "language-models");
     let has_llm_closed_beta_feature_flag = flags.iter().any(|flag| flag == "llm-closed-beta");
+    let has_predict_edits_feature_flag = flags.iter().any(|flag| flag == "predict-edits");
 
     if !session.is_staff() && !has_language_models_feature_flag {
         Err(anyhow!("permission denied"))?
@@ -4061,6 +4062,7 @@ async fn get_llm_api_token(
         session.is_staff(),
         billing_preferences,
         has_llm_closed_beta_feature_flag,
+        has_predict_edits_feature_flag,
         has_llm_subscription,
         session.current_plan(&db).await?,
         session.system_id.clone(),

crates/feature_flags/src/feature_flags.rs

@@ -59,9 +59,9 @@ impl FeatureFlag for ToolUseFeatureFlag {
     }
 }
 
-pub struct ZetaFeatureFlag;
-impl FeatureFlag for ZetaFeatureFlag {
-    const NAME: &'static str = "zeta";
+pub struct PredictEditsFeatureFlag;
+impl FeatureFlag for PredictEditsFeatureFlag {
+    const NAME: &'static str = "predict-edits";
 }
 
 pub struct GitUiFeatureFlag;

crates/inline_completion_button/src/inline_completion_button.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use copilot::{Copilot, Status};
 use editor::{scroll::Autoscroll, Editor};
-use feature_flags::{FeatureFlagAppExt, ZetaFeatureFlag};
+use feature_flags::{FeatureFlagAppExt, PredictEditsFeatureFlag};
 use fs::Fs;
 use gpui::{
     actions, div, pulsating_between, Action, Animation, AnimationExt, AppContext,
@@ -202,7 +202,7 @@ impl Render for InlineCompletionButton {
             }
 
             InlineCompletionProvider::Zed => {
-                if !cx.has_flag::<ZetaFeatureFlag>() {
+                if !cx.has_flag::<PredictEditsFeatureFlag>() {
                     return div();
                 }
 

crates/zed/src/zed/inline_completion_registry.rs

@@ -4,7 +4,7 @@ use client::Client;
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::{Editor, EditorMode};
-use feature_flags::{FeatureFlagAppExt, ZetaFeatureFlag};
+use feature_flags::{FeatureFlagAppExt, PredictEditsFeatureFlag};
 use gpui::{AnyWindowHandle, AppContext, Context, ViewContext, WeakView};
 use language::language_settings::{all_language_settings, InlineCompletionProvider};
 use settings::SettingsStore;
@@ -49,11 +49,11 @@ pub fn init(client: Arc<Client>, cx: &mut AppContext) {
         });
     }
 
-    if cx.has_flag::<ZetaFeatureFlag>() {
+    if cx.has_flag::<PredictEditsFeatureFlag>() {
         cx.on_action(clear_zeta_edit_history);
     }
 
-    cx.observe_flag::<ZetaFeatureFlag, _>({
+    cx.observe_flag::<PredictEditsFeatureFlag, _>({
         let editors = editors.clone();
         let client = client.clone();
         move |active, cx| {
@@ -164,8 +164,9 @@ fn assign_inline_completion_provider(
                 editor.set_inline_completion_provider(Some(provider), cx);
             }
         }
+
         language::language_settings::InlineCompletionProvider::Zed => {
-            if cx.has_flag::<ZetaFeatureFlag>()
+            if cx.has_flag::<PredictEditsFeatureFlag>()
                 || (cfg!(debug_assertions) && client.status().borrow().is_connected())
             {
                 let zeta = zeta::Zeta::register(client.clone(), cx);