Make plans backwards compatible (#37941)

Marshall Bowers created

This PR fixes the backwards compatibility of the new `Plan` variants.

We can't add new variants to the wire representation, as old clients
won't be able to understand them.

Release Notes:

- N/A

Change summary

crates/agent_ui/src/acp/thread_view.rs                         | 16 
crates/agent_ui/src/agent_configuration.rs                     | 12 
crates/agent_ui/src/agent_panel.rs                             | 17 
crates/agent_ui/src/debug.rs                                   |  4 
crates/agent_ui/src/message_editor.rs                          |  6 
crates/agent_ui/src/ui/end_trial_upsell.rs                     |  4 
crates/agent_ui/src/ui/preview/usage_callouts.rs               | 24 
crates/ai_onboarding/src/agent_panel_onboarding_content.rs     | 13 
crates/ai_onboarding/src/ai_onboarding.rs                      | 26 
crates/ai_onboarding/src/ai_upsell_card.rs                     | 33 
crates/ai_onboarding/src/edit_prediction_onboarding_content.rs |  6 
crates/client/src/test.rs                                      |  5 
crates/client/src/user.rs                                      | 15 
crates/cloud_api_types/src/cloud_api_types.rs                  | 11 
crates/cloud_llm_client/src/cloud_llm_client.rs                | 94 ++-
crates/language_model/src/model/cloud_model.rs                 | 14 
crates/language_models/src/provider/cloud.rs                   | 24 
crates/title_bar/src/title_bar.rs                              |  8 
18 files changed, 206 insertions(+), 126 deletions(-)

Detailed changes

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

@@ -14,6 +14,7 @@ use arrayvec::ArrayVec;
 use audio::{Audio, Sound};
 use buffer_diff::BufferDiff;
 use client::zed_urls;
+use cloud_llm_client::PlanV1;
 use collections::{HashMap, HashSet};
 use editor::scroll::Autoscroll;
 use editor::{Editor, EditorEvent, EditorMode, MultiBuffer, PathKey, SelectionEffects};
@@ -4875,7 +4876,9 @@ impl AcpThreadView {
             return None;
         }
 
-        let plan = user_store.plan().unwrap_or(cloud_llm_client::Plan::ZedFree);
+        let plan = user_store
+            .plan()
+            .unwrap_or(cloud_llm_client::Plan::V1(PlanV1::ZedFree));
 
         let usage = user_store.model_request_usage()?;
 
@@ -5134,13 +5137,12 @@ impl AcpThreadView {
         cx: &mut Context<Self>,
     ) -> Callout {
         let error_message = match plan {
-            cloud_llm_client::Plan::ZedPro => "Upgrade to usage-based billing for more prompts.",
-            cloud_llm_client::Plan::ZedProTrial | cloud_llm_client::Plan::ZedFree => {
-                "Upgrade to Zed Pro for more prompts."
+            cloud_llm_client::Plan::V1(PlanV1::ZedPro) => {
+                "Upgrade to usage-based billing for more prompts."
             }
-            cloud_llm_client::Plan::ZedProV2
-            | cloud_llm_client::Plan::ZedProTrialV2
-            | cloud_llm_client::Plan::ZedFreeV2 => "",
+            cloud_llm_client::Plan::V1(PlanV1::ZedProTrial)
+            | cloud_llm_client::Plan::V1(PlanV1::ZedFree) => "Upgrade to Zed Pro for more prompts.",
+            cloud_llm_client::Plan::V2(_) => "",
         };
 
         Callout::new()

crates/agent_ui/src/agent_configuration.rs 🔗

@@ -8,7 +8,7 @@ use std::{ops::Range, sync::Arc};
 use agent_settings::AgentSettings;
 use anyhow::Result;
 use assistant_tool::{ToolSource, ToolWorkingSet};
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1, PlanV2};
 use collections::HashMap;
 use context_server::ContextServerId;
 use editor::{Editor, SelectionEffects, scroll::Autoscroll};
@@ -515,11 +515,15 @@ impl AgentConfiguration {
                 .blend(cx.theme().colors().text_accent.opacity(0.2));
 
             let (plan_name, label_color, bg_color) = match plan {
-                Plan::ZedFree | Plan::ZedFreeV2 => ("Free", Color::Default, free_chip_bg),
-                Plan::ZedProTrial | Plan::ZedProTrialV2 => {
+                Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree) => {
+                    ("Free", Color::Default, free_chip_bg)
+                }
+                Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial) => {
                     ("Pro Trial", Color::Accent, pro_chip_bg)
                 }
-                Plan::ZedPro | Plan::ZedProV2 => ("Pro", Color::Accent, pro_chip_bg),
+                Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro) => {
+                    ("Pro", Color::Accent, pro_chip_bg)
+                }
             };
 
             Chip::new(plan_name.to_string())

crates/agent_ui/src/agent_panel.rs 🔗

@@ -51,7 +51,7 @@ use assistant_context::{AssistantContext, ContextEvent, ContextSummary};
 use assistant_slash_command::SlashCommandWorkingSet;
 use assistant_tool::ToolWorkingSet;
 use client::{UserStore, zed_urls};
-use cloud_llm_client::{CompletionIntent, Plan, UsageLimit};
+use cloud_llm_client::{CompletionIntent, Plan, PlanV1, PlanV2, UsageLimit};
 use editor::{Anchor, AnchorRangeExt as _, Editor, EditorEvent, MultiBuffer};
 use feature_flags::{self, ClaudeCodeFeatureFlag, FeatureFlagAppExt, GeminiAndNativeFeatureFlag};
 use fs::Fs;
@@ -2976,7 +2976,10 @@ impl AgentPanel {
         let plan = self.user_store.read(cx).plan();
         let has_previous_trial = self.user_store.read(cx).trial_started_at().is_some();
 
-        matches!(plan, Some(Plan::ZedFree)) && has_previous_trial
+        matches!(
+            plan,
+            Some(Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree))
+        ) && has_previous_trial
     }
 
     fn should_render_onboarding(&self, cx: &mut Context<Self>) -> bool {
@@ -2988,7 +2991,7 @@ impl AgentPanel {
 
         if user_store
             .plan()
-            .is_some_and(|plan| matches!(plan, Plan::ZedPro))
+            .is_some_and(|plan| matches!(plan, Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro)))
             && user_store
                 .subscription_period()
                 .and_then(|period| period.0.checked_add_days(chrono::Days::new(1)))
@@ -3527,9 +3530,11 @@ impl AgentPanel {
         cx: &mut Context<Self>,
     ) -> AnyElement {
         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::ZedFreeV2 => "",
+            Plan::V1(PlanV1::ZedPro) => "Upgrade to usage-based billing for more prompts.",
+            Plan::V1(PlanV1::ZedProTrial) | Plan::V1(PlanV1::ZedFree) => {
+                "Upgrade to Zed Pro for more prompts."
+            }
+            Plan::V2(_) => "",
         };
 
         Callout::new()

crates/agent_ui/src/debug.rs 🔗

@@ -1,7 +1,7 @@
 #![allow(unused, dead_code)]
 
 use client::{ModelRequestUsage, RequestUsage};
-use cloud_llm_client::{Plan, UsageLimit};
+use cloud_llm_client::{Plan, PlanV1, UsageLimit};
 use gpui::Global;
 use std::ops::{Deref, DerefMut};
 use ui::prelude::*;
@@ -75,7 +75,7 @@ impl Default for DebugAccountState {
         Self {
             enabled: false,
             trial_expired: false,
-            plan: Plan::ZedFree,
+            plan: Plan::V1(PlanV1::ZedFree),
             custom_prompt_usage: ModelRequestUsage(RequestUsage {
                 limit: UsageLimit::Unlimited,
                 amount: 0,

crates/agent_ui/src/message_editor.rs 🔗

@@ -17,7 +17,7 @@ use agent::{
 use agent_settings::{AgentProfileId, AgentSettings, CompletionMode};
 use ai_onboarding::ApiKeysWithProviders;
 use buffer_diff::BufferDiff;
-use cloud_llm_client::CompletionIntent;
+use cloud_llm_client::{CompletionIntent, PlanV1};
 use collections::{HashMap, HashSet};
 use editor::actions::{MoveUp, Paste};
 use editor::display_map::CreaseId;
@@ -1298,7 +1298,9 @@ impl MessageEditor {
             return None;
         }
 
-        let plan = user_store.plan().unwrap_or(cloud_llm_client::Plan::ZedFree);
+        let plan = user_store
+            .plan()
+            .unwrap_or(cloud_llm_client::Plan::V1(PlanV1::ZedFree));
 
         let usage = user_store.model_request_usage()?;
 

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

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use ai_onboarding::{AgentPanelOnboardingCard, PlanDefinitions};
 use client::zed_urls;
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1};
 use gpui::{AnyElement, App, IntoElement, RenderOnce, Window};
 use ui::{Divider, Tooltip, prelude::*};
 
@@ -112,7 +112,7 @@ impl Component for EndTrialUpsell {
         Some(
             v_flex()
                 .child(EndTrialUpsell {
-                    plan: Plan::ZedFree,
+                    plan: Plan::V1(PlanV1::ZedFree),
                     dismiss_upsell: Arc::new(|_, _| {}),
                 })
                 .into_any_element(),

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

@@ -1,5 +1,5 @@
 use client::{ModelRequestUsage, RequestUsage, zed_urls};
-use cloud_llm_client::{Plan, UsageLimit};
+use cloud_llm_client::{Plan, PlanV1, PlanV2, UsageLimit};
 use component::{empty_example, example_group_with_title, single_example};
 use gpui::{AnyElement, App, IntoElement, RenderOnce, Window};
 use ui::{Callout, prelude::*};
@@ -38,20 +38,20 @@ impl RenderOnce for UsageCallout {
 
         let (title, message, button_text, url) = if is_limit_reached {
             match self.plan {
-                Plan::ZedFree | Plan::ZedFreeV2 => (
+                Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree) => (
                     "Out of free prompts",
                     "Upgrade to continue, wait for the next reset, or switch to API key."
                         .to_string(),
                     "Upgrade",
                     zed_urls::account_url(cx),
                 ),
-                Plan::ZedProTrial | Plan::ZedProTrialV2 => (
+                Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial) => (
                     "Out of trial prompts",
                     "Upgrade to Zed Pro to continue, or switch to API key.".to_string(),
                     "Upgrade",
                     zed_urls::account_url(cx),
                 ),
-                Plan::ZedPro | Plan::ZedProV2 => (
+                Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro) => (
                     "Out of included prompts",
                     "Enable usage-based billing to continue.".to_string(),
                     "Manage",
@@ -60,7 +60,7 @@ impl RenderOnce for UsageCallout {
             }
         } else {
             match self.plan {
-                Plan::ZedFree => (
+                Plan::V1(PlanV1::ZedFree) => (
                     "Reaching free plan limit soon",
                     format!(
                         "{remaining} remaining - Upgrade to increase limit, or switch providers",
@@ -68,7 +68,7 @@ impl RenderOnce for UsageCallout {
                     "Upgrade",
                     zed_urls::account_url(cx),
                 ),
-                Plan::ZedProTrial => (
+                Plan::V1(PlanV1::ZedProTrial) => (
                     "Reaching trial limit soon",
                     format!(
                         "{remaining} remaining - Upgrade to increase limit, or switch providers",
@@ -76,7 +76,7 @@ impl RenderOnce for UsageCallout {
                     "Upgrade",
                     zed_urls::account_url(cx),
                 ),
-                _ => return div().into_any_element(),
+                Plan::V1(PlanV1::ZedPro) | Plan::V2(_) => return div().into_any_element(),
             }
         };
 
@@ -119,7 +119,7 @@ impl Component for UsageCallout {
                 single_example(
                     "Approaching limit (90%)",
                     UsageCallout::new(
-                        Plan::ZedFree,
+                        Plan::V1(PlanV1::ZedFree),
                         ModelRequestUsage(RequestUsage {
                             limit: UsageLimit::Limited(50),
                             amount: 45, // 90% of limit
@@ -130,7 +130,7 @@ impl Component for UsageCallout {
                 single_example(
                     "Limit reached (100%)",
                     UsageCallout::new(
-                        Plan::ZedFree,
+                        Plan::V1(PlanV1::ZedFree),
                         ModelRequestUsage(RequestUsage {
                             limit: UsageLimit::Limited(50),
                             amount: 50, // 100% of limit
@@ -147,7 +147,7 @@ impl Component for UsageCallout {
                 single_example(
                     "Approaching limit (90%)",
                     UsageCallout::new(
-                        Plan::ZedProTrial,
+                        Plan::V1(PlanV1::ZedProTrial),
                         ModelRequestUsage(RequestUsage {
                             limit: UsageLimit::Limited(150),
                             amount: 135, // 90% of limit
@@ -158,7 +158,7 @@ impl Component for UsageCallout {
                 single_example(
                     "Limit reached (100%)",
                     UsageCallout::new(
-                        Plan::ZedProTrial,
+                        Plan::V1(PlanV1::ZedProTrial),
                         ModelRequestUsage(RequestUsage {
                             limit: UsageLimit::Limited(150),
                             amount: 150, // 100% of limit
@@ -175,7 +175,7 @@ impl Component for UsageCallout {
                 single_example(
                     "Limit reached (100%)",
                     UsageCallout::new(
-                        Plan::ZedPro,
+                        Plan::V1(PlanV1::ZedPro),
                         ModelRequestUsage(RequestUsage {
                             limit: UsageLimit::Limited(500),
                             amount: 500, // 100% of limit

crates/ai_onboarding/src/agent_panel_onboarding_content.rs 🔗

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use client::{Client, UserStore};
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1, PlanV2};
 use gpui::{Entity, IntoElement, ParentElement};
 use language_model::{LanguageModelRegistry, ZED_CLOUD_PROVIDER_ID};
 use ui::prelude::*;
@@ -57,8 +57,15 @@ impl AgentPanelOnboarding {
 
 impl Render for AgentPanelOnboarding {
     fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
-        let enrolled_in_trial = self.user_store.read(cx).plan() == Some(Plan::ZedProTrial);
-        let is_pro_user = self.user_store.read(cx).plan() == Some(Plan::ZedPro);
+        let enrolled_in_trial = self.user_store.read(cx).plan().is_some_and(|plan| {
+            matches!(
+                plan,
+                Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial)
+            )
+        });
+        let is_pro_user = self.user_store.read(cx).plan().is_some_and(|plan| {
+            matches!(plan, Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro))
+        });
 
         AgentPanelOnboardingCard::new()
             .child(

crates/ai_onboarding/src/ai_onboarding.rs 🔗

@@ -10,7 +10,7 @@ pub use agent_api_keys_onboarding::{ApiKeysWithProviders, ApiKeysWithoutProvider
 pub use agent_panel_onboarding_card::AgentPanelOnboardingCard;
 pub use agent_panel_onboarding_content::AgentPanelOnboarding;
 pub use ai_upsell_card::AiUpsellCard;
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1, PlanV2};
 pub use edit_prediction_onboarding_content::EditPredictionOnboarding;
 pub use plan_definitions::PlanDefinitions;
 pub use young_account_banner::YoungAccountBanner;
@@ -308,13 +308,13 @@ impl RenderOnce for ZedAiOnboarding {
         if matches!(self.sign_in_status, SignInStatus::SignedIn) {
             match self.plan {
                 None => self.render_free_plan_state(cx.has_flag::<BillingV2FeatureFlag>(), cx),
-                Some(plan @ (Plan::ZedFree | Plan::ZedFreeV2)) => {
+                Some(plan @ (Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree))) => {
                     self.render_free_plan_state(plan.is_v2(), cx)
                 }
-                Some(plan @ (Plan::ZedProTrial | Plan::ZedProTrialV2)) => {
+                Some(plan @ (Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial))) => {
                     self.render_trial_state(plan.is_v2(), cx)
                 }
-                Some(plan @ (Plan::ZedPro | Plan::ZedProV2)) => {
+                Some(plan @ (Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro))) => {
                     self.render_pro_plan_state(plan.is_v2(), cx)
                 }
             }
@@ -370,15 +370,27 @@ impl Component for ZedAiOnboarding {
                     ),
                     single_example(
                         "Free Plan",
-                        onboarding(SignInStatus::SignedIn, Some(Plan::ZedFree), false),
+                        onboarding(
+                            SignInStatus::SignedIn,
+                            Some(Plan::V1(PlanV1::ZedFree)),
+                            false,
+                        ),
                     ),
                     single_example(
                         "Pro Trial",
-                        onboarding(SignInStatus::SignedIn, Some(Plan::ZedProTrial), false),
+                        onboarding(
+                            SignInStatus::SignedIn,
+                            Some(Plan::V1(PlanV1::ZedProTrial)),
+                            false,
+                        ),
                     ),
                     single_example(
                         "Pro Plan",
-                        onboarding(SignInStatus::SignedIn, Some(Plan::ZedPro), false),
+                        onboarding(
+                            SignInStatus::SignedIn,
+                            Some(Plan::V1(PlanV1::ZedPro)),
+                            false,
+                        ),
                     ),
                 ])
                 .into_any_element(),

crates/ai_onboarding/src/ai_upsell_card.rs 🔗

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use client::{Client, UserStore, zed_urls};
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1, PlanV2};
 use feature_flags::{BillingV2FeatureFlag, FeatureFlagAppExt};
 use gpui::{AnyElement, App, Entity, IntoElement, RenderOnce, Window};
 use ui::{CommonAnimationExt, Divider, Vector, VectorName, prelude::*};
@@ -171,7 +171,7 @@ impl RenderOnce for AiUpsellCard {
 
         match self.sign_in_status {
             SignInStatus::SignedIn => match self.user_plan {
-                None | Some(Plan::ZedFree | Plan::ZedFreeV2) => card
+                None | Some(Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree)) => card
                     .child(Label::new("Try Zed AI").size(LabelSize::Large))
                     .map(|this| {
                         if self.account_too_young {
@@ -237,16 +237,17 @@ impl RenderOnce for AiUpsellCard {
                             )
                         }
                     }),
-                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(
-                        Label::new("Here's what you get for the next 14 days:")
-                            .color(Color::Muted)
-                            .mb_2(),
-                    )
-                    .child(PlanDefinitions.pro_trial(plan.is_v2(), false)),
-                Some(plan @ (Plan::ZedPro | Plan::ZedProV2)) => card
+                Some(plan @ (Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial))) => {
+                    card.child(pro_trial_stamp)
+                        .child(Label::new("You're in the Zed Pro Trial").size(LabelSize::Large))
+                        .child(
+                            Label::new("Here's what you get for the next 14 days:")
+                                .color(Color::Muted)
+                                .mb_2(),
+                        )
+                        .child(PlanDefinitions.pro_trial(plan.is_v2(), false))
+                }
+                Some(plan @ (Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro))) => card
                     .child(certified_user_stamp)
                     .child(Label::new("You're in the Zed Pro plan").size(LabelSize::Large))
                     .child(
@@ -326,7 +327,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: false,
-                                user_plan: Some(Plan::ZedFree),
+                                user_plan: Some(Plan::V1(PlanV1::ZedFree)),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),
@@ -337,7 +338,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: true,
-                                user_plan: Some(Plan::ZedFree),
+                                user_plan: Some(Plan::V1(PlanV1::ZedFree)),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),
@@ -348,7 +349,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: false,
-                                user_plan: Some(Plan::ZedProTrial),
+                                user_plan: Some(Plan::V1(PlanV1::ZedProTrial)),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),
@@ -359,7 +360,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: false,
-                                user_plan: Some(Plan::ZedPro),
+                                user_plan: Some(Plan::V1(PlanV1::ZedPro)),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),

crates/ai_onboarding/src/edit_prediction_onboarding_content.rs 🔗

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use client::{Client, UserStore};
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1, PlanV2};
 use gpui::{Entity, IntoElement, ParentElement};
 use ui::prelude::*;
 
@@ -36,7 +36,9 @@ impl EditPredictionOnboarding {
 
 impl Render for EditPredictionOnboarding {
     fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
-        let is_free_plan = self.user_store.read(cx).plan() == Some(Plan::ZedFree);
+        let is_free_plan = self.user_store.read(cx).plan().is_some_and(|plan| {
+            matches!(plan, Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree))
+        });
 
         let github_copilot = v_flex()
             .gap_1()

crates/client/src/test.rs 🔗

@@ -1,7 +1,7 @@
 use crate::{Client, Connection, Credentials, EstablishConnectionError, UserStore};
 use anyhow::{Context as _, Result, anyhow};
 use cloud_api_client::{AuthenticatedUser, GetAuthenticatedUserResponse, PlanInfo};
-use cloud_llm_client::{CurrentUsage, Plan, UsageData, UsageLimit};
+use cloud_llm_client::{CurrentUsage, PlanV1, UsageData, UsageLimit};
 use futures::{StreamExt, stream::BoxStream};
 use gpui::{AppContext as _, BackgroundExecutor, Entity, TestAppContext};
 use http_client::{AsyncBody, Method, Request, http};
@@ -269,7 +269,8 @@ pub fn make_get_authenticated_user_response(
         },
         feature_flags: vec![],
         plan: PlanInfo {
-            plan: Plan::ZedPro,
+            plan: PlanV1::ZedPro,
+            plan_v2: None,
             subscription_period: None,
             usage: CurrentUsage {
                 model_requests: UsageData {

crates/client/src/user.rs 🔗

@@ -5,7 +5,8 @@ use cloud_api_client::websocket_protocol::MessageToClient;
 use cloud_api_client::{GetAuthenticatedUserResponse, PlanInfo};
 use cloud_llm_client::{
     EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME,
-    MODEL_REQUESTS_USAGE_AMOUNT_HEADER_NAME, MODEL_REQUESTS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
+    MODEL_REQUESTS_USAGE_AMOUNT_HEADER_NAME, MODEL_REQUESTS_USAGE_LIMIT_HEADER_NAME, Plan,
+    UsageLimit,
 };
 use collections::{HashMap, HashSet, hash_map::Entry};
 use derive_more::Deref;
@@ -692,20 +693,22 @@ impl UserStore {
         self.current_user.borrow().clone()
     }
 
-    pub fn plan(&self) -> Option<cloud_llm_client::Plan> {
+    pub fn plan(&self) -> Option<Plan> {
         #[cfg(debug_assertions)]
         if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {
+            use cloud_llm_client::PlanV1;
+
             return match plan.as_str() {
-                "free" => Some(cloud_llm_client::Plan::ZedFree),
-                "trial" => Some(cloud_llm_client::Plan::ZedProTrial),
-                "pro" => Some(cloud_llm_client::Plan::ZedPro),
+                "free" => Some(Plan::V1(PlanV1::ZedFree)),
+                "trial" => Some(Plan::V1(PlanV1::ZedProTrial)),
+                "pro" => Some(Plan::V1(PlanV1::ZedPro)),
                 _ => {
                     panic!("ZED_SIMULATE_PLAN must be one of 'free', 'trial', or 'pro'");
                 }
             };
         }
 
-        self.plan_info.as_ref().map(|info| info.plan)
+        self.plan_info.as_ref().map(|info| info.plan())
     }
 
     pub fn subscription_period(&self) -> Option<(DateTime<Utc>, DateTime<Utc>)> {

crates/cloud_api_types/src/cloud_api_types.rs 🔗

@@ -1,6 +1,7 @@
 mod timestamp;
 pub mod websocket_protocol;
 
+use cloud_llm_client::Plan;
 use serde::{Deserialize, Serialize};
 
 pub use crate::timestamp::Timestamp;
@@ -27,7 +28,9 @@ pub struct AuthenticatedUser {
 
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub struct PlanInfo {
-    pub plan: cloud_llm_client::Plan,
+    pub plan: cloud_llm_client::PlanV1,
+    #[serde(default)]
+    pub plan_v2: Option<cloud_llm_client::PlanV2>,
     pub subscription_period: Option<SubscriptionPeriod>,
     pub usage: cloud_llm_client::CurrentUsage,
     pub trial_started_at: Option<Timestamp>,
@@ -36,6 +39,12 @@ pub struct PlanInfo {
     pub has_overdue_invoices: bool,
 }
 
+impl PlanInfo {
+    pub fn plan(&self) -> Plan {
+        self.plan_v2.map(Plan::V2).unwrap_or(Plan::V1(self.plan))
+    }
+}
+
 #[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
 pub struct SubscriptionPeriod {
     pub started_at: Timestamp,

crates/cloud_llm_client/src/cloud_llm_client.rs 🔗

@@ -74,38 +74,60 @@ impl FromStr for UsageLimit {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum Plan {
+    V1(PlanV1),
+    V2(PlanV2),
+}
+
+impl Plan {
+    pub fn is_v2(&self) -> bool {
+        matches!(self, Self::V2(_))
+    }
+}
+
 #[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
-pub enum Plan {
+pub enum PlanV1 {
     #[default]
     #[serde(alias = "Free")]
     ZedFree,
-    ZedFreeV2,
     #[serde(alias = "ZedPro")]
     ZedPro,
-    ZedProV2,
     #[serde(alias = "ZedProTrial")]
     ZedProTrial,
-    ZedProTrialV2,
 }
 
-impl Plan {
-    pub fn is_v2(&self) -> bool {
-        matches!(self, Plan::ZedFreeV2 | Plan::ZedProV2 | Plan::ZedProTrialV2)
+impl FromStr for PlanV1 {
+    type Err = anyhow::Error;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value {
+            "zed_free" => Ok(Self::ZedFree),
+            "zed_pro" => Ok(Self::ZedPro),
+            "zed_pro_trial" => Ok(Self::ZedProTrial),
+            plan => Err(anyhow::anyhow!("invalid plan: {plan:?}")),
+        }
     }
 }
 
-impl FromStr for Plan {
+#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum PlanV2 {
+    #[default]
+    ZedFree,
+    ZedPro,
+    ZedProTrial,
+}
+
+impl FromStr for PlanV2 {
     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),
+            "zed_free" => Ok(Self::ZedFree),
+            "zed_pro" => Ok(Self::ZedPro),
+            "zed_pro_trial" => Ok(Self::ZedProTrial),
             plan => Err(anyhow::anyhow!("invalid plan: {plan:?}")),
         }
     }
@@ -306,7 +328,7 @@ pub struct ListModelsResponse {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct GetSubscriptionResponse {
-    pub plan: Plan,
+    pub plan: PlanV1,
     pub usage: Option<CurrentUsage>,
 }
 
@@ -330,33 +352,39 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test_plan_deserialize_snake_case() {
-        let plan = serde_json::from_value::<Plan>(json!("zed_free")).unwrap();
-        assert_eq!(plan, Plan::ZedFree);
+    fn test_plan_v1_deserialize_snake_case() {
+        let plan = serde_json::from_value::<PlanV1>(json!("zed_free")).unwrap();
+        assert_eq!(plan, PlanV1::ZedFree);
 
-        let plan = serde_json::from_value::<Plan>(json!("zed_pro")).unwrap();
-        assert_eq!(plan, Plan::ZedPro);
+        let plan = serde_json::from_value::<PlanV1>(json!("zed_pro")).unwrap();
+        assert_eq!(plan, PlanV1::ZedPro);
 
-        let plan = serde_json::from_value::<Plan>(json!("zed_pro_trial")).unwrap();
-        assert_eq!(plan, Plan::ZedProTrial);
+        let plan = serde_json::from_value::<PlanV1>(json!("zed_pro_trial")).unwrap();
+        assert_eq!(plan, PlanV1::ZedProTrial);
+    }
+
+    #[test]
+    fn test_plan_v1_deserialize_aliases() {
+        let plan = serde_json::from_value::<PlanV1>(json!("Free")).unwrap();
+        assert_eq!(plan, PlanV1::ZedFree);
 
-        let plan = serde_json::from_value::<Plan>(json!("zed_pro_v2")).unwrap();
-        assert_eq!(plan, Plan::ZedProV2);
+        let plan = serde_json::from_value::<PlanV1>(json!("ZedPro")).unwrap();
+        assert_eq!(plan, PlanV1::ZedPro);
 
-        let plan = serde_json::from_value::<Plan>(json!("zed_pro_trial_v2")).unwrap();
-        assert_eq!(plan, Plan::ZedProTrialV2);
+        let plan = serde_json::from_value::<PlanV1>(json!("ZedProTrial")).unwrap();
+        assert_eq!(plan, PlanV1::ZedProTrial);
     }
 
     #[test]
-    fn test_plan_deserialize_aliases() {
-        let plan = serde_json::from_value::<Plan>(json!("Free")).unwrap();
-        assert_eq!(plan, Plan::ZedFree);
+    fn test_plan_v2_deserialize_snake_case() {
+        let plan = serde_json::from_value::<PlanV2>(json!("zed_free")).unwrap();
+        assert_eq!(plan, PlanV2::ZedFree);
 
-        let plan = serde_json::from_value::<Plan>(json!("ZedPro")).unwrap();
-        assert_eq!(plan, Plan::ZedPro);
+        let plan = serde_json::from_value::<PlanV2>(json!("zed_pro")).unwrap();
+        assert_eq!(plan, PlanV2::ZedPro);
 
-        let plan = serde_json::from_value::<Plan>(json!("ZedProTrial")).unwrap();
-        assert_eq!(plan, Plan::ZedProTrial);
+        let plan = serde_json::from_value::<PlanV2>(json!("zed_pro_trial")).unwrap();
+        assert_eq!(plan, PlanV2::ZedProTrial);
     }
 
     #[test]

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

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use anyhow::Result;
 use client::Client;
 use cloud_api_types::websocket_protocol::MessageToClient;
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1};
 use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Global, ReadGlobal as _};
 use smol::lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};
 use thiserror::Error;
@@ -29,16 +29,16 @@ pub struct ModelRequestLimitReachedError {
 impl fmt::Display for ModelRequestLimitReachedError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let message = match self.plan {
-            Plan::ZedFree => "Model request limit reached. Upgrade to Zed Pro for more requests.",
-            Plan::ZedPro => {
+            Plan::V1(PlanV1::ZedFree) => {
+                "Model request limit reached. Upgrade to Zed Pro for more requests."
+            }
+            Plan::V1(PlanV1::ZedPro) => {
                 "Model request limit reached. Upgrade to usage-based billing for more requests."
             }
-            Plan::ZedProTrial => {
+            Plan::V1(PlanV1::ZedProTrial) => {
                 "Model request limit reached. Upgrade to Zed Pro for more requests."
             }
-            Plan::ZedFreeV2 | Plan::ZedProV2 | Plan::ZedProTrialV2 => {
-                "Model request limit reached."
-            }
+            Plan::V2(_) => "Model request limit reached.",
         };
 
         write!(f, "{message}")

crates/language_models/src/provider/cloud.rs 🔗

@@ -7,8 +7,9 @@ use cloud_llm_client::{
     CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CURRENT_PLAN_HEADER_NAME, CompletionBody,
     CompletionEvent, CompletionRequestStatus, CountTokensBody, CountTokensResponse,
     EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, MODEL_REQUESTS_RESOURCE_HEADER_VALUE, Plan,
-    SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME,
-    TOOL_USE_LIMIT_REACHED_HEADER_NAME, ZED_VERSION_HEADER_NAME,
+    PlanV1, PlanV2, SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME,
+    SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME, TOOL_USE_LIMIT_REACHED_HEADER_NAME,
+    ZED_VERSION_HEADER_NAME,
 };
 use futures::{
     AsyncBufReadExt, FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream,
@@ -480,7 +481,8 @@ impl CloudLanguageModel {
                         .headers()
                         .get(CURRENT_PLAN_HEADER_NAME)
                         .and_then(|plan| plan.to_str().ok())
-                        .and_then(|plan| cloud_llm_client::Plan::from_str(plan).ok())
+                        .and_then(|plan| cloud_llm_client::PlanV1::from_str(plan).ok())
+                        .map(Plan::V1)
                 {
                     return Err(anyhow!(ModelRequestLimitReachedError { plan }));
                 }
@@ -994,15 +996,17 @@ impl RenderOnce for ZedAiConfiguration {
     fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
         let young_account_banner = YoungAccountBanner;
 
-        let is_pro = self.plan == Some(Plan::ZedPro);
+        let is_pro = self.plan.is_some_and(|plan| {
+            matches!(plan, Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro))
+        });
         let subscription_text = match (self.plan, self.subscription_period) {
-            (Some(Plan::ZedPro), Some(_)) => {
+            (Some(Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro)), Some(_)) => {
                 "You have access to Zed's hosted models through your Pro subscription."
             }
-            (Some(Plan::ZedProTrial), Some(_)) => {
+            (Some(Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial)), Some(_)) => {
                 "You have access to Zed's hosted models through your Pro trial."
             }
-            (Some(Plan::ZedFree), Some(_)) => {
+            (Some(Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree)), Some(_)) => {
                 "You have basic access to Zed's hosted models through the Free plan."
             }
             _ => {
@@ -1161,15 +1165,15 @@ impl Component for ZedAiConfiguration {
                     ),
                     single_example(
                         "Free Plan",
-                        configuration(true, Some(Plan::ZedFree), true, false),
+                        configuration(true, Some(Plan::V1(PlanV1::ZedFree)), true, false),
                     ),
                     single_example(
                         "Zed Pro Trial Plan",
-                        configuration(true, Some(Plan::ZedProTrial), true, false),
+                        configuration(true, Some(Plan::V1(PlanV1::ZedProTrial)), true, false),
                     ),
                     single_example(
                         "Zed Pro Plan",
-                        configuration(true, Some(Plan::ZedPro), true, false),
+                        configuration(true, Some(Plan::V1(PlanV1::ZedPro)), true, false),
                     ),
                 ])
                 .into_any_element(),

crates/title_bar/src/title_bar.rs 🔗

@@ -23,7 +23,7 @@ use crate::application_menu::{
 use auto_update::AutoUpdateStatus;
 use call::ActiveCall;
 use client::{Client, UserStore, zed_urls};
-use cloud_llm_client::Plan;
+use cloud_llm_client::{Plan, PlanV1, PlanV2};
 use gpui::{
     Action, AnyElement, App, Context, Corner, Element, Entity, Focusable, InteractiveElement,
     IntoElement, MouseButton, ParentElement, Render, StatefulInteractiveElement, Styled,
@@ -668,13 +668,13 @@ impl TitleBar {
                         let user_login = user.github_login.clone();
 
                         let (plan_name, label_color, bg_color) = match plan {
-                            None | Some(Plan::ZedFree | Plan::ZedFreeV2) => {
+                            None | Some(Plan::V1(PlanV1::ZedFree) | Plan::V2(PlanV2::ZedFree)) => {
                                 ("Free", Color::Default, free_chip_bg)
                             }
-                            Some(Plan::ZedProTrial | Plan::ZedProTrialV2) => {
+                            Some(Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial)) => {
                                 ("Pro Trial", Color::Accent, pro_chip_bg)
                             }
-                            Some(Plan::ZedPro | Plan::ZedProV2) => {
+                            Some(Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro)) => {
                                 ("Pro", Color::Accent, pro_chip_bg)
                             }
                         };