cloud_api_types: Rework `Plan` type (#47784)

Marshall Bowers created

This PR reworks the `Plan` type, now that we don't need to be concerned
about the legacy plan versions.

We've also made the deserialization layer more robust, which should make
it easier to add new plan variants in the future without needing to go
through this same song and dance.

Release Notes:

- N/A

Change summary

crates/agent/src/thread.rs                                     |  5 
crates/agent_ui/src/agent_configuration.rs                     |  8 
crates/agent_ui/src/agent_panel.rs                             |  8 
crates/ai_onboarding/src/agent_panel_onboarding_content.rs     |  6 
crates/ai_onboarding/src/ai_onboarding.rs                      | 26 --
crates/ai_onboarding/src/ai_upsell_card.rs                     | 16 
crates/ai_onboarding/src/edit_prediction_onboarding_content.rs |  4 
crates/client/src/test.rs                                      |  6 
crates/client/src/user.rs                                      |  8 
crates/cloud_api_types/src/cloud_api_types.rs                  |  2 
crates/cloud_api_types/src/known_or_unknown.rs                 | 12 +
crates/cloud_api_types/src/plan.rs                             | 35 ++-
crates/language_models/src/provider/cloud.rs                   | 18 -
crates/title_bar/src/title_bar.rs                              | 12 
14 files changed, 81 insertions(+), 85 deletions(-)

Detailed changes

crates/agent/src/thread.rs 🔗

@@ -1746,10 +1746,7 @@ impl Thread {
         };
 
         let auto_retry = if model.provider_id() == ZED_CLOUD_PROVIDER_ID {
-            match plan {
-                Some(Plan::V2(_)) => true,
-                None => false,
-            }
+            plan.is_some()
         } else {
             true
         };

crates/agent_ui/src/agent_configuration.rs 🔗

@@ -9,7 +9,7 @@ use std::{ops::Range, sync::Arc};
 use agent::ContextServerRegistry;
 use anyhow::Result;
 use client::zed_urls;
-use cloud_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use collections::HashMap;
 use context_server::ContextServerId;
 use editor::{Editor, MultiBufferOffset, SelectionEffects, scroll::Autoscroll};
@@ -498,9 +498,9 @@ impl AgentConfiguration {
                 .blend(cx.theme().colors().text_accent.opacity(0.2));
 
             let (plan_name, label_color, bg_color) = match plan {
-                Plan::V2(PlanV2::ZedFree) => ("Free", Color::Default, free_chip_bg),
-                Plan::V2(PlanV2::ZedProTrial) => ("Pro Trial", Color::Accent, pro_chip_bg),
-                Plan::V2(PlanV2::ZedPro) => ("Pro", Color::Accent, pro_chip_bg),
+                Plan::ZedFree => ("Free", Color::Default, free_chip_bg),
+                Plan::ZedProTrial => ("Pro Trial", Color::Accent, pro_chip_bg),
+                Plan::ZedPro => ("Pro", Color::Accent, pro_chip_bg),
             };
 
             Chip::new(plan_name.to_string())

crates/agent_ui/src/agent_panel.rs 🔗

@@ -37,7 +37,7 @@ use anyhow::{Result, anyhow};
 use assistant_slash_command::SlashCommandWorkingSet;
 use assistant_text_thread::{TextThread, TextThreadEvent, TextThreadSummary};
 use client::UserStore;
-use cloud_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use editor::{Anchor, AnchorRangeExt as _, Editor, EditorEvent, MultiBuffer};
 use extension::ExtensionEvents;
 use extension_host::ExtensionStore;
@@ -2391,7 +2391,7 @@ impl AgentPanel {
         let plan = self.user_store.read(cx).plan();
         let has_previous_trial = self.user_store.read(cx).trial_started_at().is_some();
 
-        plan.is_some_and(|plan| plan == Plan::V2(PlanV2::ZedFree)) && has_previous_trial
+        plan.is_some_and(|plan| plan == Plan::ZedFree) && has_previous_trial
     }
 
     fn should_render_onboarding(&self, cx: &mut Context<Self>) -> bool {
@@ -2401,9 +2401,7 @@ impl AgentPanel {
 
         let user_store = self.user_store.read(cx);
 
-        if user_store
-            .plan()
-            .is_some_and(|plan| plan == Plan::V2(PlanV2::ZedPro))
+        if user_store.plan().is_some_and(|plan| plan == Plan::ZedPro)
             && user_store
                 .subscription_period()
                 .and_then(|period| period.0.checked_add_days(chrono::Days::new(1)))

crates/ai_onboarding/src/agent_panel_onboarding_content.rs 🔗

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use client::{Client, UserStore};
-use cloud_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use gpui::{Entity, IntoElement, ParentElement};
 use language_model::{LanguageModelRegistry, ZED_CLOUD_PROVIDER_ID};
 use ui::prelude::*;
@@ -58,12 +58,12 @@ impl Render for AgentPanelOnboarding {
             .user_store
             .read(cx)
             .plan()
-            .is_some_and(|plan| plan == Plan::V2(PlanV2::ZedProTrial));
+            .is_some_and(|plan| plan == Plan::ZedProTrial);
         let is_pro_user = self
             .user_store
             .read(cx)
             .plan()
-            .is_some_and(|plan| plan == Plan::V2(PlanV2::ZedPro));
+            .is_some_and(|plan| plan == Plan::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_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 pub use edit_prediction_onboarding_content::EditPredictionOnboarding;
 pub use plan_definitions::PlanDefinitions;
 pub use young_account_banner::YoungAccountBanner;
@@ -272,9 +272,9 @@ impl RenderOnce for ZedAiOnboarding {
         if matches!(self.sign_in_status, SignInStatus::SignedIn) {
             match self.plan {
                 None => self.render_free_plan_state(cx),
-                Some(Plan::V2(PlanV2::ZedFree)) => self.render_free_plan_state(cx),
-                Some(Plan::V2(PlanV2::ZedProTrial)) => self.render_trial_state(cx),
-                Some(Plan::V2(PlanV2::ZedPro)) => self.render_pro_plan_state(cx),
+                Some(Plan::ZedFree) => self.render_free_plan_state(cx),
+                Some(Plan::ZedProTrial) => self.render_trial_state(cx),
+                Some(Plan::ZedPro) => self.render_pro_plan_state(cx),
             }
         } else {
             self.render_sign_in_disclaimer(cx)
@@ -328,27 +328,15 @@ impl Component for ZedAiOnboarding {
                     ),
                     single_example(
                         "Free Plan",
-                        onboarding(
-                            SignInStatus::SignedIn,
-                            Some(Plan::V2(PlanV2::ZedFree)),
-                            false,
-                        ),
+                        onboarding(SignInStatus::SignedIn, Some(Plan::ZedFree), false),
                     ),
                     single_example(
                         "Pro Trial",
-                        onboarding(
-                            SignInStatus::SignedIn,
-                            Some(Plan::V2(PlanV2::ZedProTrial)),
-                            false,
-                        ),
+                        onboarding(SignInStatus::SignedIn, Some(Plan::ZedProTrial), false),
                     ),
                     single_example(
                         "Pro Plan",
-                        onboarding(
-                            SignInStatus::SignedIn,
-                            Some(Plan::V2(PlanV2::ZedPro)),
-                            false,
-                        ),
+                        onboarding(SignInStatus::SignedIn, Some(Plan::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_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use gpui::{AnyElement, App, Entity, IntoElement, RenderOnce, Window};
 use ui::{CommonAnimationExt, Divider, Vector, VectorName, prelude::*};
 
@@ -166,7 +166,7 @@ impl RenderOnce for AiUpsellCard {
 
         match self.sign_in_status {
             SignInStatus::SignedIn => match self.user_plan {
-                None | Some(Plan::V2(PlanV2::ZedFree)) => card
+                None | Some(Plan::ZedFree) => card
                     .child(Label::new("Try Zed AI").size(LabelSize::Large))
                     .map(|this| {
                         if self.account_too_young {
@@ -232,7 +232,7 @@ impl RenderOnce for AiUpsellCard {
                             )
                         }
                     }),
-                Some(Plan::V2(PlanV2::ZedProTrial)) => card
+                Some(Plan::ZedProTrial) => card
                     .child(pro_trial_stamp)
                     .child(Label::new("You're in the Zed Pro Trial").size(LabelSize::Large))
                     .child(
@@ -241,7 +241,7 @@ impl RenderOnce for AiUpsellCard {
                             .mb_2(),
                     )
                     .child(PlanDefinitions.pro_trial(false)),
-                Some(Plan::V2(PlanV2::ZedPro)) => card
+                Some(Plan::ZedPro) => card
                     .child(certified_user_stamp)
                     .child(Label::new("You're in the Zed Pro plan").size(LabelSize::Large))
                     .child(
@@ -321,7 +321,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: false,
-                                user_plan: Some(Plan::V2(PlanV2::ZedFree)),
+                                user_plan: Some(Plan::ZedFree),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),
@@ -332,7 +332,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: true,
-                                user_plan: Some(Plan::V2(PlanV2::ZedFree)),
+                                user_plan: Some(Plan::ZedFree),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),
@@ -343,7 +343,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: false,
-                                user_plan: Some(Plan::V2(PlanV2::ZedProTrial)),
+                                user_plan: Some(Plan::ZedProTrial),
                                 tab_index: Some(1),
                             }
                             .into_any_element(),
@@ -354,7 +354,7 @@ impl Component for AiUpsellCard {
                                 sign_in_status: SignInStatus::SignedIn,
                                 sign_in: Arc::new(|_, _| {}),
                                 account_too_young: false,
-                                user_plan: Some(Plan::V2(PlanV2::ZedPro)),
+                                user_plan: Some(Plan::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_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use gpui::{Entity, IntoElement, ParentElement};
 use ui::prelude::*;
 
@@ -40,7 +40,7 @@ impl Render for EditPredictionOnboarding {
             .user_store
             .read(cx)
             .plan()
-            .is_some_and(|plan| plan == Plan::V2(PlanV2::ZedFree));
+            .is_some_and(|plan| plan == Plan::ZedFree);
 
         let github_copilot = v_flex()
             .gap_1()

crates/client/src/test.rs 🔗

@@ -1,6 +1,8 @@
 use crate::{Client, Connection, Credentials, EstablishConnectionError, UserStore};
 use anyhow::{Context as _, Result, anyhow};
-use cloud_api_client::{AuthenticatedUser, GetAuthenticatedUserResponse, PlanInfo, PlanV2};
+use cloud_api_client::{
+    AuthenticatedUser, GetAuthenticatedUserResponse, KnownOrUnknown, Plan, PlanInfo,
+};
 use cloud_llm_client::{CurrentUsage, UsageData, UsageLimit};
 use futures::{StreamExt, stream::BoxStream};
 use gpui::{AppContext as _, Entity, TestAppContext};
@@ -264,7 +266,7 @@ pub fn make_get_authenticated_user_response(
         },
         feature_flags: vec![],
         plan: PlanInfo {
-            plan_v2: PlanV2::ZedPro,
+            plan: KnownOrUnknown::Known(Plan::ZedPro),
             subscription_period: None,
             usage: CurrentUsage {
                 edit_predictions: UsageData {

crates/client/src/user.rs 🔗

@@ -668,12 +668,12 @@ impl UserStore {
     pub fn plan(&self) -> Option<Plan> {
         #[cfg(debug_assertions)]
         if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {
-            use cloud_api_client::PlanV2;
+            use cloud_api_client::Plan;
 
             return match plan.as_str() {
-                "free" => Some(Plan::V2(PlanV2::ZedFree)),
-                "trial" => Some(Plan::V2(PlanV2::ZedProTrial)),
-                "pro" => Some(Plan::V2(PlanV2::ZedPro)),
+                "free" => Some(Plan::ZedFree),
+                "trial" => Some(Plan::ZedProTrial),
+                "pro" => Some(Plan::ZedPro),
                 _ => {
                     panic!("ZED_SIMULATE_PLAN must be one of 'free', 'trial', or 'pro'");
                 }

crates/cloud_api_types/src/cloud_api_types.rs 🔗

@@ -1,9 +1,11 @@
+mod known_or_unknown;
 mod plan;
 mod timestamp;
 pub mod websocket_protocol;
 
 use serde::{Deserialize, Serialize};
 
+pub use crate::known_or_unknown::*;
 pub use crate::plan::*;
 pub use crate::timestamp::Timestamp;
 

crates/cloud_api_types/src/known_or_unknown.rs 🔗

@@ -0,0 +1,12 @@
+use serde::{Deserialize, Serialize};
+
+/// `KnownOrUnknown` is a type that represents either a known value ([`Known`](KnownOrUnknown::Known))
+/// or an unknown value ([`Unknown`](KnownOrUnknown::Unknown)).
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum KnownOrUnknown<K, U> {
+    /// A known value.
+    Known(K),
+    /// An unknown value.
+    Unknown(U),
+}

crates/cloud_api_types/src/plan.rs 🔗

@@ -1,15 +1,10 @@
 use serde::{Deserialize, Serialize};
 
-use crate::Timestamp;
-
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum Plan {
-    V2(PlanV2),
-}
+use crate::{KnownOrUnknown, Timestamp};
 
 #[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "snake_case")]
-pub enum PlanV2 {
+pub enum Plan {
     #[default]
     ZedFree,
     ZedPro,
@@ -18,7 +13,9 @@ pub enum PlanV2 {
 
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub struct PlanInfo {
-    pub plan_v2: PlanV2,
+    /// We've named this field `plan_v3` to avoid breaking older clients when we start returning new plan variants.
+    #[serde(rename = "plan_v3")]
+    pub plan: KnownOrUnknown<Plan, String>,
     pub subscription_period: Option<SubscriptionPeriod>,
     pub usage: cloud_llm_client::CurrentUsage,
     pub trial_started_at: Option<Timestamp>,
@@ -28,7 +25,13 @@ pub struct PlanInfo {
 
 impl PlanInfo {
     pub fn plan(&self) -> Plan {
-        Plan::V2(self.plan_v2)
+        match &self.plan {
+            KnownOrUnknown::Known(plan) => *plan,
+            KnownOrUnknown::Unknown(_) => {
+                // If we get a plan that we don't recognize, fall back to the Free plan.
+                Plan::ZedFree
+            }
+        }
     }
 }
 
@@ -46,14 +49,14 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test_plan_v2_deserialize_snake_case() {
-        let plan = serde_json::from_value::<PlanV2>(json!("zed_free")).unwrap();
-        assert_eq!(plan, PlanV2::ZedFree);
+    fn test_plan_deserialize_snake_case() {
+        let plan = serde_json::from_value::<Plan>(json!("zed_free")).unwrap();
+        assert_eq!(plan, Plan::ZedFree);
 
-        let plan = serde_json::from_value::<PlanV2>(json!("zed_pro")).unwrap();
-        assert_eq!(plan, PlanV2::ZedPro);
+        let plan = serde_json::from_value::<Plan>(json!("zed_pro")).unwrap();
+        assert_eq!(plan, Plan::ZedPro);
 
-        let plan = serde_json::from_value::<PlanV2>(json!("zed_pro_trial")).unwrap();
-        assert_eq!(plan, PlanV2::ZedProTrial);
+        let plan = serde_json::from_value::<Plan>(json!("zed_pro_trial")).unwrap();
+        assert_eq!(plan, Plan::ZedProTrial);
     }
 }

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

@@ -3,7 +3,7 @@ use anthropic::AnthropicModelMode;
 use anyhow::{Context as _, Result, anyhow};
 use chrono::{DateTime, Utc};
 use client::{Client, UserStore, zed_urls};
-use cloud_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use cloud_llm_client::{
     CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CLIENT_SUPPORTS_X_AI_HEADER_NAME, CompletionBody,
     CompletionEvent, CountTokensBody, CountTokensResponse, ListModelsResponse,
@@ -984,17 +984,15 @@ struct ZedAiConfiguration {
 
 impl RenderOnce for ZedAiConfiguration {
     fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
-        let is_pro = self
-            .plan
-            .is_some_and(|plan| plan == Plan::V2(PlanV2::ZedPro));
+        let is_pro = self.plan.is_some_and(|plan| plan == Plan::ZedPro);
         let subscription_text = match (self.plan, self.subscription_period) {
-            (Some(Plan::V2(PlanV2::ZedPro)), Some(_)) => {
+            (Some(Plan::ZedPro), Some(_)) => {
                 "You have access to Zed's hosted models through your Pro subscription."
             }
-            (Some(Plan::V2(PlanV2::ZedProTrial)), Some(_)) => {
+            (Some(Plan::ZedProTrial), Some(_)) => {
                 "You have access to Zed's hosted models through your Pro trial."
             }
-            (Some(Plan::V2(PlanV2::ZedFree)), Some(_)) => {
+            (Some(Plan::ZedFree), Some(_)) => {
                 if self.eligible_for_trial {
                     "Subscribe for access to Zed's hosted models. Start with a 14 day free trial."
                 } else {
@@ -1157,15 +1155,15 @@ impl Component for ZedAiConfiguration {
                     ),
                     single_example(
                         "Free Plan",
-                        configuration(true, Some(Plan::V2(PlanV2::ZedFree)), true, false),
+                        configuration(true, Some(Plan::ZedFree), true, false),
                     ),
                     single_example(
                         "Zed Pro Trial Plan",
-                        configuration(true, Some(Plan::V2(PlanV2::ZedProTrial)), true, false),
+                        configuration(true, Some(Plan::ZedProTrial), true, false),
                     ),
                     single_example(
                         "Zed Pro Plan",
-                        configuration(true, Some(Plan::V2(PlanV2::ZedPro)), true, false),
+                        configuration(true, Some(Plan::ZedPro), true, false),
                     ),
                 ])
                 .into_any_element(),

crates/title_bar/src/title_bar.rs 🔗

@@ -21,7 +21,7 @@ use crate::application_menu::{
 use auto_update::AutoUpdateStatus;
 use call::ActiveCall;
 use client::{Client, UserStore, zed_urls};
-use cloud_api_types::{Plan, PlanV2};
+use cloud_api_types::Plan;
 use gpui::{
     Action, AnyElement, App, Context, Corner, Element, Entity, FocusHandle, Focusable,
     InteractiveElement, IntoElement, MouseButton, ParentElement, Render,
@@ -962,13 +962,9 @@ impl TitleBar {
                     let user_login = user_login.clone();
 
                     let (plan_name, label_color, bg_color) = match plan {
-                        None | Some(Plan::V2(PlanV2::ZedFree)) => {
-                            ("Free", Color::Default, free_chip_bg)
-                        }
-                        Some(Plan::V2(PlanV2::ZedProTrial)) => {
-                            ("Pro Trial", Color::Accent, pro_chip_bg)
-                        }
-                        Some(Plan::V2(PlanV2::ZedPro)) => ("Pro", Color::Accent, pro_chip_bg),
+                        None | Some(Plan::ZedFree) => ("Free", Color::Default, free_chip_bg),
+                        Some(Plan::ZedProTrial) => ("Pro Trial", Color::Accent, pro_chip_bg),
+                        Some(Plan::ZedPro) => ("Pro", Color::Accent, pro_chip_bg),
                     };
 
                     menu.when(is_signed_in, |this| {