From 4a7630204a56b20262b80c16e0201664c48b8dd5 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Wed, 15 Jan 2025 10:52:10 -0300 Subject: [PATCH] Check for `predict-edits` feature flag, remove `is_staff` check (#23165) Release Notes: - N/A --------- Co-authored-by: Thorsten Ball --- 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 +++--- .../src/inline_completion_button.rs | 4 ++-- crates/zed/src/zed/inline_completion_registry.rs | 9 +++++---- 6 files changed, 21 insertions(+), 11 deletions(-) diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index 830ea3673a..8bf3fa6443 100644 --- a/crates/collab/src/llm.rs +++ b/crates/collab/src/llm.rs @@ -440,8 +440,11 @@ async fn predict_edits( _country_code_header: Option>, Json(params): Json, ) -> Result { - 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 diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs index bc0a14cda1..5629beb322 100644 --- a/crates/collab/src/llm/token.rs +++ b/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, @@ -37,6 +39,7 @@ impl LlmTokenClaims { is_staff: bool, billing_preferences: Option, has_llm_closed_beta_feature_flag: bool, + has_predict_edits_feature_flag: bool, has_llm_subscription: bool, plan: rpc::proto::Plan, system_id: Option, @@ -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| { diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 5cd635b4c3..37d3d15df7 100644 --- a/crates/collab/src/rpc.rs +++ b/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(), diff --git a/crates/feature_flags/src/feature_flags.rs b/crates/feature_flags/src/feature_flags.rs index b32e56aa5e..f8767aafa5 100644 --- a/crates/feature_flags/src/feature_flags.rs +++ b/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; diff --git a/crates/inline_completion_button/src/inline_completion_button.rs b/crates/inline_completion_button/src/inline_completion_button.rs index 0147a98e82..b210a163e7 100644 --- a/crates/inline_completion_button/src/inline_completion_button.rs +++ b/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::() { + if !cx.has_flag::() { return div(); } diff --git a/crates/zed/src/zed/inline_completion_registry.rs b/crates/zed/src/zed/inline_completion_registry.rs index efd3737c9f..475e871d09 100644 --- a/crates/zed/src/zed/inline_completion_registry.rs +++ b/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, cx: &mut AppContext) { }); } - if cx.has_flag::() { + if cx.has_flag::() { cx.on_action(clear_zeta_edit_history); } - cx.observe_flag::({ + cx.observe_flag::({ 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::() + if cx.has_flag::() || (cfg!(debug_assertions) && client.status().borrow().is_connected()) { let zeta = zeta::Zeta::register(client.clone(), cx);