zeta: Refresh LLM token in case it expired (#21796)

Release Notes: - N/A --------- Co-authored-by: Antonio <[email protected]> Co-authored-by: Bennet <[email protected]>
zed-industries · Dec 10, 2024 · 96499b7 · 96499b7
1 parent 09006aa
commit 96499b7
Show file tree

Hide file tree

Showing 4 changed files with 190 additions and 113 deletions.
diff --git a/crates/feature_flags/src/feature_flags.rs b/crates/feature_flags/src/feature_flags.rs
@@ -59,6 +59,11 @@ impl FeatureFlag for ToolUseFeatureFlag {
     }
 }
 
+pub struct ZetaFeatureFlag;
+impl FeatureFlag for ZetaFeatureFlag {
+    const NAME: &'static str = "zeta";
+}
+
 pub struct Remoting {}
 impl FeatureFlag for Remoting {
     const NAME: &'static str = "remoting";

diff --git 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;
+use feature_flags::{FeatureFlagAppExt, ZetaFeatureFlag};
 use fs::Fs;
 use gpui::{
     div, Action, AnchorCorner, AppContext, AsyncWindowContext, Entity, IntoElement, ParentElement,
@@ -199,7 +199,7 @@ impl Render for InlineCompletionButton {
             }
 
             InlineCompletionProvider::Zeta => {
-                if !cx.is_staff() {
+                if !cx.has_flag::<ZetaFeatureFlag>() {
                     return div();
                 }
 

diff --git a/crates/zed/src/zed/inline_completion_registry.rs b/crates/zed/src/zed/inline_completion_registry.rs
@@ -4,9 +4,9 @@ use client::Client;
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::{Editor, EditorMode};
-use feature_flags::FeatureFlagAppExt;
+use feature_flags::{FeatureFlagAppExt, ZetaFeatureFlag};
 use gpui::{AnyWindowHandle, AppContext, Context, ViewContext, WeakView};
-use language::language_settings::all_language_settings;
+use language::language_settings::{all_language_settings, InlineCompletionProvider};
 use settings::SettingsStore;
 use supermaven::{Supermaven, SupermavenCompletionProvider};
 
@@ -49,22 +49,45 @@ pub fn init(client: Arc<Client>, cx: &mut AppContext) {
         });
     }
 
-    cx.observe_global::<SettingsStore>(move |cx| {
-        let new_provider = all_language_settings(None, cx).inline_completions.provider;
-        if new_provider != provider {
-            provider = new_provider;
-            for (editor, window) in editors.borrow().iter() {
-                _ = window.update(cx, |_window, cx| {
-                    _ = editor.update(cx, |editor, cx| {
-                        assign_inline_completion_provider(editor, provider, &client, cx);
-                    })
-                });
+    cx.observe_flag::<ZetaFeatureFlag, _>({
+        let editors = editors.clone();
+        let client = client.clone();
+        move |_flag, cx| {
+            let provider = all_language_settings(None, cx).inline_completions.provider;
+            assign_inline_completion_providers(&editors, provider, &client, cx)
+        }
+    })
+    .detach();
+
+    cx.observe_global::<SettingsStore>({
+        let editors = editors.clone();
+        let client = client.clone();
+        move |cx| {
+            let new_provider = all_language_settings(None, cx).inline_completions.provider;
+            if new_provider != provider {
+                provider = new_provider;
+                assign_inline_completion_providers(&editors, provider, &client, cx)
             }
         }
     })
     .detach();
 }
 
+fn assign_inline_completion_providers(
+    editors: &Rc<RefCell<HashMap<WeakView<Editor>, AnyWindowHandle>>>,
+    provider: InlineCompletionProvider,
+    client: &Arc<Client>,
+    cx: &mut AppContext,
+) {
+    for (editor, window) in editors.borrow().iter() {
+        _ = window.update(cx, |_window, cx| {
+            _ = editor.update(cx, |editor, cx| {
+                assign_inline_completion_provider(editor, provider, &client, cx);
+            })
+        });
+    }
+}
+
 fn register_backward_compatible_actions(editor: &mut Editor, cx: &ViewContext<Editor>) {
     // We renamed some of these actions to not be copilot-specific, but that
     // would have not been backwards-compatible. So here we are re-registering
@@ -129,7 +152,7 @@ fn assign_inline_completion_provider(
             }
         }
         language::language_settings::InlineCompletionProvider::Zeta => {
-            if cx.is_staff() {
+            if cx.has_flag::<ZetaFeatureFlag>() {
                 let zeta = zeta::Zeta::register(client.clone(), cx);
                 if let Some(buffer) = editor.buffer().read(cx).as_singleton() {
                     if buffer.read(cx).file().is_some() {