From 71867096c88a70dc1609005baa4389bba9e5f548 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 12 Feb 2025 18:35:25 -0800
Subject: [PATCH] Migrate edit_prediction_provider setting before updating its
 value to 'zed' during onboarding (#24781)

This fixes a bug where we'd update your settings to an invalid state if
you were using the old `inline_completion_provider` setting, then
onboarded to Zeta, then migrated your settings.

Release Notes:

- N/A

Co-authored-by: Michael Sloan <mgsloan@gmail.com>
Co-authored-by: Agus Zubiaga <hi@aguz.me>
---
 Cargo.lock                          |  2 ++
 crates/migrator/src/migrator.rs     | 48 +++++++++++++++++------------
 crates/zeta/Cargo.toml              |  2 ++
 crates/zeta/src/onboarding_modal.rs | 16 ++++++++++
 4 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9653529d7e..a5661b189d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -17073,6 +17073,8 @@ dependencies = [
  "language_models",
  "log",
  "menu",
+ "migrator",
+ "paths",
  "postage",
  "project",
  "regex",
diff --git a/crates/migrator/src/migrator.rs b/crates/migrator/src/migrator.rs
index c4eca93009..72576241f9 100644
--- a/crates/migrator/src/migrator.rs
+++ b/crates/migrator/src/migrator.rs
@@ -68,6 +68,17 @@ pub fn migrate_settings(text: &str) -> Result<Option<String>> {
     )
 }
 
+pub fn migrate_edit_prediction_provider_settings(text: &str) -> Result<Option<String>> {
+    migrate(
+        &text,
+        &[(
+            SETTINGS_REPLACE_NESTED_KEY,
+            replace_edit_prediction_provider_setting,
+        )],
+        &EDIT_PREDICTION_SETTINGS_MIGRATION_QUERY,
+    )
+}
+
 type MigrationPatterns = &'static [(
     &'static str,
     fn(&str, &QueryMatch, &Query) -> Option<(Range<usize>, String)>,
@@ -550,7 +561,10 @@ pub static CONTEXT_REPLACE: LazyLock<HashMap<&str, &str>> = LazyLock::new(|| {
 
 const SETTINGS_MIGRATION_PATTERNS: MigrationPatterns = &[
     (SETTINGS_STRING_REPLACE_QUERY, replace_setting_name),
-    (SETTINGS_REPLACE_NESTED_KEY, replace_setting_nested_key),
+    (
+        SETTINGS_REPLACE_NESTED_KEY,
+        replace_edit_prediction_provider_setting,
+    ),
     (
         SETTINGS_REPLACE_IN_LANGUAGES_QUERY,
         replace_setting_in_languages,
@@ -568,6 +582,14 @@ static SETTINGS_MIGRATION_QUERY: LazyLock<Query> = LazyLock::new(|| {
     .unwrap()
 });
 
+static EDIT_PREDICTION_SETTINGS_MIGRATION_QUERY: LazyLock<Query> = LazyLock::new(|| {
+    Query::new(
+        &tree_sitter_json::LANGUAGE.into(),
+        SETTINGS_REPLACE_NESTED_KEY,
+    )
+    .unwrap()
+});
+
 const SETTINGS_STRING_REPLACE_QUERY: &str = r#"(document
     (object
         (pair
@@ -622,7 +644,7 @@ const SETTINGS_REPLACE_NESTED_KEY: &str = r#"
 )
 "#;
 
-fn replace_setting_nested_key(
+fn replace_edit_prediction_provider_setting(
     contents: &str,
     mat: &QueryMatch,
     query: &Query,
@@ -641,27 +663,13 @@ fn replace_setting_nested_key(
         .byte_range();
     let setting_name = contents.get(setting_range.clone())?;
 
-    let new_setting_name = SETTINGS_NESTED_STRING_REPLACE
-        .get(&parent_object_name)?
-        .get(setting_name)?;
+    if parent_object_name == "features" && setting_name == "inline_completion_provider" {
+        return Some((setting_range, "edit_prediction_provider".into()));
+    }
 
-    Some((setting_range, new_setting_name.to_string()))
+    None
 }
 
-/// ```json
-/// "features": {
-///   "inline_completion_provider": "copilot"
-/// },
-/// ```
-pub static SETTINGS_NESTED_STRING_REPLACE: LazyLock<
-    HashMap<&'static str, HashMap<&'static str, &'static str>>,
-> = LazyLock::new(|| {
-    HashMap::from_iter([(
-        "features",
-        HashMap::from_iter([("inline_completion_provider", "edit_prediction_provider")]),
-    )])
-});
-
 const SETTINGS_REPLACE_IN_LANGUAGES_QUERY: &str = r#"
 (object
   (pair
diff --git a/crates/zeta/Cargo.toml b/crates/zeta/Cargo.toml
index 7e1f46c5fe..f8f0bb4da3 100644
--- a/crates/zeta/Cargo.toml
+++ b/crates/zeta/Cargo.toml
@@ -36,6 +36,8 @@ language.workspace = true
 language_models.workspace = true
 log.workspace = true
 menu.workspace = true
+migrator.workspace = true
+paths.workspace = true
 postage.workspace = true
 project.workspace = true
 regex.workspace = true
diff --git a/crates/zeta/src/onboarding_modal.rs b/crates/zeta/src/onboarding_modal.rs
index 63aaa2b672..a9ad7469e0 100644
--- a/crates/zeta/src/onboarding_modal.rs
+++ b/crates/zeta/src/onboarding_modal.rs
@@ -1,6 +1,7 @@
 use std::{sync::Arc, time::Duration};
 
 use crate::{onboarding_event, ZED_PREDICT_DATA_COLLECTION_CHOICE};
+use anyhow::Context as _;
 use client::{Client, UserStore};
 use db::kvp::KEY_VALUE_STORE;
 use feature_flags::FeatureFlagAppExt as _;
@@ -83,6 +84,7 @@ impl ZedPredictModal {
         let task = self
             .user_store
             .update(cx, |this, cx| this.accept_terms_of_service(cx));
+        let fs = self.fs.clone();
 
         cx.spawn(|this, mut cx| async move {
             task.await?;
@@ -101,6 +103,20 @@ impl ZedPredictModal {
                 .await
                 .log_err();
 
+            // Make sure edit prediction provider setting is using the new key
+            let settings_path = paths::settings_file().as_path();
+            let settings_path = fs.canonicalize(settings_path).await.with_context(|| {
+                format!("Failed to canonicalize settings path {:?}", settings_path)
+            })?;
+
+            if let Some(settings) = fs.load(&settings_path).await.log_err() {
+                if let Some(new_settings) =
+                    migrator::migrate_edit_prediction_provider_settings(&settings)?
+                {
+                    fs.atomic_write(settings_path, new_settings).await?;
+                }
+            }
+
             this.update(&mut cx, |this, cx| {
                 update_settings_file::<AllLanguageSettings>(this.fs.clone(), cx, move |file, _| {
                     file.features