Fix data collection permission asked multiple times for same worktree (#24016)

João Marcos created

After the user confirmation, only the current instance of the
completions provider had the answer stored.

In this PR, the changes are propagated by having each provider have an
`Entity<choice>`, and having a lookup map with one `Entity<choice>` for
each worktree that `Zeta` has seen.

Release Notes:

- N/A

Change summary

crates/zeta/src/persistence.rs |  10 --
crates/zeta/src/zeta.rs        | 111 ++++++++++++++++++++---------------
2 files changed, 64 insertions(+), 57 deletions(-)

Detailed changes

crates/zeta/src/persistence.rs 🔗

@@ -1,5 +1,3 @@
-use anyhow::Result;
-use collections::HashMap;
 use std::path::{Path, PathBuf};
 use workspace::WorkspaceDb;
 
@@ -18,12 +16,8 @@ define_connection!(
 );
 
 impl ZetaDb {
-    pub fn get_all_zeta_preferences(&self) -> Result<HashMap<PathBuf, bool>> {
-        Ok(self.get_all_zeta_preferences_query()?.into_iter().collect())
-    }
-
     query! {
-        fn get_all_zeta_preferences_query() -> Result<Vec<(PathBuf, bool)>> {
+        pub fn get_all_data_collection_preferences() -> Result<Vec<(PathBuf, bool)>> {
             SELECT worktree_path, accepted_data_collection FROM zeta_preferences
         }
     }
@@ -36,7 +30,7 @@ impl ZetaDb {
     }
 
     query! {
-        pub async fn save_accepted_data_collection(worktree_path: PathBuf, accepted_data_collection: bool) -> Result<()> {
+        pub async fn save_data_collection_choice(worktree_path: PathBuf, accepted_data_collection: bool) -> Result<()> {
             INSERT INTO zeta_preferences
                 (worktree_path, accepted_data_collection)
             VALUES

crates/zeta/src/zeta.rs 🔗

@@ -10,6 +10,7 @@ pub use rate_completion_modal::*;
 use anyhow::{anyhow, Context as _, Result};
 use arrayvec::ArrayVec;
 use client::{Client, UserStore};
+use collections::hash_map::Entry;
 use collections::{HashMap, HashSet, VecDeque};
 use feature_flags::FeatureFlagAppExt as _;
 use futures::AsyncReadExt;
@@ -24,7 +25,6 @@ use language::{
 };
 use language_models::LlmApiToken;
 use rpc::{PredictEditsParams, PredictEditsResponse, EXPIRED_LLM_TOKEN_HEADER_NAME};
-use serde::{Deserialize, Serialize};
 use std::{
     borrow::Cow,
     cmp, env,
@@ -903,28 +903,24 @@ and then another
         new_snapshot
     }
 
-    pub fn data_collection_choice_at(&self, path: &Path) -> DataCollectionChoice {
-        match self.data_collection_preferences.per_worktree.get(path) {
-            Some(true) => DataCollectionChoice::Enabled,
-            Some(false) => DataCollectionChoice::Disabled,
-            None => DataCollectionChoice::NotAnswered,
-        }
-    }
-
-    fn update_data_collection_choice_for_worktree(
+    /// Creates a `Entity<DataCollectionChoice>` for each unique worktree abs path it sees.
+    pub fn data_collection_choice_at(
         &mut self,
-        absolute_path_of_project_worktree: PathBuf,
-        can_collect_data: bool,
+        worktree_abs_path: PathBuf,
         cx: &mut Context<Self>,
-    ) {
-        self.data_collection_preferences
+    ) -> Entity<DataCollectionChoice> {
+        match self
+            .data_collection_preferences
             .per_worktree
-            .insert(absolute_path_of_project_worktree.clone(), can_collect_data);
-
-        db::write_and_log(cx, move || {
-            persistence::DB
-                .save_accepted_data_collection(absolute_path_of_project_worktree, can_collect_data)
-        });
+            .entry(worktree_abs_path)
+        {
+            Entry::Vacant(entry) => {
+                let choice = cx.new(|_| DataCollectionChoice::NotAnswered);
+                entry.insert(choice.clone());
+                choice
+            }
+            Entry::Occupied(entry) => entry.get().clone(),
+        }
     }
 
     fn set_never_ask_again_for_data_collection(&mut self, cx: &mut Context<Self>) {
@@ -959,23 +955,32 @@ and then another
             .map(|value| value == "true")
             .unwrap_or(false);
 
-        let preferences_per_project = persistence::DB
-            .get_all_zeta_preferences()
+        let preferences_per_worktree = persistence::DB
+            .get_all_data_collection_preferences()
             .log_err()
-            .unwrap_or_else(HashMap::default);
+            .into_iter()
+            .flatten()
+            .map(|(path, choice)| {
+                let choice = cx.new(|_| DataCollectionChoice::from(choice));
+                (path, choice)
+            })
+            .collect();
 
         DataCollectionPreferences {
             never_ask_again,
-            per_worktree: preferences_per_project,
+            per_worktree: preferences_per_worktree,
         }
     }
 }
 
-#[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Default, Debug)]
 struct DataCollectionPreferences {
     /// Set when a user clicks on "Never Ask Again", can never be unset.
     never_ask_again: bool,
-    per_worktree: HashMap<PathBuf, bool>,
+    /// The choices for each worktree.
+    ///
+    /// This is filled when loading from database, or when querying if no matching path is found.
+    per_worktree: HashMap<PathBuf, Entity<DataCollectionChoice>>,
 }
 
 fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
@@ -1281,7 +1286,7 @@ struct PendingCompletion {
     _task: Task<()>,
 }
 
-#[derive(Clone, Copy)]
+#[derive(Debug, Clone, Copy)]
 pub enum DataCollectionChoice {
     NotAnswered,
     Enabled,
@@ -1289,21 +1294,21 @@ pub enum DataCollectionChoice {
 }
 
 impl DataCollectionChoice {
-    pub fn is_enabled(&self) -> bool {
+    pub fn is_enabled(self) -> bool {
         match self {
             Self::Enabled => true,
             Self::NotAnswered | Self::Disabled => false,
         }
     }
 
-    pub fn is_answered(&self) -> bool {
+    pub fn is_answered(self) -> bool {
         match self {
             Self::Enabled | Self::Disabled => true,
             Self::NotAnswered => false,
         }
     }
 
-    pub fn toggle(&self) -> DataCollectionChoice {
+    pub fn toggle(self) -> DataCollectionChoice {
         match self {
             Self::Enabled => Self::Disabled,
             Self::Disabled => Self::Enabled,
@@ -1312,6 +1317,15 @@ impl DataCollectionChoice {
     }
 }
 
+impl From<bool> for DataCollectionChoice {
+    fn from(value: bool) -> Self {
+        match value {
+            true => DataCollectionChoice::Enabled,
+            false => DataCollectionChoice::Disabled,
+        }
+    }
+}
+
 pub struct ZetaInlineCompletionProvider {
     zeta: Entity<Zeta>,
     pending_completions: ArrayVec<PendingCompletion, 2>,
@@ -1323,7 +1337,7 @@ pub struct ZetaInlineCompletionProvider {
 pub struct ProviderDataCollection {
     workspace: WeakEntity<Workspace>,
     worktree_root_path: PathBuf,
-    choice: DataCollectionChoice,
+    choice: Entity<DataCollectionChoice>,
 }
 
 impl ProviderDataCollection {
@@ -1351,7 +1365,9 @@ impl ProviderDataCollection {
             })
         })?;
 
-        let choice = zeta.read(cx).data_collection_choice_at(&worktree_root_path);
+        let choice = zeta.update(cx, |zeta, cx| {
+            zeta.data_collection_choice_at(worktree_root_path.clone(), cx)
+        });
 
         Some(ProviderDataCollection {
             workspace: workspace.downgrade(),
@@ -1360,22 +1376,18 @@ impl ProviderDataCollection {
         })
     }
 
-    fn set_choice(&mut self, choice: DataCollectionChoice, zeta: &Entity<Zeta>, cx: &mut App) {
-        self.choice = choice;
+    fn set_choice(&mut self, choice: DataCollectionChoice, cx: &mut App) {
+        self.choice.update(cx, |this, _| *this = choice);
 
         let worktree_root_path = self.worktree_root_path.clone();
 
-        zeta.update(cx, |zeta, cx| {
-            zeta.update_data_collection_choice_for_worktree(
-                worktree_root_path,
-                choice.is_enabled(),
-                cx,
-            )
+        db::write_and_log(cx, move || {
+            persistence::DB.save_data_collection_choice(worktree_root_path, choice.is_enabled())
         });
     }
 
-    fn toggle_choice(&mut self, zeta: &Entity<Zeta>, cx: &mut App) {
-        self.set_choice(self.choice.toggle(), zeta, cx);
+    fn toggle_choice(&mut self, cx: &mut App) {
+        self.set_choice(self.choice.read(cx).toggle(), cx);
     }
 }
 
@@ -1394,7 +1406,7 @@ impl ZetaInlineCompletionProvider {
 
     fn set_data_collection_choice(&mut self, choice: DataCollectionChoice, cx: &mut App) {
         if let Some(data_collection) = self.data_collection.as_mut() {
-            data_collection.set_choice(choice, &self.zeta, cx);
+            data_collection.set_choice(choice, cx);
         }
     }
 }
@@ -1420,12 +1432,12 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
         true
     }
 
-    fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
+    fn data_collection_state(&self, cx: &App) -> DataCollectionState {
         let Some(data_collection) = self.data_collection.as_ref() else {
             return DataCollectionState::Unknown;
         };
 
-        if data_collection.choice.is_enabled() {
+        if data_collection.choice.read(cx).is_enabled() {
             DataCollectionState::Enabled
         } else {
             DataCollectionState::Disabled
@@ -1434,7 +1446,7 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
 
     fn toggle_data_collection(&mut self, cx: &mut App) {
         if let Some(data_collection) = self.data_collection.as_mut() {
-            data_collection.toggle_choice(&self.zeta, cx);
+            data_collection.toggle_choice(cx);
         }
     }
 
@@ -1486,7 +1498,9 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
         let can_collect_data = self
             .data_collection
             .as_ref()
-            .map_or(false, |data_collection| data_collection.choice.is_enabled());
+            .map_or(false, |data_collection| {
+                data_collection.choice.read(cx).is_enabled()
+            });
 
         let task = cx.spawn(|this, mut cx| async move {
             if debounce {
@@ -1579,7 +1593,7 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
             return;
         };
 
-        if data_collection.choice.is_answered()
+        if data_collection.choice.read(cx).is_answered()
             || self
                 .zeta
                 .read(cx)
@@ -1633,7 +1647,6 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
                             })
                             .with_tertiary_click_message("Never Ask Again")
                             .on_tertiary_click({
-                                let zeta = zeta.clone();
                                 move |_window, cx| {
                                     zeta.update(cx, |zeta, cx| {
                                         zeta.set_never_ask_again_for_data_collection(cx);