codex_protocol/
account.rs

1use schemars::JsonSchema;
2use serde::Deserialize;
3use serde::Serialize;
4use ts_rs::TS;
5
6use crate::auth::KnownPlan;
7use crate::auth::PlanType as AuthPlanType;
8
9#[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)]
10#[serde(rename_all = "lowercase")]
11#[ts(rename_all = "lowercase")]
12pub enum PlanType {
13    #[default]
14    Free,
15    Go,
16    Plus,
17    Pro,
18    ProLite,
19    Team,
20    #[serde(rename = "self_serve_business_usage_based")]
21    #[ts(rename = "self_serve_business_usage_based")]
22    SelfServeBusinessUsageBased,
23    Business,
24    #[serde(rename = "enterprise_cbp_usage_based")]
25    #[ts(rename = "enterprise_cbp_usage_based")]
26    EnterpriseCbpUsageBased,
27    Enterprise,
28    Edu,
29    #[serde(other)]
30    Unknown,
31}
32
33/// Account state returned by a model provider before it is adapted to an app-facing wire type.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub enum ProviderAccount {
36    ApiKey,
37    Chatgpt {
38        email: Option<String>,
39        plan_type: PlanType,
40    },
41    AmazonBedrock {
42        uses_codex_managed_credentials: bool,
43    },
44}
45
46impl PlanType {
47    pub fn is_team_like(self) -> bool {
48        matches!(self, Self::Team | Self::SelfServeBusinessUsageBased)
49    }
50
51    pub fn is_business_like(self) -> bool {
52        matches!(self, Self::Business | Self::EnterpriseCbpUsageBased)
53    }
54
55    pub fn is_workspace_account(self) -> bool {
56        matches!(
57            self,
58            Self::Team
59                | Self::SelfServeBusinessUsageBased
60                | Self::Business
61                | Self::EnterpriseCbpUsageBased
62                | Self::Enterprise
63                | Self::Edu
64        )
65    }
66}
67
68impl From<AuthPlanType> for PlanType {
69    fn from(plan_type: AuthPlanType) -> Self {
70        match plan_type {
71            AuthPlanType::Known(plan) => plan.into(),
72            AuthPlanType::Unknown(_) => Self::Unknown,
73        }
74    }
75}
76
77impl From<KnownPlan> for PlanType {
78    fn from(plan: KnownPlan) -> Self {
79        match plan {
80            KnownPlan::Free => Self::Free,
81            KnownPlan::Go => Self::Go,
82            KnownPlan::Plus => Self::Plus,
83            KnownPlan::Pro => Self::Pro,
84            KnownPlan::ProLite => Self::ProLite,
85            KnownPlan::Team => Self::Team,
86            KnownPlan::SelfServeBusinessUsageBased => Self::SelfServeBusinessUsageBased,
87            KnownPlan::Business => Self::Business,
88            KnownPlan::EnterpriseCbpUsageBased => Self::EnterpriseCbpUsageBased,
89            KnownPlan::Enterprise => Self::Enterprise,
90            KnownPlan::Edu => Self::Edu,
91        }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::PlanType;
98    use crate::auth::KnownPlan;
99    use crate::auth::PlanType as AuthPlanType;
100    use pretty_assertions::assert_eq;
101
102    #[test]
103    fn usage_based_plan_types_use_expected_wire_names() {
104        assert_eq!(
105            serde_json::to_string(&PlanType::SelfServeBusinessUsageBased)
106                .expect("self-serve business usage based should serialize"),
107            "\"self_serve_business_usage_based\""
108        );
109        assert_eq!(
110            serde_json::to_string(&PlanType::EnterpriseCbpUsageBased)
111                .expect("enterprise cbp usage based should serialize"),
112            "\"enterprise_cbp_usage_based\""
113        );
114        assert_eq!(
115            serde_json::to_string(&PlanType::ProLite).expect("prolite should serialize"),
116            "\"prolite\""
117        );
118        assert_eq!(
119            serde_json::from_str::<PlanType>("\"self_serve_business_usage_based\"")
120                .expect("self-serve business usage based should deserialize"),
121            PlanType::SelfServeBusinessUsageBased
122        );
123        assert_eq!(
124            serde_json::from_str::<PlanType>("\"prolite\"").expect("prolite should deserialize"),
125            PlanType::ProLite
126        );
127        assert_eq!(
128            serde_json::from_str::<PlanType>("\"enterprise_cbp_usage_based\"")
129                .expect("enterprise cbp usage based should deserialize"),
130            PlanType::EnterpriseCbpUsageBased
131        );
132    }
133
134    #[test]
135    fn plan_family_helpers_group_usage_based_variants_with_existing_plans() {
136        assert_eq!(PlanType::Team.is_team_like(), true);
137        assert_eq!(PlanType::SelfServeBusinessUsageBased.is_team_like(), true);
138        assert_eq!(PlanType::Business.is_team_like(), false);
139
140        assert_eq!(PlanType::Business.is_business_like(), true);
141        assert_eq!(PlanType::EnterpriseCbpUsageBased.is_business_like(), true);
142        assert_eq!(PlanType::Team.is_business_like(), false);
143    }
144
145    #[test]
146    fn workspace_account_helper_includes_usage_based_workspace_plans() {
147        assert_eq!(PlanType::Team.is_workspace_account(), true);
148        assert_eq!(
149            PlanType::SelfServeBusinessUsageBased.is_workspace_account(),
150            true
151        );
152        assert_eq!(PlanType::Business.is_workspace_account(), true);
153        assert_eq!(
154            PlanType::EnterpriseCbpUsageBased.is_workspace_account(),
155            true
156        );
157        assert_eq!(PlanType::Enterprise.is_workspace_account(), true);
158        assert_eq!(PlanType::Edu.is_workspace_account(), true);
159        assert_eq!(PlanType::Pro.is_workspace_account(), false);
160    }
161
162    #[test]
163    fn auth_plan_type_converts_to_account_plan_type() {
164        assert_eq!(
165            PlanType::from(AuthPlanType::Known(KnownPlan::EnterpriseCbpUsageBased)),
166            PlanType::EnterpriseCbpUsageBased
167        );
168        assert_eq!(
169            PlanType::from(AuthPlanType::Known(KnownPlan::Enterprise)),
170            PlanType::Enterprise
171        );
172        assert_eq!(
173            PlanType::from(AuthPlanType::Unknown("mystery-tier".to_string())),
174            PlanType::Unknown
175        );
176    }
177}
codex_protocol/account.rs

codex_protocol/
account.rs
