edit_prediction: Expose allow_data_collection in settings (#51389)

Closes #48394

Moves the data collection preference for Zed's Edit Predictions out of
the internal KV store and into `settings.json` as a proper
`allow_data_collection` setting under `edit_predictions`.

**Migration:** Existing users' choices are preserved. When
`allow_data_collection` is absent from `settings.json`, the resolved
value falls back to the legacy KV entry
(`zed_predict_data_collection_choice`). Once the user toggles the
setting or sets it explicitly, the new setting takes precedence and the
KV entry is ignored.

**Bug fixed:** The original implementation of `toggle_data_collection`
read the raw (unresolved) settings content to determine the current
state. When `allow_data_collection` was absent from `settings.json` but
the KV store held `"true"`, the raw read returned `None → false`,
causing the first toggle click to write `Some(true)` (re-enabling)
instead of `Some(false)` (disabling). The fix reads the resolved
`is_data_collection_enabled()` value before entering the
`update_settings_file` closure.

## Manual testing

**Setting takes effect:**
1. Open settings (`cmd+,`) and add `"allow_data_collection": true` under
`edit_predictions`. Save.
2. Open a file — the data collection indicator in the editor should
reflect the enabled state.
3. Flip to `false` and confirm it updates.

**Toggle correctly disables from KV-enabled state (migration bug fix):**
1. Remove `allow_data_collection` from `settings.json`.
2. Write the legacy KV entry directly:
   ```
   sqlite3 ~/Library/Application\ Support/Zed/db/0-dev/db.sqlite \
"INSERT OR REPLACE INTO kv_store(key,value)
VALUES('zed_predict_data_collection_choice','true');"
   ```
3. Restart Zed. The data collection toggle should show as **enabled**
(reading from KV store).
4. Click the toggle once to disable. `allow_data_collection` should
appear as `false` in `settings.json` — not `true`, which was the pre-fix
behaviour.

**Upsell modal still appears for new users:**
1. Clear both KV keys and restart:
   ```
   sqlite3 ~/Library/Application\ Support/Zed/db/0-dev/db.sqlite \
"DELETE FROM kv_store WHERE key IN
('zed_predict_data_collection_choice','dismissed-edit-predict-upsell');"
   ```
2. Open any file so the status bar is visible.
3. Click the edit prediction button (bottom-right status bar) — it
should have a muted dot indicator.
4. The upsell modal should appear. Dismissing it should prevent it from
reappearing.

## Release Notes:

- `allow_data_collection` for Zed's Edit Predictions can now be set
explicitly in `settings.json` under `edit_predictions`. Existing
preferences stored in the internal database are preserved as a fallback.

---------

Co-authored-by: Ben Kunkle <ben.kunkle@gmail.com>

assets/settings/default.json

@@ -1683,6 +1683,14 @@
       "prompt_format": "infer",
       "max_output_tokens": 64,
     },
+    // Controls whether Zed may collect training data when using Zed's Edit Predictions.
+    // Data is only captured when the project is detected as open source.
+    // Possible values:
+    //   - "default": use the preference previously set via the status-bar toggle,
+    //     or false if no preference has been stored.
+    //   - "yes": allow data collection for files in open-source projects.
+    //   - "no": never allow data collection.
+    "allow_data_collection": "default",
   },
   // Settings specific to journaling
   "journal": {

crates/client/src/user.rs

@@ -740,6 +740,21 @@ impl UserStore {
             .get(&current_organization.id)
     }
 
+    #[cfg(any(test, feature = "test-support"))]
+    pub fn set_current_organization_configuration_for_test(
+        &mut self,
+        organization: Arc<Organization>,
+        configuration: OrganizationConfiguration,
+        cx: &mut Context<Self>,
+    ) {
+        self.current_organization = Some(organization.clone());
+        self.organizations = vec![organization.clone()];
+        self.configuration_by_organization
+            .insert(organization.id.clone(), configuration);
+        cx.emit(Event::OrganizationChanged);
+        cx.notify();
+    }
+
     pub fn plan(&self) -> Option<Plan> {
         #[cfg(debug_assertions)]
         if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {

crates/edit_prediction/src/edit_prediction.rs

@@ -40,7 +40,8 @@ use release_channel::AppVersion;
 use semver::Version;
 use serde::de::DeserializeOwned;
 use settings::{
-    EditPredictionPromptFormat, EditPredictionProvider, Settings as _, update_settings_file,
+    EditPredictionDataCollectionChoice, EditPredictionPromptFormat, EditPredictionProvider,
+    Settings as _, update_settings_file,
 };
 use std::collections::{VecDeque, hash_map};
 use std::env;
@@ -151,7 +152,7 @@ pub struct EditPredictionStore {
     preferred_experiment: Option<String>,
     available_experiments: Vec<String>,
     pub mercury: Mercury,
-    data_collection_choice: DataCollectionChoice,
+    legacy_data_collection_enabled: bool,
     reject_predictions_tx: mpsc::UnboundedSender<EditPredictionRejectionPayload>,
     settled_predictions_tx: mpsc::UnboundedSender<Instant>,
     shown_predictions: VecDeque<EditPrediction>,
@@ -757,9 +758,8 @@ impl EditPredictionStore {
     }
 
     pub fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
-        let data_collection_choice = Self::load_data_collection_choice(cx);
-
         let llm_token = global_llm_token(cx);
+        let legacy_data_collection_enabled = Self::load_legacy_data_collection_enabled(cx);
 
         let (reject_tx, reject_rx) = mpsc::unbounded();
         cx.background_spawn({
@@ -814,8 +814,8 @@ impl EditPredictionStore {
             preferred_experiment: None,
             available_experiments: Vec::new(),
             mercury: Mercury::new(cx),
+            legacy_data_collection_enabled,
 
-            data_collection_choice,
             reject_predictions_tx: reject_tx,
             settled_predictions_tx,
             rated_predictions: Default::default(),
@@ -2770,38 +2770,45 @@ impl EditPredictionStore {
     }
 
     pub(crate) fn is_data_collection_enabled(&self, cx: &App) -> bool {
-        self.data_collection_choice.is_enabled(cx)
-    }
+        if !self.is_data_collection_allowed_by_organization(cx) {
+            return false;
+        }
 
-    fn load_data_collection_choice(cx: &App) -> DataCollectionChoice {
-        let choice = KeyValueStore::global(cx)
-            .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
-            .log_err()
-            .flatten();
+        if cx.is_staff() {
+            return true;
+        }
 
-        match choice.as_deref() {
-            Some("true") => DataCollectionChoice::Enabled,
-            Some("false") => DataCollectionChoice::Disabled,
-            Some(_) => {
-                log::error!("unknown value in '{ZED_PREDICT_DATA_COLLECTION_CHOICE}'");
-                DataCollectionChoice::NotAnswered
-            }
-            None => DataCollectionChoice::NotAnswered,
+        match all_language_settings(None, cx)
+            .edit_predictions
+            .allow_data_collection
+        {
+            EditPredictionDataCollectionChoice::Yes => true,
+            EditPredictionDataCollectionChoice::No => false,
+            // Fall back to the legacy KV entry captured when the store was
+            // created, preserving existing users' choices without per-request
+            // database reads.
+            EditPredictionDataCollectionChoice::Default => self.legacy_data_collection_enabled,
         }
     }
 
-    fn toggle_data_collection_choice(&mut self, cx: &mut Context<Self>) {
-        self.data_collection_choice = self.data_collection_choice.toggle();
-        let new_choice = self.data_collection_choice;
-        let is_enabled = new_choice.is_enabled(cx);
-        let kvp = KeyValueStore::global(cx);
-        db::write_and_log(cx, move || async move {
-            kvp.write_kvp(
-                ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
-                is_enabled.to_string(),
-            )
-            .await
-        });
+    fn load_legacy_data_collection_enabled(cx: &App) -> bool {
+        KeyValueStore::global(cx)
+            .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
+            .log_err()
+            .flatten()
+            .as_deref()
+            == Some("true")
+    }
+
+    pub(crate) fn is_data_collection_allowed_by_organization(&self, cx: &App) -> bool {
+        self.user_store
+            .read(cx)
+            .current_organization_configuration()
+            .is_none_or(|organization_configuration| {
+                organization_configuration
+                    .edit_prediction
+                    .is_feedback_enabled
+            })
     }
 
     pub fn shown_predictions(&self) -> impl DoubleEndedIterator<Item = &EditPrediction> {
@@ -2994,70 +3001,37 @@ pub struct ZedUpdateRequiredError {
     minimum_version: Version,
 }
 
-#[derive(Debug, Clone, Copy)]
-pub enum DataCollectionChoice {
-    NotAnswered,
-    Enabled,
-    Disabled,
-}
-
-impl DataCollectionChoice {
-    pub fn is_enabled(self, cx: &App) -> bool {
-        if cx.is_staff() {
-            return true;
-        }
-        match self {
-            Self::Enabled => true,
-            Self::NotAnswered | Self::Disabled => false,
-        }
-    }
-
-    #[must_use]
-    pub fn toggle(&self) -> DataCollectionChoice {
-        match self {
-            Self::Enabled => Self::Disabled,
-            Self::Disabled => Self::Enabled,
-            Self::NotAnswered => Self::Enabled,
-        }
-    }
-}
-
-impl From<bool> for DataCollectionChoice {
-    fn from(value: bool) -> Self {
-        match value {
-            true => DataCollectionChoice::Enabled,
-            false => DataCollectionChoice::Disabled,
-        }
-    }
-}
-
 struct ZedPredictUpsell;
 
+fn is_upsell_dismissed(cx: &App) -> bool {
+    // To make this backwards compatible with older versions of Zed, we
+    // check if the user has seen the previous Edit Prediction Onboarding
+    // before, by checking the data collection choice which was written to
+    // the database once the user clicked on "Accept and Enable"
+    let kvp = KeyValueStore::global(cx);
+    if kvp
+        .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
+        .log_err()
+        .is_some_and(|s| s.is_some())
+    {
+        return true;
+    }
+
+    kvp.read_kvp(ZedPredictUpsell::KEY)
+        .log_err()
+        .is_some_and(|s| s.is_some())
+}
+
 impl Dismissable for ZedPredictUpsell {
     const KEY: &'static str = "dismissed-edit-predict-upsell";
 
     fn dismissed(cx: &App) -> bool {
-        // To make this backwards compatible with older versions of Zed, we
-        // check if the user has seen the previous Edit Prediction Onboarding
-        // before, by checking the data collection choice which was written to
-        // the database once the user clicked on "Accept and Enable"
-        let kvp = KeyValueStore::global(cx);
-        if kvp
-            .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
-            .log_err()
-            .is_some_and(|s| s.is_some())
-        {
-            return true;
-        }
-
-        kvp.read_kvp(Self::KEY)
-            .log_err()
-            .is_some_and(|s| s.is_some())
+        is_upsell_dismissed(cx)
     }
 }
 
 pub fn should_show_upsell_modal(cx: &App) -> bool {
-    !ZedPredictUpsell::dismissed(cx)
+    !is_upsell_dismissed(cx)
 }
 
 pub fn init(cx: &mut App) {

crates/edit_prediction/src/edit_prediction_tests.rs

@@ -3,11 +3,16 @@ use crate::udiff::apply_diff_to_string;
 use client::{RefreshLlmTokenListener, UserStore, test::FakeServer};
 use clock::FakeSystemClock;
 use clock::ReplicaId;
-use cloud_api_types::{CreateLlmTokenResponse, LlmToken};
+use cloud_api_types::{
+    CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration,
+    OrganizationEditPredictionConfiguration, OrganizationId,
+};
 use cloud_llm_client::{
     EditPredictionRejectReason, EditPredictionRejection, RejectEditPredictionsBody,
     predict_edits_v3::{PredictEditsV3Request, PredictEditsV3Response},
 };
+use db::AppDatabase;
+use settings::EditPredictionDataCollectionChoice;
 
 use futures::{
     AsyncReadExt, FutureExt, StreamExt,
@@ -2390,12 +2395,31 @@ struct RequestChannels {
 
 fn init_test_with_fake_client(
     cx: &mut TestAppContext,
+) -> (Entity<EditPredictionStore>, RequestChannels) {
+    init_test_with_fake_client_and_legacy_data_collection(cx, None)
+}
+
+fn init_test_with_fake_client_and_legacy_data_collection(
+    cx: &mut TestAppContext,
+    legacy_data_collection_choice: Option<&str>,
 ) -> (Entity<EditPredictionStore>, RequestChannels) {
     cx.update(move |cx| {
+        cx.set_global(AppDatabase::test_new());
         let settings_store = SettingsStore::test(cx);
         cx.set_global(settings_store);
         zlog::init_test();
 
+        if let Some(legacy_data_collection_choice) = legacy_data_collection_choice {
+            KeyValueStore::global(cx)
+                .write_kvp(
+                    ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
+                    legacy_data_collection_choice.to_string(),
+                )
+                .now_or_never()
+                .expect("legacy data collection write should complete immediately")
+                .expect("legacy data collection write should succeed");
+        }
+
         let (predict_req_tx, predict_req_rx) = mpsc::unbounded();
         let (reject_req_tx, reject_req_rx) = mpsc::unbounded();
 
@@ -2772,6 +2796,7 @@ async fn test_v3_prediction_strips_cursor_marker_from_edit_text(cx: &mut TestApp
 
 fn init_test(cx: &mut TestAppContext) {
     cx.update(|cx| {
+        cx.set_global(AppDatabase::test_new());
         let settings_store = SettingsStore::test(cx);
         cx.set_global(settings_store);
     });
@@ -3392,6 +3417,252 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) {
     }
 }
 
+#[gpui::test]
+async fn test_data_collection_disabled_by_default(cx: &mut TestAppContext) {
+    let (ep_store, _channels) = init_test_with_fake_client(cx);
+
+    cx.update(|cx| {
+        assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+#[gpui::test]
+async fn test_data_collection_enabled_via_legacy_kv_store(cx: &mut TestAppContext) {
+    let (ep_store, _channels) =
+        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
+
+    cx.update(|cx| {
+        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+#[gpui::test]
+async fn test_data_collection_default_uses_cached_legacy_value(cx: &mut TestAppContext) {
+    let (ep_store, _channels) =
+        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
+
+    cx.update(|cx| {
+        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+
+    cx.update(|cx| KeyValueStore::global(cx))
+        .delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
+        .await
+        .unwrap();
+
+    cx.update(|cx| {
+        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+#[gpui::test]
+async fn test_data_collection_setting_overrides_kv_store(cx: &mut TestAppContext) {
+    let (ep_store, _channels) =
+        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
+
+    // An explicit false in settings.json wins over the KV store.
+    cx.update_global::<SettingsStore, _>(|settings, cx| {
+        settings.update_user_settings(cx, |content| {
+            content
+                .project
+                .all_languages
+                .edit_predictions
+                .get_or_insert_default()
+                .allow_data_collection = Some(EditPredictionDataCollectionChoice::No);
+        });
+    });
+
+    cx.update(|cx| {
+        assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+#[gpui::test]
+async fn test_data_collection_enabled_via_setting(cx: &mut TestAppContext) {
+    let (ep_store, _channels) = init_test_with_fake_client(cx);
+
+    cx.update_global::<SettingsStore, _>(|settings, cx| {
+        settings.update_user_settings(cx, |content| {
+            content
+                .project
+                .all_languages
+                .edit_predictions
+                .get_or_insert_default()
+                .allow_data_collection = Some(EditPredictionDataCollectionChoice::Yes);
+        });
+    });
+
+    cx.update(|cx| {
+        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+#[gpui::test]
+async fn test_data_collection_always_enabled_for_staff(cx: &mut TestAppContext) {
+    let (ep_store, _channels) = init_test_with_fake_client(cx);
+
+    cx.update(|cx| {
+        cx.set_staff(true);
+        assert!(ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+#[gpui::test]
+async fn test_data_collection_disabled_by_organization_configuration(cx: &mut TestAppContext) {
+    let (ep_store, _channels) = init_test_with_fake_client(cx);
+
+    cx.update_global::<SettingsStore, _>(|settings, cx| {
+        settings.update_user_settings(cx, |content| {
+            content
+                .project
+                .all_languages
+                .edit_predictions
+                .get_or_insert_default()
+                .allow_data_collection = Some(EditPredictionDataCollectionChoice::Yes);
+        });
+    });
+
+    let user_store = cx.update(|cx| ep_store.read(cx).user_store.clone());
+    cx.update(|cx| {
+        user_store.update(cx, |user_store, cx| {
+            user_store.set_current_organization_configuration_for_test(
+                Arc::new(Organization {
+                    id: OrganizationId("org-1".into()),
+                    name: "Org 1".into(),
+                    is_personal: false,
+                }),
+                OrganizationConfiguration {
+                    is_zed_model_provider_enabled: true,
+                    is_agent_thread_feedback_enabled: true,
+                    is_collaboration_enabled: true,
+                    edit_prediction: OrganizationEditPredictionConfiguration {
+                        is_enabled: true,
+                        is_feedback_enabled: false,
+                    },
+                },
+                cx,
+            );
+        });
+
+        assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
+    });
+}
+
+// When a user had data collection enabled via the legacy KV store (with no explicit
+// setting in settings.json), toggle_data_collection must read the *resolved* state
+// (true) and write Some(false).
+#[gpui::test]
+async fn test_toggle_data_collection_from_kv_enabled_state(cx: &mut TestAppContext) {
+    let (ep_store, _channels) =
+        init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
+
+    cx.update(|cx| {
+        assert!(
+            ep_store.read(cx).is_data_collection_enabled(cx),
+            "data collection should be enabled via KV store before toggle"
+        );
+    });
+
+    // Simulate what toggle_data_collection does: capture the resolved current
+    // state, then write its inverse.
+    let is_currently_enabled = cx.update(|cx| ep_store.read(cx).is_data_collection_enabled(cx));
+    cx.update_global::<SettingsStore, _>(|settings, cx| {
+        settings.update_user_settings(cx, |content| {
+            content
+                .project
+                .all_languages
+                .edit_predictions
+                .get_or_insert_default()
+                .allow_data_collection = Some(if is_currently_enabled {
+                EditPredictionDataCollectionChoice::No
+            } else {
+                EditPredictionDataCollectionChoice::Yes
+            });
+        });
+    });
+
+    cx.update(|cx| {
+        assert!(
+            !ep_store.read(cx).is_data_collection_enabled(cx),
+            "data collection should be disabled after toggling off from KV-enabled state"
+        );
+    });
+}
+
+#[gpui::test]
+async fn test_upsell_shown_by_default(cx: &mut TestAppContext) {
+    init_test(cx);
+    let kvp = cx.update(|cx| KeyValueStore::global(cx));
+    kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
+        .await
+        .ok();
+    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.ok();
+
+    cx.update(|cx| assert!(should_show_upsell_modal(cx)));
+}
+
+#[gpui::test]
+async fn test_upsell_dismissed_when_data_collection_choice_in_kv_store(cx: &mut TestAppContext) {
+    init_test(cx);
+
+    // Any value for the data collection key means the old upsell was already
+    // shown, regardless of whether data collection was accepted or declined.
+    for value in &["true", "false"] {
+        cx.update(|cx| KeyValueStore::global(cx))
+            .write_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into(), value.to_string())
+            .await
+            .unwrap();
+
+        cx.update(|cx| {
+            assert!(
+                !should_show_upsell_modal(cx),
+                "upsell should be suppressed when data collection choice is '{value}'"
+            );
+        });
+    }
+
+    cx.update(|cx| KeyValueStore::global(cx))
+        .delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
+        .await
+        .unwrap();
+}
+
+#[gpui::test]
+async fn test_upsell_dismissed_when_dismissed_key_set(cx: &mut TestAppContext) {
+    init_test(cx);
+    let kvp = cx.update(|cx| KeyValueStore::global(cx));
+    kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
+        .await
+        .ok();
+    kvp.write_kvp(ZedPredictUpsell::KEY.into(), "1".into())
+        .await
+        .unwrap();
+
+    cx.update(|cx| assert!(!should_show_upsell_modal(cx)));
+
+    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.unwrap();
+}
+
+#[gpui::test]
+async fn test_upsell_dismissed_via_dismissable_api(cx: &mut TestAppContext) {
+    init_test(cx);
+    let kvp = cx.update(|cx| KeyValueStore::global(cx));
+    kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
+        .await
+        .ok();
+    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.ok();
+
+    cx.update(|cx| {
+        assert!(should_show_upsell_modal(cx));
+        ZedPredictUpsell::set_dismissed(true, cx);
+    });
+    cx.run_until_parked();
+
+    cx.update(|cx| assert!(!should_show_upsell_modal(cx)));
+
+    kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.unwrap();
+}
+
 #[ctor::ctor]
 fn init_logger() {
     zlog::init_test();

crates/edit_prediction/src/zed_edit_prediction_delegate.rs

@@ -7,9 +7,11 @@ use edit_prediction_types::{
     EditPredictionIconSet, SuggestionDisplayType,
 };
 use feature_flags::FeatureFlagAppExt;
+use fs::Fs;
 use gpui::{App, Entity, prelude::*};
 use language::{Buffer, ToPoint as _};
 use project::Project;
+use settings::{EditPredictionDataCollectionChoice, update_settings_file};
 
 use crate::{BufferEditPrediction, EditPredictionStore};
 
@@ -75,24 +77,7 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate {
                     .read(cx)
                     .is_file_open_source(&self.project, file, cx);
 
-            if let Some(organization_configuration) = self
-                .store
-                .read(cx)
-                .user_store
-                .read(cx)
-                .current_organization_configuration()
-            {
-                if !organization_configuration
-                    .edit_prediction
-                    .is_feedback_enabled
-                {
-                    return DataCollectionState::Disabled {
-                        is_project_open_source,
-                    };
-                }
-            }
-
-            if self.store.read(cx).data_collection_choice.is_enabled(cx) {
+            if self.store.read(cx).is_data_collection_enabled(cx) {
                 DataCollectionState::Enabled {
                     is_project_open_source,
                 }
@@ -102,9 +87,9 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate {
                 }
             }
         } else {
-            return DataCollectionState::Disabled {
+            DataCollectionState::Disabled {
                 is_project_open_source: false,
-            };
+            }
         }
     }
 
@@ -113,27 +98,26 @@ impl EditPredictionDelegate for ZedEditPredictionDelegate {
             return false;
         }
 
-        if let Some(organization_configuration) = self
-            .store
+        self.store
             .read(cx)
-            .user_store
-            .read(cx)
-            .current_organization_configuration()
-        {
-            if !organization_configuration
-                .edit_prediction
-                .is_feedback_enabled
-            {
-                return false;
-            }
-        }
-
-        true
+            .is_data_collection_allowed_by_organization(cx)
     }
 
     fn toggle_data_collection(&mut self, cx: &mut App) {
-        self.store.update(cx, |store, cx| {
-            store.toggle_data_collection_choice(cx);
+        let fs = <dyn Fs>::global(cx);
+        let is_currently_enabled = self.store.read(cx).is_data_collection_enabled(cx);
+        update_settings_file(fs, cx, move |settings, _| {
+            let edit_predictions = settings
+                .project
+                .all_languages
+                .edit_predictions
+                .get_or_insert_default();
+
+            edit_predictions.allow_data_collection = Some(if is_currently_enabled {
+                EditPredictionDataCollectionChoice::No
+            } else {
+                EditPredictionDataCollectionChoice::Yes
+            });
         });
     }
 

crates/language/src/language_settings.rs

@@ -16,10 +16,10 @@ use itertools::{Either, Itertools};
 use settings::{DocumentFoldingRanges, DocumentSymbols, IntoGpui, SemanticTokens};
 
 pub use settings::{
-    AutoIndentMode, CompletionSettingsContent, EditPredictionPromptFormat, EditPredictionProvider,
-    EditPredictionsMode, FormatOnSave, Formatter, FormatterList, InlayHintKind,
-    LanguageSettingsContent, LineEndingSetting, LspInsertMode, RewrapBehavior,
-    ShowWhitespaceSetting, SoftWrap, WordsCompletionMode,
+    AutoIndentMode, CompletionSettingsContent, EditPredictionDataCollectionChoice,
+    EditPredictionPromptFormat, EditPredictionProvider, EditPredictionsMode, FormatOnSave,
+    Formatter, FormatterList, InlayHintKind, LanguageSettingsContent, LineEndingSetting,
+    LspInsertMode, RewrapBehavior, ShowWhitespaceSetting, SoftWrap, WordsCompletionMode,
 };
 use settings::{RegisterSetting, Settings, SettingsLocation, SettingsStore, merge_from::MergeFrom};
 use shellexpand;
@@ -478,6 +478,11 @@ pub struct EditPredictionSettings {
     pub ollama: Option<OpenAiCompatibleEditPredictionSettings>,
     pub open_ai_compatible_api: Option<OpenAiCompatibleEditPredictionSettings>,
     pub examples_dir: Option<Arc<Path>>,
+    /// Controls whether training data collection is enabled.
+    ///
+    /// `Default` means the value stored in the legacy KV store is used as a fallback,
+    /// preserving existing users' choices without a migration.
+    pub allow_data_collection: EditPredictionDataCollectionChoice,
 }
 
 impl EditPredictionSettings {
@@ -867,6 +872,7 @@ impl settings::Settings for AllLanguageSettings {
                 ollama: ollama_settings,
                 open_ai_compatible_api: openai_compatible_settings,
                 examples_dir: edit_predictions.examples_dir,
+                allow_data_collection: edit_predictions.allow_data_collection.unwrap_or_default(),
             },
             defaults: default_language_settings,
             languages,

crates/settings_content/src/language.rs

@@ -182,6 +182,14 @@ pub struct EditPredictionSettingsContent {
     pub open_ai_compatible_api: Option<CustomEditPredictionProviderSettingsContent>,
     /// The directory where manually captured edit prediction examples are stored.
     pub examples_dir: Option<Arc<Path>>,
+    /// Controls whether Zed may collect training data when using Zed's Edit Predictions.
+    /// Data is only ever captured for files in projects that are detected as open source.
+    ///
+    /// - `"default"`: use the preference previously set via the status-bar toggle,
+    ///   or false if no preference has been stored.
+    /// - `"yes"`: allow data collection for files in open-source projects.
+    /// - `"no"`: never allow data collection.
+    pub allow_data_collection: Option<EditPredictionDataCollectionChoice>,
 }
 
 #[with_fallible_options]
@@ -318,6 +326,33 @@ pub struct OllamaEditPredictionSettingsContent {
     pub prompt_format: Option<EditPredictionPromptFormat>,
 }
 
+/// Controls whether Zed collects training data when using Zed's Edit Predictions.
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    Default,
+    Eq,
+    PartialEq,
+    Serialize,
+    Deserialize,
+    JsonSchema,
+    MergeFrom,
+    strum::VariantArray,
+    strum::VariantNames,
+)]
+#[serde(rename_all = "snake_case")]
+pub enum EditPredictionDataCollectionChoice {
+    /// Use the preference previously set via the status-bar toggle, or false
+    /// if no preference has been stored.
+    #[default]
+    Default,
+    /// Allow Zed to collect training data from open-source projects.
+    Yes,
+    /// Never allow training data collection.
+    No,
+}
+
 /// The mode in which edit predictions should be displayed.
 #[derive(
     Copy,

crates/settings_ui/src/page_data.rs

@@ -423,7 +423,7 @@ fn appearance_page() -> SettingsPage {
                                     .as_ref()?
                                     .discriminant() as usize])
                         },
-                        write: |settings_content, value, _app: &App| {
+                        write: |settings_content, value, app: &App| {
                             let Some(value) = value else {
                                 settings_content.theme.theme = None;
                                 return;
@@ -438,7 +438,7 @@ fn appearance_page() -> SettingsPage {
                                                 theme_settings::ThemeAppearanceMode::Light => light.clone(),
                                                 theme_settings::ThemeAppearanceMode::Dark => dark.clone(),
                                                 theme_settings::ThemeAppearanceMode::System => {
-                                                    if SystemAppearance::global(_app).is_light() {
+                                                    if SystemAppearance::global(app).is_light() {
                                                         light.clone()
                                                     } else {
                                                         dark.clone()
@@ -9419,7 +9419,7 @@ fn non_editor_language_settings_data() -> Box<[SettingsPageItem]> {
     )
 }
 
-fn edit_prediction_language_settings_section() -> [SettingsPageItem; 4] {
+fn edit_prediction_language_settings_section() -> [SettingsPageItem; 5] {
     [
         SettingsPageItem::SectionHeader("Edit Predictions"),
         SettingsPageItem::SubPageLink(SubPageLink {
@@ -9431,6 +9431,32 @@ fn edit_prediction_language_settings_section() -> [SettingsPageItem; 4] {
             files: USER,
             render: render_edit_prediction_setup_page
         }),
+        SettingsPageItem::SettingItem(SettingItem {
+            title: "Data Collection",
+            description: "Controls whether Zed may collect training data when using Zed's Edit Predictions. Data is only collected for files in projects detected as open source. The default value uses the preference previously set via the status-bar toggle, or false if no preference has been stored.",
+            field: Box::new(SettingField {
+                json_path: Some("edit_predictions.allow_data_collection"),
+                pick: |settings_content| {
+                    settings_content
+                        .project
+                        .all_languages
+                        .edit_predictions
+                        .as_ref()?
+                        .allow_data_collection
+                        .as_ref()
+                },
+                write: |settings_content, value, _app| {
+                    settings_content
+                        .project
+                        .all_languages
+                        .edit_predictions
+                        .get_or_insert_default()
+                        .allow_data_collection = value;
+                },
+            }),
+            metadata: None,
+            files: USER,
+        }),
         SettingsPageItem::SettingItem(SettingItem {
             title: "Show Edit Predictions",
             description: "Controls whether edit predictions are shown immediately or manually.",

crates/settings_ui/src/settings_ui.rs

@@ -504,6 +504,7 @@ fn init_renderers(cx: &mut App) {
         .add_basic_renderer::<settings::TerminalBlink>(render_dropdown)
         .add_basic_renderer::<settings::CursorShapeContent>(render_dropdown)
         .add_basic_renderer::<settings::EditPredictionPromptFormat>(render_dropdown)
+        .add_basic_renderer::<settings::EditPredictionDataCollectionChoice>(render_dropdown)
         .add_basic_renderer::<f32>(render_editable_number_field)
         .add_basic_renderer::<u32>(render_editable_number_field)
         .add_basic_renderer::<u64>(render_editable_number_field)