zeta: Revised data-collection onboarding experience (#24031)

Agus Zubiaga , Danilo , Danilo Leal , and João Marcos created

Release Notes:

- N/A

---------

Co-authored-by: Danilo <danilo@zed.dev>
Co-authored-by: Danilo Leal <daniloleal09@gmail.com>
Co-authored-by: João Marcos <marcospb19@hotmail.com>

Change summary

Cargo.lock                                                      |  31 
Cargo.toml                                                      |   2 
assets/icons/zed_predict_bg.svg                                 |   2 
crates/editor/Cargo.toml                                        |   2 
crates/editor/src/editor.rs                                     |  19 
crates/inline_completion/src/inline_completion.rs               |   6 
crates/inline_completion_button/Cargo.toml                      |   1 
crates/inline_completion_button/src/inline_completion_button.rs |  94 
crates/rpc/src/llm.rs                                           |   2 
crates/title_bar/Cargo.toml                                     |   4 
crates/title_bar/src/title_bar.rs                               |  14 
crates/ui/src/components/button/button_like.rs                  |   5 
crates/ui/src/components/popover_menu.rs                        |  30 
crates/ui/src/components/toggle.rs                              |  39 
crates/worktree/src/worktree.rs                                 |   2 
crates/zed/Cargo.toml                                           |   1 
crates/zed/src/main.rs                                          |   1 
crates/zed/src/zed.rs                                           |   1 
crates/zed/src/zed/inline_completion_registry.rs                |  51 
crates/zed_actions/src/lib.rs                                   |   2 
crates/zed_predict_onboarding/Cargo.toml                        |  31 
crates/zed_predict_onboarding/LICENSE-GPL                       |   1 
crates/zed_predict_onboarding/src/lib.rs                        |   5 
crates/zeta/Cargo.toml                                          |   7 
crates/zeta/src/init.rs                                         |  60 
crates/zeta/src/license_detection.rs                            | 210 +
crates/zeta/src/onboarding_banner.rs                            |  48 
crates/zeta/src/onboarding_modal.rs                             | 208 +
crates/zeta/src/persistence.rs                                  |  48 
crates/zeta/src/rate_completion_modal.rs                        |  38 
crates/zeta/src/zeta.rs                                         | 372 +-
31 files changed, 748 insertions(+), 589 deletions(-)

Detailed changes

Cargo.lock 🔗

@@ -4059,7 +4059,7 @@ dependencies = [
  "util",
  "uuid",
  "workspace",
- "zed_predict_onboarding",
+ "zed_actions",
 ]
 
 [[package]]
@@ -6454,7 +6454,6 @@ dependencies = [
  "ui",
  "workspace",
  "zed_actions",
- "zed_predict_onboarding",
  "zeta",
 ]
 
@@ -13590,7 +13589,7 @@ dependencies = [
  "windows 0.58.0",
  "workspace",
  "zed_actions",
- "zed_predict_onboarding",
+ "zeta",
 ]
 
 [[package]]
@@ -16588,7 +16587,6 @@ dependencies = [
  "winresource",
  "workspace",
  "zed_actions",
- "zed_predict_onboarding",
  "zeta",
 ]
 
@@ -16702,25 +16700,6 @@ dependencies = [
  "zed_extension_api 0.1.0",
 ]
 
-[[package]]
-name = "zed_predict_onboarding"
-version = "0.1.0"
-dependencies = [
- "chrono",
- "client",
- "db",
- "feature_flags",
- "fs",
- "gpui",
- "language",
- "menu",
- "settings",
- "theme",
- "ui",
- "util",
- "workspace",
-]
-
 [[package]]
 name = "zed_proto"
 version = "0.2.1"
@@ -16906,6 +16885,7 @@ dependencies = [
  "anyhow",
  "arrayvec",
  "call",
+ "chrono",
  "client",
  "clock",
  "collections",
@@ -16915,6 +16895,7 @@ dependencies = [
  "editor",
  "env_logger 0.11.6",
  "feature_flags",
+ "fs",
  "futures 0.3.31",
  "gpui",
  "http_client",
@@ -16924,6 +16905,8 @@ dependencies = [
  "language_models",
  "log",
  "menu",
+ "postage",
+ "regex",
  "reqwest_client",
  "rpc",
  "serde",
@@ -16936,10 +16919,12 @@ dependencies = [
  "tree-sitter-go",
  "tree-sitter-rust",
  "ui",
+ "unindent",
  "util",
  "uuid",
  "workspace",
  "worktree",
+ "zed_actions",
 ]
 
 [[package]]

Cargo.toml 🔗

@@ -152,7 +152,6 @@ members = [
     "crates/worktree",
     "crates/zed",
     "crates/zed_actions",
-    "crates/zed_predict_onboarding",
     "crates/zeta",
 
     #
@@ -348,7 +347,6 @@ workspace = { path = "crates/workspace" }
 worktree = { path = "crates/worktree" }
 zed = { path = "crates/zed" }
 zed_actions = { path = "crates/zed_actions" }
-zed_predict_onboarding = { path = "crates/zed_predict_onboarding" }
 zeta = { path = "crates/zeta" }
 
 #

assets/icons/zed_predict_bg.svg 🔗

@@ -1,4 +1,4 @@
-<svg width="420" height="128" xmlns="http://www.w3.org/2000/svg">
+<svg width="440" height="128" xmlns="http://www.w3.org/2000/svg">
   <defs>
     <pattern id="tilePattern" width="22" height="22" patternUnits="userSpaceOnUse">
       <svg width="16" height="16" viewBox="0 0 16 16" fill="none" xmlns="http://www.w3.org/2000/svg">

crates/editor/Cargo.toml 🔗

@@ -87,7 +87,7 @@ url.workspace = true
 util.workspace = true
 uuid.workspace = true
 workspace.workspace = true
-zed_predict_onboarding.workspace = true
+zed_actions.workspace = true
 
 [dev-dependencies]
 ctor.workspace = true

crates/editor/src/editor.rs 🔗

@@ -69,7 +69,6 @@ pub use element::{
 };
 use futures::{future, FutureExt};
 use fuzzy::StringMatchCandidate;
-use zed_predict_onboarding::ZedPredictModal;
 
 use code_context_menus::{
     AvailableCodeAction, CodeActionContents, CodeActionsItem, CodeActionsMenu, CodeContextMenu,
@@ -617,7 +616,8 @@ pub struct Editor {
     active_diagnostics: Option<ActiveDiagnosticGroup>,
     soft_wrap_mode_override: Option<language_settings::SoftWrap>,
 
-    project: Option<Entity<Project>>,
+    // TODO: make this a access method
+    pub project: Option<Entity<Project>>,
     semantics_provider: Option<Rc<dyn SemanticsProvider>>,
     completion_provider: Option<Box<dyn CompletionProvider>>,
     collaboration_hub: Option<Box<dyn CollaborationHub>>,
@@ -3944,20 +3944,7 @@ impl Editor {
     }
 
     fn toggle_zed_predict_onboarding(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        let (Some(workspace), Some(project)) = (self.workspace(), self.project.as_ref()) else {
-            return;
-        };
-
-        let project = project.read(cx);
-
-        ZedPredictModal::toggle(
-            workspace,
-            project.user_store().clone(),
-            project.client().clone(),
-            project.fs().clone(),
-            window,
-            cx,
-        );
+        window.dispatch_action(zed_actions::OpenZedPredictOnboarding.boxed_clone(), cx);
     }
 
     fn do_completion(

crates/inline_completion/src/inline_completion.rs 🔗

@@ -21,8 +21,6 @@ pub struct InlineCompletion {
 pub enum DataCollectionState {
     /// The provider doesn't support data collection.
     Unsupported,
-    /// When there's a file not saved yet. In this case, we can't tell to which project it belongs.
-    Unknown,
     /// Data collection is enabled
     Enabled,
     /// Data collection is disabled or unanswered.
@@ -34,10 +32,6 @@ impl DataCollectionState {
         !matches!(self, DataCollectionState::Unsupported)
     }
 
-    pub fn is_unknown(&self) -> bool {
-        matches!(self, DataCollectionState::Unknown)
-    }
-
     pub fn is_enabled(&self) -> bool {
         matches!(self, DataCollectionState::Enabled)
     }

crates/inline_completion_button/Cargo.toml 🔗

@@ -29,7 +29,6 @@ workspace.workspace = true
 zed_actions.workspace = true
 zeta.workspace = true
 client.workspace = true
-zed_predict_onboarding.workspace = true
 
 [dev-dependencies]
 copilot = { workspace = true, features = ["test-support"] }

crates/inline_completion_button/src/inline_completion_button.rs 🔗

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use client::{Client, UserStore};
+use client::UserStore;
 use copilot::{Copilot, Status};
 use editor::{actions::ShowInlineCompletion, scroll::Autoscroll, Editor};
 use feature_flags::{
@@ -21,15 +21,14 @@ use settings::{update_settings_file, Settings, SettingsStore};
 use std::{path::Path, sync::Arc, time::Duration};
 use supermaven::{AccountStatus, Supermaven};
 use ui::{
-    prelude::*, ButtonLike, Clickable, ContextMenu, ContextMenuEntry, IconButton,
-    IconWithIndicator, Indicator, PopoverMenu, PopoverMenuHandle, Tooltip,
+    prelude::*, Clickable, ContextMenu, ContextMenuEntry, IconButton, IconButtonShape, PopoverMenu,
+    PopoverMenuHandle, Tooltip,
 };
 use workspace::{
     create_and_open_local_file, item::ItemHandle, notifications::NotificationId, StatusItemView,
     Toast, Workspace,
 };
 use zed_actions::OpenBrowser;
-use zed_predict_onboarding::ZedPredictModal;
 use zeta::RateCompletionModal;
 
 actions!(zeta, [RateCompletions]);
@@ -46,7 +45,6 @@ pub struct InlineCompletionButton {
     language: Option<Arc<Language>>,
     file: Option<Arc<dyn File>>,
     inline_completion_provider: Option<Arc<dyn inline_completion::InlineCompletionProviderHandle>>,
-    client: Arc<Client>,
     fs: Arc<dyn Fs>,
     workspace: WeakEntity<Workspace>,
     user_store: Entity<UserStore>,
@@ -230,71 +228,49 @@ impl Render for InlineCompletionButton {
                     return div();
                 }
 
+                fn icon_button() -> IconButton {
+                    IconButton::new("zed-predict-pending-button", IconName::ZedPredict)
+                        .shape(IconButtonShape::Square)
+                }
+
                 let current_user_terms_accepted =
                     self.user_store.read(cx).current_user_has_accepted_terms();
 
                 if !current_user_terms_accepted.unwrap_or(false) {
-                    let workspace = self.workspace.clone();
-                    let user_store = self.user_store.clone();
-                    let client = self.client.clone();
-                    let fs = self.fs.clone();
-
                     let signed_in = current_user_terms_accepted.is_some();
+                    let tooltip_meta = if signed_in {
+                        "Read Terms of Service"
+                    } else {
+                        "Sign in to use"
+                    };
 
                     return div().child(
-                        ButtonLike::new("zeta-pending-tos-icon")
-                            .child(
-                                IconWithIndicator::new(
-                                    Icon::new(IconName::ZedPredict),
-                                    Some(Indicator::dot().color(Color::Error)),
-                                )
-                                .indicator_border_color(Some(
-                                    cx.theme().colors().status_bar_background,
-                                ))
-                                .into_any_element(),
-                            )
+                        icon_button()
                             .tooltip(move |window, cx| {
                                 Tooltip::with_meta(
                                     "Edit Predictions",
                                     None,
-                                    if signed_in {
-                                        "Read Terms of Service"
-                                    } else {
-                                        "Sign in to use"
-                                    },
+                                    tooltip_meta,
                                     window,
                                     cx,
                                 )
                             })
                             .on_click(cx.listener(move |_, _, window, cx| {
-                                if let Some(workspace) = workspace.upgrade() {
-                                    ZedPredictModal::toggle(
-                                        workspace,
-                                        user_store.clone(),
-                                        client.clone(),
-                                        fs.clone(),
-                                        window,
-                                        cx,
-                                    );
-                                }
+                                window.dispatch_action(
+                                    zed_actions::OpenZedPredictOnboarding.boxed_clone(),
+                                    cx,
+                                );
                             })),
                     );
                 }
 
                 let this = cx.entity().clone();
-                let button = IconButton::new("zeta", IconName::ZedPredict).when(
-                    !self.popover_menu_handle.is_deployed(),
-                    |button| {
-                        button.tooltip(|window, cx| {
-                            Tooltip::for_action("Edit Prediction", &ToggleMenu, window, cx)
-                        })
-                    },
-                );
 
-                let is_refreshing = self
-                    .inline_completion_provider
-                    .as_ref()
-                    .map_or(false, |provider| provider.is_refreshing(cx));
+                if !self.popover_menu_handle.is_deployed() {
+                    icon_button().tooltip(|window, cx| {
+                        Tooltip::for_action("Edit Prediction", &ToggleMenu, window, cx)
+                    });
+                }
 
                 let mut popover_menu = PopoverMenu::new("zeta")
                     .menu(move |window, cx| {
@@ -303,9 +279,14 @@ impl Render for InlineCompletionButton {
                     .anchor(Corner::BottomRight)
                     .with_handle(self.popover_menu_handle.clone());
 
+                let is_refreshing = self
+                    .inline_completion_provider
+                    .as_ref()
+                    .map_or(false, |provider| provider.is_refreshing(cx));
+
                 if is_refreshing {
                     popover_menu = popover_menu.trigger(
-                        button.with_animation(
+                        icon_button().with_animation(
                             "pulsating-label",
                             Animation::new(Duration::from_secs(2))
                                 .repeat()
@@ -314,7 +295,7 @@ impl Render for InlineCompletionButton {
                         ),
                     );
                 } else {
-                    popover_menu = popover_menu.trigger(button);
+                    popover_menu = popover_menu.trigger(icon_button());
                 }
 
                 div().child(popover_menu.into_any_element())
@@ -328,7 +309,6 @@ impl InlineCompletionButton {
         workspace: WeakEntity<Workspace>,
         fs: Arc<dyn Fs>,
         user_store: Entity<UserStore>,
-        client: Arc<Client>,
         popover_menu_handle: PopoverMenuHandle<ContextMenu>,
         cx: &mut Context<Self>,
     ) -> Self {
@@ -348,7 +328,6 @@ impl InlineCompletionButton {
             inline_completion_provider: None,
             popover_menu_handle,
             workspace,
-            client,
             fs,
             user_store,
         }
@@ -447,10 +426,15 @@ impl InlineCompletionButton {
 
             if data_collection.is_supported() {
                 let provider = provider.clone();
-                menu = menu.separator().item(
-                    ContextMenuEntry::new("Data Collection")
+                menu = menu
+                    .separator()
+                    .header("Help Improve The Model")
+                    .header("For OSS Projects Only");
+                menu = menu.item(
+                    // TODO: We want to add something later that communicates whether
+                    // the current project is open-source.
+                    ContextMenuEntry::new("Share Training Data")
                         .toggleable(IconPosition::Start, data_collection.is_enabled())
-                        .disabled(data_collection.is_unknown())
                         .handler(move |_, cx| {
                             provider.toggle_data_collection(cx);
                         }),

crates/rpc/src/llm.rs 🔗

@@ -41,7 +41,7 @@ pub struct PredictEditsParams {
     pub input_excerpt: String,
     /// Whether the user provided consent for sampling this interaction.
     #[serde(default)]
-    pub can_collect_data: bool,
+    pub data_collection_permission: bool,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

crates/title_bar/Cargo.toml 🔗

@@ -41,13 +41,13 @@ serde.workspace = true
 settings.workspace = true
 smallvec.workspace = true
 story = { workspace = true, optional = true }
+telemetry.workspace = true
 theme.workspace = true
 ui.workspace = true
 util.workspace = true
-telemetry.workspace = true
 workspace.workspace = true
 zed_actions.workspace = true
-zed_predict_onboarding.workspace = true
+zeta.workspace = true
 
 [target.'cfg(windows)'.dependencies]
 windows.workspace = true

crates/title_bar/src/title_bar.rs 🔗

@@ -34,7 +34,7 @@ use ui::{
 use util::ResultExt;
 use workspace::{notifications::NotifyResultExt, Workspace};
 use zed_actions::{OpenBrowser, OpenRecent, OpenRemote};
-use zed_predict_onboarding::ZedPredictBanner;
+use zeta::ZedPredictBanner;
 
 #[cfg(feature = "stories")]
 pub use stories::*;
@@ -162,6 +162,7 @@ impl Render for TitleBar {
                     .id("titlebar-content")
                     .flex()
                     .flex_row()
+                    .items_center()
                     .justify_between()
                     .w_full()
                     // Note: On Windows the title bar behavior is handled by the platform implementation.
@@ -268,7 +269,6 @@ impl TitleBar {
         let project = workspace.project().clone();
         let user_store = workspace.app_state().user_store.clone();
         let client = workspace.app_state().client.clone();
-        let fs = workspace.app_state().fs.clone();
         let active_call = ActiveCall::global(cx);
 
         let platform_style = PlatformStyle::platform();
@@ -296,15 +296,7 @@ impl TitleBar {
         subscriptions.push(cx.observe_window_activation(window, Self::window_activation_changed));
         subscriptions.push(cx.observe(&user_store, |_, _, cx| cx.notify()));
 
-        let zed_predict_banner = cx.new(|cx| {
-            ZedPredictBanner::new(
-                workspace.weak_handle(),
-                user_store.clone(),
-                client.clone(),
-                fs.clone(),
-                cx,
-            )
-        });
+        let zed_predict_banner = cx.new(ZedPredictBanner::new);
 
         Self {
             platform_style,

crates/ui/src/components/button/button_like.rs 🔗

@@ -385,6 +385,11 @@ impl ButtonLike {
         Self::new(id).rounding(ButtonLikeRounding::Right)
     }
 
+    pub fn opacity(mut self, opacity: f32) -> Self {
+        self.base = self.base.opacity(opacity);
+        self
+    }
+
     pub(crate) fn height(mut self, height: DefiniteLength) -> Self {
         self.height = Some(height);
         self

crates/ui/src/components/popover_menu.rs 🔗

@@ -57,12 +57,19 @@ impl<M> Default for PopoverMenuHandle<M> {
 struct PopoverMenuHandleState<M> {
     menu_builder: Rc<dyn Fn(&mut Window, &mut App) -> Option<Entity<M>>>,
     menu: Rc<RefCell<Option<Entity<M>>>>,
+    on_open: Option<Rc<dyn Fn(&mut Window, &mut App)>>,
 }
 
 impl<M: ManagedView> PopoverMenuHandle<M> {
     pub fn show(&self, window: &mut Window, cx: &mut App) {
         if let Some(state) = self.0.borrow().as_ref() {
-            show_menu(&state.menu_builder, &state.menu, window, cx);
+            show_menu(
+                &state.menu_builder,
+                &state.menu,
+                state.on_open.clone(),
+                window,
+                cx,
+            );
         }
     }
 
@@ -118,6 +125,7 @@ pub struct PopoverMenu<M: ManagedView> {
     attach: Option<Corner>,
     offset: Option<Point<Pixels>>,
     trigger_handle: Option<PopoverMenuHandle<M>>,
+    on_open: Option<Rc<dyn Fn(&mut Window, &mut App)>>,
     full_width: bool,
 }
 
@@ -132,6 +140,7 @@ impl<M: ManagedView> PopoverMenu<M> {
             attach: None,
             offset: None,
             trigger_handle: None,
+            on_open: None,
             full_width: false,
         }
     }
@@ -155,11 +164,14 @@ impl<M: ManagedView> PopoverMenu<M> {
     }
 
     pub fn trigger<T: PopoverTrigger>(mut self, t: T) -> Self {
-        self.child_builder = Some(Box::new(|menu, builder| {
+        let on_open = self.on_open.clone();
+        self.child_builder = Some(Box::new(move |menu, builder| {
             let open = menu.borrow().is_some();
             t.toggle_state(open)
                 .when_some(builder, |el, builder| {
-                    el.on_click(move |_event, window, cx| show_menu(&builder, &menu, window, cx))
+                    el.on_click(move |_event, window, cx| {
+                        show_menu(&builder, &menu, on_open.clone(), window, cx)
+                    })
                 })
                 .into_any_element()
         }));
@@ -185,6 +197,12 @@ impl<M: ManagedView> PopoverMenu<M> {
         self
     }
 
+    /// attach something upon opening the menu
+    pub fn on_open(mut self, on_open: Rc<dyn Fn(&mut Window, &mut App)>) -> Self {
+        self.on_open = Some(on_open);
+        self
+    }
+
     fn resolved_attach(&self) -> Corner {
         self.attach.unwrap_or(match self.anchor {
             Corner::TopLeft => Corner::BottomLeft,
@@ -209,6 +227,7 @@ impl<M: ManagedView> PopoverMenu<M> {
 fn show_menu<M: ManagedView>(
     builder: &Rc<dyn Fn(&mut Window, &mut App) -> Option<Entity<M>>>,
     menu: &Rc<RefCell<Option<Entity<M>>>>,
+    on_open: Option<Rc<dyn Fn(&mut Window, &mut App)>>,
     window: &mut Window,
     cx: &mut App,
 ) {
@@ -232,6 +251,10 @@ fn show_menu<M: ManagedView>(
     window.focus(&new_menu.focus_handle(cx));
     *menu.borrow_mut() = Some(new_menu);
     window.refresh();
+
+    if let Some(on_open) = on_open {
+        on_open(window, cx);
+    }
 }
 
 pub struct PopoverMenuElementState<M> {
@@ -311,6 +334,7 @@ impl<M: ManagedView> Element for PopoverMenu<M> {
                         *trigger_handle.0.borrow_mut() = Some(PopoverMenuHandleState {
                             menu_builder,
                             menu: element_state.menu.clone(),
+                            on_open: self.on_open.clone(),
                         });
                     }
                 }

crates/ui/src/components/toggle.rs 🔗

@@ -1,4 +1,6 @@
-use gpui::{div, hsla, prelude::*, AnyView, ElementId, Hsla, IntoElement, Styled, Window};
+use gpui::{
+    div, hsla, prelude::*, AnyView, CursorStyle, ElementId, Hsla, IntoElement, Styled, Window,
+};
 use std::sync::Arc;
 
 use crate::utils::is_light;
@@ -45,6 +47,7 @@ pub struct Checkbox {
     filled: bool,
     style: ToggleStyle,
     tooltip: Option<Box<dyn Fn(&mut Window, &mut App) -> AnyView>>,
+    label: Option<SharedString>,
 }
 
 impl Checkbox {
@@ -58,6 +61,7 @@ impl Checkbox {
             filled: false,
             style: ToggleStyle::default(),
             tooltip: None,
+            label: None,
         }
     }
 
@@ -99,6 +103,12 @@ impl Checkbox {
         self.tooltip = Some(Box::new(tooltip));
         self
     }
+
+    /// Set the label for the checkbox.
+    pub fn label(mut self, label: impl Into<SharedString>) -> Self {
+        self.label = Some(label.into());
+        self
+    }
 }
 
 impl Checkbox {
@@ -116,11 +126,11 @@ impl Checkbox {
 
     fn border_color(&self, cx: &App) -> Hsla {
         if self.disabled {
-            return cx.theme().colors().border_disabled;
+            return cx.theme().colors().border_variant;
         }
 
         match self.style.clone() {
-            ToggleStyle::Ghost => cx.theme().colors().border_variant,
+            ToggleStyle::Ghost => cx.theme().colors().border,
             ToggleStyle::ElevationBased(elevation) => elevation.on_elevation_bg(cx),
             ToggleStyle::Custom(color) => color.opacity(0.3),
         }
@@ -153,10 +163,8 @@ impl RenderOnce for Checkbox {
         let bg_color = self.bg_color(cx);
         let border_color = self.border_color(cx);
 
-        h_flex()
-            .id(self.id)
+        let checkbox = h_flex()
             .justify_center()
-            .items_center()
             .size(DynamicSpacing::Base20.rems(cx))
             .group(group_id.clone())
             .child(
@@ -171,13 +179,24 @@ impl RenderOnce for Checkbox {
                     .bg(bg_color)
                     .border_1()
                     .border_color(border_color)
+                    .when(self.disabled, |this| {
+                        this.cursor(CursorStyle::OperationNotAllowed)
+                    })
+                    .when(self.disabled, |this| {
+                        this.bg(cx.theme().colors().element_disabled.opacity(0.6))
+                    })
                     .when(!self.disabled, |this| {
                         this.group_hover(group_id.clone(), |el| {
                             el.bg(cx.theme().colors().element_hover)
                         })
                     })
                     .children(icon),
-            )
+            );
+
+        h_flex()
+            .id(self.id)
+            .gap(DynamicSpacing::Base06.rems(cx))
+            .child(checkbox)
             .when_some(
                 self.on_click.filter(|_| !self.disabled),
                 |this, on_click| {
@@ -186,6 +205,11 @@ impl RenderOnce for Checkbox {
                     })
                 },
             )
+            // TODO: Allow label size to be different from default.
+            // TODO: Allow label color to be different from muted.
+            .when_some(self.label, |this, label| {
+                this.child(Label::new(label).color(Color::Muted))
+            })
             .when_some(self.tooltip, |this, tooltip| {
                 this.tooltip(move |window, cx| tooltip(window, cx))
             })
@@ -203,6 +227,7 @@ pub struct CheckboxWithLabel {
     style: ToggleStyle,
 }
 
+// TODO: Remove `CheckboxWithLabel` now that `label` is a method of `Checkbox`.
 impl CheckboxWithLabel {
     /// Creates a checkbox with an attached label.
     pub fn new(

crates/worktree/src/worktree.rs 🔗

@@ -87,7 +87,7 @@ pub const FS_WATCH_LATENCY: Duration = Duration::from_millis(100);
 /// May correspond to a directory or a single file.
 /// Possible examples:
 /// * a drag and dropped file — may be added as an invisible, "ephemeral" entry to the current worktree
-/// * a directory opened in Zed — may be added as a visible entry to the current worktree
+/// * a directory opened in Zed — may be added as a visible entry to the current worktree
 ///
 /// Uses [`Entry`] to track the state of each file/directory, can look up absolute paths for entries.
 pub enum Worktree {

crates/zed/Cargo.toml 🔗

@@ -16,7 +16,6 @@ path = "src/main.rs"
 
 [dependencies]
 activity_indicator.workspace = true
-zed_predict_onboarding.workspace = true
 anyhow.workspace = true
 assets.workspace = true
 assistant.workspace = true

crates/zed/src/main.rs 🔗

@@ -439,7 +439,6 @@ fn main() {
         inline_completion_registry::init(
             app_state.client.clone(),
             app_state.user_store.clone(),
-            app_state.fs.clone(),
             cx,
         );
         let prompt_builder = PromptBuilder::load(app_state.fs.clone(), stdout_is_a_pty(), cx);

crates/zed/src/zed.rs 🔗

@@ -176,7 +176,6 @@ pub fn initialize_workspace(
                 workspace.weak_handle(),
                 app_state.fs.clone(),
                 app_state.user_store.clone(),
-                app_state.client.clone(),
                 popover_menu_handle.clone(),
                 cx,
             )

crates/zed/src/zed/inline_completion_registry.rs 🔗

@@ -1,21 +1,17 @@
-use std::{cell::RefCell, rc::Rc, sync::Arc};
-
 use client::{Client, UserStore};
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::{Editor, EditorMode};
 use feature_flags::{FeatureFlagAppExt, PredictEditsFeatureFlag};
-use fs::Fs;
 use gpui::{AnyWindowHandle, App, AppContext, Context, Entity, WeakEntity};
 use language::language_settings::{all_language_settings, InlineCompletionProvider};
 use settings::SettingsStore;
+use std::{cell::RefCell, rc::Rc, sync::Arc};
 use supermaven::{Supermaven, SupermavenCompletionProvider};
 use ui::Window;
-use workspace::Workspace;
-use zed_predict_onboarding::ZedPredictModal;
 use zeta::ProviderDataCollection;
 
-pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, fs: Arc<dyn Fs>, cx: &mut App) {
+pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
     cx.observe_new({
         let editors = editors.clone();
@@ -96,7 +92,6 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, fs: Arc<dyn Fs>,
         let editors = editors.clone();
         let client = client.clone();
         let user_store = user_store.clone();
-        let fs = fs.clone();
         move |cx| {
             let new_provider = all_language_settings(None, cx).inline_completions.provider;
             if new_provider != provider {
@@ -120,21 +115,10 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, fs: Arc<dyn Fs>,
                                 return;
                             };
 
-                            let Some(Some(workspace)) = window
-                                .update(cx, |_, window, _| window.root().flatten())
-                                .ok()
-                            else {
-                                return;
-                            };
-
                             window
                                 .update(cx, |_, window, cx| {
-                                    ZedPredictModal::toggle(
-                                        workspace,
-                                        user_store.clone(),
-                                        client.clone(),
-                                        fs.clone(),
-                                        window,
+                                    window.dispatch_action(
+                                        Box::new(zed_actions::OpenZedPredictOnboarding),
                                         cx,
                                     );
                                 })
@@ -228,6 +212,7 @@ fn assign_inline_completion_provider(
     window: &mut Window,
     cx: &mut Context<Editor>,
 ) {
+    // TODO: Do we really want to collect data only for singleton buffers?
     let singleton_buffer = editor.buffer().read(cx).as_singleton();
 
     match provider {
@@ -255,7 +240,23 @@ fn assign_inline_completion_provider(
             if cx.has_flag::<PredictEditsFeatureFlag>()
                 || (cfg!(debug_assertions) && client.status().borrow().is_connected())
             {
-                let zeta = zeta::Zeta::register(client.clone(), user_store, cx);
+                let mut worktree = None;
+
+                if let Some(buffer) = &singleton_buffer {
+                    if let Some(file) = buffer.read(cx).file() {
+                        let id = file.worktree_id(cx);
+                        if let Some(inner_worktree) = editor
+                            .project
+                            .as_ref()
+                            .and_then(|project| project.read(cx).worktree_for_id(id, cx))
+                        {
+                            worktree = Some(inner_worktree);
+                        }
+                    }
+                }
+
+                let zeta = zeta::Zeta::register(worktree, client.clone(), user_store, cx);
+
                 if let Some(buffer) = &singleton_buffer {
                     if buffer.read(cx).file().is_some() {
                         zeta.update(cx, |zeta, cx| {
@@ -264,12 +265,8 @@ fn assign_inline_completion_provider(
                     }
                 }
 
-                let data_collection = ProviderDataCollection::new(
-                    zeta.clone(),
-                    window.root::<Workspace>().flatten(),
-                    singleton_buffer,
-                    cx,
-                );
+                let data_collection =
+                    ProviderDataCollection::new(zeta.clone(), singleton_buffer, cx);
 
                 let provider =
                     cx.new(|_| zeta::ZetaInlineCompletionProvider::new(zeta, data_collection));

crates/zed_actions/src/lib.rs 🔗

@@ -186,3 +186,5 @@ pub mod outline {
     /// A pointer to outline::toggle function, exposed here to sewer the breadcrumbs <-> outline dependency.
     pub static TOGGLE_OUTLINE: OnceLock<fn(AnyView, &mut Window, &mut App)> = OnceLock::new();
 }
+
+actions!(zed_predict_onboarding, [OpenZedPredictOnboarding]);

crates/zed_predict_onboarding/Cargo.toml 🔗

@@ -1,31 +0,0 @@
-[package]
-name = "zed_predict_onboarding"
-version = "0.1.0"
-edition = "2021"
-publish = false
-license = "GPL-3.0-or-later"
-
-[lints]
-workspace = true
-
-[lib]
-path = "src/lib.rs"
-doctest = false
-
-[features]
-test-support = []
-
-[dependencies]
-chrono.workspace = true
-client.workspace = true
-db.workspace = true
-feature_flags.workspace = true
-fs.workspace = true
-gpui.workspace = true
-language.workspace = true
-menu.workspace = true
-settings.workspace = true
-theme.workspace = true
-ui.workspace = true
-util.workspace = true
-workspace.workspace = true

crates/zeta/Cargo.toml 🔗

@@ -19,12 +19,14 @@ test-support = []
 [dependencies]
 anyhow.workspace = true
 arrayvec.workspace = true
+chrono.workspace = true
 client.workspace = true
 collections.workspace = true
 command_palette_hooks.workspace = true
 db.workspace = true
 editor.workspace = true
 feature_flags.workspace = true
+fs.workspace = true
 futures.workspace = true
 gpui.workspace = true
 http_client.workspace = true
@@ -34,6 +36,8 @@ language.workspace = true
 language_models.workspace = true
 log.workspace = true
 menu.workspace = true
+postage.workspace = true
+regex.workspace = true
 rpc.workspace = true
 serde.workspace = true
 serde_json.workspace = true
@@ -46,6 +50,8 @@ ui.workspace = true
 util.workspace = true
 uuid.workspace = true
 workspace.workspace = true
+worktree.workspace = true
+zed_actions.workspace = true
 
 [dev-dependencies]
 collections = { workspace = true, features = ["test-support"] }
@@ -64,6 +70,7 @@ settings = { workspace = true, features = ["test-support"] }
 theme = { workspace = true, features = ["test-support"] }
 tree-sitter-go.workspace = true
 tree-sitter-rust.workspace = true
+unindent.workspace = true
 workspace = { workspace = true, features = ["test-support"] }
 worktree = { workspace = true, features = ["test-support"] }
 call = { workspace = true, features = ["test-support"] }

crates/zeta/src/init.rs 🔗

@@ -0,0 +1,60 @@
+use std::any::{Any, TypeId};
+
+use command_palette_hooks::CommandPaletteFilter;
+use feature_flags::{
+    FeatureFlagAppExt as _, PredictEditsFeatureFlag, PredictEditsRateCompletionsFeatureFlag,
+};
+use ui::App;
+use workspace::Workspace;
+
+use crate::{onboarding_modal::ZedPredictModal, RateCompletionModal, RateCompletions};
+
+pub fn init(cx: &mut App) {
+    cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
+        workspace.register_action(|workspace, _: &RateCompletions, window, cx| {
+            if cx.has_flag::<PredictEditsRateCompletionsFeatureFlag>() {
+                RateCompletionModal::toggle(workspace, window, cx);
+            }
+        });
+
+        workspace.register_action(
+            move |workspace, _: &zed_actions::OpenZedPredictOnboarding, window, cx| {
+                if cx.has_flag::<PredictEditsFeatureFlag>() {
+                    ZedPredictModal::toggle(
+                        workspace,
+                        workspace.user_store().clone(),
+                        workspace.client().clone(),
+                        workspace.app_state().fs.clone(),
+                        window,
+                        cx,
+                    )
+                }
+            },
+        );
+    })
+    .detach();
+
+    feature_gate_predict_edits_rating_actions(cx);
+}
+
+fn feature_gate_predict_edits_rating_actions(cx: &mut App) {
+    let rate_completion_action_types = [TypeId::of::<RateCompletions>()];
+
+    CommandPaletteFilter::update_global(cx, |filter, _cx| {
+        filter.hide_action_types(&rate_completion_action_types);
+        filter.hide_action_types(&[zed_actions::OpenZedPredictOnboarding.type_id()]);
+    });
+
+    cx.observe_flag::<PredictEditsRateCompletionsFeatureFlag, _>(move |is_enabled, cx| {
+        if is_enabled {
+            CommandPaletteFilter::update_global(cx, |filter, _cx| {
+                filter.show_action_types(rate_completion_action_types.iter());
+            });
+        } else {
+            CommandPaletteFilter::update_global(cx, |filter, _cx| {
+                filter.hide_action_types(&rate_completion_action_types);
+            });
+        }
+    })
+    .detach();
+}

crates/zeta/src/license_detection.rs 🔗

@@ -0,0 +1,210 @@
+use regex::Regex;
+
+pub fn is_license_eligible_for_data_collection(license: &str) -> bool {
+    // TODO: Include more licenses later (namely, Apache)
+    for pattern in [MIT_LICENSE_REGEX, ISC_LICENSE_REGEX] {
+        let regex = Regex::new(pattern.trim()).unwrap();
+        if regex.is_match(license.trim()) {
+            return true;
+        }
+    }
+    false
+}
+
+const MIT_LICENSE_REGEX: &str = r#"
+^.*MIT License.*
+
+Copyright.*?
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files \(the "Software"\), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software\.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT\. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE\.$
+"#;
+
+const ISC_LICENSE_REGEX: &str = r#"
+^ISC License
+
+Copyright.*?
+
+Permission to use, copy, modify, and/or distribute this software for any
+purpose with or without fee is hereby granted, provided that the above
+copyright notice and this permission notice appear in all copies\.
+
+THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS\. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE\.$
+"#;
+
+#[cfg(test)]
+mod tests {
+    use unindent::unindent;
+
+    use crate::is_license_eligible_for_data_collection;
+
+    #[test]
+    fn test_mit_positive_detection() {
+        let example_license = unindent(
+            r#"
+                MIT License
+
+                Copyright (c) 2024 John Doe
+
+                Permission is hereby granted, free of charge, to any person obtaining a copy
+                of this software and associated documentation files (the "Software"), to deal
+                in the Software without restriction, including without limitation the rights
+                to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+                copies of the Software, and to permit persons to whom the Software is
+                furnished to do so, subject to the following conditions:
+
+                The above copyright notice and this permission notice shall be included in all
+                copies or substantial portions of the Software.
+
+                THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+                IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+                FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+                AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+                LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+                OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+                SOFTWARE.
+            "#
+            .trim(),
+        );
+
+        assert!(is_license_eligible_for_data_collection(&example_license));
+
+        let example_license = unindent(
+            r#"
+                The MIT License (MIT)
+
+                Copyright (c) 2019 John Doe
+
+                Permission is hereby granted, free of charge, to any person obtaining a copy
+                of this software and associated documentation files (the "Software"), to deal
+                in the Software without restriction, including without limitation the rights
+                to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+                copies of the Software, and to permit persons to whom the Software is
+                furnished to do so, subject to the following conditions:
+
+                The above copyright notice and this permission notice shall be included in all
+                copies or substantial portions of the Software.
+
+                THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+                IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+                FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+                AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+                LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+                OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+                SOFTWARE.
+            "#
+            .trim(),
+        );
+
+        assert!(is_license_eligible_for_data_collection(&example_license));
+    }
+
+    #[test]
+    fn test_mit_negative_detection() {
+        let example_license = unindent(
+            r#"
+                MIT License
+
+                Copyright (c) 2024 John Doe
+
+                Permission is hereby granted, free of charge, to any person obtaining a copy
+                of this software and associated documentation files (the "Software"), to deal
+                in the Software without restriction, including without limitation the rights
+                to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+                copies of the Software, and to permit persons to whom the Software is
+                furnished to do so, subject to the following conditions:
+
+                The above copyright notice and this permission notice shall be included in all
+                copies or substantial portions of the Software.
+
+                THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+                IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+                FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+                AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+                LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+                OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+                SOFTWARE.
+
+                This project is dual licensed under the MIT License and the Apache License, Version 2.0.
+            "#
+            .trim(),
+        );
+
+        assert!(!is_license_eligible_for_data_collection(&example_license));
+    }
+
+    #[test]
+    fn test_isc_positive_detection() {
+        let example_license = unindent(
+            r#"
+                ISC License
+
+                Copyright (c) 2024, John Doe
+
+                Permission to use, copy, modify, and/or distribute this software for any
+                purpose with or without fee is hereby granted, provided that the above
+                copyright notice and this permission notice appear in all copies.
+
+                THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+                WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+                MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+                ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+                WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+                ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+                OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+            "#
+            .trim(),
+        );
+
+        assert!(is_license_eligible_for_data_collection(&example_license));
+    }
+
+    #[test]
+    fn test_isc_negative_detection() {
+        let example_license = unindent(
+            r#"
+                ISC License
+
+                Copyright (c) 2024, John Doe
+
+                Permission to use, copy, modify, and/or distribute this software for any
+                purpose with or without fee is hereby granted, provided that the above
+                copyright notice and this permission notice appear in all copies.
+
+                THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+                WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+                MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+                ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+                WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+                ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+                OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+                This project is dual licensed under the ISC License and the MIT License.
+            "#
+            .trim(),
+        );
+
+        assert!(!is_license_eligible_for_data_collection(&example_license));
+    }
+}

crates/zed_predict_onboarding/src/banner.rs → crates/zeta/src/onboarding_banner.rs 🔗

@@ -1,40 +1,20 @@
-use std::sync::Arc;
-
-use crate::ZedPredictModal;
 use chrono::Utc;
-use client::{Client, UserStore};
 use feature_flags::{FeatureFlagAppExt as _, PredictEditsFeatureFlag};
-use fs::Fs;
-use gpui::{Entity, Subscription, WeakEntity};
+use gpui::Subscription;
 use language::language_settings::{all_language_settings, InlineCompletionProvider};
 use settings::SettingsStore;
 use ui::{prelude::*, ButtonLike, Tooltip};
 use util::ResultExt;
-use workspace::Workspace;
 
-/// Prompts user to try AI inline prediction feature
+/// Prompts the user to try Zed's Edit Prediction feature
 pub struct ZedPredictBanner {
-    workspace: WeakEntity<Workspace>,
-    user_store: Entity<UserStore>,
-    client: Arc<Client>,
-    fs: Arc<dyn Fs>,
     dismissed: bool,
     _subscription: Subscription,
 }
 
 impl ZedPredictBanner {
-    pub fn new(
-        workspace: WeakEntity<Workspace>,
-        user_store: Entity<UserStore>,
-        client: Arc<Client>,
-        fs: Arc<dyn Fs>,
-        cx: &mut Context<Self>,
-    ) -> Self {
+    pub fn new(cx: &mut Context<Self>) -> Self {
         Self {
-            workspace,
-            user_store,
-            client,
-            fs,
             dismissed: get_dismissed(),
             _subscription: cx.observe_global::<SettingsStore>(Self::handle_settings_changed),
         }
@@ -126,24 +106,8 @@ impl Render for ZedPredictBanner {
                                     .child(Label::new("Edit Prediction").size(LabelSize::Small)),
                             ),
                     )
-                    .on_click({
-                        let workspace = self.workspace.clone();
-                        let user_store = self.user_store.clone();
-                        let client = self.client.clone();
-                        let fs = self.fs.clone();
-                        move |_, window, cx| {
-                            let Some(workspace) = workspace.upgrade() else {
-                                return;
-                            };
-                            ZedPredictModal::toggle(
-                                workspace,
-                                user_store.clone(),
-                                client.clone(),
-                                fs.clone(),
-                                window,
-                                cx,
-                            );
-                        }
+                    .on_click(|_, window, cx| {
+                        window.dispatch_action(Box::new(zed_actions::OpenZedPredictOnboarding), cx)
                     }),
             )
             .child(
@@ -163,6 +127,6 @@ impl Render for ZedPredictBanner {
                 ),
             );
 
-        div().pr_1().child(banner)
+        div().pr_2().child(banner)
     }
 }

crates/zed_predict_onboarding/src/modal.rs → crates/zeta/src/onboarding_modal.rs 🔗

@@ -1,6 +1,8 @@
 use std::{sync::Arc, time::Duration};
 
+use crate::{Zeta, ZED_PREDICT_DATA_COLLECTION_CHOICE};
 use client::{Client, UserStore};
+use db::kvp::KEY_VALUE_STORE;
 use feature_flags::FeatureFlagAppExt as _;
 use fs::Fs;
 use gpui::{
@@ -9,10 +11,12 @@ use gpui::{
 };
 use language::language_settings::{AllLanguageSettings, InlineCompletionProvider};
 use settings::{update_settings_file, Settings};
-use ui::{prelude::*, CheckboxWithLabel, TintColor};
+use ui::{prelude::*, Checkbox, TintColor, Tooltip};
+use util::ResultExt;
 use workspace::{notifications::NotifyTaskExt, ModalView, Workspace};
+use worktree::Worktree;
 
-/// Introduces user to AI inline prediction feature and terms of service
+/// Introduces user to Zed's Edit Prediction feature and terms of service
 pub struct ZedPredictModal {
     user_store: Entity<UserStore>,
     client: Arc<Client>,
@@ -20,6 +24,9 @@ pub struct ZedPredictModal {
     focus_handle: FocusHandle,
     sign_in_status: SignInStatus,
     terms_of_service: bool,
+    data_collection_expanded: bool,
+    data_collection_opted_in: bool,
+    worktrees: Vec<Entity<Worktree>>,
 }
 
 #[derive(PartialEq, Eq)]
@@ -33,34 +40,26 @@ enum SignInStatus {
 }
 
 impl ZedPredictModal {
-    fn new(
+    pub fn toggle(
+        workspace: &mut Workspace,
         user_store: Entity<UserStore>,
         client: Arc<Client>,
         fs: Arc<dyn Fs>,
-        cx: &mut Context<Self>,
-    ) -> Self {
-        ZedPredictModal {
+        window: &mut Window,
+        cx: &mut Context<Workspace>,
+    ) {
+        let worktrees = workspace.visible_worktrees(cx).collect();
+
+        workspace.toggle_modal(window, cx, |_window, cx| Self {
             user_store,
             client,
             fs,
             focus_handle: cx.focus_handle(),
             sign_in_status: SignInStatus::Idle,
             terms_of_service: false,
-        }
-    }
-
-    pub fn toggle(
-        workspace: Entity<Workspace>,
-        user_store: Entity<UserStore>,
-        client: Arc<Client>,
-        fs: Arc<dyn Fs>,
-        window: &mut Window,
-        cx: &mut App,
-    ) {
-        workspace.update(cx, |this, cx| {
-            this.toggle_modal(window, cx, |_window, cx| {
-                ZedPredictModal::new(user_store, client, fs, cx)
-            });
+            data_collection_expanded: false,
+            data_collection_opted_in: false,
+            worktrees,
         });
     }
 
@@ -74,6 +73,11 @@ impl ZedPredictModal {
         cx.notify();
     }
 
+    fn inline_completions_doc(&mut self, _: &ClickEvent, _: &mut Window, cx: &mut Context<Self>) {
+        cx.open_url("https://zed.dev/docs/configuring-zed#inline-completions");
+        cx.notify();
+    }
+
     fn accept_and_enable(&mut self, _: &ClickEvent, window: &mut Window, cx: &mut Context<Self>) {
         let task = self
             .user_store
@@ -82,6 +86,20 @@ impl ZedPredictModal {
         cx.spawn(|this, mut cx| async move {
             task.await?;
 
+            let mut data_collection_opted_in = false;
+            this.update(&mut cx, |this, _cx| {
+                data_collection_opted_in = this.data_collection_opted_in;
+            })
+            .ok();
+
+            KEY_VALUE_STORE
+                .write_kvp(
+                    ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
+                    data_collection_opted_in.to_string(),
+                )
+                .await
+                .log_err();
+
             this.update(&mut cx, |this, cx| {
                 update_settings_file::<AllLanguageSettings>(this.fs.clone(), cx, move |file, _| {
                     file.features
@@ -89,6 +107,13 @@ impl ZedPredictModal {
                         .inline_completion_provider = Some(InlineCompletionProvider::Zed);
                 });
 
+                if this.worktrees.is_empty() {
+                    cx.emit(DismissEvent);
+                    return;
+                }
+
+                Zeta::register(None, this.client.clone(), this.user_store.clone(), cx);
+
                 cx.emit(DismissEvent);
             })
         })
@@ -135,16 +160,16 @@ impl ModalView for ZedPredictModal {}
 impl Render for ZedPredictModal {
     fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         let base = v_flex()
-            .w(px(420.))
+            .id("zed predict tos")
+            .key_context("ZedPredictModal")
+            .w(px(440.))
             .p_4()
             .relative()
             .gap_2()
             .overflow_hidden()
             .elevation_3(cx)
-            .id("zed predict tos")
             .track_focus(&self.focus_handle(cx))
             .on_action(cx.listener(Self::cancel))
-            .key_context("ZedPredictModal")
             .on_action(cx.listener(|_, _: &menu::Cancel, _window, cx| {
                 cx.emit(DismissEvent);
             }))
@@ -155,15 +180,15 @@ impl Render for ZedPredictModal {
                 div()
                     .p_1p5()
                     .absolute()
-                    .top_0()
-                    .left_0()
+                    .top_1()
+                    .left_1p5()
                     .right_0()
                     .h(px(200.))
                     .child(
                         svg()
                             .path("icons/zed_predict_bg.svg")
                             .text_color(cx.theme().colors().icon_disabled)
-                            .w(px(416.))
+                            .w(px(418.))
                             .h(px(128.))
                             .overflow_hidden(),
                     ),
@@ -249,24 +274,49 @@ impl Render for ZedPredictModal {
 
         if self.user_store.read(cx).current_user().is_some() {
             let copy = match self.sign_in_status {
-                SignInStatus::Idle => "Get accurate and helpful edit predictions at every keystroke. To set Zed as your inline completions provider, ensure you:",
+                SignInStatus::Idle => "Get accurate and instant edit predictions at every keystroke. Before setting Zed as your inline completions provider:",
                 SignInStatus::SignedIn => "Almost there! Ensure you:",
                 SignInStatus::Waiting => unreachable!(),
             };
 
+            let accordion_icons = if self.data_collection_expanded {
+                (IconName::ChevronUp, IconName::ChevronDown)
+            } else {
+                (IconName::ChevronDown, IconName::ChevronUp)
+            };
+
+            fn label_item(label_text: impl Into<SharedString>) -> impl Element {
+                Label::new(label_text).color(Color::Muted).into_element()
+            }
+
+            fn info_item(label_text: impl Into<SharedString>) -> impl Element {
+                h_flex()
+                    .gap_2()
+                    .child(Icon::new(IconName::Check).size(IconSize::XSmall))
+                    .child(label_item(label_text))
+            }
+
+            fn multiline_info_item<E1: Into<SharedString>, E2: IntoElement>(
+                first_line: E1,
+                second_line: E2,
+            ) -> impl Element {
+                v_flex()
+                    .child(info_item(first_line))
+                    .child(div().pl_5().child(second_line))
+            }
+
             base.child(Label::new(copy).color(Color::Muted))
                 .child(
                     h_flex()
-                        .gap_0p5()
-                        .child(CheckboxWithLabel::new(
-                            "tos-checkbox",
-                            Label::new("Have read and accepted the").color(Color::Muted),
-                            self.terms_of_service.into(),
-                            cx.listener(move |this, state, _window, cx| {
-                                this.terms_of_service = *state == ToggleState::Selected;
-                                cx.notify()
-                            }),
-                        ))
+                        .child(
+                            Checkbox::new("tos-checkbox", self.terms_of_service.into())
+                                .fill()
+                                .label("Read and accept the")
+                                .on_click(cx.listener(move |this, state, _window, cx| {
+                                    this.terms_of_service = *state == ToggleState::Selected;
+                                    cx.notify()
+                                })),
+                        )
                         .child(
                             Button::new("view-tos", "Terms of Service")
                                 .icon(IconName::ArrowUpRight)
@@ -275,6 +325,88 @@ impl Render for ZedPredictModal {
                                 .on_click(cx.listener(Self::view_terms)),
                         ),
                 )
+                .child(
+                    v_flex()
+                        .child(
+                            h_flex()
+                                .child(
+                                    Checkbox::new(
+                                        "training-data-checkbox",
+                                        self.data_collection_opted_in.into(),
+                                    )
+                                    .label("Optionally share training data (OSS-only).")
+                                    .fill()
+                                    .when(self.worktrees.is_empty(), |element| {
+                                        element.disabled(true).tooltip(move |window, cx| {
+                                            Tooltip::with_meta(
+                                                "No Project Open",
+                                                None,
+                                                "Open a project to enable this option.",
+                                                window,
+                                                cx,
+                                            )
+                                        })
+                                    })
+                                    .on_click(cx.listener(
+                                        move |this, state, _window, cx| {
+                                            this.data_collection_opted_in =
+                                                *state == ToggleState::Selected;
+                                            cx.notify()
+                                        },
+                                    )),
+                                )
+                                // TODO: show each worktree if more than 1
+                                .child(
+                                    Button::new("learn-more", "Learn More")
+                                        .icon(accordion_icons.0)
+                                        .icon_size(IconSize::Indicator)
+                                        .icon_color(Color::Muted)
+                                        .on_click(cx.listener(|this, _, _, cx| {
+                                            this.data_collection_expanded =
+                                                !this.data_collection_expanded;
+                                            cx.notify()
+                                        })),
+                                ),
+                        )
+                        .when(self.data_collection_expanded, |element| {
+                            element.child(
+                                v_flex()
+                                    .mt_2()
+                                    .p_2()
+                                    .rounded_md()
+                                    .bg(cx.theme().colors().editor_background.opacity(0.5))
+                                    .border_1()
+                                    .border_color(cx.theme().colors().border_variant)
+                                    .child(
+                                        div().child(
+                                            Label::new("To improve edit predictions, help fine-tune Zed's model by sharing data from the open-source projects you work on.")
+                                                .mb_1()
+                                        )
+                                    )
+                                    .child(info_item(
+                                        "We ask this exclusively for open-source projects.",
+                                    ))
+                                    .child(info_item(
+                                        "Zed automatically detects if your project is open-source.",
+                                    ))
+                                    .child(info_item(
+                                        "This setting is valid for all OSS projects you open in Zed.",
+                                    ))
+                                    .child(info_item("Toggle it anytime via the status bar menu."))
+                                    .child(multiline_info_item(
+                                        "Files that can contain sensitive data, like `.env`, are",
+                                        h_flex()
+                                            .child(label_item("excluded by default via the"))
+                                            .child(
+                                                Button::new("doc-link", "disabled_globs").on_click(
+                                                    cx.listener(Self::inline_completions_doc),
+                                                ),
+                                            )
+                                            .child(label_item("setting.")),
+                                    )),
+                            )
+                        }),
+                )
                 .child(
                     v_flex()
                         .mt_2()

crates/zeta/src/persistence.rs 🔗

@@ -1,48 +0,0 @@
-use std::path::{Path, PathBuf};
-use workspace::WorkspaceDb;
-
-use db::sqlez_macros::sql;
-use db::{define_connection, query};
-
-define_connection!(
-    pub static ref DB: ZetaDb<WorkspaceDb> = &[
-        sql! (
-            CREATE TABLE zeta_preferences(
-                worktree_path BLOB NOT NULL PRIMARY KEY,
-                accepted_data_collection INTEGER
-            ) STRICT;
-        ),
-    ];
-);
-
-impl ZetaDb {
-    query! {
-        pub fn get_all_data_collection_preferences() -> Result<Vec<(PathBuf, bool)>> {
-            SELECT worktree_path, accepted_data_collection FROM zeta_preferences
-        }
-    }
-
-    query! {
-        pub fn get_accepted_data_collection(worktree_path: &Path) -> Result<Option<bool>> {
-            SELECT accepted_data_collection FROM zeta_preferences
-            WHERE worktree_path = ?
-        }
-    }
-
-    query! {
-        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
-                (?1, ?2)
-            ON CONFLICT (worktree_path) DO UPDATE SET
-                accepted_data_collection = ?2
-        }
-    }
-
-    query! {
-        pub async fn clear_all_zeta_preferences() -> Result<()> {
-            DELETE FROM zeta_preferences
-        }
-    }
-}

crates/zeta/src/rate_completion_modal.rs 🔗

@@ -1,10 +1,8 @@
 use crate::{CompletionDiffElement, InlineCompletion, InlineCompletionRating, Zeta};
-use command_palette_hooks::CommandPaletteFilter;
 use editor::Editor;
-use feature_flags::{FeatureFlagAppExt as _, PredictEditsRateCompletionsFeatureFlag};
 use gpui::{actions, prelude::*, App, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable};
 use language::language_settings;
-use std::{any::TypeId, time::Duration};
+use std::time::Duration;
 use ui::{prelude::*, KeyBinding, List, ListItem, ListItemSpacing, Tooltip};
 use workspace::{ModalView, Workspace};
 
@@ -21,40 +19,6 @@ actions!(
     ]
 );
 
-pub fn init(cx: &mut App) {
-    cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
-        workspace.register_action(|workspace, _: &RateCompletions, window, cx| {
-            if cx.has_flag::<PredictEditsRateCompletionsFeatureFlag>() {
-                RateCompletionModal::toggle(workspace, window, cx);
-            }
-        });
-    })
-    .detach();
-
-    feature_gate_predict_edits_rating_actions(cx);
-}
-
-fn feature_gate_predict_edits_rating_actions(cx: &mut App) {
-    let rate_completion_action_types = [TypeId::of::<RateCompletions>()];
-
-    CommandPaletteFilter::update_global(cx, |filter, _cx| {
-        filter.hide_action_types(&rate_completion_action_types);
-    });
-
-    cx.observe_flag::<PredictEditsRateCompletionsFeatureFlag, _>(move |is_enabled, cx| {
-        if is_enabled {
-            CommandPaletteFilter::update_global(cx, |filter, _cx| {
-                filter.show_action_types(rate_completion_action_types.iter());
-            });
-        } else {
-            CommandPaletteFilter::update_global(cx, |filter, _cx| {
-                filter.hide_action_types(&rate_completion_action_types);
-            });
-        }
-    })
-    .detach();
-}
-
 pub struct RateCompletionModal {
     zeta: Entity<Zeta>,
     active_completion: Option<ActiveCompletion>,

crates/zeta/src/zeta.rs 🔗

@@ -1,22 +1,26 @@
 mod completion_diff_element;
-mod persistence;
+mod init;
+mod license_detection;
+mod onboarding_banner;
+mod onboarding_modal;
 mod rate_completion_modal;
 
 pub(crate) use completion_diff_element::*;
 use db::kvp::KEY_VALUE_STORE;
+pub use init::*;
 use inline_completion::DataCollectionState;
+pub use license_detection::is_license_eligible_for_data_collection;
+pub use onboarding_banner::*;
 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;
 use gpui::{
     actions, App, AppContext as _, AsyncApp, Context, Entity, EntityId, Global, Subscription, Task,
-    WeakEntity,
 };
 use http_client::{HttpClient, Method};
 use language::{
@@ -24,33 +28,32 @@ use language::{
     OffsetRangeExt, Point, ToOffset, ToPoint,
 };
 use language_models::LlmApiToken;
+use postage::watch;
 use rpc::{PredictEditsParams, PredictEditsResponse, EXPIRED_LLM_TOKEN_HEADER_NAME};
+use settings::WorktreeId;
 use std::{
     borrow::Cow,
-    cmp, env,
+    cmp,
     fmt::Write,
     future::Future,
     mem,
     ops::Range,
-    path::{Path, PathBuf},
+    path::Path,
+    rc::Rc,
     sync::Arc,
     time::{Duration, Instant},
 };
 use telemetry_events::InlineCompletionRating;
 use util::ResultExt;
 use uuid::Uuid;
-use workspace::{
-    notifications::{simple_message_notification::MessageNotification, NotificationId},
-    Workspace,
-};
+use worktree::Worktree;
 
 const CURSOR_MARKER: &'static str = "<|user_cursor_is_here|>";
 const START_OF_FILE_MARKER: &'static str = "<|start_of_file|>";
 const EDITABLE_REGION_START_MARKER: &'static str = "<|editable_region_start|>";
 const EDITABLE_REGION_END_MARKER: &'static str = "<|editable_region_end|>";
 const BUFFER_CHANGE_GROUPING_INTERVAL: Duration = Duration::from_secs(1);
-const ZED_PREDICT_DATA_COLLECTION_NEVER_ASK_AGAIN_KEY: &'static str =
-    "zed_predict_data_collection_never_ask_again";
+const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice";
 
 // TODO(mgsloan): more systematic way to choose or tune these fairly arbitrary constants?
 
@@ -206,11 +209,12 @@ pub struct Zeta {
     registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
     shown_completions: VecDeque<InlineCompletion>,
     rated_completions: HashSet<InlineCompletionId>,
-    data_collection_preferences: DataCollectionPreferences,
+    data_collection_choice: Entity<DataCollectionChoice>,
     llm_token: LlmApiToken,
     _llm_token_subscription: Subscription,
     tos_accepted: bool, // Terms of service accepted
     _user_store_subscription: Subscription,
+    license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
 }
 
 impl Zeta {
@@ -219,15 +223,28 @@ impl Zeta {
     }
 
     pub fn register(
+        worktree: Option<Entity<Worktree>>,
         client: Arc<Client>,
         user_store: Entity<UserStore>,
         cx: &mut App,
     ) -> Entity<Self> {
-        Self::global(cx).unwrap_or_else(|| {
+        let this = Self::global(cx).unwrap_or_else(|| {
             let model = cx.new(|cx| Self::new(client, user_store, cx));
             cx.set_global(ZetaGlobal(model.clone()));
             model
-        })
+        });
+
+        this.update(cx, move |this, cx| {
+            if let Some(worktree) = worktree {
+                worktree.update(cx, |worktree, cx| {
+                    this.license_detection_watchers
+                        .entry(worktree.id())
+                        .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(worktree, cx)));
+                });
+            }
+        });
+
+        this
     }
 
     pub fn clear_history(&mut self) {
@@ -236,13 +253,17 @@ impl Zeta {
 
     fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
         let refresh_llm_token_listener = language_models::RefreshLlmTokenListener::global(cx);
+
+        let data_collection_choice = Self::load_data_collection_choices();
+        let data_collection_choice = cx.new(|_| data_collection_choice);
+
         Self {
             client,
             events: VecDeque::new(),
             shown_completions: VecDeque::new(),
             rated_completions: HashSet::default(),
             registered_buffers: HashMap::default(),
-            data_collection_preferences: Self::load_data_collection_preferences(cx),
+            data_collection_choice,
             llm_token: LlmApiToken::default(),
             _llm_token_subscription: cx.subscribe(
                 &refresh_llm_token_listener,
@@ -271,6 +292,7 @@ impl Zeta {
                     _ => {}
                 }
             }),
+            license_detection_watchers: HashMap::default(),
         }
     }
 
@@ -342,7 +364,7 @@ impl Zeta {
         &mut self,
         buffer: &Entity<Buffer>,
         cursor: language::Anchor,
-        can_collect_data: bool,
+        data_collection_permission: bool,
         cx: &mut Context<Self>,
         perform_predict_edits: F,
     ) -> Task<Result<Option<InlineCompletion>>>
@@ -407,7 +429,7 @@ impl Zeta {
                 input_events: input_events.clone(),
                 input_excerpt: input_excerpt.clone(),
                 outline: Some(input_outline.clone()),
-                can_collect_data,
+                data_collection_permission,
             };
 
             let response = perform_predict_edits(client, llm_token, is_staff, body).await?;
@@ -587,13 +609,13 @@ and then another
         &mut self,
         buffer: &Entity<Buffer>,
         position: language::Anchor,
-        can_collect_data: bool,
+        data_collection_permission: bool,
         cx: &mut Context<Self>,
     ) -> Task<Result<Option<InlineCompletion>>> {
         self.request_completion_impl(
             buffer,
             position,
-            can_collect_data,
+            data_collection_permission,
             cx,
             Self::perform_predict_edits,
         )
@@ -903,84 +925,55 @@ and then another
         new_snapshot
     }
 
-    /// Creates a `Entity<DataCollectionChoice>` for each unique worktree abs path it sees.
-    pub fn data_collection_choice_at(
-        &mut self,
-        worktree_abs_path: PathBuf,
-        cx: &mut Context<Self>,
-    ) -> Entity<DataCollectionChoice> {
-        match self
-            .data_collection_preferences
-            .per_worktree
-            .entry(worktree_abs_path)
-        {
-            Entry::Vacant(entry) => {
-                let choice = cx.new(|_| DataCollectionChoice::NotAnswered);
-                entry.insert(choice.clone());
-                choice
+    fn load_data_collection_choices() -> DataCollectionChoice {
+        let choice = KEY_VALUE_STORE
+            .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
+            .log_err()
+            .flatten();
+
+        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
             }
-            Entry::Occupied(entry) => entry.get().clone(),
+            None => DataCollectionChoice::NotAnswered,
         }
     }
+}
 
-    fn set_never_ask_again_for_data_collection(&mut self, cx: &mut Context<Self>) {
-        self.data_collection_preferences.never_ask_again = true;
+struct LicenseDetectionWatcher {
+    is_open_source_rx: watch::Receiver<bool>,
+    _is_open_source_task: Task<()>,
+}
 
-        // persist choice
-        db::write_and_log(cx, move || {
-            KEY_VALUE_STORE.write_kvp(
-                ZED_PREDICT_DATA_COLLECTION_NEVER_ASK_AGAIN_KEY.into(),
-                "true".to_string(),
-            )
-        });
-    }
+impl LicenseDetectionWatcher {
+    pub fn new(worktree: &Worktree, cx: &mut Context<Worktree>) -> Self {
+        let (mut is_open_source_tx, is_open_source_rx) = watch::channel_with::<bool>(false);
 
-    fn load_data_collection_preferences(cx: &mut Context<Self>) -> DataCollectionPreferences {
-        if env::var("ZED_PREDICT_CLEAR_DATA_COLLECTION_PREFERENCES").is_ok() {
-            db::write_and_log(cx, move || async move {
-                KEY_VALUE_STORE
-                    .delete_kvp(ZED_PREDICT_DATA_COLLECTION_NEVER_ASK_AGAIN_KEY.into())
-                    .await
-                    .log_err();
+        let loaded_file_fut = worktree.load_file(Path::new("LICENSE"), false, cx);
 
-                persistence::DB.clear_all_zeta_preferences().await
-            });
-            return DataCollectionPreferences::default();
-        }
-
-        let never_ask_again = KEY_VALUE_STORE
-            .read_kvp(ZED_PREDICT_DATA_COLLECTION_NEVER_ASK_AGAIN_KEY)
-            .log_err()
-            .flatten()
-            .map(|value| value == "true")
-            .unwrap_or(false);
+        Self {
+            is_open_source_rx,
+            _is_open_source_task: cx.spawn(|_, _| async move {
+                // TODO: Don't display error if file not found
+                let Some(loaded_file) = loaded_file_fut.await.log_err() else {
+                    return;
+                };
 
-        let preferences_per_worktree = persistence::DB
-            .get_all_data_collection_preferences()
-            .log_err()
-            .into_iter()
-            .flatten()
-            .map(|(path, choice)| {
-                let choice = cx.new(|_| DataCollectionChoice::from(choice));
-                (path, choice)
-            })
-            .collect();
+                let is_loaded_file_open_source_thing: bool =
+                    is_license_eligible_for_data_collection(&loaded_file.text);
 
-        DataCollectionPreferences {
-            never_ask_again,
-            per_worktree: preferences_per_worktree,
+                *is_open_source_tx.borrow_mut() = is_loaded_file_open_source_thing;
+            }),
         }
     }
-}
 
-#[derive(Default, Debug)]
-struct DataCollectionPreferences {
-    /// Set when a user clicks on "Never Ask Again", can never be unset.
-    never_ask_again: 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>>,
+    /// Answers false until we find out it's open source
+    pub fn is_open_source(&self) -> bool {
+        *self.is_open_source_rx.borrow()
+    }
 }
 
 fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
@@ -1308,7 +1301,7 @@ impl DataCollectionChoice {
         }
     }
 
-    pub fn toggle(self) -> DataCollectionChoice {
+    pub fn toggle(&self) -> DataCollectionChoice {
         match self {
             Self::Enabled => Self::Disabled,
             Self::Disabled => Self::Enabled,
@@ -1326,87 +1319,93 @@ impl From<bool> for DataCollectionChoice {
     }
 }
 
-pub struct ZetaInlineCompletionProvider {
-    zeta: Entity<Zeta>,
-    pending_completions: ArrayVec<PendingCompletion, 2>,
-    next_pending_completion_id: usize,
-    current_completion: Option<CurrentInlineCompletion>,
-    data_collection: Option<ProviderDataCollection>,
-}
-
 pub struct ProviderDataCollection {
-    workspace: WeakEntity<Workspace>,
-    worktree_root_path: PathBuf,
-    choice: Entity<DataCollectionChoice>,
+    /// When set to None, data collection is not possible in the provider buffer
+    choice: Option<Entity<DataCollectionChoice>>,
+    license_detection_watcher: Option<Rc<LicenseDetectionWatcher>>,
 }
 
 impl ProviderDataCollection {
-    pub fn new(
-        zeta: Entity<Zeta>,
-        workspace: Option<Entity<Workspace>>,
-        buffer: Option<Entity<Buffer>>,
-        cx: &mut App,
-    ) -> Option<ProviderDataCollection> {
-        let workspace = workspace?;
-
-        let worktree_root_path = buffer?.update(cx, |buffer, cx| {
-            let file = buffer.file()?;
+    pub fn new(zeta: Entity<Zeta>, buffer: Option<Entity<Buffer>>, cx: &mut App) -> Self {
+        let choice_and_watcher = buffer.and_then(|buffer| {
+            let file = buffer.read(cx).file()?;
 
             if !file.is_local() || file.is_private() {
                 return None;
             }
 
-            workspace.update(cx, |workspace, cx| {
-                Some(
-                    workspace
-                        .absolute_path_of_worktree(file.worktree_id(cx), cx)?
-                        .to_path_buf(),
-                )
-            })
-        })?;
+            let zeta = zeta.read(cx);
+            let choice = zeta.data_collection_choice.clone();
+
+            // Unwrap safety: there should be a watcher for each worktree
+            let license_detection_watcher = zeta
+                .license_detection_watchers
+                .get(&file.worktree_id(cx))
+                .cloned()?;
 
-        let choice = zeta.update(cx, |zeta, cx| {
-            zeta.data_collection_choice_at(worktree_root_path.clone(), cx)
+            Some((choice, license_detection_watcher))
         });
 
-        Some(ProviderDataCollection {
-            workspace: workspace.downgrade(),
-            worktree_root_path,
-            choice,
-        })
+        if let Some((choice, watcher)) = choice_and_watcher {
+            ProviderDataCollection {
+                choice: Some(choice),
+                license_detection_watcher: Some(watcher),
+            }
+        } else {
+            ProviderDataCollection {
+                choice: None,
+                license_detection_watcher: None,
+            }
+        }
     }
 
-    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();
+    pub fn data_collection_permission(&self, cx: &App) -> bool {
+        self.choice
+            .as_ref()
+            .is_some_and(|choice| choice.read(cx).is_enabled())
+            && self
+                .license_detection_watcher
+                .as_ref()
+                .is_some_and(|watcher| watcher.is_open_source())
+    }
+
+    pub fn toggle(&mut self, cx: &mut App) {
+        if let Some(choice) = self.choice.as_mut() {
+            let new_choice = choice.update(cx, |choice, _cx| {
+                let new_choice = choice.toggle();
+                *choice = new_choice;
+                new_choice
+            });
 
-        db::write_and_log(cx, move || {
-            persistence::DB.save_data_collection_choice(worktree_root_path, choice.is_enabled())
-        });
+            db::write_and_log(cx, move || {
+                KEY_VALUE_STORE.write_kvp(
+                    ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
+                    new_choice.is_enabled().to_string(),
+                )
+            });
+        }
     }
+}
 
-    fn toggle_choice(&mut self, cx: &mut App) {
-        self.set_choice(self.choice.read(cx).toggle(), cx);
-    }
+pub struct ZetaInlineCompletionProvider {
+    zeta: Entity<Zeta>,
+    pending_completions: ArrayVec<PendingCompletion, 2>,
+    next_pending_completion_id: usize,
+    current_completion: Option<CurrentInlineCompletion>,
+    /// None if this is entirely disabled for this provider
+    provider_data_collection: ProviderDataCollection,
 }
 
 impl ZetaInlineCompletionProvider {
     pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(8);
 
-    pub fn new(zeta: Entity<Zeta>, data_collection: Option<ProviderDataCollection>) -> Self {
+    pub fn new(zeta: Entity<Zeta>, provider_data_collection: ProviderDataCollection) -> Self {
         Self {
             zeta,
             pending_completions: ArrayVec::new(),
             next_pending_completion_id: 0,
             current_completion: None,
-            data_collection,
-        }
-    }
-
-    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, cx);
+            provider_data_collection,
         }
     }
 }
@@ -1433,11 +1432,7 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
     }
 
     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.read(cx).is_enabled() {
+        if self.provider_data_collection.data_collection_permission(cx) {
             DataCollectionState::Enabled
         } else {
             DataCollectionState::Disabled
@@ -1445,9 +1440,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(cx);
-        }
+        self.provider_data_collection.toggle(cx);
     }
 
     fn is_enabled(
@@ -1495,12 +1488,8 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
 
         let pending_completion_id = self.next_pending_completion_id;
         self.next_pending_completion_id += 1;
-        let can_collect_data = self
-            .data_collection
-            .as_ref()
-            .map_or(false, |data_collection| {
-                data_collection.choice.read(cx).is_enabled()
-            });
+        let data_collection_permission =
+            self.provider_data_collection.data_collection_permission(cx);
 
         let task = cx.spawn(|this, mut cx| async move {
             if debounce {
@@ -1509,7 +1498,7 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
 
             let completion_request = this.update(&mut cx, |this, cx| {
                 this.zeta.update(cx, |zeta, cx| {
-                    zeta.request_completion(&buffer, position, can_collect_data, cx)
+                    zeta.request_completion(&buffer, position, data_collection_permission, cx)
                 })
             });
 
@@ -1596,79 +1585,8 @@ impl inline_completion::InlineCompletionProvider for ZetaInlineCompletionProvide
         // Right now we don't support cycling.
     }
 
-    fn accept(&mut self, cx: &mut Context<Self>) {
+    fn accept(&mut self, _cx: &mut Context<Self>) {
         self.pending_completions.clear();
-
-        let Some(data_collection) = self.data_collection.as_mut() else {
-            return;
-        };
-
-        if data_collection.choice.read(cx).is_answered()
-            || self
-                .zeta
-                .read(cx)
-                .data_collection_preferences
-                .never_ask_again
-        {
-            return;
-        }
-
-        struct ZetaDataCollectionNotification;
-        let notification_id = NotificationId::unique::<ZetaDataCollectionNotification>();
-
-        const DATA_COLLECTION_INFO_URL: &str = "https://zed.dev/terms-of-service"; // TODO: Replace for a link that's dedicated to Edit Predictions data collection
-
-        let this = cx.entity();
-        data_collection
-            .workspace
-            .update(cx, |workspace, cx| {
-                workspace.show_notification(notification_id, cx, |cx| {
-                    let zeta = self.zeta.clone();
-
-                    cx.new(move |_cx| {
-                        let message =
-                            "To allow Zed to suggest better edits, turn on data collection. You \
-                            can turn off at any time via the status bar menu.";
-                        MessageNotification::new(message)
-                            .with_title("Per-Project Data Collection Program")
-                            .show_close_button(false)
-                            .with_click_message("Turn On")
-                            .on_click({
-                                let this = this.clone();
-                                move |_window, cx| {
-                                    this.update(cx, |this, cx| {
-                                        this.set_data_collection_choice(
-                                            DataCollectionChoice::Enabled,
-                                            cx,
-                                        )
-                                    });
-                                }
-                            })
-                            .with_secondary_click_message("Turn Off")
-                            .on_secondary_click({
-                                move |_window, cx| {
-                                    this.update(cx, |this, cx| {
-                                        this.set_data_collection_choice(
-                                            DataCollectionChoice::Disabled,
-                                            cx,
-                                        )
-                                    });
-                                }
-                            })
-                            .with_tertiary_click_message("Never Ask Again")
-                            .on_tertiary_click({
-                                move |_window, cx| {
-                                    zeta.update(cx, |zeta, cx| {
-                                        zeta.set_never_ask_again_for_data_collection(cx);
-                                    });
-                                }
-                            })
-                            .more_info_message("Learn More")
-                            .more_info_url(DATA_COLLECTION_INFO_URL)
-                    })
-                });
-            })
-            .log_err();
     }
 
     fn discard(&mut self, _cx: &mut Context<Self>) {