cloud_llm_client: Add another `Plan` variant (#37852)

Marshall Bowers created

This PR adds a corresponding `FreeV2` variant to the `Plan`.

Release Notes:

- N/A

Change summary

crates/agent_ui/src/acp/thread_view.rs           |  4 ++
crates/agent_ui/src/agent_configuration.rs       |  2 
crates/agent_ui/src/agent_panel.rs               | 25 ++++++++++------
crates/agent_ui/src/ui/end_trial_upsell.rs       | 15 ++++++---
crates/agent_ui/src/ui/preview/usage_callouts.rs |  2 
crates/ai_onboarding/src/ai_onboarding.rs        | 27 +++++++++--------
crates/ai_onboarding/src/ai_upsell_card.rs       | 23 ++++++++-------
crates/cloud_llm_client/src/cloud_llm_client.rs  | 10 ++++++
crates/language_model/src/model/cloud_model.rs   |  4 ++
crates/title_bar/src/title_bar.rs                |  4 ++
10 files changed, 72 insertions(+), 44 deletions(-)

Detailed changes

crates/agent_ui/src/acp/thread_view.rs 🔗

@@ -5010,7 +5010,9 @@ impl AcpThreadView {
             cloud_llm_client::Plan::ZedProTrial | cloud_llm_client::Plan::ZedFree => {
                 "Upgrade to Zed Pro for more prompts."
             }
-            cloud_llm_client::Plan::ZedProV2 | cloud_llm_client::Plan::ZedProTrialV2 => "",
+            cloud_llm_client::Plan::ZedProV2
+            | cloud_llm_client::Plan::ZedProTrialV2
+            | cloud_llm_client::Plan::ZedFreeV2 => "",
         };
 
         Callout::new()

crates/agent_ui/src/agent_configuration.rs 🔗

@@ -515,7 +515,7 @@ impl AgentConfiguration {
                 .blend(cx.theme().colors().text_accent.opacity(0.2));
 
             let (plan_name, label_color, bg_color) = match plan {
-                Plan::ZedFree => ("Free", Color::Default, free_chip_bg),
+                Plan::ZedFree | Plan::ZedFreeV2 => ("Free", Color::Default, free_chip_bg),
                 Plan::ZedProTrial | Plan::ZedProTrialV2 => {
                     ("Pro Trial", Color::Accent, pro_chip_bg)
                 }

crates/agent_ui/src/agent_panel.rs 🔗

@@ -3066,6 +3066,8 @@ impl AgentPanel {
             return None;
         }
 
+        let plan = self.user_store.read(cx).plan()?;
+
         Some(
             v_flex()
                 .absolute()
@@ -3074,15 +3076,18 @@ impl AgentPanel {
                 .bg(cx.theme().colors().panel_background)
                 .opacity(0.85)
                 .block_mouse_except_scroll()
-                .child(EndTrialUpsell::new(Arc::new({
-                    let this = cx.entity();
-                    move |_, cx| {
-                        this.update(cx, |_this, cx| {
-                            TrialEndUpsell::set_dismissed(true, cx);
-                            cx.notify();
-                        });
-                    }
-                }))),
+                .child(EndTrialUpsell::new(
+                    plan,
+                    Arc::new({
+                        let this = cx.entity();
+                        move |_, cx| {
+                            this.update(cx, |_this, cx| {
+                                TrialEndUpsell::set_dismissed(true, cx);
+                                cx.notify();
+                            });
+                        }
+                    }),
+                )),
         )
     }
 
@@ -3518,7 +3523,7 @@ impl AgentPanel {
         let error_message = match plan {
             Plan::ZedPro => "Upgrade to usage-based billing for more prompts.",
             Plan::ZedProTrial | Plan::ZedFree => "Upgrade to Zed Pro for more prompts.",
-            Plan::ZedProV2 | Plan::ZedProTrialV2 => "",
+            Plan::ZedProV2 | Plan::ZedProTrialV2 | Plan::ZedFreeV2 => "",
         };
 
         Callout::new()

crates/agent_ui/src/ui/end_trial_upsell.rs 🔗

@@ -2,18 +2,22 @@ use std::sync::Arc;
 
 use ai_onboarding::{AgentPanelOnboardingCard, PlanDefinitions};
 use client::zed_urls;
-use feature_flags::{BillingV2FeatureFlag, FeatureFlagAppExt as _};
+use cloud_llm_client::Plan;
 use gpui::{AnyElement, App, IntoElement, RenderOnce, Window};
 use ui::{Divider, Tooltip, prelude::*};
 
 #[derive(IntoElement, RegisterComponent)]
 pub struct EndTrialUpsell {
+    plan: Plan,
     dismiss_upsell: Arc<dyn Fn(&mut Window, &mut App)>,
 }
 
 impl EndTrialUpsell {
-    pub fn new(dismiss_upsell: Arc<dyn Fn(&mut Window, &mut App)>) -> Self {
-        Self { dismiss_upsell }
+    pub fn new(plan: Plan, dismiss_upsell: Arc<dyn Fn(&mut Window, &mut App)>) -> Self {
+        Self {
+            plan,
+            dismiss_upsell,
+        }
     }
 }
 
@@ -32,7 +36,7 @@ impl RenderOnce for EndTrialUpsell {
                     )
                     .child(Divider::horizontal()),
             )
-            .child(PlanDefinitions.pro_plan(cx.has_flag::<BillingV2FeatureFlag>(), false))
+            .child(PlanDefinitions.pro_plan(self.plan.is_v2(), false))
             .child(
                 Button::new("cta-button", "Upgrade to Zed Pro")
                     .full_width()
@@ -63,7 +67,7 @@ impl RenderOnce for EndTrialUpsell {
                     )
                     .child(Divider::horizontal()),
             )
-            .child(PlanDefinitions.free_plan(cx.has_flag::<BillingV2FeatureFlag>()));
+            .child(PlanDefinitions.free_plan(self.plan.is_v2()));
 
         AgentPanelOnboardingCard::new()
             .child(Headline::new("Your Zed Pro Trial has expired"))
@@ -108,6 +112,7 @@ impl Component for EndTrialUpsell {
         Some(
             v_flex()
                 .child(EndTrialUpsell {
+                    plan: Plan::ZedFree,
                     dismiss_upsell: Arc::new(|_, _| {}),
                 })
                 .into_any_element(),

crates/agent_ui/src/ui/preview/usage_callouts.rs 🔗

@@ -38,7 +38,7 @@ impl RenderOnce for UsageCallout {
 
         let (title, message, button_text, url) = if is_limit_reached {
             match self.plan {
-                Plan::ZedFree => (
+                Plan::ZedFree | Plan::ZedFreeV2 => (
                     "Out of free prompts",
                     "Upgrade to continue, wait for the next reset, or switch to API key."
                         .to_string(),

crates/ai_onboarding/src/ai_onboarding.rs 🔗

@@ -113,7 +113,7 @@ impl ZedAiOnboarding {
             .into_any_element()
     }
 
-    fn render_free_plan_state(&self, cx: &mut App) -> AnyElement {
+    fn render_free_plan_state(&self, is_v2: bool, cx: &mut App) -> AnyElement {
         if self.account_too_young {
             v_flex()
                 .relative()
@@ -136,9 +136,7 @@ impl ZedAiOnboarding {
                                 )
                                 .child(Divider::horizontal()),
                         )
-                        .child(
-                            PlanDefinitions.pro_plan(cx.has_flag::<BillingV2FeatureFlag>(), true),
-                        )
+                        .child(PlanDefinitions.pro_plan(is_v2, true))
                         .child(
                             Button::new("pro", "Get Started")
                                 .full_width()
@@ -181,7 +179,7 @@ impl ZedAiOnboarding {
                                 )
                                 .child(Divider::horizontal()),
                         )
-                        .child(PlanDefinitions.free_plan(cx.has_flag::<BillingV2FeatureFlag>())),
+                        .child(PlanDefinitions.free_plan(is_v2)),
                 )
                 .when_some(
                     self.dismiss_onboarding.as_ref(),
@@ -219,9 +217,7 @@ impl ZedAiOnboarding {
                                 )
                                 .child(Divider::horizontal()),
                         )
-                        .child(
-                            PlanDefinitions.pro_trial(cx.has_flag::<BillingV2FeatureFlag>(), true),
-                        )
+                        .child(PlanDefinitions.pro_trial(is_v2, true))
                         .child(
                             Button::new("pro", "Start Free Trial")
                                 .full_width()
@@ -311,11 +307,16 @@ impl RenderOnce for ZedAiOnboarding {
     fn render(self, _window: &mut ui::Window, cx: &mut App) -> impl IntoElement {
         if matches!(self.sign_in_status, SignInStatus::SignedIn) {
             match self.plan {
-                None | Some(Plan::ZedFree) => self.render_free_plan_state(cx),
-                Some(Plan::ZedProTrial) => self.render_trial_state(false, cx),
-                Some(Plan::ZedProTrialV2) => self.render_trial_state(true, cx),
-                Some(Plan::ZedPro) => self.render_pro_plan_state(false, cx),
-                Some(Plan::ZedProV2) => self.render_pro_plan_state(true, cx),
+                None => self.render_free_plan_state(cx.has_flag::<BillingV2FeatureFlag>(), cx),
+                Some(plan @ (Plan::ZedFree | Plan::ZedFreeV2)) => {
+                    self.render_free_plan_state(plan.is_v2(), cx)
+                }
+                Some(plan @ (Plan::ZedProTrial | Plan::ZedProTrialV2)) => {
+                    self.render_trial_state(plan.is_v2(), cx)
+                }
+                Some(plan @ (Plan::ZedPro | Plan::ZedProV2)) => {
+                    self.render_pro_plan_state(plan.is_v2(), cx)
+                }
             }
         } else {
             self.render_sign_in_disclaimer(cx)

crates/ai_onboarding/src/ai_upsell_card.rs 🔗

@@ -50,6 +50,10 @@ impl AiUpsellCard {
 
 impl RenderOnce for AiUpsellCard {
     fn render(self, _window: &mut Window, cx: &mut App) -> impl IntoElement {
+        let is_v2_plan = self
+            .user_plan
+            .map_or(cx.has_flag::<BillingV2FeatureFlag>(), |plan| plan.is_v2());
+
         let pro_section = v_flex()
             .flex_grow()
             .w_full()
@@ -65,7 +69,7 @@ impl RenderOnce for AiUpsellCard {
                     )
                     .child(Divider::horizontal()),
             )
-            .child(PlanDefinitions.pro_plan(cx.has_flag::<BillingV2FeatureFlag>(), false));
+            .child(PlanDefinitions.pro_plan(is_v2_plan, false));
 
         let free_section = v_flex()
             .flex_grow()
@@ -82,7 +86,7 @@ impl RenderOnce for AiUpsellCard {
                     )
                     .child(Divider::horizontal()),
             )
-            .child(PlanDefinitions.free_plan(cx.has_flag::<BillingV2FeatureFlag>()));
+            .child(PlanDefinitions.free_plan(is_v2_plan));
 
         let grid_bg = h_flex()
             .absolute()
@@ -167,7 +171,7 @@ impl RenderOnce for AiUpsellCard {
 
         match self.sign_in_status {
             SignInStatus::SignedIn => match self.user_plan {
-                None | Some(Plan::ZedFree) => card
+                None | Some(Plan::ZedFree | Plan::ZedFreeV2) => card
                     .child(Label::new("Try Zed AI").size(LabelSize::Large))
                     .map(|this| {
                         if self.account_too_young {
@@ -186,10 +190,7 @@ impl RenderOnce for AiUpsellCard {
                                             )
                                             .child(Divider::horizontal()),
                                     )
-                                    .child(
-                                        PlanDefinitions
-                                            .pro_plan(cx.has_flag::<BillingV2FeatureFlag>(), true),
-                                    )
+                                    .child(PlanDefinitions.pro_plan(is_v2_plan, true))
                                     .child(
                                         Button::new("pro", "Get Started")
                                             .full_width()
@@ -236,7 +237,7 @@ impl RenderOnce for AiUpsellCard {
                             )
                         }
                     }),
-                Some(plan @ Plan::ZedProTrial | plan @ Plan::ZedProTrialV2) => card
+                Some(plan @ (Plan::ZedProTrial | Plan::ZedProTrialV2)) => card
                     .child(pro_trial_stamp)
                     .child(Label::new("You're in the Zed Pro Trial").size(LabelSize::Large))
                     .child(
@@ -244,8 +245,8 @@ impl RenderOnce for AiUpsellCard {
                             .color(Color::Muted)
                             .mb_2(),
                     )
-                    .child(PlanDefinitions.pro_trial(plan == Plan::ZedProTrialV2, false)),
-                Some(plan @ Plan::ZedPro | plan @ Plan::ZedProV2) => card
+                    .child(PlanDefinitions.pro_trial(plan.is_v2(), false)),
+                Some(plan @ (Plan::ZedPro | Plan::ZedProV2)) => card
                     .child(certified_user_stamp)
                     .child(Label::new("You're in the Zed Pro plan").size(LabelSize::Large))
                     .child(
@@ -253,7 +254,7 @@ impl RenderOnce for AiUpsellCard {
                             .color(Color::Muted)
                             .mb_2(),
                     )
-                    .child(PlanDefinitions.pro_plan(plan == Plan::ZedProV2, false)),
+                    .child(PlanDefinitions.pro_plan(plan.is_v2(), false)),
             },
             // Signed Out State
             _ => card

crates/cloud_llm_client/src/cloud_llm_client.rs 🔗

@@ -80,6 +80,7 @@ pub enum Plan {
     #[default]
     #[serde(alias = "Free")]
     ZedFree,
+    ZedFreeV2,
     #[serde(alias = "ZedPro")]
     ZedPro,
     ZedProV2,
@@ -88,14 +89,23 @@ pub enum Plan {
     ZedProTrialV2,
 }
 
+impl Plan {
+    pub fn is_v2(&self) -> bool {
+        matches!(self, Plan::ZedFreeV2 | Plan::ZedProV2 | Plan::ZedProTrialV2)
+    }
+}
+
 impl FromStr for Plan {
     type Err = anyhow::Error;
 
     fn from_str(value: &str) -> Result<Self, Self::Err> {
         match value {
             "zed_free" => Ok(Plan::ZedFree),
+            "zed_free_v2" => Ok(Plan::ZedFreeV2),
             "zed_pro" => Ok(Plan::ZedPro),
+            "zed_pro_v2" => Ok(Plan::ZedProV2),
             "zed_pro_trial" => Ok(Plan::ZedProTrial),
+            "zed_pro_trial_v2" => Ok(Plan::ZedProTrialV2),
             plan => Err(anyhow::anyhow!("invalid plan: {plan:?}")),
         }
     }

crates/language_model/src/model/cloud_model.rs 🔗

@@ -36,7 +36,9 @@ impl fmt::Display for ModelRequestLimitReachedError {
             Plan::ZedProTrial => {
                 "Model request limit reached. Upgrade to Zed Pro for more requests."
             }
-            Plan::ZedProV2 | Plan::ZedProTrialV2 => "Model request limit reached.",
+            Plan::ZedFreeV2 | Plan::ZedProV2 | Plan::ZedProTrialV2 => {
+                "Model request limit reached."
+            }
         };
 
         write!(f, "{message}")

crates/title_bar/src/title_bar.rs 🔗

@@ -659,7 +659,9 @@ impl TitleBar {
                         let user_login = user.github_login.clone();
 
                         let (plan_name, label_color, bg_color) = match plan {
-                            None | Some(Plan::ZedFree) => ("Free", Color::Default, free_chip_bg),
+                            None | Some(Plan::ZedFree | Plan::ZedFreeV2) => {
+                                ("Free", Color::Default, free_chip_bg)
+                            }
                             Some(Plan::ZedProTrial | Plan::ZedProTrialV2) => {
                                 ("Pro Trial", Color::Accent, pro_chip_bg)
                             }