Ensure Sweep and Mercury keys are loaded for Edit Prediction button (#44894)

Ben Kunkle created 1 week ago

Follow up for #44505 

Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/edit_prediction/src/mercury.rs                   |  6 ++
crates/edit_prediction/src/sweep_ai.rs                  |  6 ++
crates/edit_prediction_ui/src/edit_prediction_button.rs | 31 +++++++---
3 files changed, 34 insertions(+), 9 deletions(-)

Detailed changes

crates/edit_prediction/src/mercury.rs

@@ -309,3 +309,9 @@ pub fn mercury_api_token(cx: &mut App) -> Entity<ApiKeyState> {
         })
         .clone()
 }
+
+pub fn load_mercury_api_token(cx: &mut App) -> Task<Result<(), language_model::AuthenticateError>> {
+    mercury_api_token(cx).update(cx, |key_state, cx| {
+        key_state.load_if_needed(MERCURY_CREDENTIALS_URL, |s| s, cx)
+    })
+}

crates/edit_prediction/src/sweep_ai.rs

@@ -282,6 +282,12 @@ pub fn sweep_api_token(cx: &mut App) -> Entity<ApiKeyState> {
         .clone()
 }
 
+pub fn load_sweep_api_token(cx: &mut App) -> Task<Result<(), language_model::AuthenticateError>> {
+    sweep_api_token(cx).update(cx, |key_state, cx| {
+        key_state.load_if_needed(SWEEP_CREDENTIALS_URL, |s| s, cx)
+    })
+}
+
 #[derive(Debug, Clone, Serialize)]
 struct AutocompleteRequest {
     pub debug_info: Arc<str>,

crates/edit_prediction_ui/src/edit_prediction_button.rs

@@ -487,6 +487,21 @@ impl EditPredictionButton {
         cx.observe_global::<SettingsStore>(move |_, cx| cx.notify())
             .detach();
 
+        cx.observe_global::<EditPredictionStore>(move |_, cx| cx.notify())
+            .detach();
+
+        let sweep_api_token_task = edit_prediction::sweep_ai::load_sweep_api_token(cx);
+        let mercury_api_token_task = edit_prediction::mercury::load_mercury_api_token(cx);
+
+        cx.spawn(async move |this, cx| {
+            _ = futures::join!(sweep_api_token_task, mercury_api_token_task);
+            this.update(cx, |_, cx| {
+                cx.notify();
+            })
+            .ok();
+        })
+        .detach();
+
         CodestralEditPredictionDelegate::ensure_api_key_loaded(client.http_client(), cx);
 
         Self {
@@ -503,7 +518,7 @@ impl EditPredictionButton {
         }
     }
 
-    fn get_available_providers(&self, cx: &App) -> Vec<EditPredictionProvider> {
+    fn get_available_providers(&self, cx: &mut App) -> Vec<EditPredictionProvider> {
         let mut providers = Vec::new();
 
         providers.push(EditPredictionProvider::Zed);
@@ -532,12 +547,10 @@ impl EditPredictionButton {
             providers.push(EditPredictionProvider::Codestral);
         }
 
-        let ep_store = EditPredictionStore::try_global(cx);
-
         if cx.has_flag::<SweepFeatureFlag>()
-            && ep_store
-                .as_ref()
-                .is_some_and(|ep_store| ep_store.read(cx).has_sweep_api_token(cx))
+            && edit_prediction::sweep_ai::sweep_api_token(cx)
+                .read(cx)
+                .has_key()
         {
             providers.push(EditPredictionProvider::Experimental(
                 EXPERIMENTAL_SWEEP_EDIT_PREDICTION_PROVIDER_NAME,
@@ -545,9 +558,9 @@ impl EditPredictionButton {
         }
 
         if cx.has_flag::<MercuryFeatureFlag>()
-            && ep_store
-                .as_ref()
-                .is_some_and(|ep_store| ep_store.read(cx).has_mercury_api_token(cx))
+            && edit_prediction::mercury::mercury_api_token(cx)
+                .read(cx)
+                .has_key()
         {
             providers.push(EditPredictionProvider::Experimental(
                 EXPERIMENTAL_MERCURY_EDIT_PREDICTION_PROVIDER_NAME,