edit_predictions_ui: Enforce org configuration to disable Zed provider

Tom Houlé and Marshall Bowers created 1 day ago

The Zed provider is now disabled in the selection dropdown when the current organization has disabled edit predictions. It is also unselected when you switch organizations to one such organization.

This is also already enforced server side.

Co-authored-by: Marshall Bowers <[e-mail]>

Change summary

crates/edit_prediction_ui/src/edit_prediction_button.rs | 29 ++++++++++
crates/zed/src/zed/edit_prediction_registry.rs          | 18 ++++++
2 files changed, 44 insertions(+), 3 deletions(-)

Detailed changes

crates/edit_prediction_ui/src/edit_prediction_button.rs

@@ -573,6 +573,14 @@ impl EditPredictionButton {
         current_provider: EditPredictionProvider,
         cx: &mut App,
     ) -> ContextMenu {
+        let organization_configuration = self
+            .user_store
+            .read(cx)
+            .current_organization_configuration();
+
+        let is_zed_provider_disabled = organization_configuration
+            .is_some_and(|configuration| !configuration.edit_prediction.is_enabled);
+
         let available_providers = get_available_providers(cx);
 
         let providers: Vec<_> = available_providers
@@ -592,7 +600,26 @@ impl EditPredictionButton {
 
                 menu = menu.item(
                     ContextMenuEntry::new(name)
-                        .toggleable(IconPosition::Start, is_current)
+                        .toggleable(
+                            IconPosition::Start,
+                            is_current
+                                && (provider == EditPredictionProvider::Zed
+                                    && !is_zed_provider_disabled),
+                        )
+                        .disabled(
+                            provider == EditPredictionProvider::Zed && is_zed_provider_disabled,
+                        )
+                        .when(
+                            provider == EditPredictionProvider::Zed && is_zed_provider_disabled,
+                            |item| {
+                                item.documentation_aside(DocumentationSide::Left, move |_cx| {
+                                    Label::new(
+                                        "Edit predictions are disabled for this organization.",
+                                    )
+                                    .into_any_element()
+                                })
+                            },
+                        )
                         .handler(move |_, cx| {
                             set_completion_provider(fs.clone(), cx, provider);
                         }),

crates/zed/src/zed/edit_prediction_registry.rs

@@ -62,8 +62,9 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
         let editors = editors.clone();
         let client = client.clone();
 
-        move |user_store, event, cx| {
-            if let client::user::Event::PrivateUserInfoUpdated = event {
+        move |user_store, event, cx| match event {
+            client::user::Event::PrivateUserInfoUpdated
+            | client::user::Event::OrganizationChanged => {
                 let provider_config = edit_prediction_provider_config_for_settings(cx);
                 assign_edit_prediction_providers(
                     &editors,
@@ -73,6 +74,7 @@ pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
                     cx,
                 );
             }
+            _ => {}
         }
     })
     .detach();
@@ -275,6 +277,18 @@ fn assign_edit_prediction_provider(
         Some(EditPredictionProviderConfig::Zed(model)) => {
             let ep_store = edit_prediction::EditPredictionStore::global(client, &user_store, cx);
 
+            if let Some(organization_configuration) =
+                user_store.read(cx).current_organization_configuration()
+            {
+                if !organization_configuration.edit_prediction.is_enabled {
+                    editor.set_edit_prediction_provider::<ZedEditPredictionDelegate>(
+                        None, window, cx,
+                    );
+
+                    return;
+                }
+            }
+
             if let Some(project) = editor.project() {
                 ep_store.update(cx, |ep_store, cx| {
                     ep_store.set_edit_prediction_model(model);