Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
This commit is contained in:
João Marcos 2025-01-31 07:22:31 -03:00 • committed by GitHub
parent be4c3cfbd2
commit 6f467281e0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 64 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

10
crates/zeta/src/persistence.rs
View file

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

111
crates/zeta/src/zeta.rs
View file

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