edit_prediction_ui: Enforce org configuration to disable Zed provider (#53655)

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.

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

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);