From bf403ca7882a668a2753b8244dc6ef1d7db1bcdb Mon Sep 17 00:00:00 2001
From: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
Date: Wed, 7 Aug 2024 12:46:49 +0200
Subject: [PATCH] chore: Remove assistant labels for users without feature flag
 enabled (#15899)

Release Notes:

- N/A
---
 crates/assistant/src/model_selector.rs | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/crates/assistant/src/model_selector.rs b/crates/assistant/src/model_selector.rs
index 341aa512c575f019e067e1ee4a6d26af365c1626..1b936133fd871471fe6f1fce580de6205a646ddb 100644
--- a/crates/assistant/src/model_selector.rs
+++ b/crates/assistant/src/model_selector.rs
@@ -1,4 +1,5 @@
 use feature_flags::LanguageModels;
+use feature_flags::ZedPro;
 use language_model::{LanguageModel, LanguageModelAvailability, LanguageModelRegistry};
 use proto::Plan;
 
@@ -143,8 +144,9 @@ impl PickerDelegate for ModelPickerDelegate {
         selected: bool,
         cx: &mut ViewContext<Picker<Self>>,
     ) -> Option<Self::ListItem> {
+        use feature_flags::FeatureFlagAppExt;
         let model_info = self.filtered_models.get(ix)?;
-
+        let show_badges = cx.has_flag::<ZedPro>();
         Some(
             ListItem::new(ix)
                 .inset(true)
@@ -170,11 +172,13 @@ impl PickerDelegate for ModelPickerDelegate {
                                 .children(match model_info.availability {
                                     LanguageModelAvailability::Public => None,
                                     LanguageModelAvailability::RequiresPlan(Plan::Free) => None,
-                                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro) => Some(
-                                        Label::new("Pro")
-                                            .size(LabelSize::XSmall)
-                                            .color(Color::Muted),
-                                    ),
+                                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro) => {
+                                        show_badges.then(|| {
+                                            Label::new("Pro")
+                                                .size(LabelSize::XSmall)
+                                                .color(Color::Muted)
+                                        })
+                                    }
                                 }),
                         )
                         .child(div().when(model_info.is_selected, |this| {