zeta2: Feature flag (#40505)

Agus Zubiaga created 2 weeks ago

Release Notes:

- N/A

Change summary

Cargo.lock                                     |  2 +
crates/zed/src/zed/edit_prediction_registry.rs |  4 ++
crates/zeta2/Cargo.toml                        |  1 
crates/zeta2/src/zeta2.rs                      | 11 +++++++++
crates/zeta2_tools/Cargo.toml                  |  1 
crates/zeta2_tools/src/zeta2_tools.rs          | 23 +++++++++++++------
6 files changed, 34 insertions(+), 8 deletions(-)

Detailed changes

Cargo.lock

@@ -21781,6 +21781,7 @@ dependencies = [
  "cloud_zeta2_prompt",
  "edit_prediction",
  "edit_prediction_context",
+ "feature_flags",
  "futures 0.3.31",
  "gpui",
  "indoc",
@@ -21813,6 +21814,7 @@ dependencies = [
  "collections",
  "edit_prediction_context",
  "editor",
+ "feature_flags",
  "futures 0.3.31",
  "gpui",
  "indoc",

crates/zed/src/zed/edit_prediction_registry.rs

@@ -3,6 +3,7 @@ use codestral::CodestralCompletionProvider;
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::Editor;
+use feature_flags::FeatureFlagAppExt;
 use gpui::{AnyWindowHandle, App, AppContext as _, Context, Entity, WeakEntity};
 use language::language_settings::{EditPredictionProvider, all_language_settings};
 use language_models::MistralLanguageModelProvider;
@@ -11,6 +12,7 @@ use std::{cell::RefCell, rc::Rc, sync::Arc};
 use supermaven::{Supermaven, SupermavenCompletionProvider};
 use ui::Window;
 use zeta::ZetaEditPredictionProvider;
+use zeta2::Zeta2FeatureFlag;
 
 pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
@@ -217,7 +219,7 @@ fn assign_edit_prediction_provider(
                 }
 
                 if let Some(project) = editor.project() {
-                    if std::env::var("ZED_ZETA2").is_ok() {
+                    if cx.has_flag::<Zeta2FeatureFlag>() {
                         let zeta = zeta2::Zeta::global(client, &user_store, cx);
                         let provider = cx.new(|cx| {
                             zeta2::ZetaEditPredictionProvider::new(

crates/zeta2/Cargo.toml

@@ -20,6 +20,7 @@ cloud_llm_client.workspace = true
 cloud_zeta2_prompt.workspace = true
 edit_prediction.workspace = true
 edit_prediction_context.workspace = true
+feature_flags.workspace = true
 futures.workspace = true
 gpui.workspace = true
 indoc.workspace = true

crates/zeta2/src/zeta2.rs

@@ -11,6 +11,7 @@ use edit_prediction_context::{
     DeclarationId, DeclarationStyle, EditPredictionContext, EditPredictionContextOptions,
     EditPredictionExcerptOptions, EditPredictionScoreOptions, SyntaxIndex, SyntaxIndexState,
 };
+use feature_flags::FeatureFlag;
 use futures::AsyncReadExt as _;
 use futures::channel::{mpsc, oneshot};
 use gpui::http_client::{AsyncBody, Method};
@@ -66,6 +67,16 @@ pub const DEFAULT_OPTIONS: ZetaOptions = ZetaOptions {
     file_indexing_parallelism: 1,
 };
 
+pub struct Zeta2FeatureFlag;
+
+impl FeatureFlag for Zeta2FeatureFlag {
+    const NAME: &'static str = "zeta2";
+
+    fn enabled_for_staff() -> bool {
+        false
+    }
+}
+
 #[derive(Clone)]
 struct ZetaGlobal(Entity<Zeta>);
 

crates/zeta2_tools/Cargo.toml

@@ -18,6 +18,7 @@ cloud_llm_client.workspace = true
 collections.workspace = true
 edit_prediction_context.workspace = true
 editor.workspace = true
+feature_flags.workspace = true
 futures.workspace = true
 gpui.workspace = true
 language.workspace = true

crates/zeta2_tools/src/zeta2_tools.rs

@@ -8,6 +8,7 @@ use client::{Client, UserStore};
 use cloud_llm_client::predict_edits_v3::{DeclarationScoreComponents, PromptFormat};
 use collections::HashMap;
 use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
+use feature_flags::FeatureFlagAppExt as _;
 use futures::{StreamExt as _, channel::oneshot};
 use gpui::{
     CursorStyle, Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task, WeakEntity,
@@ -20,7 +21,7 @@ use ui::{ContextMenu, ContextMenuEntry, DropdownMenu, prelude::*};
 use ui_input::SingleLineInput;
 use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
 use workspace::{Item, SplitDirection, Workspace};
-use zeta2::{PredictionDebugInfo, Zeta, ZetaOptions};
+use zeta2::{PredictionDebugInfo, Zeta, Zeta2FeatureFlag, ZetaOptions};
 
 use edit_prediction_context::{
     DeclarationStyle, EditPredictionContextOptions, EditPredictionExcerptOptions,
@@ -676,17 +677,25 @@ impl Zeta2Inspector {
     }
 
     fn render_content(&self, cx: &mut Context<Self>) -> AnyElement {
+        if !cx.has_flag::<Zeta2FeatureFlag>() {
+            return Self::render_message("`zeta2` feature flag is not enabled");
+        }
+
         match self.last_prediction.as_ref() {
-            None => v_flex()
-                .size_full()
-                .justify_center()
-                .items_center()
-                .child(Label::new("No prediction").size(LabelSize::Large))
-                .into_any(),
+            None => Self::render_message("No prediction"),
             Some(prediction) => self.render_last_prediction(prediction, cx).into_any(),
         }
     }
 
+    fn render_message(message: impl Into<SharedString>) -> AnyElement {
+        v_flex()
+            .size_full()
+            .justify_center()
+            .items_center()
+            .child(Label::new(message).size(LabelSize::Large))
+            .into_any()
+    }
+
     fn render_last_prediction(&self, prediction: &LastPrediction, cx: &mut Context<Self>) -> Div {
         match &self.active_view {
             ActiveView::Context => div().size_full().child(prediction.context_editor.clone()),