tenant.openagents/omega

1use serde::{Deserialize, Serialize};
2
3use crate::{KnownOrUnknown, Timestamp};
4
5#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
6#[serde(rename_all = "snake_case")]
7pub enum Plan {
8    #[default]
9    ZedFree,
10    ZedPro,
11    ZedProTrial,
12    ZedBusiness,
13    ZedVip,
14    ZedStudent,
15}
16
17#[derive(Debug, PartialEq, Serialize, Deserialize)]
18pub struct PlanInfo {
19    /// We've named this field `plan_v3` to avoid breaking older clients when we start returning new plan variants.
20    #[serde(rename = "plan_v3")]
21    pub plan: KnownOrUnknown<Plan, String>,
22    pub subscription_period: Option<SubscriptionPeriod>,
23    pub usage: cloud_llm_client::CurrentUsage,
24    pub trial_started_at: Option<Timestamp>,
25    pub is_account_too_young: bool,
26    pub has_overdue_invoices: bool,
27}
28
29impl PlanInfo {
30    pub fn plan(&self) -> Plan {
31        match &self.plan {
32            KnownOrUnknown::Known(plan) => *plan,
33            KnownOrUnknown::Unknown(_) => {
34                // If we get a plan that we don't recognize, fall back to the Free plan.
35                Plan::ZedFree
36            }
37        }
38    }
39}
40
41#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
42pub struct SubscriptionPeriod {
43    pub started_at: Timestamp,
44    pub ended_at: Timestamp,
45}
46
47#[cfg(test)]
48mod tests {
49    use pretty_assertions::assert_eq;
50    use serde_json::json;
51
52    use super::*;
53
54    #[test]
55    fn test_plan_deserialize_snake_case() {
56        let plan = serde_json::from_value::<Plan>(json!("zed_free")).unwrap();
57        assert_eq!(plan, Plan::ZedFree);
58
59        let plan = serde_json::from_value::<Plan>(json!("zed_pro")).unwrap();
60        assert_eq!(plan, Plan::ZedPro);
61
62        let plan = serde_json::from_value::<Plan>(json!("zed_pro_trial")).unwrap();
63        assert_eq!(plan, Plan::ZedProTrial);
64
65        let plan = serde_json::from_value::<Plan>(json!("zed_student")).unwrap();
66        assert_eq!(plan, Plan::ZedStudent);
67    }
68}
69