use std::fmt;
use base64::Engine;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct RedactedString(String);
impl RedactedString {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn expose_secret(&self) -> &str {
&self.0
}
pub fn into_inner(self) -> String {
self.0
}
}
impl fmt::Debug for RedactedString {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("<redacted>")
}
}
impl From<String> for RedactedString {
fn from(value: String) -> Self {
Self::new(value)
}
}
impl From<&str> for RedactedString {
fn from(value: &str) -> Self {
Self::new(value)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccountsStore {
pub version: u32,
pub accounts: Vec<StoredAccount>,
#[serde(default)]
pub masked_account_ids: Vec<String>,
}
impl Default for AccountsStore {
fn default() -> Self {
Self {
version: 1,
accounts: Vec::new(),
masked_account_ids: Vec::new(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredAccount {
pub id: String,
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
pub chatgpt_user_id: Option<String>,
pub chatgpt_account_is_fedramp: bool,
pub token_last_refresh_at: Option<DateTime<Utc>>,
#[serde(default)]
pub subscription_expires_at: Option<DateTime<Utc>>,
pub auth_mode: AuthMode,
pub auth_data: AuthData,
pub created_at: DateTime<Utc>,
pub last_used_at: Option<DateTime<Utc>>,
}
impl StoredAccount {
pub fn new_api_key(name: String, api_key: String) -> Self {
Self {
id: Uuid::new_v4().to_string(),
name,
email: None,
plan_type: None,
chatgpt_user_id: None,
chatgpt_account_is_fedramp: false,
token_last_refresh_at: None,
subscription_expires_at: None,
auth_mode: AuthMode::ApiKey,
auth_data: AuthData::ApiKey {
key: RedactedString::new(api_key),
},
created_at: Utc::now(),
last_used_at: None,
}
}
pub fn new_chatgpt(account: NewChatGptAccount) -> Self {
Self {
id: Uuid::new_v4().to_string(),
name: account.name,
email: account.email,
plan_type: account.plan_type,
chatgpt_user_id: account.chatgpt_user_id,
chatgpt_account_is_fedramp: account.chatgpt_account_is_fedramp,
token_last_refresh_at: Some(account.token_last_refresh_at),
subscription_expires_at: account.subscription_expires_at,
auth_mode: AuthMode::ChatGPT,
auth_data: AuthData::ChatGPT {
id_token: account.id_token,
access_token: account.access_token,
refresh_token: account.refresh_token,
account_id: account.account_id,
},
created_at: Utc::now(),
last_used_at: None,
}
}
}
#[derive(Debug, Clone)]
pub struct NewChatGptAccount {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
pub chatgpt_user_id: Option<String>,
pub chatgpt_account_is_fedramp: bool,
pub token_last_refresh_at: DateTime<Utc>,
pub subscription_expires_at: Option<DateTime<Utc>>,
pub id_token: RedactedString,
pub access_token: RedactedString,
pub refresh_token: RedactedString,
pub account_id: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuthMode {
ApiKey,
ChatGPT,
}
impl fmt::Display for AuthMode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ApiKey => f.write_str("api_key"),
Self::ChatGPT => f.write_str("chatgpt"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AuthData {
ApiKey {
key: RedactedString,
},
ChatGPT {
id_token: RedactedString,
access_token: RedactedString,
refresh_token: RedactedString,
account_id: Option<String>,
},
}
#[derive(Debug, Clone, Default)]
pub struct ChatGptIdTokenClaims {
pub email: Option<String>,
pub plan_type: Option<String>,
pub user_id: Option<String>,
pub account_id: Option<String>,
pub account_is_fedramp: bool,
pub subscription_expires_at: Option<DateTime<Utc>>,
}
pub fn parse_chatgpt_id_token_claims(id_token: &str) -> ChatGptIdTokenClaims {
let parts: Vec<&str> = id_token.split('.').collect();
if parts.len() != 3 {
return ChatGptIdTokenClaims::default();
}
let payload = match base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(parts[1]) {
Ok(bytes) => bytes,
Err(_) => return ChatGptIdTokenClaims::default(),
};
let json: serde_json::Value = match serde_json::from_slice(&payload) {
Ok(value) => value,
Err(_) => return ChatGptIdTokenClaims::default(),
};
let profile_claims = json.get("https://api.openai.com/profile");
let auth_claims = json.get("https://api.openai.com/auth");
ChatGptIdTokenClaims {
email: json
.get("email")
.and_then(|v| v.as_str())
.or_else(|| {
profile_claims
.and_then(|profile| profile.get("email"))
.and_then(|v| v.as_str())
})
.map(String::from),
plan_type: auth_claims
.and_then(|auth| auth.get("chatgpt_plan_type"))
.and_then(|v| v.as_str())
.map(String::from),
user_id: auth_claims
.and_then(|auth| auth.get("chatgpt_user_id").or_else(|| auth.get("user_id")))
.and_then(|v| v.as_str())
.map(String::from),
account_id: auth_claims
.and_then(|auth| auth.get("chatgpt_account_id"))
.and_then(|v| v.as_str())
.map(String::from),
account_is_fedramp: auth_claims
.and_then(|auth| auth.get("chatgpt_account_is_fedramp"))
.and_then(|v| v.as_bool())
.unwrap_or(false),
subscription_expires_at: auth_claims
.and_then(|auth| auth.get("chatgpt_subscription_active_until"))
.and_then(|v| v.as_str())
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
.map(|value| value.with_timezone(&Utc)),
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthDotJson {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub auth_mode: Option<String>,
#[serde(rename = "OPENAI_API_KEY", skip_serializing_if = "Option::is_none")]
pub openai_api_key: Option<RedactedString>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tokens: Option<TokenData>,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_refresh: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenData {
pub id_token: RedactedString,
pub access_token: RedactedString,
pub refresh_token: RedactedString,
#[serde(skip_serializing_if = "Option::is_none")]
pub account_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageInfo {
pub account_id: String,
pub limit_id: Option<String>,
pub limit_name: Option<String>,
pub plan_type: Option<String>,
pub primary_used_percent: Option<f64>,
pub primary_window_minutes: Option<i64>,
pub primary_resets_at: Option<i64>,
pub secondary_used_percent: Option<f64>,
pub secondary_window_minutes: Option<i64>,
pub secondary_resets_at: Option<i64>,
pub has_credits: Option<bool>,
pub unlimited_credits: Option<bool>,
pub credits_balance: Option<String>,
pub rate_limit_reached_type: Option<String>,
#[serde(default)]
pub additional_limits: Vec<UsageLimitInfo>,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageLimitInfo {
pub limit_id: Option<String>,
pub limit_name: Option<String>,
pub primary_used_percent: Option<f64>,
pub primary_window_minutes: Option<i64>,
pub primary_resets_at: Option<i64>,
pub secondary_used_percent: Option<f64>,
pub secondary_window_minutes: Option<i64>,
pub secondary_resets_at: Option<i64>,
}
impl UsageInfo {
pub fn error(account_id: String, error: String) -> Self {
Self {
account_id,
limit_id: None,
limit_name: None,
plan_type: None,
primary_used_percent: None,
primary_window_minutes: None,
primary_resets_at: None,
secondary_used_percent: None,
secondary_window_minutes: None,
secondary_resets_at: None,
has_credits: None,
unlimited_credits: None,
credits_balance: None,
rate_limit_reached_type: None,
additional_limits: Vec::new(),
error: Some(error),
}
}
pub fn unsupported(account_id: String) -> Self {
Self {
account_id,
limit_id: None,
limit_name: None,
plan_type: Some("api_key".to_string()),
primary_used_percent: None,
primary_window_minutes: None,
primary_resets_at: None,
secondary_used_percent: None,
secondary_window_minutes: None,
secondary_resets_at: None,
has_credits: None,
unlimited_credits: None,
credits_balance: None,
rate_limit_reached_type: None,
additional_limits: Vec::new(),
error: Some("usage unsupported".to_string()),
}
}
}
#[derive(Debug, Clone, Deserialize)]
pub struct RateLimitStatusPayload {
#[serde(rename = "plan_type")]
pub plan_type: PlanType,
#[serde(rename = "rate_limit", default)]
pub rate_limit: Option<Option<Box<RateLimitStatusDetails>>>,
#[serde(rename = "credits", default)]
pub credits: Option<Option<Box<CreditStatusDetails>>>,
#[serde(rename = "additional_rate_limits", default)]
pub additional_rate_limits: Option<Option<Vec<AdditionalRateLimitDetails>>>,
#[serde(rename = "rate_limit_reached_type", default)]
pub rate_limit_reached_type: Option<Option<RateLimitReachedType>>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct AdditionalRateLimitDetails {
#[serde(rename = "limit_name")]
pub limit_name: String,
#[serde(rename = "metered_feature")]
pub metered_feature: String,
#[serde(rename = "rate_limit", default)]
pub rate_limit: Option<Option<Box<RateLimitStatusDetails>>>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct RateLimitReachedType {
#[serde(rename = "type")]
pub kind: RateLimitReachedKind,
}
#[derive(Debug, Clone, Copy, Deserialize)]
pub enum RateLimitReachedKind {
#[serde(rename = "rate_limit_reached")]
RateLimitReached,
#[serde(rename = "workspace_owner_credits_depleted")]
WorkspaceOwnerCreditsDepleted,
#[serde(rename = "workspace_member_credits_depleted")]
WorkspaceMemberCreditsDepleted,
#[serde(rename = "workspace_owner_usage_limit_reached")]
WorkspaceOwnerUsageLimitReached,
#[serde(rename = "workspace_member_usage_limit_reached")]
WorkspaceMemberUsageLimitReached,
#[serde(rename = "unknown", other)]
Unknown,
}
impl RateLimitReachedKind {
pub fn as_str(self) -> Option<&'static str> {
match self {
Self::RateLimitReached => Some("rate_limit_reached"),
Self::WorkspaceOwnerCreditsDepleted => Some("workspace_owner_credits_depleted"),
Self::WorkspaceMemberCreditsDepleted => Some("workspace_member_credits_depleted"),
Self::WorkspaceOwnerUsageLimitReached => Some("workspace_owner_usage_limit_reached"),
Self::WorkspaceMemberUsageLimitReached => Some("workspace_member_usage_limit_reached"),
Self::Unknown => None,
}
}
}
#[allow(dead_code)]
#[derive(Debug, Clone, Deserialize)]
pub struct RateLimitStatusDetails {
#[serde(rename = "allowed")]
pub allowed: bool,
#[serde(rename = "limit_reached")]
pub limit_reached: bool,
#[serde(rename = "primary_window", default)]
pub primary_window: Option<Option<Box<RateLimitWindowSnapshot>>>,
#[serde(rename = "secondary_window", default)]
pub secondary_window: Option<Option<Box<RateLimitWindowSnapshot>>>,
}
#[allow(dead_code)]
#[derive(Debug, Clone, Deserialize)]
pub struct RateLimitWindowSnapshot {
#[serde(rename = "used_percent")]
pub used_percent: i32,
#[serde(rename = "limit_window_seconds")]
pub limit_window_seconds: i32,
#[serde(rename = "reset_after_seconds")]
pub reset_after_seconds: i32,
#[serde(rename = "reset_at")]
pub reset_at: i32,
}
#[derive(Debug, Clone, Deserialize)]
pub struct CreditStatusDetails {
#[serde(rename = "has_credits")]
pub has_credits: bool,
#[serde(rename = "unlimited")]
pub unlimited: bool,
#[serde(rename = "balance", default)]
pub balance: Option<Option<String>>,
}
#[derive(Debug, Clone, Copy, Deserialize)]
pub enum PlanType {
#[serde(rename = "guest")]
Guest,
#[serde(rename = "free")]
Free,
#[serde(rename = "go")]
Go,
#[serde(rename = "plus")]
Plus,
#[serde(rename = "pro")]
Pro,
#[serde(rename = "prolite")]
ProLite,
#[serde(rename = "free_workspace")]
FreeWorkspace,
#[serde(rename = "team")]
Team,
#[serde(rename = "self_serve_business_usage_based")]
SelfServeBusinessUsageBased,
#[serde(rename = "business")]
Business,
#[serde(rename = "enterprise_cbp_usage_based")]
EnterpriseCbpUsageBased,
#[serde(rename = "education")]
Education,
#[serde(rename = "quorum")]
Quorum,
#[serde(rename = "k12")]
K12,
#[serde(rename = "enterprise")]
Enterprise,
#[serde(rename = "edu")]
Edu,
#[serde(rename = "unknown", other)]
Unknown,
}
impl PlanType {
pub fn as_str(self) -> &'static str {
match self {
Self::Guest => "guest",
Self::Free => "free",
Self::Go => "go",
Self::Plus => "plus",
Self::Pro => "pro",
Self::ProLite => "prolite",
Self::FreeWorkspace => "free_workspace",
Self::Team => "team",
Self::SelfServeBusinessUsageBased => "self_serve_business_usage_based",
Self::Business => "business",
Self::EnterpriseCbpUsageBased => "enterprise_cbp_usage_based",
Self::Education => "education",
Self::Quorum => "quorum",
Self::K12 => "k12",
Self::Enterprise => "enterprise",
Self::Edu => "edu",
Self::Unknown => "unknown",
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn redacted_string_debug_redacts_secret_but_json_preserves_value() {
let secret = "secret-value-that-must-stay-private";
let value = RedactedString::new(secret);
assert_eq!(format!("{value:?}"), "<redacted>");
let json = serde_json::to_string(&value).expect("serialize redacted string");
assert!(json.contains(secret));
let decoded: RedactedString =
serde_json::from_str(&json).expect("deserialize redacted string");
assert_eq!(decoded.expose_secret(), secret);
}
#[test]
fn account_debug_redacts_api_key_and_chatgpt_tokens() {
let api_secret = "sk-codex-switch-test-secret";
let id_secret = "id-token-codex-switch-test-secret";
let access_secret = "access-token-codex-switch-test-secret";
let refresh_secret = "refresh-token-codex-switch-test-secret";
let api_auth = AuthData::ApiKey {
key: api_secret.into(),
};
let chatgpt_auth = AuthData::ChatGPT {
id_token: id_secret.into(),
access_token: access_secret.into(),
refresh_token: refresh_secret.into(),
account_id: Some("account-id".to_string()),
};
let api_account = StoredAccount::new_api_key("api".to_string(), api_secret.to_string());
let chatgpt_account = StoredAccount::new_chatgpt(NewChatGptAccount {
name: "chatgpt".to_string(),
email: Some("user@example.com".to_string()),
plan_type: Some("pro".to_string()),
chatgpt_user_id: Some("user-id".to_string()),
chatgpt_account_is_fedramp: false,
token_last_refresh_at: Utc::now(),
subscription_expires_at: None,
id_token: id_secret.into(),
access_token: access_secret.into(),
refresh_token: refresh_secret.into(),
account_id: Some("account-id".to_string()),
});
let debug = format!("{api_auth:?} {chatgpt_auth:?} {api_account:?} {chatgpt_account:?}");
for secret in [api_secret, id_secret, access_secret, refresh_secret] {
assert!(!debug.contains(secret), "debug output leaked {secret}");
}
assert!(debug.contains("<redacted>"));
}
#[test]
fn accounts_and_auth_json_keep_raw_secret_json_values() {
let api_secret = "sk-json-compat-secret";
let id_secret = "id-token-json-compat-secret";
let access_secret = "access-token-json-compat-secret";
let refresh_secret = "refresh-token-json-compat-secret";
let api_account = StoredAccount::new_api_key("api".to_string(), api_secret.to_string());
let chatgpt_account = StoredAccount::new_chatgpt(NewChatGptAccount {
name: "chatgpt".to_string(),
email: None,
plan_type: Some("pro".to_string()),
chatgpt_user_id: None,
chatgpt_account_is_fedramp: false,
token_last_refresh_at: Utc::now(),
subscription_expires_at: None,
id_token: id_secret.into(),
access_token: access_secret.into(),
refresh_token: refresh_secret.into(),
account_id: Some("account-id".to_string()),
});
let store = AccountsStore {
version: 1,
accounts: vec![api_account, chatgpt_account],
masked_account_ids: Vec::new(),
};
let auth_json = AuthDotJson {
auth_mode: Some("chatgpt".to_string()),
openai_api_key: Some(api_secret.into()),
tokens: Some(TokenData {
id_token: id_secret.into(),
access_token: access_secret.into(),
refresh_token: refresh_secret.into(),
account_id: Some("account-id".to_string()),
}),
last_refresh: None,
};
let auth_debug = format!(
"{auth_json:?} {:?}",
auth_json.tokens.as_ref().expect("tokens")
);
for secret in [api_secret, id_secret, access_secret, refresh_secret] {
assert!(
!auth_debug.contains(secret),
"auth debug output leaked {secret}"
);
}
assert!(auth_debug.contains("<redacted>"));
let store_json = serde_json::to_string(&store).expect("serialize accounts store");
let auth_file_json = serde_json::to_string(&auth_json).expect("serialize auth json");
for secret in [api_secret, id_secret, access_secret, refresh_secret] {
assert!(store_json.contains(secret), "accounts json lost {secret}");
assert!(auth_file_json.contains(secret), "auth json lost {secret}");
}
let decoded_auth: AuthDotJson =
serde_json::from_str(&auth_file_json).expect("deserialize auth json");
assert_eq!(
decoded_auth
.openai_api_key
.as_ref()
.map(RedactedString::expose_secret),
Some(api_secret)
);
let decoded_tokens = decoded_auth.tokens.expect("tokens");
assert_eq!(decoded_tokens.id_token.expose_secret(), id_secret);
assert_eq!(decoded_tokens.access_token.expose_secret(), access_secret);
assert_eq!(decoded_tokens.refresh_token.expose_secret(), refresh_secret);
}
#[test]
fn accounts_store_does_not_serialize_active_account_id() {
let store = AccountsStore::default();
let value = serde_json::to_value(&store).expect("serialize accounts store");
assert!(value.get("active_account_id").is_none());
}
#[test]
fn accounts_store_ignores_legacy_active_account_id() {
let store: AccountsStore = serde_json::from_value(serde_json::json!({
"version": 1,
"accounts": [],
"active_account_id": "legacy-account-id",
"masked_account_ids": []
}))
.expect("deserialize accounts store");
assert!(store.accounts.is_empty());
assert!(store.masked_account_ids.is_empty());
}
}