diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 097fef2af1..87c723aa12 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1590,6 +1590,34 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "codex-auth" +version = "0.0.0" +dependencies = [ + "anyhow", + "async-trait", + "base64 0.22.1", + "chrono", + "codex-app-server-protocol", + "codex-client", + "codex-keyring-store", + "codex-otel", + "codex-protocol", + "keyring", + "once_cell", + "pretty_assertions", + "reqwest", + "schemars 0.8.22", + "serde", + "serde_json", + "serial_test", + "sha2", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "codex-backend-client" version = "0.0.0" @@ -1824,6 +1852,7 @@ dependencies = [ "codex-arg0", "codex-artifacts", "codex-async-utils", + "codex-auth", "codex-client", "codex-config", "codex-connectors", @@ -1831,7 +1860,6 @@ dependencies = [ "codex-file-search", "codex-git", "codex-hooks", - "codex-keyring-store", "codex-network-proxy", "codex-otel", "codex-protocol", @@ -1867,7 +1895,6 @@ dependencies = [ "image", "indexmap 2.13.0", "insta", - "keyring", "landlock", "libc", "maplit", @@ -1890,7 +1917,6 @@ dependencies = [ "serde_yaml", "serial_test", "sha1", - "sha2", "shlex", "similar", "tempfile", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 77ffb61204..cc8692379c 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -68,6 +68,7 @@ members = [ "test-macros", "package-manager", "artifacts", + "auth", ] resolver = "2" @@ -84,6 +85,7 @@ license = "Apache-2.0" # Internal app_test_support = { path = "app-server/tests/common" } codex-ansi-escape = { path = "ansi-escape" } +codex-auth = { path = "auth" } codex-api = { path = "codex-api" } codex-artifacts = { path = "artifacts" } codex-package-manager = { path = "package-manager" } diff --git a/codex-rs/auth/Cargo.toml b/codex-rs/auth/Cargo.toml new file mode 100644 index 0000000000..4c8a0e2a82 --- /dev/null +++ b/codex-rs/auth/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "codex-auth" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dependencies] +async-trait = { workspace = true } +base64 = { workspace = true } +chrono = { workspace = true, features = ["serde"] } +codex-app-server-protocol = { workspace = true } +codex-client = { workspace = true } +codex-keyring-store = { workspace = true } +codex-otel = { workspace = true } +codex-protocol = { workspace = true } +once_cell = { workspace = true } +reqwest = { workspace = true, features = ["json", "stream"] } +schemars = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +sha2 = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +tracing = { workspace = true, features = ["log"] } + +[dev-dependencies] +anyhow = { workspace = true } +keyring = { workspace = true, features = ["crypto-rust"] } +pretty_assertions = { workspace = true } +serial_test = { workspace = true } +tempfile = { workspace = true } + +[lints] +workspace = true diff --git a/codex-rs/auth/src/lib.rs b/codex-rs/auth/src/lib.rs new file mode 100644 index 0000000000..72fcff57d7 --- /dev/null +++ b/codex-rs/auth/src/lib.rs @@ -0,0 +1,1332 @@ +mod storage; +pub mod token_data; + +use async_trait::async_trait; +use chrono::Utc; +use reqwest::StatusCode; +use serde::Deserialize; +use serde::Serialize; +use std::env; +use std::fmt::Debug; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use std::sync::RwLock; + +use codex_app_server_protocol::AuthMode as ApiAuthMode; +use codex_otel::TelemetryAuthMode; + +pub use crate::storage::AuthCredentialsStoreMode; +pub use crate::storage::AuthDotJson; +use crate::storage::AuthStorageBackend; +use crate::storage::create_auth_storage; +use crate::token_data::KnownPlan as InternalKnownPlan; +use crate::token_data::PlanType as InternalPlanType; +use crate::token_data::TokenData; +use crate::token_data::parse_chatgpt_jwt_claims; +use codex_client::CodexHttpClient; +use codex_protocol::account::PlanType as AccountPlanType; +use serde_json::Value; +use thiserror::Error; + +/// Account type for the current user. +/// +/// This is used internally to determine the base URL for generating responses, +/// and to gate ChatGPT-only behaviors like rate limits and available models (as +/// opposed to API key-based auth). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AuthMode { + ApiKey, + Chatgpt, +} + +impl From for TelemetryAuthMode { + fn from(mode: AuthMode) -> Self { + match mode { + AuthMode::ApiKey => TelemetryAuthMode::ApiKey, + AuthMode::Chatgpt => TelemetryAuthMode::Chatgpt, + } + } +} + +/// Authentication mechanism used by the current user. +#[derive(Debug, Clone)] +pub enum CodexAuth { + ApiKey(ApiKeyAuth), + Chatgpt(ChatgptAuth), + ChatgptAuthTokens(ChatgptAuthTokens), +} + +#[derive(Debug, Clone)] +pub struct ApiKeyAuth { + api_key: String, +} + +#[derive(Debug, Clone)] +pub struct ChatgptAuth { + state: ChatgptAuthState, + storage: Arc, +} + +#[derive(Debug, Clone)] +pub struct ChatgptAuthTokens { + state: ChatgptAuthState, +} + +#[derive(Debug, Clone)] +struct ChatgptAuthState { + auth_dot_json: Arc>>, + client: CodexHttpClient, +} + +impl PartialEq for CodexAuth { + fn eq(&self, other: &Self) -> bool { + self.api_auth_mode() == other.api_auth_mode() + } +} + +// TODO(pakrym): use token exp field to check for expiration instead +const TOKEN_REFRESH_INTERVAL: i64 = 8; + +const REFRESH_TOKEN_EXPIRED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token has expired. Please log out and sign in again."; +const REFRESH_TOKEN_REUSED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token was already used. Please log out and sign in again."; +const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token was revoked. Please log out and sign in again."; +const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str = + "Your access token could not be refreshed. Please log out and sign in again."; +const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again."; +const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; +pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE"; +static DEFAULT_CLIENT_FACTORY: OnceLock CodexHttpClient> = OnceLock::new(); + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +#[error("{message}")] +pub struct RefreshTokenFailedError { + pub reason: RefreshTokenFailedReason, + pub message: String, +} + +impl RefreshTokenFailedError { + pub fn new(reason: RefreshTokenFailedReason, message: impl Into) -> Self { + Self { + reason, + message: message.into(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RefreshTokenFailedReason { + Expired, + Exhausted, + Revoked, + Other, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[error("default client factory already initialized")] +pub struct SetDefaultClientFactoryError; + +pub fn set_default_client_factory( + factory: fn() -> CodexHttpClient, +) -> Result<(), SetDefaultClientFactoryError> { + DEFAULT_CLIENT_FACTORY + .set(factory) + .map_err(|_| SetDefaultClientFactoryError) +} + +#[derive(Debug, Error)] +pub enum RefreshTokenError { + #[error("{0}")] + Permanent(#[from] RefreshTokenFailedError), + #[error(transparent)] + Transient(#[from] std::io::Error), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalAuthTokens { + pub access_token: String, + pub chatgpt_account_id: String, + pub chatgpt_plan_type: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExternalAuthRefreshReason { + Unauthorized, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ExternalAuthRefreshContext { + pub reason: ExternalAuthRefreshReason, + pub previous_account_id: Option, +} + +#[async_trait] +pub trait ExternalAuthRefresher: Send + Sync { + async fn refresh( + &self, + context: ExternalAuthRefreshContext, + ) -> std::io::Result; +} + +impl RefreshTokenError { + pub fn failed_reason(&self) -> Option { + match self { + Self::Permanent(error) => Some(error.reason), + Self::Transient(_) => None, + } + } +} + +impl From for std::io::Error { + fn from(err: RefreshTokenError) -> Self { + match err { + RefreshTokenError::Permanent(failed) => std::io::Error::other(failed), + RefreshTokenError::Transient(inner) => inner, + } + } +} + +impl CodexAuth { + fn from_auth_dot_json( + codex_home: &Path, + auth_dot_json: AuthDotJson, + auth_credentials_store_mode: AuthCredentialsStoreMode, + client: CodexHttpClient, + ) -> std::io::Result { + let auth_mode = auth_dot_json.resolved_mode(); + if auth_mode == ApiAuthMode::ApiKey { + let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else { + return Err(std::io::Error::other("API key auth is missing a key.")); + }; + return Ok(CodexAuth::from_api_key_with_client(api_key, client)); + } + + let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); + let state = ChatgptAuthState { + auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), + client, + }; + + match auth_mode { + ApiAuthMode::Chatgpt => { + let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); + Ok(Self::Chatgpt(ChatgptAuth { state, storage })) + } + ApiAuthMode::ChatgptAuthTokens => { + Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) + } + ApiAuthMode::ApiKey => unreachable!("api key mode is handled above"), + } + } + + /// Loads the available auth information from auth storage. + pub fn from_auth_storage( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + ) -> std::io::Result> { + load_auth(codex_home, false, auth_credentials_store_mode) + } + + pub fn auth_mode(&self) -> AuthMode { + match self { + Self::ApiKey(_) => AuthMode::ApiKey, + Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => AuthMode::Chatgpt, + } + } + + pub fn api_auth_mode(&self) -> ApiAuthMode { + match self { + Self::ApiKey(_) => ApiAuthMode::ApiKey, + Self::Chatgpt(_) => ApiAuthMode::Chatgpt, + Self::ChatgptAuthTokens(_) => ApiAuthMode::ChatgptAuthTokens, + } + } + + pub fn is_api_key_auth(&self) -> bool { + self.auth_mode() == AuthMode::ApiKey + } + + pub fn is_chatgpt_auth(&self) -> bool { + self.auth_mode() == AuthMode::Chatgpt + } + + pub fn is_external_chatgpt_tokens(&self) -> bool { + matches!(self, Self::ChatgptAuthTokens(_)) + } + + /// Returns `None` if `auth_mode() != AuthMode::ApiKey`. + pub fn api_key(&self) -> Option<&str> { + match self { + Self::ApiKey(auth) => Some(auth.api_key.as_str()), + Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => None, + } + } + + /// Returns `Err` if `is_chatgpt_auth()` is false. + pub fn get_token_data(&self) -> Result { + let auth_dot_json: Option = self.get_current_auth_json(); + match auth_dot_json { + Some(AuthDotJson { + tokens: Some(tokens), + last_refresh: Some(_), + .. + }) => Ok(tokens), + _ => Err(std::io::Error::other("Token data is not available.")), + } + } + + /// Returns the token string used for bearer authentication. + pub fn get_token(&self) -> Result { + match self { + Self::ApiKey(auth) => Ok(auth.api_key.clone()), + Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => { + let access_token = self.get_token_data()?.access_token; + Ok(access_token) + } + } + } + + /// Returns `None` if `is_chatgpt_auth()` is false. + pub fn get_account_id(&self) -> Option { + self.get_current_token_data().and_then(|t| t.account_id) + } + + /// Returns `None` if `is_chatgpt_auth()` is false. + pub fn get_account_email(&self) -> Option { + self.get_current_token_data().and_then(|t| t.id_token.email) + } + + /// Returns `None` if `is_chatgpt_auth()` is false. + pub fn get_chatgpt_user_id(&self) -> Option { + self.get_current_token_data() + .and_then(|t| t.id_token.chatgpt_user_id) + } + + /// Account-facing plan classification derived from the current token. + /// Returns a high-level `AccountPlanType` (e.g., Free/Plus/Pro/Team/…) + /// mapped from the ID token's internal plan value. Prefer this when you + /// need to make UI or product decisions based on the user's subscription. + /// When ChatGPT auth is active but the token omits the plan claim, report + /// `Unknown` instead of treating the account as invalid. + pub fn account_plan_type(&self) -> Option { + let map_known = |kp: &InternalKnownPlan| match kp { + InternalKnownPlan::Free => AccountPlanType::Free, + InternalKnownPlan::Go => AccountPlanType::Go, + InternalKnownPlan::Plus => AccountPlanType::Plus, + InternalKnownPlan::Pro => AccountPlanType::Pro, + InternalKnownPlan::Team => AccountPlanType::Team, + InternalKnownPlan::Business => AccountPlanType::Business, + InternalKnownPlan::Enterprise => AccountPlanType::Enterprise, + InternalKnownPlan::Edu => AccountPlanType::Edu, + }; + + self.get_current_token_data().map(|t| { + t.id_token + .chatgpt_plan_type + .map(|pt| match pt { + InternalPlanType::Known(k) => map_known(&k), + InternalPlanType::Unknown(_) => AccountPlanType::Unknown, + }) + .unwrap_or(AccountPlanType::Unknown) + }) + } + + /// Returns `None` if `is_chatgpt_auth()` is false. + fn get_current_auth_json(&self) -> Option { + let state = match self { + Self::Chatgpt(auth) => &auth.state, + Self::ChatgptAuthTokens(auth) => &auth.state, + Self::ApiKey(_) => return None, + }; + #[expect(clippy::unwrap_used)] + state.auth_dot_json.lock().unwrap().clone() + } + + /// Returns `None` if `is_chatgpt_auth()` is false. + fn get_current_token_data(&self) -> Option { + self.get_current_auth_json().and_then(|t| t.tokens) + } + + /// Consider this private to integration tests. + pub fn create_dummy_chatgpt_auth_for_testing() -> Self { + let auth_dot_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: Default::default(), + access_token: "Access Token".to_string(), + refresh_token: "test".to_string(), + account_id: Some("account_id".to_string()), + }), + last_refresh: Some(Utc::now()), + }; + + let client = create_client(); + let state = ChatgptAuthState { + auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), + client, + }; + let storage = create_auth_storage(PathBuf::new(), AuthCredentialsStoreMode::File); + Self::Chatgpt(ChatgptAuth { state, storage }) + } + + fn from_api_key_with_client(api_key: &str, _client: CodexHttpClient) -> Self { + Self::ApiKey(ApiKeyAuth { + api_key: api_key.to_owned(), + }) + } + + pub fn from_api_key(api_key: &str) -> Self { + Self::from_api_key_with_client(api_key, create_client()) + } +} + +impl ChatgptAuth { + fn current_auth_json(&self) -> Option { + #[expect(clippy::unwrap_used)] + self.state.auth_dot_json.lock().unwrap().clone() + } + + fn current_token_data(&self) -> Option { + self.current_auth_json().and_then(|auth| auth.tokens) + } + + fn storage(&self) -> &Arc { + &self.storage + } + + fn client(&self) -> &CodexHttpClient { + &self.state.client + } +} + +pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; +pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; + +pub fn read_openai_api_key_from_env() -> Option { + env::var(OPENAI_API_KEY_ENV_VAR) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +pub fn read_codex_api_key_from_env() -> Option { + env::var(CODEX_API_KEY_ENV_VAR) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +/// Delete the auth.json file inside `codex_home` if it exists. Returns `Ok(true)` +/// if a file was removed, `Ok(false)` if no auth file was present. +pub fn logout( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result { + let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); + storage.delete() +} + +/// Writes an `auth.json` that contains only the API key. +pub fn login_with_api_key( + codex_home: &Path, + api_key: &str, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some(api_key.to_string()), + tokens: None, + last_refresh: None, + }; + save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) +} + +/// Writes an in-memory auth payload for externally managed ChatGPT tokens. +pub fn login_with_chatgpt_auth_tokens( + codex_home: &Path, + access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, +) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson::from_external_access_token( + access_token, + chatgpt_account_id, + chatgpt_plan_type, + )?; + save_auth( + codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + ) +} + +/// Persist the provided auth payload using the specified backend. +pub fn save_auth( + codex_home: &Path, + auth: &AuthDotJson, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result<()> { + let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); + storage.save(auth) +} + +/// Load CLI auth data using the configured credential store backend. +/// Returns `None` when no credentials are stored. This function is +/// provided only for tests. Production code should not directly load +/// from the auth.json storage. It should use the AuthManager abstraction +/// instead. +pub fn load_auth_dot_json( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result> { + let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); + storage.load() +} + +pub fn load_auth( + codex_home: &Path, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result> { + let build_auth = |auth_dot_json: AuthDotJson, storage_mode| { + let client = create_client(); + CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, storage_mode, client) + }; + + // API key via env var takes precedence over any other auth method. + if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() { + let client = create_client(); + return Ok(Some(CodexAuth::from_api_key_with_client( + api_key.as_str(), + client, + ))); + } + + // External ChatGPT auth tokens live in the in-memory (ephemeral) store. Always check this + // first so external auth takes precedence over any persisted credentials. + let ephemeral_storage = create_auth_storage( + codex_home.to_path_buf(), + AuthCredentialsStoreMode::Ephemeral, + ); + if let Some(auth_dot_json) = ephemeral_storage.load()? { + let auth = build_auth(auth_dot_json, AuthCredentialsStoreMode::Ephemeral)?; + return Ok(Some(auth)); + } + + // If the caller explicitly requested ephemeral auth, there is no persisted fallback. + if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { + return Ok(None); + } + + // Fall back to the configured persistent store (file/keyring/auto) for managed auth. + let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); + let auth_dot_json = match storage.load()? { + Some(auth) => auth, + None => return Ok(None), + }; + + let auth = build_auth(auth_dot_json, auth_credentials_store_mode)?; + Ok(Some(auth)) +} + +fn logout_all_stores( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result { + if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { + return logout(codex_home, AuthCredentialsStoreMode::Ephemeral); + } + let removed_ephemeral = logout(codex_home, AuthCredentialsStoreMode::Ephemeral)?; + let removed_managed = logout(codex_home, auth_credentials_store_mode)?; + Ok(removed_ephemeral || removed_managed) +} + +fn create_client() -> CodexHttpClient { + DEFAULT_CLIENT_FACTORY.get().copied().map_or_else( + || CodexHttpClient::new(reqwest::Client::new()), + |factory| factory(), + ) +} + +fn try_parse_error_message(text: &str) -> String { + tracing::debug!("Parsing server error response: {text}"); + let json = serde_json::from_str::(text).unwrap_or_default(); + if let Some(error) = json.get("error") + && let Some(message) = error.get("message") + && let Some(message_str) = message.as_str() + { + return message_str.to_string(); + } + if text.is_empty() { + return "Unknown error".to_string(); + } + text.to_string() +} + +// Persist refreshed tokens into auth storage and update last_refresh. +fn persist_tokens( + storage: &Arc, + id_token: Option, + access_token: Option, + refresh_token: Option, +) -> std::io::Result { + let mut auth_dot_json = storage + .load()? + .ok_or(std::io::Error::other("Token data is not available."))?; + + let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default); + if let Some(id_token) = id_token { + tokens.id_token = parse_chatgpt_jwt_claims(&id_token).map_err(std::io::Error::other)?; + } + if let Some(access_token) = access_token { + tokens.access_token = access_token; + } + if let Some(refresh_token) = refresh_token { + tokens.refresh_token = refresh_token; + } + auth_dot_json.last_refresh = Some(Utc::now()); + storage.save(&auth_dot_json)?; + Ok(auth_dot_json) +} + +// Requests refreshed ChatGPT OAuth tokens from the auth service using a refresh token. +// The caller is responsible for persisting any returned tokens. +async fn request_chatgpt_token_refresh( + refresh_token: String, + client: &CodexHttpClient, +) -> Result { + let refresh_request = RefreshRequest { + client_id: CLIENT_ID, + grant_type: "refresh_token", + refresh_token, + }; + + let endpoint = refresh_token_endpoint(); + + // Use shared client factory to include standard headers + let response = client + .post(endpoint.as_str()) + .header("Content-Type", "application/json") + .json(&refresh_request) + .send() + .await + .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))?; + + let status = response.status(); + if status.is_success() { + let refresh_response = response + .json::() + .await + .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))?; + Ok(refresh_response) + } else { + let body = response.text().await.unwrap_or_default(); + tracing::error!("Failed to refresh token: {status}: {body}"); + if status == StatusCode::UNAUTHORIZED { + let failed = classify_refresh_token_failure(&body); + Err(RefreshTokenError::Permanent(failed)) + } else { + let message = try_parse_error_message(&body); + Err(RefreshTokenError::Transient(std::io::Error::other( + format!("Failed to refresh token: {status}: {message}"), + ))) + } + } +} + +fn classify_refresh_token_failure(body: &str) -> RefreshTokenFailedError { + let code = extract_refresh_token_error_code(body); + + let normalized_code = code.as_deref().map(str::to_ascii_lowercase); + let reason = match normalized_code.as_deref() { + Some("refresh_token_expired") => RefreshTokenFailedReason::Expired, + Some("refresh_token_reused") => RefreshTokenFailedReason::Exhausted, + Some("refresh_token_invalidated") => RefreshTokenFailedReason::Revoked, + _ => RefreshTokenFailedReason::Other, + }; + + if reason == RefreshTokenFailedReason::Other { + tracing::warn!( + backend_code = normalized_code.as_deref(), + backend_body = body, + "Encountered unknown 401 response while refreshing token" + ); + } + + let message = match reason { + RefreshTokenFailedReason::Expired => REFRESH_TOKEN_EXPIRED_MESSAGE.to_string(), + RefreshTokenFailedReason::Exhausted => REFRESH_TOKEN_REUSED_MESSAGE.to_string(), + RefreshTokenFailedReason::Revoked => REFRESH_TOKEN_INVALIDATED_MESSAGE.to_string(), + RefreshTokenFailedReason::Other => REFRESH_TOKEN_UNKNOWN_MESSAGE.to_string(), + }; + + RefreshTokenFailedError::new(reason, message) +} + +fn extract_refresh_token_error_code(body: &str) -> Option { + if body.trim().is_empty() { + return None; + } + + let Value::Object(map) = serde_json::from_str::(body).ok()? else { + return None; + }; + + if let Some(error_value) = map.get("error") { + match error_value { + Value::Object(obj) => { + if let Some(code) = obj.get("code").and_then(Value::as_str) { + return Some(code.to_string()); + } + } + Value::String(code) => { + return Some(code.to_string()); + } + _ => {} + } + } + + map.get("code").and_then(Value::as_str).map(str::to_string) +} + +#[derive(Serialize)] +struct RefreshRequest { + client_id: &'static str, + grant_type: &'static str, + refresh_token: String, +} + +#[derive(Deserialize, Clone)] +struct RefreshResponse { + id_token: Option, + access_token: Option, + refresh_token: Option, +} + +// Shared constant for token refresh (client id used for oauth token refresh flow) +pub const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; + +fn refresh_token_endpoint() -> String { + std::env::var(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR) + .unwrap_or_else(|_| REFRESH_TOKEN_URL.to_string()) +} + +impl AuthDotJson { + fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { + let mut token_info = + parse_chatgpt_jwt_claims(&external.access_token).map_err(std::io::Error::other)?; + token_info.chatgpt_account_id = Some(external.chatgpt_account_id.clone()); + token_info.chatgpt_plan_type = external + .chatgpt_plan_type + .as_deref() + .map(InternalPlanType::from_raw_value) + .or(token_info.chatgpt_plan_type) + .or(Some(InternalPlanType::Unknown("unknown".to_string()))); + let tokens = TokenData { + id_token: token_info, + access_token: external.access_token.clone(), + refresh_token: String::new(), + account_id: Some(external.chatgpt_account_id.clone()), + }; + + Ok(Self { + auth_mode: Some(ApiAuthMode::ChatgptAuthTokens), + openai_api_key: None, + tokens: Some(tokens), + last_refresh: Some(Utc::now()), + }) + } + + fn from_external_access_token( + access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, + ) -> std::io::Result { + let external = ExternalAuthTokens { + access_token: access_token.to_string(), + chatgpt_account_id: chatgpt_account_id.to_string(), + chatgpt_plan_type: chatgpt_plan_type.map(str::to_string), + }; + Self::from_external_tokens(&external) + } + + fn resolved_mode(&self) -> ApiAuthMode { + if let Some(mode) = self.auth_mode { + return mode; + } + if self.openai_api_key.is_some() { + return ApiAuthMode::ApiKey; + } + ApiAuthMode::Chatgpt + } + + fn storage_mode( + &self, + auth_credentials_store_mode: AuthCredentialsStoreMode, + ) -> AuthCredentialsStoreMode { + if self.resolved_mode() == ApiAuthMode::ChatgptAuthTokens { + AuthCredentialsStoreMode::Ephemeral + } else { + auth_credentials_store_mode + } + } +} + +/// Internal cached auth state. +#[derive(Clone)] +struct CachedAuth { + auth: Option, + /// Callback used to refresh external auth by asking the parent app for new tokens. + external_refresher: Option>, +} + +impl Debug for CachedAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachedAuth") + .field( + "auth_mode", + &self.auth.as_ref().map(CodexAuth::api_auth_mode), + ) + .field( + "external_refresher", + &self.external_refresher.as_ref().map(|_| "present"), + ) + .finish() + } +} + +enum UnauthorizedRecoveryStep { + Reload, + RefreshToken, + ExternalRefresh, + Done, +} + +enum ReloadOutcome { + /// Reload was performed and the cached auth changed + ReloadedChanged, + /// Reload was performed and the cached auth remained the same + ReloadedNoChange, + /// Reload was skipped (missing or mismatched account id) + Skipped, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum UnauthorizedRecoveryMode { + Managed, + External, +} + +// UnauthorizedRecovery is a state machine that handles an attempt to refresh the authentication when requests +// to API fail with 401 status code. +// The client calls next() every time it encounters a 401 error, one time per retry. +// For API key based authentication, we don't do anything and let the error bubble to the user. +// +// For ChatGPT based authentication, we: +// 1. Attempt to reload the auth data from disk. We only reload if the account id matches the one the current process is running as. +// 2. Attempt to refresh the token using OAuth token refresh flow. +// If after both steps the server still responds with 401 we let the error bubble to the user. +// +// For external ChatGPT auth tokens (chatgptAuthTokens), UnauthorizedRecovery does not touch disk or refresh +// tokens locally. Instead it calls the ExternalAuthRefresher (account/chatgptAuthTokens/refresh) to ask the +// parent app for new tokens, stores them in the ephemeral auth store, and retries once. +pub struct UnauthorizedRecovery { + manager: Arc, + step: UnauthorizedRecoveryStep, + expected_account_id: Option, + mode: UnauthorizedRecoveryMode, +} + +impl UnauthorizedRecovery { + fn new(manager: Arc) -> Self { + let cached_auth = manager.auth_cached(); + let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id); + let mode = if cached_auth + .as_ref() + .is_some_and(CodexAuth::is_external_chatgpt_tokens) + { + UnauthorizedRecoveryMode::External + } else { + UnauthorizedRecoveryMode::Managed + }; + let step = match mode { + UnauthorizedRecoveryMode::Managed => UnauthorizedRecoveryStep::Reload, + UnauthorizedRecoveryMode::External => UnauthorizedRecoveryStep::ExternalRefresh, + }; + Self { + manager, + step, + expected_account_id, + mode, + } + } + + pub fn has_next(&self) -> bool { + if !self + .manager + .auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_chatgpt_auth) + { + return false; + } + + if self.mode == UnauthorizedRecoveryMode::External + && !self.manager.has_external_auth_refresher() + { + return false; + } + + !matches!(self.step, UnauthorizedRecoveryStep::Done) + } + + pub async fn next(&mut self) -> Result<(), RefreshTokenError> { + if !self.has_next() { + return Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + "No more recovery steps available.", + ))); + } + + match self.step { + UnauthorizedRecoveryStep::Reload => { + match self + .manager + .reload_if_account_id_matches(self.expected_account_id.as_deref()) + { + ReloadOutcome::ReloadedChanged | ReloadOutcome::ReloadedNoChange => { + self.step = UnauthorizedRecoveryStep::RefreshToken; + } + ReloadOutcome::Skipped => { + self.step = UnauthorizedRecoveryStep::Done; + return Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), + ))); + } + } + } + UnauthorizedRecoveryStep::RefreshToken => { + self.manager.refresh_token_from_authority().await?; + self.step = UnauthorizedRecoveryStep::Done; + } + UnauthorizedRecoveryStep::ExternalRefresh => { + self.manager + .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await?; + self.step = UnauthorizedRecoveryStep::Done; + } + UnauthorizedRecoveryStep::Done => {} + } + Ok(()) + } +} + +/// Central manager providing a single source of truth for auth.json derived +/// authentication data. It loads once (or on preference change) and then +/// hands out cloned `CodexAuth` values so the rest of the program has a +/// consistent snapshot. +/// +/// External modifications to `auth.json` will NOT be observed until +/// `reload()` is called explicitly. This matches the design goal of avoiding +/// different parts of the program seeing inconsistent auth data mid‑run. +#[derive(Debug)] +pub struct AuthManager { + codex_home: PathBuf, + inner: RwLock, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: RwLock>, +} + +impl AuthManager { + /// Create a new manager loading the initial auth using the provided + /// preferred auth method. Errors loading auth are swallowed; `auth()` will + /// simply return `None` in that case so callers can treat it as an + /// unauthenticated state. + pub fn new( + codex_home: PathBuf, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, + ) -> Self { + let managed_auth = load_auth( + &codex_home, + enable_codex_api_key_env, + auth_credentials_store_mode, + ) + .ok() + .flatten(); + Self { + codex_home, + inner: RwLock::new(CachedAuth { + auth: managed_auth, + external_refresher: None, + }), + enable_codex_api_key_env, + auth_credentials_store_mode, + forced_chatgpt_workspace_id: RwLock::new(None), + } + } + + /// Create an AuthManager with a specific CodexAuth, for testing only. + pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { + let cached = CachedAuth { + auth: Some(auth), + external_refresher: None, + }; + + Arc::new(Self { + codex_home: PathBuf::from("non-existent"), + inner: RwLock::new(cached), + enable_codex_api_key_env: false, + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + forced_chatgpt_workspace_id: RwLock::new(None), + }) + } + + /// Create an AuthManager with a specific CodexAuth and codex home, for testing only. + pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { + let cached = CachedAuth { + auth: Some(auth), + external_refresher: None, + }; + Arc::new(Self { + codex_home, + inner: RwLock::new(cached), + enable_codex_api_key_env: false, + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + forced_chatgpt_workspace_id: RwLock::new(None), + }) + } + + /// Current cached auth (clone) without attempting a refresh. + pub fn auth_cached(&self) -> Option { + self.inner.read().ok().and_then(|c| c.auth.clone()) + } + + /// Current cached auth (clone). May be `None` if not logged in or load failed. + /// Refreshes cached ChatGPT tokens if they are stale before returning. + pub async fn auth(&self) -> Option { + let auth = self.auth_cached()?; + if let Err(err) = self.refresh_if_stale(&auth).await { + tracing::error!("Failed to refresh token: {}", err); + return Some(auth); + } + self.auth_cached() + } + + /// Force a reload of the auth information from auth.json. Returns + /// whether the auth value changed. + pub fn reload(&self) -> bool { + tracing::info!("Reloading auth"); + let new_auth = self.load_auth_from_storage(); + self.set_cached_auth(new_auth) + } + + fn reload_if_account_id_matches(&self, expected_account_id: Option<&str>) -> ReloadOutcome { + let expected_account_id = match expected_account_id { + Some(account_id) => account_id, + None => { + tracing::info!("Skipping auth reload because no account id is available."); + return ReloadOutcome::Skipped; + } + }; + + let new_auth = self.load_auth_from_storage(); + let new_account_id = new_auth.as_ref().and_then(CodexAuth::get_account_id); + + if new_account_id.as_deref() != Some(expected_account_id) { + let found_account_id = new_account_id.as_deref().unwrap_or("unknown"); + tracing::info!( + "Skipping auth reload due to account id mismatch (expected: {expected_account_id}, found: {found_account_id})" + ); + return ReloadOutcome::Skipped; + } + + tracing::info!("Reloading auth for account {expected_account_id}"); + let cached_before_reload = self.auth_cached(); + let auth_changed = + !Self::auths_equal_for_refresh(cached_before_reload.as_ref(), new_auth.as_ref()); + self.set_cached_auth(new_auth); + if auth_changed { + ReloadOutcome::ReloadedChanged + } else { + ReloadOutcome::ReloadedNoChange + } + } + + fn auths_equal_for_refresh(a: Option<&CodexAuth>, b: Option<&CodexAuth>) -> bool { + match (a, b) { + (None, None) => true, + (Some(a), Some(b)) => match (a.api_auth_mode(), b.api_auth_mode()) { + (ApiAuthMode::ApiKey, ApiAuthMode::ApiKey) => a.api_key() == b.api_key(), + (ApiAuthMode::Chatgpt, ApiAuthMode::Chatgpt) + | (ApiAuthMode::ChatgptAuthTokens, ApiAuthMode::ChatgptAuthTokens) => { + a.get_current_auth_json() == b.get_current_auth_json() + } + _ => false, + }, + _ => false, + } + } + + fn auths_equal(a: Option<&CodexAuth>, b: Option<&CodexAuth>) -> bool { + match (a, b) { + (None, None) => true, + (Some(a), Some(b)) => a == b, + _ => false, + } + } + + fn load_auth_from_storage(&self) -> Option { + load_auth( + &self.codex_home, + self.enable_codex_api_key_env, + self.auth_credentials_store_mode, + ) + .ok() + .flatten() + } + + fn set_cached_auth(&self, new_auth: Option) -> bool { + if let Ok(mut guard) = self.inner.write() { + let previous = guard.auth.as_ref(); + let changed = !AuthManager::auths_equal(previous, new_auth.as_ref()); + tracing::info!("Reloaded auth, changed: {changed}"); + guard.auth = new_auth; + changed + } else { + false + } + } + + pub fn set_external_auth_refresher(&self, refresher: Arc) { + if let Ok(mut guard) = self.inner.write() { + guard.external_refresher = Some(refresher); + } + } + + pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { + if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() { + *guard = workspace_id; + } + } + + pub fn forced_chatgpt_workspace_id(&self) -> Option { + self.forced_chatgpt_workspace_id + .read() + .ok() + .and_then(|guard| guard.clone()) + } + + pub fn has_external_auth_refresher(&self) -> bool { + self.inner + .read() + .ok() + .map(|guard| guard.external_refresher.is_some()) + .unwrap_or(false) + } + + pub fn is_external_auth_active(&self) -> bool { + self.auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_external_chatgpt_tokens) + } + + /// Convenience constructor returning an `Arc` wrapper. + pub fn shared( + codex_home: PathBuf, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, + ) -> Arc { + Arc::new(Self::new( + codex_home, + enable_codex_api_key_env, + auth_credentials_store_mode, + )) + } + + pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { + UnauthorizedRecovery::new(Arc::clone(self)) + } + + /// Attempt to refresh the token by first performing a guarded reload. Auth + /// is reloaded from storage only when the account id matches the currently + /// cached account id. If the persisted token differs from the cached token, we + /// can assume that some other instance already refreshed it. If the persisted + /// token is the same as the cached, then ask the token authority to refresh. + pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> { + let auth_before_reload = self.auth_cached(); + let expected_account_id = auth_before_reload + .as_ref() + .and_then(CodexAuth::get_account_id); + + match self.reload_if_account_id_matches(expected_account_id.as_deref()) { + ReloadOutcome::ReloadedChanged => { + tracing::info!("Skipping token refresh because auth changed after guarded reload."); + Ok(()) + } + ReloadOutcome::ReloadedNoChange => self.refresh_token_from_authority().await, + ReloadOutcome::Skipped => { + Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), + ))) + } + } + } + + /// Attempt to refresh the current auth token from the authority that issued + /// the token. On success, reloads the auth state from disk so other components + /// observe refreshed token. If the token refresh fails, returns the error to + /// the caller. + pub async fn refresh_token_from_authority(&self) -> Result<(), RefreshTokenError> { + tracing::info!("Refreshing token"); + + let auth = match self.auth_cached() { + Some(auth) => auth, + None => return Ok(()), + }; + match auth { + CodexAuth::ChatgptAuthTokens(_) => { + self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await + } + CodexAuth::Chatgpt(chatgpt_auth) => { + let token_data = chatgpt_auth.current_token_data().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other( + "Token data is not available.", + )) + })?; + self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token) + .await?; + Ok(()) + } + CodexAuth::ApiKey(_) => Ok(()), + } + } + + /// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true) + /// if a file was removed, Ok(false) if no auth file existed. On success, + /// reloads the in‑memory auth cache so callers immediately observe the + /// unauthenticated state. + pub fn logout(&self) -> std::io::Result { + let removed = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; + // Always reload to clear any cached auth (even if file absent). + self.reload(); + Ok(removed) + } + + pub fn get_api_auth_mode(&self) -> Option { + self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) + } + + pub fn auth_mode(&self) -> Option { + self.auth_cached().as_ref().map(CodexAuth::auth_mode) + } + + async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result { + let chatgpt_auth = match auth { + CodexAuth::Chatgpt(chatgpt_auth) => chatgpt_auth, + _ => return Ok(false), + }; + + let auth_dot_json = match chatgpt_auth.current_auth_json() { + Some(auth_dot_json) => auth_dot_json, + None => return Ok(false), + }; + let tokens = match auth_dot_json.tokens { + Some(tokens) => tokens, + None => return Ok(false), + }; + let last_refresh = match auth_dot_json.last_refresh { + Some(last_refresh) => last_refresh, + None => return Ok(false), + }; + if last_refresh >= Utc::now() - chrono::Duration::days(TOKEN_REFRESH_INTERVAL) { + return Ok(false); + } + self.refresh_and_persist_chatgpt_token(chatgpt_auth, tokens.refresh_token) + .await?; + Ok(true) + } + + async fn refresh_external_auth( + &self, + reason: ExternalAuthRefreshReason, + ) -> Result<(), RefreshTokenError> { + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); + let refresher = match self.inner.read() { + Ok(guard) => guard.external_refresher.clone(), + Err(_) => { + return Err(RefreshTokenError::Transient(std::io::Error::other( + "failed to read external auth state", + ))); + } + }; + + let Some(refresher) = refresher else { + return Err(RefreshTokenError::Transient(std::io::Error::other( + "external auth refresher is not configured", + ))); + }; + + let previous_account_id = self + .auth_cached() + .as_ref() + .and_then(CodexAuth::get_account_id); + let context = ExternalAuthRefreshContext { + reason, + previous_account_id, + }; + + let refreshed = refresher.refresh(context).await?; + if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref() + && refreshed.chatgpt_account_id != expected_workspace_id + { + return Err(RefreshTokenError::Transient(std::io::Error::other( + format!( + "external auth refresh returned workspace {:?}, expected {expected_workspace_id:?}", + refreshed.chatgpt_account_id, + ), + ))); + } + let auth_dot_json = + AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; + save_auth( + &self.codex_home, + &auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + ) + .map_err(RefreshTokenError::Transient)?; + self.reload(); + Ok(()) + } + + // Refreshes ChatGPT OAuth tokens, persists the updated auth state, and + // reloads the in-memory cache so callers immediately observe new tokens. + async fn refresh_and_persist_chatgpt_token( + &self, + auth: &ChatgptAuth, + refresh_token: String, + ) -> Result<(), RefreshTokenError> { + let refresh_response = request_chatgpt_token_refresh(refresh_token, auth.client()).await?; + + persist_tokens( + auth.storage(), + refresh_response.id_token, + refresh_response.access_token, + refresh_response.refresh_token, + ) + .map_err(RefreshTokenError::from)?; + self.reload(); + + Ok(()) + } +} + +#[cfg(test)] +mod tests; diff --git a/codex-rs/core/src/auth/storage.rs b/codex-rs/auth/src/storage.rs similarity index 100% rename from codex-rs/core/src/auth/storage.rs rename to codex-rs/auth/src/storage.rs diff --git a/codex-rs/core/src/auth/storage_tests.rs b/codex-rs/auth/src/storage_tests.rs similarity index 100% rename from codex-rs/core/src/auth/storage_tests.rs rename to codex-rs/auth/src/storage_tests.rs diff --git a/codex-rs/auth/src/tests.rs b/codex-rs/auth/src/tests.rs new file mode 100644 index 0000000000..6cd60377ba --- /dev/null +++ b/codex-rs/auth/src/tests.rs @@ -0,0 +1,285 @@ +use super::*; +use crate::storage::FileAuthStorage; +use crate::storage::get_auth_file; +use crate::token_data::IdTokenInfo; +use crate::token_data::KnownPlan as InternalKnownPlan; +use crate::token_data::PlanType as InternalPlanType; +use base64::Engine; +use codex_protocol::account::PlanType as AccountPlanType; +use pretty_assertions::assert_eq; +use serde::Serialize; +use serde_json::json; +use serial_test::serial; +use tempfile::tempdir; + +#[tokio::test] +async fn refresh_without_id_token() { + let codex_home = tempdir().unwrap(); + let fake_jwt = write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: None, + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let storage = create_auth_storage( + codex_home.path().to_path_buf(), + AuthCredentialsStoreMode::File, + ); + let updated = super::persist_tokens( + &storage, + None, + Some("new-access-token".to_string()), + Some("new-refresh-token".to_string()), + ) + .expect("update_tokens should succeed"); + + let tokens = updated.tokens.expect("tokens should exist"); + assert_eq!(tokens.id_token.raw_jwt, fake_jwt); + assert_eq!(tokens.access_token, "new-access-token"); + assert_eq!(tokens.refresh_token, "new-refresh-token"); +} + +#[test] +fn login_with_api_key_overwrites_existing_auth_json() { + let dir = tempdir().unwrap(); + let auth_path = dir.path().join("auth.json"); + let stale_auth = json!({ + "OPENAI_API_KEY": "sk-old", + "tokens": { + "id_token": "stale.header.payload", + "access_token": "stale-access", + "refresh_token": "stale-refresh", + "account_id": "stale-acc" + } + }); + std::fs::write( + &auth_path, + serde_json::to_string_pretty(&stale_auth).unwrap(), + ) + .unwrap(); + + super::login_with_api_key(dir.path(), "sk-new", AuthCredentialsStoreMode::File) + .expect("login_with_api_key should succeed"); + + let storage = FileAuthStorage::new(dir.path().to_path_buf()); + let auth = storage + .try_read_auth_json(&auth_path) + .expect("auth.json should parse"); + assert_eq!(auth.openai_api_key.as_deref(), Some("sk-new")); + assert!(auth.tokens.is_none(), "tokens should be cleared"); +} + +#[test] +fn missing_auth_json_returns_none() { + let dir = tempdir().unwrap(); + let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File) + .expect("call should succeed"); + assert_eq!(auth, None); +} + +#[tokio::test] +#[serial(codex_api_key)] +async fn pro_account_with_no_api_key_uses_chatgpt_auth() { + let codex_home = tempdir().unwrap(); + let fake_jwt = write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: None, + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) + .unwrap() + .unwrap(); + assert_eq!(None, auth.api_key()); + assert_eq!(AuthMode::Chatgpt, auth.auth_mode()); + assert_eq!(auth.get_chatgpt_user_id().as_deref(), Some("user-12345")); + + let auth_dot_json = auth + .get_current_auth_json() + .expect("AuthDotJson should exist"); + let last_refresh = auth_dot_json + .last_refresh + .expect("last_refresh should be recorded"); + + assert_eq!( + AuthDotJson { + auth_mode: None, + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: Some("user@example.com".to_string()), + chatgpt_plan_type: Some(InternalPlanType::Known(InternalKnownPlan::Pro)), + chatgpt_user_id: Some("user-12345".to_string()), + chatgpt_account_id: None, + raw_jwt: fake_jwt, + }, + access_token: "test-access-token".to_string(), + refresh_token: "test-refresh-token".to_string(), + account_id: None, + }), + last_refresh: Some(last_refresh), + }, + auth_dot_json + ); +} + +#[tokio::test] +#[serial(codex_api_key)] +async fn loads_api_key_from_auth_json() { + let dir = tempdir().unwrap(); + let auth_file = dir.path().join("auth.json"); + std::fs::write( + auth_file, + r#"{"OPENAI_API_KEY":"sk-test-key","tokens":null,"last_refresh":null}"#, + ) + .unwrap(); + + let auth = super::load_auth(dir.path(), false, AuthCredentialsStoreMode::File) + .unwrap() + .unwrap(); + assert_eq!(auth.auth_mode(), AuthMode::ApiKey); + assert_eq!(auth.api_key(), Some("sk-test-key")); + + assert!(auth.get_token_data().is_err()); +} + +#[test] +fn logout_removes_auth_file() -> Result<(), std::io::Error> { + let dir = tempdir()?; + let auth_dot_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-test-key".to_string()), + tokens: None, + last_refresh: None, + }; + super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?; + let auth_file = get_auth_file(dir.path()); + assert!(auth_file.exists()); + assert!(logout(dir.path(), AuthCredentialsStoreMode::File)?); + assert!(!auth_file.exists()); + Ok(()) +} + +struct AuthFileParams { + openai_api_key: Option, + chatgpt_plan_type: Option, + chatgpt_account_id: Option, +} + +fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result { + let auth_file = get_auth_file(codex_home); + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + let header = Header { + alg: "none", + typ: "JWT", + }; + let mut auth_payload = serde_json::json!({ + "chatgpt_user_id": "user-12345", + "user_id": "user-12345", + }); + + if let Some(chatgpt_plan_type) = params.chatgpt_plan_type { + auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type); + } + + if let Some(chatgpt_account_id) = params.chatgpt_account_id { + let org_value = serde_json::Value::String(chatgpt_account_id); + auth_payload["chatgpt_account_id"] = org_value; + } + + let payload = serde_json::json!({ + "email": "user@example.com", + "email_verified": true, + "https://api.openai.com/auth": auth_payload, + }); + let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b); + let header_b64 = b64(&serde_json::to_vec(&header)?); + let payload_b64 = b64(&serde_json::to_vec(&payload)?); + let signature_b64 = b64(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + let auth_json_data = json!({ + "OPENAI_API_KEY": params.openai_api_key, + "tokens": { + "id_token": fake_jwt, + "access_token": "test-access-token", + "refresh_token": "test-refresh-token" + }, + "last_refresh": Utc::now(), + }); + let auth_json = serde_json::to_string_pretty(&auth_json_data)?; + std::fs::write(auth_file, auth_json)?; + Ok(fake_jwt) +} + +#[test] +fn plan_type_maps_known_plan() { + let codex_home = tempdir().unwrap(); + let _jwt = write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: None, + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth available"); + + pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Pro)); +} + +#[test] +fn plan_type_maps_unknown_to_unknown() { + let codex_home = tempdir().unwrap(); + let _jwt = write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("mystery-tier".to_string()), + chatgpt_account_id: None, + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth available"); + + pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown)); +} + +#[test] +fn missing_plan_type_maps_to_unknown() { + let codex_home = tempdir().unwrap(); + let _jwt = write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: None, + chatgpt_account_id: None, + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth available"); + + pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown)); +} diff --git a/codex-rs/auth/src/token_data.rs b/codex-rs/auth/src/token_data.rs new file mode 100644 index 0000000000..4e7c624f1d --- /dev/null +++ b/codex-rs/auth/src/token_data.rs @@ -0,0 +1,179 @@ +use base64::Engine; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; + +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)] +pub struct TokenData { + /// Flat info parsed from the JWT in auth.json. + #[serde( + deserialize_with = "deserialize_id_token", + serialize_with = "serialize_id_token" + )] + pub id_token: IdTokenInfo, + + /// This is a JWT. + pub access_token: String, + + pub refresh_token: String, + + pub account_id: Option, +} + +/// Flat subset of useful claims in id_token from auth.json. +#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] +pub struct IdTokenInfo { + pub email: Option, + /// The ChatGPT subscription plan type + /// (e.g., "free", "plus", "pro", "business", "enterprise", "edu"). + /// (Note: values may vary by backend.) + pub(crate) chatgpt_plan_type: Option, + /// ChatGPT user identifier associated with the token, if present. + pub chatgpt_user_id: Option, + /// Organization/workspace identifier associated with the token, if present. + pub chatgpt_account_id: Option, + pub raw_jwt: String, +} + +impl IdTokenInfo { + pub fn get_chatgpt_plan_type(&self) -> Option { + self.chatgpt_plan_type.as_ref().map(|t| match t { + PlanType::Known(plan) => format!("{plan:?}"), + PlanType::Unknown(s) => s.clone(), + }) + } + + pub fn is_workspace_account(&self) -> bool { + matches!( + self.chatgpt_plan_type, + Some(PlanType::Known( + KnownPlan::Team | KnownPlan::Business | KnownPlan::Enterprise | KnownPlan::Edu + )) + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum PlanType { + Known(KnownPlan), + Unknown(String), +} + +impl PlanType { + pub(crate) fn from_raw_value(raw: &str) -> Self { + match raw.to_ascii_lowercase().as_str() { + "free" => Self::Known(KnownPlan::Free), + "go" => Self::Known(KnownPlan::Go), + "plus" => Self::Known(KnownPlan::Plus), + "pro" => Self::Known(KnownPlan::Pro), + "team" => Self::Known(KnownPlan::Team), + "business" => Self::Known(KnownPlan::Business), + "enterprise" => Self::Known(KnownPlan::Enterprise), + "education" | "edu" => Self::Known(KnownPlan::Edu), + _ => Self::Unknown(raw.to_string()), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum KnownPlan { + Free, + Go, + Plus, + Pro, + Team, + Business, + Enterprise, + Edu, +} + +#[derive(Deserialize)] +struct IdClaims { + #[serde(default)] + email: Option, + #[serde(rename = "https://api.openai.com/profile", default)] + profile: Option, + #[serde(rename = "https://api.openai.com/auth", default)] + auth: Option, +} + +#[derive(Deserialize)] +struct ProfileClaims { + #[serde(default)] + email: Option, +} + +#[derive(Deserialize)] +struct AuthClaims { + #[serde(default)] + chatgpt_plan_type: Option, + #[serde(default)] + chatgpt_user_id: Option, + #[serde(default)] + user_id: Option, + #[serde(default)] + chatgpt_account_id: Option, +} + +#[derive(Debug, Error)] +pub enum IdTokenInfoError { + #[error("invalid ID token format")] + InvalidFormat, + #[error(transparent)] + Base64(#[from] base64::DecodeError), + #[error(transparent)] + Json(#[from] serde_json::Error), +} + +pub fn parse_chatgpt_jwt_claims(jwt: &str) -> Result { + // JWT format: header.payload.signature + let mut parts = jwt.split('.'); + let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) { + (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s), + _ => return Err(IdTokenInfoError::InvalidFormat), + }; + + let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?; + let claims: IdClaims = serde_json::from_slice(&payload_bytes)?; + let email = claims + .email + .or_else(|| claims.profile.and_then(|profile| profile.email)); + + match claims.auth { + Some(auth) => Ok(IdTokenInfo { + email, + raw_jwt: jwt.to_string(), + chatgpt_plan_type: auth.chatgpt_plan_type, + chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id), + chatgpt_account_id: auth.chatgpt_account_id, + }), + None => Ok(IdTokenInfo { + email, + raw_jwt: jwt.to_string(), + chatgpt_plan_type: None, + chatgpt_user_id: None, + chatgpt_account_id: None, + }), + } +} + +fn deserialize_id_token<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let s = String::deserialize(deserializer)?; + parse_chatgpt_jwt_claims(&s).map_err(serde::de::Error::custom) +} + +fn serialize_id_token(id_token: &IdTokenInfo, serializer: S) -> Result +where + S: serde::Serializer, +{ + serializer.serialize_str(&id_token.raw_jwt) +} + +#[cfg(test)] +#[path = "token_data_tests.rs"] +mod tests; diff --git a/codex-rs/auth/src/token_data_tests.rs b/codex-rs/auth/src/token_data_tests.rs new file mode 100644 index 0000000000..a4f594f6fd --- /dev/null +++ b/codex-rs/auth/src/token_data_tests.rs @@ -0,0 +1,110 @@ +use super::*; +use base64::Engine; +use pretty_assertions::assert_eq; +use serde::Serialize; + +#[test] +fn id_token_info_parses_email_and_plan() { + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + let header = Header { + alg: "none", + typ: "JWT", + }; + let payload = serde_json::json!({ + "email": "user@example.com", + "https://api.openai.com/auth": { + "chatgpt_plan_type": "pro" + } + }); + + fn b64url_no_pad(bytes: &[u8]) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes) + } + + let header_b64 = b64url_no_pad(&serde_json::to_vec(&header).unwrap()); + let payload_b64 = b64url_no_pad(&serde_json::to_vec(&payload).unwrap()); + let signature_b64 = b64url_no_pad(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + let info = parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse"); + assert_eq!(info.email.as_deref(), Some("user@example.com")); + assert_eq!(info.get_chatgpt_plan_type().as_deref(), Some("Pro")); +} + +#[test] +fn id_token_info_parses_go_plan() { + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + let header = Header { + alg: "none", + typ: "JWT", + }; + let payload = serde_json::json!({ + "email": "user@example.com", + "https://api.openai.com/auth": { + "chatgpt_plan_type": "go" + } + }); + + fn b64url_no_pad(bytes: &[u8]) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes) + } + + let header_b64 = b64url_no_pad(&serde_json::to_vec(&header).unwrap()); + let payload_b64 = b64url_no_pad(&serde_json::to_vec(&payload).unwrap()); + let signature_b64 = b64url_no_pad(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + let info = parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse"); + assert_eq!(info.email.as_deref(), Some("user@example.com")); + assert_eq!(info.get_chatgpt_plan_type().as_deref(), Some("Go")); +} + +#[test] +fn id_token_info_handles_missing_fields() { + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + let header = Header { + alg: "none", + typ: "JWT", + }; + let payload = serde_json::json!({ "sub": "123" }); + + fn b64url_no_pad(bytes: &[u8]) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes) + } + + let header_b64 = b64url_no_pad(&serde_json::to_vec(&header).unwrap()); + let payload_b64 = b64url_no_pad(&serde_json::to_vec(&payload).unwrap()); + let signature_b64 = b64url_no_pad(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + let info = parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse"); + assert!(info.email.is_none()); + assert!(info.get_chatgpt_plan_type().is_none()); +} + +#[test] +fn workspace_account_detection_matches_workspace_plans() { + let workspace = IdTokenInfo { + chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Business)), + ..IdTokenInfo::default() + }; + assert_eq!(workspace.is_workspace_account(), true); + + let personal = IdTokenInfo { + chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Pro)), + ..IdTokenInfo::default() + }; + assert_eq!(personal.is_workspace_account(), false); +} diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index ef6b8a0132..130bdf2064 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -28,6 +28,7 @@ chardetng = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-api = { workspace = true } +codex-auth = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } @@ -40,7 +41,6 @@ codex-execpolicy = { workspace = true } codex-file-search = { workspace = true } codex-git = { workspace = true } codex-hooks = { workspace = true } -codex-keyring-store = { workspace = true } codex-network-proxy = { workspace = true } codex-otel = { workspace = true } codex-artifacts = { workspace = true } @@ -58,6 +58,7 @@ codex-utils-string = { workspace = true } codex-utils-stream-parser = { workspace = true } codex-windows-sandbox = { package = "codex-windows-sandbox", path = "../windows-sandbox-rs" } csv = { workspace = true } +ctor = { workspace = true } dirs = { workspace = true } dunce = { workspace = true } encoding_rs = { workspace = true } @@ -68,7 +69,6 @@ http = { workspace = true } iana-time-zone = { workspace = true } image = { workspace = true, features = ["jpeg", "png", "webp"] } indexmap = { workspace = true } -keyring = { workspace = true, features = ["crypto-rust"] } libc = { workspace = true } notify = { workspace = true } once_cell = { workspace = true } @@ -87,7 +87,6 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } serde_yaml = { workspace = true } sha1 = { workspace = true } -sha2 = { workspace = true } shlex = { workspace = true } similar = { workspace = true } tempfile = { workspace = true } @@ -118,13 +117,11 @@ wildmatch = { workspace = true } zip = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] -keyring = { workspace = true, features = ["linux-native-async-persistent"] } landlock = { workspace = true } seccompiler = { workspace = true } [target.'cfg(target_os = "macos")'.dependencies] core-foundation = "0.9" -keyring = { workspace = true, features = ["apple-native"] } # Build OpenSSL from source for musl builds. [target.x86_64-unknown-linux-musl.dependencies] @@ -135,16 +132,12 @@ openssl-sys = { workspace = true, features = ["vendored"] } openssl-sys = { workspace = true, features = ["vendored"] } [target.'cfg(target_os = "windows")'.dependencies] -keyring = { workspace = true, features = ["windows-native"] } windows-sys = { version = "0.52", features = [ "Win32_Foundation", "Win32_System_Com", "Win32_UI_Shell", ] } -[target.'cfg(any(target_os = "freebsd", target_os = "openbsd"))'.dependencies] -keyring = { workspace = true, features = ["sync-secret-service"] } - [target.'cfg(unix)'.dependencies] codex-shell-escalation = { workspace = true } @@ -152,13 +145,13 @@ codex-shell-escalation = { workspace = true } assert_cmd = { workspace = true } assert_matches = { workspace = true } codex-arg0 = { workspace = true } +codex-auth = { workspace = true } codex-otel = { workspace = true, features = [ "disable-default-metrics-exporter", ] } codex-test-macros = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { workspace = true } -ctor = { workspace = true } insta = { workspace = true } maplit = { workspace = true } opentelemetry = { workspace = true } diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index 8bb2b23d87..08e93954aa 100644 --- a/codex-rs/core/src/auth.rs +++ b/codex-rs/core/src/auth.rs @@ -1,461 +1,16 @@ -mod storage; +pub use codex_auth::*; -use async_trait::async_trait; -use chrono::Utc; -use reqwest::StatusCode; -use serde::Deserialize; -use serde::Serialize; -#[cfg(test)] -use serial_test::serial; -use std::env; -use std::fmt::Debug; -use std::path::Path; -use std::path::PathBuf; -use std::sync::Arc; -use std::sync::Mutex; -use std::sync::RwLock; - -use codex_app_server_protocol::AuthMode as ApiAuthMode; -use codex_otel::TelemetryAuthMode; -use codex_protocol::config_types::ForcedLoginMethod; - -pub use crate::auth::storage::AuthCredentialsStoreMode; -pub use crate::auth::storage::AuthDotJson; -use crate::auth::storage::AuthStorageBackend; -use crate::auth::storage::create_auth_storage; use crate::config::Config; -use crate::error::RefreshTokenFailedError; -use crate::error::RefreshTokenFailedReason; -use crate::token_data::KnownPlan as InternalKnownPlan; -use crate::token_data::PlanType as InternalPlanType; -use crate::token_data::TokenData; -use crate::token_data::parse_chatgpt_jwt_claims; -use crate::util::try_parse_error_message; -use codex_client::CodexHttpClient; -use codex_protocol::account::PlanType as AccountPlanType; -use serde_json::Value; -use thiserror::Error; +use codex_protocol::config_types::ForcedLoginMethod; +use std::path::Path; -/// Account type for the current user. -/// -/// This is used internally to determine the base URL for generating responses, -/// and to gate ChatGPT-only behaviors like rate limits and available models (as -/// opposed to API key-based auth). -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum AuthMode { - ApiKey, - Chatgpt, -} - -impl From for TelemetryAuthMode { - fn from(mode: AuthMode) -> Self { - match mode { - AuthMode::ApiKey => TelemetryAuthMode::ApiKey, - AuthMode::Chatgpt => TelemetryAuthMode::Chatgpt, - } - } -} - -/// Authentication mechanism used by the current user. -#[derive(Debug, Clone)] -pub enum CodexAuth { - ApiKey(ApiKeyAuth), - Chatgpt(ChatgptAuth), - ChatgptAuthTokens(ChatgptAuthTokens), -} - -#[derive(Debug, Clone)] -pub struct ApiKeyAuth { - api_key: String, -} - -#[derive(Debug, Clone)] -pub struct ChatgptAuth { - state: ChatgptAuthState, - storage: Arc, -} - -#[derive(Debug, Clone)] -pub struct ChatgptAuthTokens { - state: ChatgptAuthState, -} - -#[derive(Debug, Clone)] -struct ChatgptAuthState { - auth_dot_json: Arc>>, - client: CodexHttpClient, -} - -impl PartialEq for CodexAuth { - fn eq(&self, other: &Self) -> bool { - self.api_auth_mode() == other.api_auth_mode() - } -} - -// TODO(pakrym): use token exp field to check for expiration instead -const TOKEN_REFRESH_INTERVAL: i64 = 8; - -const REFRESH_TOKEN_EXPIRED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token has expired. Please log out and sign in again."; -const REFRESH_TOKEN_REUSED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token was already used. Please log out and sign in again."; -const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be refreshed because your refresh token was revoked. Please log out and sign in again."; -const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str = - "Your access token could not be refreshed. Please log out and sign in again."; -const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again."; -const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; -pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE"; - -#[derive(Debug, Error)] -pub enum RefreshTokenError { - #[error("{0}")] - Permanent(#[from] RefreshTokenFailedError), - #[error(transparent)] - Transient(#[from] std::io::Error), -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ExternalAuthTokens { - pub access_token: String, - pub chatgpt_account_id: String, - pub chatgpt_plan_type: Option, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum ExternalAuthRefreshReason { - Unauthorized, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ExternalAuthRefreshContext { - pub reason: ExternalAuthRefreshReason, - pub previous_account_id: Option, -} - -#[async_trait] -pub trait ExternalAuthRefresher: Send + Sync { - async fn refresh( - &self, - context: ExternalAuthRefreshContext, - ) -> std::io::Result; -} - -impl RefreshTokenError { - pub fn failed_reason(&self) -> Option { - match self { - Self::Permanent(error) => Some(error.reason), - Self::Transient(_) => None, - } - } -} - -impl From for std::io::Error { - fn from(err: RefreshTokenError) -> Self { - match err { - RefreshTokenError::Permanent(failed) => std::io::Error::other(failed), - RefreshTokenError::Transient(inner) => inner, - } - } -} - -impl CodexAuth { - fn from_auth_dot_json( - codex_home: &Path, - auth_dot_json: AuthDotJson, - auth_credentials_store_mode: AuthCredentialsStoreMode, - client: CodexHttpClient, - ) -> std::io::Result { - let auth_mode = auth_dot_json.resolved_mode(); - if auth_mode == ApiAuthMode::ApiKey { - let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else { - return Err(std::io::Error::other("API key auth is missing a key.")); - }; - return Ok(CodexAuth::from_api_key_with_client(api_key, client)); - } - - let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); - let state = ChatgptAuthState { - auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, - }; - - match auth_mode { - ApiAuthMode::Chatgpt => { - let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); - Ok(Self::Chatgpt(ChatgptAuth { state, storage })) - } - ApiAuthMode::ChatgptAuthTokens => { - Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) - } - ApiAuthMode::ApiKey => unreachable!("api key mode is handled above"), - } - } - - /// Loads the available auth information from auth storage. - pub fn from_auth_storage( - codex_home: &Path, - auth_credentials_store_mode: AuthCredentialsStoreMode, - ) -> std::io::Result> { - load_auth(codex_home, false, auth_credentials_store_mode) - } - - pub fn auth_mode(&self) -> AuthMode { - match self { - Self::ApiKey(_) => AuthMode::ApiKey, - Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => AuthMode::Chatgpt, - } - } - - pub fn api_auth_mode(&self) -> ApiAuthMode { - match self { - Self::ApiKey(_) => ApiAuthMode::ApiKey, - Self::Chatgpt(_) => ApiAuthMode::Chatgpt, - Self::ChatgptAuthTokens(_) => ApiAuthMode::ChatgptAuthTokens, - } - } - - pub fn is_api_key_auth(&self) -> bool { - self.auth_mode() == AuthMode::ApiKey - } - - pub fn is_chatgpt_auth(&self) -> bool { - self.auth_mode() == AuthMode::Chatgpt - } - - pub fn is_external_chatgpt_tokens(&self) -> bool { - matches!(self, Self::ChatgptAuthTokens(_)) - } - - /// Returns `None` if `auth_mode() != AuthMode::ApiKey`. - pub fn api_key(&self) -> Option<&str> { - match self { - Self::ApiKey(auth) => Some(auth.api_key.as_str()), - Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => None, - } - } - - /// Returns `Err` if `is_chatgpt_auth()` is false. - pub fn get_token_data(&self) -> Result { - let auth_dot_json: Option = self.get_current_auth_json(); - match auth_dot_json { - Some(AuthDotJson { - tokens: Some(tokens), - last_refresh: Some(_), - .. - }) => Ok(tokens), - _ => Err(std::io::Error::other("Token data is not available.")), - } - } - - /// Returns the token string used for bearer authentication. - pub fn get_token(&self) -> Result { - match self { - Self::ApiKey(auth) => Ok(auth.api_key.clone()), - Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => { - let access_token = self.get_token_data()?.access_token; - Ok(access_token) - } - } - } - - /// Returns `None` if `is_chatgpt_auth()` is false. - pub fn get_account_id(&self) -> Option { - self.get_current_token_data().and_then(|t| t.account_id) - } - - /// Returns `None` if `is_chatgpt_auth()` is false. - pub fn get_account_email(&self) -> Option { - self.get_current_token_data().and_then(|t| t.id_token.email) - } - - /// Returns `None` if `is_chatgpt_auth()` is false. - pub fn get_chatgpt_user_id(&self) -> Option { - self.get_current_token_data() - .and_then(|t| t.id_token.chatgpt_user_id) - } - - /// Account-facing plan classification derived from the current token. - /// Returns a high-level `AccountPlanType` (e.g., Free/Plus/Pro/Team/…) - /// mapped from the ID token's internal plan value. Prefer this when you - /// need to make UI or product decisions based on the user's subscription. - /// When ChatGPT auth is active but the token omits the plan claim, report - /// `Unknown` instead of treating the account as invalid. - pub fn account_plan_type(&self) -> Option { - let map_known = |kp: &InternalKnownPlan| match kp { - InternalKnownPlan::Free => AccountPlanType::Free, - InternalKnownPlan::Go => AccountPlanType::Go, - InternalKnownPlan::Plus => AccountPlanType::Plus, - InternalKnownPlan::Pro => AccountPlanType::Pro, - InternalKnownPlan::Team => AccountPlanType::Team, - InternalKnownPlan::Business => AccountPlanType::Business, - InternalKnownPlan::Enterprise => AccountPlanType::Enterprise, - InternalKnownPlan::Edu => AccountPlanType::Edu, - }; - - self.get_current_token_data().map(|t| { - t.id_token - .chatgpt_plan_type - .map(|pt| match pt { - InternalPlanType::Known(k) => map_known(&k), - InternalPlanType::Unknown(_) => AccountPlanType::Unknown, - }) - .unwrap_or(AccountPlanType::Unknown) - }) - } - - /// Returns `None` if `is_chatgpt_auth()` is false. - fn get_current_auth_json(&self) -> Option { - let state = match self { - Self::Chatgpt(auth) => &auth.state, - Self::ChatgptAuthTokens(auth) => &auth.state, - Self::ApiKey(_) => return None, - }; - #[expect(clippy::unwrap_used)] - state.auth_dot_json.lock().unwrap().clone() - } - - /// Returns `None` if `is_chatgpt_auth()` is false. - fn get_current_token_data(&self) -> Option { - self.get_current_auth_json().and_then(|t| t.tokens) - } - - /// Consider this private to integration tests. - pub fn create_dummy_chatgpt_auth_for_testing() -> Self { - let auth_dot_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: Default::default(), - access_token: "Access Token".to_string(), - refresh_token: "test".to_string(), - account_id: Some("account_id".to_string()), - }), - last_refresh: Some(Utc::now()), - }; - - let client = crate::default_client::create_client(); - let state = ChatgptAuthState { - auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, - }; - let storage = create_auth_storage(PathBuf::new(), AuthCredentialsStoreMode::File); - Self::Chatgpt(ChatgptAuth { state, storage }) - } - - fn from_api_key_with_client(api_key: &str, _client: CodexHttpClient) -> Self { - Self::ApiKey(ApiKeyAuth { - api_key: api_key.to_owned(), - }) - } - - pub fn from_api_key(api_key: &str) -> Self { - Self::from_api_key_with_client(api_key, crate::default_client::create_client()) - } -} - -impl ChatgptAuth { - fn current_auth_json(&self) -> Option { - #[expect(clippy::unwrap_used)] - self.state.auth_dot_json.lock().unwrap().clone() - } - - fn current_token_data(&self) -> Option { - self.current_auth_json().and_then(|auth| auth.tokens) - } - - fn storage(&self) -> &Arc { - &self.storage - } - - fn client(&self) -> &CodexHttpClient { - &self.state.client - } -} - -pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; -pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; - -pub fn read_openai_api_key_from_env() -> Option { - env::var(OPENAI_API_KEY_ENV_VAR) - .ok() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -pub fn read_codex_api_key_from_env() -> Option { - env::var(CODEX_API_KEY_ENV_VAR) - .ok() - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -/// Delete the auth.json file inside `codex_home` if it exists. Returns `Ok(true)` -/// if a file was removed, `Ok(false)` if no auth file was present. -pub fn logout( - codex_home: &Path, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result { - let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); - storage.delete() -} - -/// Writes an `auth.json` that contains only the API key. -pub fn login_with_api_key( - codex_home: &Path, - api_key: &str, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result<()> { - let auth_dot_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::ApiKey), - openai_api_key: Some(api_key.to_string()), - tokens: None, - last_refresh: None, - }; - save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) -} - -/// Writes an in-memory auth payload for externally managed ChatGPT tokens. -pub fn login_with_chatgpt_auth_tokens( - codex_home: &Path, - access_token: &str, - chatgpt_account_id: &str, - chatgpt_plan_type: Option<&str>, -) -> std::io::Result<()> { - let auth_dot_json = AuthDotJson::from_external_access_token( - access_token, - chatgpt_account_id, - chatgpt_plan_type, - )?; - save_auth( - codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - ) -} - -/// Persist the provided auth payload using the specified backend. -pub fn save_auth( - codex_home: &Path, - auth: &AuthDotJson, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result<()> { - let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); - storage.save(auth) -} - -/// Load CLI auth data using the configured credential store backend. -/// Returns `None` when no credentials are stored. This function is -/// provided only for tests. Production code should not directly load -/// from the auth.json storage. It should use the AuthManager abstraction -/// instead. -pub fn load_auth_dot_json( - codex_home: &Path, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result> { - let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); - storage.load() +#[ctor::ctor] +fn init_default_auth_client_factory() { + let _ = set_default_client_factory(crate::default_client::create_client); } pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { - let Some(auth) = load_auth( + let Some(auth) = codex_auth::load_auth( &config.codex_home, true, config.cli_auth_credentials_store_mode, @@ -505,7 +60,6 @@ pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { } }; - // workspace is the external identifier for account id. let chatgpt_account_id = token_data.id_token.chatgpt_account_id.as_deref(); if chatgpt_account_id != Some(expected_account_id) { let message = match chatgpt_account_id { @@ -532,8 +86,6 @@ fn logout_with_message( message: String, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result<()> { - // External auth tokens live in the ephemeral store, but persistent auth may still exist - // from earlier logins. Clear both so a forced logout truly removes all active auth. let removal_result = logout_all_stores(codex_home, auth_credentials_store_mode); let error_message = match removal_result { Ok(_) => message, @@ -554,817 +106,6 @@ fn logout_all_stores( Ok(removed_ephemeral || removed_managed) } -fn load_auth( - codex_home: &Path, - enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result> { - let build_auth = |auth_dot_json: AuthDotJson, storage_mode| { - let client = crate::default_client::create_client(); - CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, storage_mode, client) - }; - - // API key via env var takes precedence over any other auth method. - if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() { - let client = crate::default_client::create_client(); - return Ok(Some(CodexAuth::from_api_key_with_client( - api_key.as_str(), - client, - ))); - } - - // External ChatGPT auth tokens live in the in-memory (ephemeral) store. Always check this - // first so external auth takes precedence over any persisted credentials. - let ephemeral_storage = create_auth_storage( - codex_home.to_path_buf(), - AuthCredentialsStoreMode::Ephemeral, - ); - if let Some(auth_dot_json) = ephemeral_storage.load()? { - let auth = build_auth(auth_dot_json, AuthCredentialsStoreMode::Ephemeral)?; - return Ok(Some(auth)); - } - - // If the caller explicitly requested ephemeral auth, there is no persisted fallback. - if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { - return Ok(None); - } - - // Fall back to the configured persistent store (file/keyring/auto) for managed auth. - let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); - let auth_dot_json = match storage.load()? { - Some(auth) => auth, - None => return Ok(None), - }; - - let auth = build_auth(auth_dot_json, auth_credentials_store_mode)?; - Ok(Some(auth)) -} - -// Persist refreshed tokens into auth storage and update last_refresh. -fn persist_tokens( - storage: &Arc, - id_token: Option, - access_token: Option, - refresh_token: Option, -) -> std::io::Result { - let mut auth_dot_json = storage - .load()? - .ok_or(std::io::Error::other("Token data is not available."))?; - - let tokens = auth_dot_json.tokens.get_or_insert_with(TokenData::default); - if let Some(id_token) = id_token { - tokens.id_token = parse_chatgpt_jwt_claims(&id_token).map_err(std::io::Error::other)?; - } - if let Some(access_token) = access_token { - tokens.access_token = access_token; - } - if let Some(refresh_token) = refresh_token { - tokens.refresh_token = refresh_token; - } - auth_dot_json.last_refresh = Some(Utc::now()); - storage.save(&auth_dot_json)?; - Ok(auth_dot_json) -} - -// Requests refreshed ChatGPT OAuth tokens from the auth service using a refresh token. -// The caller is responsible for persisting any returned tokens. -async fn request_chatgpt_token_refresh( - refresh_token: String, - client: &CodexHttpClient, -) -> Result { - let refresh_request = RefreshRequest { - client_id: CLIENT_ID, - grant_type: "refresh_token", - refresh_token, - }; - - let endpoint = refresh_token_endpoint(); - - // Use shared client factory to include standard headers - let response = client - .post(endpoint.as_str()) - .header("Content-Type", "application/json") - .json(&refresh_request) - .send() - .await - .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))?; - - let status = response.status(); - if status.is_success() { - let refresh_response = response - .json::() - .await - .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))?; - Ok(refresh_response) - } else { - let body = response.text().await.unwrap_or_default(); - tracing::error!("Failed to refresh token: {status}: {body}"); - if status == StatusCode::UNAUTHORIZED { - let failed = classify_refresh_token_failure(&body); - Err(RefreshTokenError::Permanent(failed)) - } else { - let message = try_parse_error_message(&body); - Err(RefreshTokenError::Transient(std::io::Error::other( - format!("Failed to refresh token: {status}: {message}"), - ))) - } - } -} - -fn classify_refresh_token_failure(body: &str) -> RefreshTokenFailedError { - let code = extract_refresh_token_error_code(body); - - let normalized_code = code.as_deref().map(str::to_ascii_lowercase); - let reason = match normalized_code.as_deref() { - Some("refresh_token_expired") => RefreshTokenFailedReason::Expired, - Some("refresh_token_reused") => RefreshTokenFailedReason::Exhausted, - Some("refresh_token_invalidated") => RefreshTokenFailedReason::Revoked, - _ => RefreshTokenFailedReason::Other, - }; - - if reason == RefreshTokenFailedReason::Other { - tracing::warn!( - backend_code = normalized_code.as_deref(), - backend_body = body, - "Encountered unknown 401 response while refreshing token" - ); - } - - let message = match reason { - RefreshTokenFailedReason::Expired => REFRESH_TOKEN_EXPIRED_MESSAGE.to_string(), - RefreshTokenFailedReason::Exhausted => REFRESH_TOKEN_REUSED_MESSAGE.to_string(), - RefreshTokenFailedReason::Revoked => REFRESH_TOKEN_INVALIDATED_MESSAGE.to_string(), - RefreshTokenFailedReason::Other => REFRESH_TOKEN_UNKNOWN_MESSAGE.to_string(), - }; - - RefreshTokenFailedError::new(reason, message) -} - -fn extract_refresh_token_error_code(body: &str) -> Option { - if body.trim().is_empty() { - return None; - } - - let Value::Object(map) = serde_json::from_str::(body).ok()? else { - return None; - }; - - if let Some(error_value) = map.get("error") { - match error_value { - Value::Object(obj) => { - if let Some(code) = obj.get("code").and_then(Value::as_str) { - return Some(code.to_string()); - } - } - Value::String(code) => { - return Some(code.to_string()); - } - _ => {} - } - } - - map.get("code").and_then(Value::as_str).map(str::to_string) -} - -#[derive(Serialize)] -struct RefreshRequest { - client_id: &'static str, - grant_type: &'static str, - refresh_token: String, -} - -#[derive(Deserialize, Clone)] -struct RefreshResponse { - id_token: Option, - access_token: Option, - refresh_token: Option, -} - -// Shared constant for token refresh (client id used for oauth token refresh flow) -pub const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; - -fn refresh_token_endpoint() -> String { - std::env::var(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR) - .unwrap_or_else(|_| REFRESH_TOKEN_URL.to_string()) -} - -impl AuthDotJson { - fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { - let mut token_info = - parse_chatgpt_jwt_claims(&external.access_token).map_err(std::io::Error::other)?; - token_info.chatgpt_account_id = Some(external.chatgpt_account_id.clone()); - token_info.chatgpt_plan_type = external - .chatgpt_plan_type - .as_deref() - .map(InternalPlanType::from_raw_value) - .or(token_info.chatgpt_plan_type) - .or(Some(InternalPlanType::Unknown("unknown".to_string()))); - let tokens = TokenData { - id_token: token_info, - access_token: external.access_token.clone(), - refresh_token: String::new(), - account_id: Some(external.chatgpt_account_id.clone()), - }; - - Ok(Self { - auth_mode: Some(ApiAuthMode::ChatgptAuthTokens), - openai_api_key: None, - tokens: Some(tokens), - last_refresh: Some(Utc::now()), - }) - } - - fn from_external_access_token( - access_token: &str, - chatgpt_account_id: &str, - chatgpt_plan_type: Option<&str>, - ) -> std::io::Result { - let external = ExternalAuthTokens { - access_token: access_token.to_string(), - chatgpt_account_id: chatgpt_account_id.to_string(), - chatgpt_plan_type: chatgpt_plan_type.map(str::to_string), - }; - Self::from_external_tokens(&external) - } - - fn resolved_mode(&self) -> ApiAuthMode { - if let Some(mode) = self.auth_mode { - return mode; - } - if self.openai_api_key.is_some() { - return ApiAuthMode::ApiKey; - } - ApiAuthMode::Chatgpt - } - - fn storage_mode( - &self, - auth_credentials_store_mode: AuthCredentialsStoreMode, - ) -> AuthCredentialsStoreMode { - if self.resolved_mode() == ApiAuthMode::ChatgptAuthTokens { - AuthCredentialsStoreMode::Ephemeral - } else { - auth_credentials_store_mode - } - } -} - -/// Internal cached auth state. -#[derive(Clone)] -struct CachedAuth { - auth: Option, - /// Callback used to refresh external auth by asking the parent app for new tokens. - external_refresher: Option>, -} - -impl Debug for CachedAuth { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("CachedAuth") - .field( - "auth_mode", - &self.auth.as_ref().map(CodexAuth::api_auth_mode), - ) - .field( - "external_refresher", - &self.external_refresher.as_ref().map(|_| "present"), - ) - .finish() - } -} - -enum UnauthorizedRecoveryStep { - Reload, - RefreshToken, - ExternalRefresh, - Done, -} - -enum ReloadOutcome { - /// Reload was performed and the cached auth changed - ReloadedChanged, - /// Reload was performed and the cached auth remained the same - ReloadedNoChange, - /// Reload was skipped (missing or mismatched account id) - Skipped, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum UnauthorizedRecoveryMode { - Managed, - External, -} - -// UnauthorizedRecovery is a state machine that handles an attempt to refresh the authentication when requests -// to API fail with 401 status code. -// The client calls next() every time it encounters a 401 error, one time per retry. -// For API key based authentication, we don't do anything and let the error bubble to the user. -// -// For ChatGPT based authentication, we: -// 1. Attempt to reload the auth data from disk. We only reload if the account id matches the one the current process is running as. -// 2. Attempt to refresh the token using OAuth token refresh flow. -// If after both steps the server still responds with 401 we let the error bubble to the user. -// -// For external ChatGPT auth tokens (chatgptAuthTokens), UnauthorizedRecovery does not touch disk or refresh -// tokens locally. Instead it calls the ExternalAuthRefresher (account/chatgptAuthTokens/refresh) to ask the -// parent app for new tokens, stores them in the ephemeral auth store, and retries once. -pub struct UnauthorizedRecovery { - manager: Arc, - step: UnauthorizedRecoveryStep, - expected_account_id: Option, - mode: UnauthorizedRecoveryMode, -} - -impl UnauthorizedRecovery { - fn new(manager: Arc) -> Self { - let cached_auth = manager.auth_cached(); - let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id); - let mode = if cached_auth - .as_ref() - .is_some_and(CodexAuth::is_external_chatgpt_tokens) - { - UnauthorizedRecoveryMode::External - } else { - UnauthorizedRecoveryMode::Managed - }; - let step = match mode { - UnauthorizedRecoveryMode::Managed => UnauthorizedRecoveryStep::Reload, - UnauthorizedRecoveryMode::External => UnauthorizedRecoveryStep::ExternalRefresh, - }; - Self { - manager, - step, - expected_account_id, - mode, - } - } - - pub fn has_next(&self) -> bool { - if !self - .manager - .auth_cached() - .as_ref() - .is_some_and(CodexAuth::is_chatgpt_auth) - { - return false; - } - - if self.mode == UnauthorizedRecoveryMode::External - && !self.manager.has_external_auth_refresher() - { - return false; - } - - !matches!(self.step, UnauthorizedRecoveryStep::Done) - } - - pub async fn next(&mut self) -> Result<(), RefreshTokenError> { - if !self.has_next() { - return Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( - RefreshTokenFailedReason::Other, - "No more recovery steps available.", - ))); - } - - match self.step { - UnauthorizedRecoveryStep::Reload => { - match self - .manager - .reload_if_account_id_matches(self.expected_account_id.as_deref()) - { - ReloadOutcome::ReloadedChanged | ReloadOutcome::ReloadedNoChange => { - self.step = UnauthorizedRecoveryStep::RefreshToken; - } - ReloadOutcome::Skipped => { - self.step = UnauthorizedRecoveryStep::Done; - return Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( - RefreshTokenFailedReason::Other, - REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), - ))); - } - } - } - UnauthorizedRecoveryStep::RefreshToken => { - self.manager.refresh_token_from_authority().await?; - self.step = UnauthorizedRecoveryStep::Done; - } - UnauthorizedRecoveryStep::ExternalRefresh => { - self.manager - .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) - .await?; - self.step = UnauthorizedRecoveryStep::Done; - } - UnauthorizedRecoveryStep::Done => {} - } - Ok(()) - } -} - -/// Central manager providing a single source of truth for auth.json derived -/// authentication data. It loads once (or on preference change) and then -/// hands out cloned `CodexAuth` values so the rest of the program has a -/// consistent snapshot. -/// -/// External modifications to `auth.json` will NOT be observed until -/// `reload()` is called explicitly. This matches the design goal of avoiding -/// different parts of the program seeing inconsistent auth data mid‑run. -#[derive(Debug)] -pub struct AuthManager { - codex_home: PathBuf, - inner: RwLock, - enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - forced_chatgpt_workspace_id: RwLock>, -} - -impl AuthManager { - /// Create a new manager loading the initial auth using the provided - /// preferred auth method. Errors loading auth are swallowed; `auth()` will - /// simply return `None` in that case so callers can treat it as an - /// unauthenticated state. - pub fn new( - codex_home: PathBuf, - enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - ) -> Self { - let managed_auth = load_auth( - &codex_home, - enable_codex_api_key_env, - auth_credentials_store_mode, - ) - .ok() - .flatten(); - Self { - codex_home, - inner: RwLock::new(CachedAuth { - auth: managed_auth, - external_refresher: None, - }), - enable_codex_api_key_env, - auth_credentials_store_mode, - forced_chatgpt_workspace_id: RwLock::new(None), - } - } - - /// Create an AuthManager with a specific CodexAuth, for testing only. - pub(crate) fn from_auth_for_testing(auth: CodexAuth) -> Arc { - let cached = CachedAuth { - auth: Some(auth), - external_refresher: None, - }; - - Arc::new(Self { - codex_home: PathBuf::from("non-existent"), - inner: RwLock::new(cached), - enable_codex_api_key_env: false, - auth_credentials_store_mode: AuthCredentialsStoreMode::File, - forced_chatgpt_workspace_id: RwLock::new(None), - }) - } - - /// Create an AuthManager with a specific CodexAuth and codex home, for testing only. - pub(crate) fn from_auth_for_testing_with_home( - auth: CodexAuth, - codex_home: PathBuf, - ) -> Arc { - let cached = CachedAuth { - auth: Some(auth), - external_refresher: None, - }; - Arc::new(Self { - codex_home, - inner: RwLock::new(cached), - enable_codex_api_key_env: false, - auth_credentials_store_mode: AuthCredentialsStoreMode::File, - forced_chatgpt_workspace_id: RwLock::new(None), - }) - } - - /// Current cached auth (clone) without attempting a refresh. - pub fn auth_cached(&self) -> Option { - self.inner.read().ok().and_then(|c| c.auth.clone()) - } - - /// Current cached auth (clone). May be `None` if not logged in or load failed. - /// Refreshes cached ChatGPT tokens if they are stale before returning. - pub async fn auth(&self) -> Option { - let auth = self.auth_cached()?; - if let Err(err) = self.refresh_if_stale(&auth).await { - tracing::error!("Failed to refresh token: {}", err); - return Some(auth); - } - self.auth_cached() - } - - /// Force a reload of the auth information from auth.json. Returns - /// whether the auth value changed. - pub fn reload(&self) -> bool { - tracing::info!("Reloading auth"); - let new_auth = self.load_auth_from_storage(); - self.set_cached_auth(new_auth) - } - - fn reload_if_account_id_matches(&self, expected_account_id: Option<&str>) -> ReloadOutcome { - let expected_account_id = match expected_account_id { - Some(account_id) => account_id, - None => { - tracing::info!("Skipping auth reload because no account id is available."); - return ReloadOutcome::Skipped; - } - }; - - let new_auth = self.load_auth_from_storage(); - let new_account_id = new_auth.as_ref().and_then(CodexAuth::get_account_id); - - if new_account_id.as_deref() != Some(expected_account_id) { - let found_account_id = new_account_id.as_deref().unwrap_or("unknown"); - tracing::info!( - "Skipping auth reload due to account id mismatch (expected: {expected_account_id}, found: {found_account_id})" - ); - return ReloadOutcome::Skipped; - } - - tracing::info!("Reloading auth for account {expected_account_id}"); - let cached_before_reload = self.auth_cached(); - let auth_changed = - !Self::auths_equal_for_refresh(cached_before_reload.as_ref(), new_auth.as_ref()); - self.set_cached_auth(new_auth); - if auth_changed { - ReloadOutcome::ReloadedChanged - } else { - ReloadOutcome::ReloadedNoChange - } - } - - fn auths_equal_for_refresh(a: Option<&CodexAuth>, b: Option<&CodexAuth>) -> bool { - match (a, b) { - (None, None) => true, - (Some(a), Some(b)) => match (a.api_auth_mode(), b.api_auth_mode()) { - (ApiAuthMode::ApiKey, ApiAuthMode::ApiKey) => a.api_key() == b.api_key(), - (ApiAuthMode::Chatgpt, ApiAuthMode::Chatgpt) - | (ApiAuthMode::ChatgptAuthTokens, ApiAuthMode::ChatgptAuthTokens) => { - a.get_current_auth_json() == b.get_current_auth_json() - } - _ => false, - }, - _ => false, - } - } - - fn auths_equal(a: Option<&CodexAuth>, b: Option<&CodexAuth>) -> bool { - match (a, b) { - (None, None) => true, - (Some(a), Some(b)) => a == b, - _ => false, - } - } - - fn load_auth_from_storage(&self) -> Option { - load_auth( - &self.codex_home, - self.enable_codex_api_key_env, - self.auth_credentials_store_mode, - ) - .ok() - .flatten() - } - - fn set_cached_auth(&self, new_auth: Option) -> bool { - if let Ok(mut guard) = self.inner.write() { - let previous = guard.auth.as_ref(); - let changed = !AuthManager::auths_equal(previous, new_auth.as_ref()); - tracing::info!("Reloaded auth, changed: {changed}"); - guard.auth = new_auth; - changed - } else { - false - } - } - - pub fn set_external_auth_refresher(&self, refresher: Arc) { - if let Ok(mut guard) = self.inner.write() { - guard.external_refresher = Some(refresher); - } - } - - pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { - if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() { - *guard = workspace_id; - } - } - - pub fn forced_chatgpt_workspace_id(&self) -> Option { - self.forced_chatgpt_workspace_id - .read() - .ok() - .and_then(|guard| guard.clone()) - } - - pub fn has_external_auth_refresher(&self) -> bool { - self.inner - .read() - .ok() - .map(|guard| guard.external_refresher.is_some()) - .unwrap_or(false) - } - - pub fn is_external_auth_active(&self) -> bool { - self.auth_cached() - .as_ref() - .is_some_and(CodexAuth::is_external_chatgpt_tokens) - } - - /// Convenience constructor returning an `Arc` wrapper. - pub fn shared( - codex_home: PathBuf, - enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - ) -> Arc { - Arc::new(Self::new( - codex_home, - enable_codex_api_key_env, - auth_credentials_store_mode, - )) - } - - pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { - UnauthorizedRecovery::new(Arc::clone(self)) - } - - /// Attempt to refresh the token by first performing a guarded reload. Auth - /// is reloaded from storage only when the account id matches the currently - /// cached account id. If the persisted token differs from the cached token, we - /// can assume that some other instance already refreshed it. If the persisted - /// token is the same as the cached, then ask the token authority to refresh. - pub async fn refresh_token(&self) -> Result<(), RefreshTokenError> { - let auth_before_reload = self.auth_cached(); - let expected_account_id = auth_before_reload - .as_ref() - .and_then(CodexAuth::get_account_id); - - match self.reload_if_account_id_matches(expected_account_id.as_deref()) { - ReloadOutcome::ReloadedChanged => { - tracing::info!("Skipping token refresh because auth changed after guarded reload."); - Ok(()) - } - ReloadOutcome::ReloadedNoChange => self.refresh_token_from_authority().await, - ReloadOutcome::Skipped => { - Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( - RefreshTokenFailedReason::Other, - REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE.to_string(), - ))) - } - } - } - - /// Attempt to refresh the current auth token from the authority that issued - /// the token. On success, reloads the auth state from disk so other components - /// observe refreshed token. If the token refresh fails, returns the error to - /// the caller. - pub async fn refresh_token_from_authority(&self) -> Result<(), RefreshTokenError> { - tracing::info!("Refreshing token"); - - let auth = match self.auth_cached() { - Some(auth) => auth, - None => return Ok(()), - }; - match auth { - CodexAuth::ChatgptAuthTokens(_) => { - self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) - .await - } - CodexAuth::Chatgpt(chatgpt_auth) => { - let token_data = chatgpt_auth.current_token_data().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other( - "Token data is not available.", - )) - })?; - self.refresh_and_persist_chatgpt_token(&chatgpt_auth, token_data.refresh_token) - .await?; - Ok(()) - } - CodexAuth::ApiKey(_) => Ok(()), - } - } - - /// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true) - /// if a file was removed, Ok(false) if no auth file existed. On success, - /// reloads the in‑memory auth cache so callers immediately observe the - /// unauthenticated state. - pub fn logout(&self) -> std::io::Result { - let removed = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; - // Always reload to clear any cached auth (even if file absent). - self.reload(); - Ok(removed) - } - - pub fn get_api_auth_mode(&self) -> Option { - self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) - } - - pub fn auth_mode(&self) -> Option { - self.auth_cached().as_ref().map(CodexAuth::auth_mode) - } - - async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result { - let chatgpt_auth = match auth { - CodexAuth::Chatgpt(chatgpt_auth) => chatgpt_auth, - _ => return Ok(false), - }; - - let auth_dot_json = match chatgpt_auth.current_auth_json() { - Some(auth_dot_json) => auth_dot_json, - None => return Ok(false), - }; - let tokens = match auth_dot_json.tokens { - Some(tokens) => tokens, - None => return Ok(false), - }; - let last_refresh = match auth_dot_json.last_refresh { - Some(last_refresh) => last_refresh, - None => return Ok(false), - }; - if last_refresh >= Utc::now() - chrono::Duration::days(TOKEN_REFRESH_INTERVAL) { - return Ok(false); - } - self.refresh_and_persist_chatgpt_token(chatgpt_auth, tokens.refresh_token) - .await?; - Ok(true) - } - - async fn refresh_external_auth( - &self, - reason: ExternalAuthRefreshReason, - ) -> Result<(), RefreshTokenError> { - let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); - let refresher = match self.inner.read() { - Ok(guard) => guard.external_refresher.clone(), - Err(_) => { - return Err(RefreshTokenError::Transient(std::io::Error::other( - "failed to read external auth state", - ))); - } - }; - - let Some(refresher) = refresher else { - return Err(RefreshTokenError::Transient(std::io::Error::other( - "external auth refresher is not configured", - ))); - }; - - let previous_account_id = self - .auth_cached() - .as_ref() - .and_then(CodexAuth::get_account_id); - let context = ExternalAuthRefreshContext { - reason, - previous_account_id, - }; - - let refreshed = refresher.refresh(context).await?; - if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref() - && refreshed.chatgpt_account_id != expected_workspace_id - { - return Err(RefreshTokenError::Transient(std::io::Error::other( - format!( - "external auth refresh returned workspace {:?}, expected {expected_workspace_id:?}", - refreshed.chatgpt_account_id, - ), - ))); - } - let auth_dot_json = - AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; - save_auth( - &self.codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - ) - .map_err(RefreshTokenError::Transient)?; - self.reload(); - Ok(()) - } - - // Refreshes ChatGPT OAuth tokens, persists the updated auth state, and - // reloads the in-memory cache so callers immediately observe new tokens. - async fn refresh_and_persist_chatgpt_token( - &self, - auth: &ChatgptAuth, - refresh_token: String, - ) -> Result<(), RefreshTokenError> { - let refresh_response = request_chatgpt_token_refresh(refresh_token, auth.client()).await?; - - persist_tokens( - auth.storage(), - refresh_response.id_token, - refresh_response.access_token, - refresh_response.refresh_token, - ) - .map_err(RefreshTokenError::from)?; - self.reload(); - - Ok(()) - } -} - #[cfg(test)] #[path = "auth_tests.rs"] mod tests; diff --git a/codex-rs/core/src/auth_tests.rs b/codex-rs/core/src/auth_tests.rs index 0c4a574f34..86eae54949 100644 --- a/codex-rs/core/src/auth_tests.rs +++ b/codex-rs/core/src/auth_tests.rs @@ -1,233 +1,12 @@ use super::*; -use crate::auth::storage::FileAuthStorage; -use crate::auth::storage::get_auth_file; use crate::config::Config; use crate::config::ConfigBuilder; -use crate::token_data::IdTokenInfo; -use crate::token_data::KnownPlan as InternalKnownPlan; -use crate::token_data::PlanType as InternalPlanType; -use codex_protocol::account::PlanType as AccountPlanType; - use base64::Engine; -use codex_protocol::config_types::ForcedLoginMethod; -use pretty_assertions::assert_eq; use serde::Serialize; use serde_json::json; +use serial_test::serial; use tempfile::tempdir; -#[tokio::test] -async fn refresh_without_id_token() { - let codex_home = tempdir().unwrap(); - let fake_jwt = write_auth_file( - AuthFileParams { - openai_api_key: None, - chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: None, - }, - codex_home.path(), - ) - .expect("failed to write auth file"); - - let storage = create_auth_storage( - codex_home.path().to_path_buf(), - AuthCredentialsStoreMode::File, - ); - let updated = super::persist_tokens( - &storage, - None, - Some("new-access-token".to_string()), - Some("new-refresh-token".to_string()), - ) - .expect("update_tokens should succeed"); - - let tokens = updated.tokens.expect("tokens should exist"); - assert_eq!(tokens.id_token.raw_jwt, fake_jwt); - assert_eq!(tokens.access_token, "new-access-token"); - assert_eq!(tokens.refresh_token, "new-refresh-token"); -} - -#[test] -fn login_with_api_key_overwrites_existing_auth_json() { - let dir = tempdir().unwrap(); - let auth_path = dir.path().join("auth.json"); - let stale_auth = json!({ - "OPENAI_API_KEY": "sk-old", - "tokens": { - "id_token": "stale.header.payload", - "access_token": "stale-access", - "refresh_token": "stale-refresh", - "account_id": "stale-acc" - } - }); - std::fs::write( - &auth_path, - serde_json::to_string_pretty(&stale_auth).unwrap(), - ) - .unwrap(); - - super::login_with_api_key(dir.path(), "sk-new", AuthCredentialsStoreMode::File) - .expect("login_with_api_key should succeed"); - - let storage = FileAuthStorage::new(dir.path().to_path_buf()); - let auth = storage - .try_read_auth_json(&auth_path) - .expect("auth.json should parse"); - assert_eq!(auth.openai_api_key.as_deref(), Some("sk-new")); - assert!(auth.tokens.is_none(), "tokens should be cleared"); -} - -#[test] -fn missing_auth_json_returns_none() { - let dir = tempdir().unwrap(); - let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File) - .expect("call should succeed"); - assert_eq!(auth, None); -} - -#[tokio::test] -#[serial(codex_api_key)] -async fn pro_account_with_no_api_key_uses_chatgpt_auth() { - let codex_home = tempdir().unwrap(); - let fake_jwt = write_auth_file( - AuthFileParams { - openai_api_key: None, - chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: None, - }, - codex_home.path(), - ) - .expect("failed to write auth file"); - - let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) - .unwrap() - .unwrap(); - assert_eq!(None, auth.api_key()); - assert_eq!(AuthMode::Chatgpt, auth.auth_mode()); - assert_eq!(auth.get_chatgpt_user_id().as_deref(), Some("user-12345")); - - let auth_dot_json = auth - .get_current_auth_json() - .expect("AuthDotJson should exist"); - let last_refresh = auth_dot_json - .last_refresh - .expect("last_refresh should be recorded"); - - assert_eq!( - AuthDotJson { - auth_mode: None, - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: Some("user@example.com".to_string()), - chatgpt_plan_type: Some(InternalPlanType::Known(InternalKnownPlan::Pro)), - chatgpt_user_id: Some("user-12345".to_string()), - chatgpt_account_id: None, - raw_jwt: fake_jwt, - }, - access_token: "test-access-token".to_string(), - refresh_token: "test-refresh-token".to_string(), - account_id: None, - }), - last_refresh: Some(last_refresh), - }, - auth_dot_json - ); -} - -#[tokio::test] -#[serial(codex_api_key)] -async fn loads_api_key_from_auth_json() { - let dir = tempdir().unwrap(); - let auth_file = dir.path().join("auth.json"); - std::fs::write( - auth_file, - r#"{"OPENAI_API_KEY":"sk-test-key","tokens":null,"last_refresh":null}"#, - ) - .unwrap(); - - let auth = super::load_auth(dir.path(), false, AuthCredentialsStoreMode::File) - .unwrap() - .unwrap(); - assert_eq!(auth.auth_mode(), AuthMode::ApiKey); - assert_eq!(auth.api_key(), Some("sk-test-key")); - - assert!(auth.get_token_data().is_err()); -} - -#[test] -fn logout_removes_auth_file() -> Result<(), std::io::Error> { - let dir = tempdir()?; - let auth_dot_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::ApiKey), - openai_api_key: Some("sk-test-key".to_string()), - tokens: None, - last_refresh: None, - }; - super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?; - let auth_file = get_auth_file(dir.path()); - assert!(auth_file.exists()); - assert!(logout(dir.path(), AuthCredentialsStoreMode::File)?); - assert!(!auth_file.exists()); - Ok(()) -} - -struct AuthFileParams { - openai_api_key: Option, - chatgpt_plan_type: Option, - chatgpt_account_id: Option, -} - -fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result { - let auth_file = get_auth_file(codex_home); - // Create a minimal valid JWT for the id_token field. - #[derive(Serialize)] - struct Header { - alg: &'static str, - typ: &'static str, - } - let header = Header { - alg: "none", - typ: "JWT", - }; - let mut auth_payload = serde_json::json!({ - "chatgpt_user_id": "user-12345", - "user_id": "user-12345", - }); - - if let Some(chatgpt_plan_type) = params.chatgpt_plan_type { - auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type); - } - - if let Some(chatgpt_account_id) = params.chatgpt_account_id { - let org_value = serde_json::Value::String(chatgpt_account_id); - auth_payload["chatgpt_account_id"] = org_value; - } - - let payload = serde_json::json!({ - "email": "user@example.com", - "email_verified": true, - "https://api.openai.com/auth": auth_payload, - }); - let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b); - let header_b64 = b64(&serde_json::to_vec(&header)?); - let payload_b64 = b64(&serde_json::to_vec(&payload)?); - let signature_b64 = b64(b"sig"); - let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - - let auth_json_data = json!({ - "OPENAI_API_KEY": params.openai_api_key, - "tokens": { - "id_token": fake_jwt, - "access_token": "test-access-token", - "refresh_token": "test-refresh-token" - }, - "last_refresh": Utc::now(), - }); - let auth_json = serde_json::to_string_pretty(&auth_json_data)?; - std::fs::write(auth_file, auth_json)?; - Ok(fake_jwt) -} - async fn build_config( codex_home: &Path, forced_login_method: Option, @@ -243,40 +22,106 @@ async fn build_config( config } -/// Use sparingly. -/// TODO (gpeal): replace this with an injectable env var provider. -#[cfg(test)] struct EnvVarGuard { key: &'static str, original: Option, } -#[cfg(test)] impl EnvVarGuard { fn set(key: &'static str, value: &str) -> Self { - let original = env::var_os(key); + let original = std::env::var_os(key); unsafe { - env::set_var(key, value); + std::env::set_var(key, value); } Self { key, original } } } -#[cfg(test)] impl Drop for EnvVarGuard { fn drop(&mut self) { unsafe { match &self.original { - Some(value) => env::set_var(self.key, value), - None => env::remove_var(self.key), + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), } } } } +struct ResidencyRequirementGuard; + +impl ResidencyRequirementGuard { + fn set(requirement: Option) -> Self { + crate::default_client::set_default_client_residency_requirement(requirement); + Self + } +} + +impl Drop for ResidencyRequirementGuard { + fn drop(&mut self) { + crate::default_client::set_default_client_residency_requirement(None); + } +} + +struct AuthFileParams { + openai_api_key: Option, + chatgpt_plan_type: Option, + chatgpt_account_id: Option, +} + +fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result { + let auth_file = codex_home.join("auth.json"); + + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + + let header = Header { + alg: "none", + typ: "JWT", + }; + let mut auth_payload = serde_json::json!({ + "chatgpt_user_id": "user-12345", + "user_id": "user-12345", + }); + + if let Some(chatgpt_plan_type) = params.chatgpt_plan_type { + auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type); + } + + if let Some(chatgpt_account_id) = params.chatgpt_account_id { + auth_payload["chatgpt_account_id"] = serde_json::Value::String(chatgpt_account_id); + } + + let payload = serde_json::json!({ + "email": "user@example.com", + "email_verified": true, + "https://api.openai.com/auth": auth_payload, + }); + let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); + let header_b64 = encode(&serde_json::to_vec(&header)?); + let payload_b64 = encode(&serde_json::to_vec(&payload)?); + let signature_b64 = encode(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + let auth_json_data = json!({ + "OPENAI_API_KEY": params.openai_api_key, + "tokens": { + "id_token": fake_jwt, + "access_token": "test-access-token", + "refresh_token": "test-refresh-token" + }, + "last_refresh": chrono::Utc::now(), + }); + std::fs::write(auth_file, serde_json::to_string_pretty(&auth_json_data)?)?; + Ok(fake_jwt) +} + #[tokio::test] async fn enforce_login_restrictions_logs_out_for_method_mismatch() { - let codex_home = tempdir().unwrap(); + let codex_home = tempdir().expect("tempdir"); login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File) .expect("seed api key"); @@ -294,7 +139,7 @@ async fn enforce_login_restrictions_logs_out_for_method_mismatch() { #[tokio::test] #[serial(codex_api_key)] async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { - let codex_home = tempdir().unwrap(); + let codex_home = tempdir().expect("tempdir"); let _jwt = write_auth_file( AuthFileParams { openai_api_key: None, @@ -319,7 +164,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { #[tokio::test] #[serial(codex_api_key)] async fn enforce_login_restrictions_allows_matching_workspace() { - let codex_home = tempdir().unwrap(); + let codex_home = tempdir().expect("tempdir"); let _jwt = write_auth_file( AuthFileParams { openai_api_key: None, @@ -342,7 +187,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() { #[tokio::test] async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_forced_chatgpt_workspace_id_is_set() { - let codex_home = tempdir().unwrap(); + let codex_home = tempdir().expect("tempdir"); login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File) .expect("seed api key"); @@ -359,7 +204,7 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f #[serial(codex_api_key)] async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() { let _guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env"); - let codex_home = tempdir().unwrap(); + let codex_home = tempdir().expect("tempdir"); let config = build_config(codex_home.path(), Some(ForcedLoginMethod::Chatgpt), None).await; @@ -371,10 +216,18 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() { ); } -#[test] -fn plan_type_maps_known_plan() { - let codex_home = tempdir().unwrap(); - let _jwt = write_auth_file( +#[tokio::test] +#[serial(codex_auth_refresh_env)] +async fn auth_refresh_uses_core_default_http_client_factory() { + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + let codex_home = tempdir().expect("tempdir"); + write_auth_file( AuthFileParams { openai_api_key: None, chatgpt_plan_type: Some("pro".to_string()), @@ -384,49 +237,42 @@ fn plan_type_maps_known_plan() { ) .expect("failed to write auth file"); - let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) - .expect("load auth") - .expect("auth available"); + let server = MockServer::start().await; + let expected_originator = crate::default_client::originator().value; + let expected_user_agent = crate::default_client::get_codex_user_agent(); + let _refresh_url_guard = EnvVarGuard::set( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + &format!("{}/oauth/token", server.uri()), + ); + let _residency_guard = + ResidencyRequirementGuard::set(Some(crate::config_loader::ResidencyRequirement::Us)); - pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Pro)); -} - -#[test] -fn plan_type_maps_unknown_to_unknown() { - let codex_home = tempdir().unwrap(); - let _jwt = write_auth_file( - AuthFileParams { - openai_api_key: None, - chatgpt_plan_type: Some("mystery-tier".to_string()), - chatgpt_account_id: None, - }, - codex_home.path(), - ) - .expect("failed to write auth file"); - - let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) - .expect("load auth") - .expect("auth available"); - - pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown)); -} - -#[test] -fn missing_plan_type_maps_to_unknown() { - let codex_home = tempdir().unwrap(); - let _jwt = write_auth_file( - AuthFileParams { - openai_api_key: None, - chatgpt_plan_type: None, - chatgpt_account_id: None, - }, - codex_home.path(), - ) - .expect("failed to write auth file"); - - let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) - .expect("load auth") - .expect("auth available"); - - pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown)); + Mock::given(method("POST")) + .and(path("/oauth/token")) + .and(header("originator", expected_originator.as_str())) + .and(header("user-agent", expected_user_agent.as_str())) + .and(header(crate::default_client::RESIDENCY_HEADER_NAME, "us")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "new-access-token", + "refresh_token": "new-refresh-token" + }))) + .mount(&server) + .await; + + let auth_manager = AuthManager::shared( + codex_home.path().to_path_buf(), + false, + AuthCredentialsStoreMode::File, + ); + auth_manager + .refresh_token_from_authority() + .await + .expect("refresh should succeed"); + + let refreshed = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File) + .expect("load auth.json") + .expect("stored auth"); + let tokens = refreshed.tokens.expect("token data"); + assert_eq!(tokens.access_token, "new-access-token"); + assert_eq!(tokens.refresh_token, "new-refresh-token"); } diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index f3bb4dc8e5..6e2da2767a 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -21,6 +21,8 @@ use thiserror::Error; use tokio::task::JoinError; pub type Result = std::result::Result; +pub use codex_auth::RefreshTokenFailedError; +pub use codex_auth::RefreshTokenFailedReason; /// Limit UI error messages to a reasonable size while keeping useful context. const ERROR_MESSAGE_UI_MAX_BYTES: usize = 2 * 1024; // 2 KiB @@ -261,30 +263,6 @@ impl std::fmt::Display for ResponseStreamFailed { } } -#[derive(Debug, Clone, PartialEq, Eq, Error)] -#[error("{message}")] -pub struct RefreshTokenFailedError { - pub reason: RefreshTokenFailedReason, - pub message: String, -} - -impl RefreshTokenFailedError { - pub fn new(reason: RefreshTokenFailedReason, message: impl Into) -> Self { - Self { - reason, - message: message.into(), - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum RefreshTokenFailedReason { - Expired, - Exhausted, - Revoked, - Other, -} - #[derive(Debug)] pub struct UnexpectedResponseError { pub status: StatusCode, @@ -655,5 +633,492 @@ pub fn get_error_message_ui(e: &CodexErr) -> String { } #[cfg(test)] -#[path = "error_tests.rs"] -mod tests; +mod tests { + use super::*; + use crate::exec::StreamOutput; + use chrono::DateTime; + use chrono::Duration as ChronoDuration; + use chrono::TimeZone; + use chrono::Utc; + use codex_protocol::protocol::RateLimitWindow; + use pretty_assertions::assert_eq; + use reqwest::Response; + use reqwest::ResponseBuilderExt; + use reqwest::StatusCode; + use reqwest::Url; + + fn rate_limit_snapshot() -> RateLimitSnapshot { + let primary_reset_at = Utc + .with_ymd_and_hms(2024, 1, 1, 1, 0, 0) + .unwrap() + .timestamp(); + let secondary_reset_at = Utc + .with_ymd_and_hms(2024, 1, 1, 2, 0, 0) + .unwrap() + .timestamp(); + RateLimitSnapshot { + limit_id: None, + limit_name: None, + primary: Some(RateLimitWindow { + used_percent: 50.0, + window_minutes: Some(60), + resets_at: Some(primary_reset_at), + }), + secondary: Some(RateLimitWindow { + used_percent: 30.0, + window_minutes: Some(120), + resets_at: Some(secondary_reset_at), + }), + credits: None, + plan_type: None, + } + } + + fn with_now_override(now: DateTime, f: impl FnOnce() -> T) -> T { + NOW_OVERRIDE.with(|cell| { + *cell.borrow_mut() = Some(now); + let result = f(); + *cell.borrow_mut() = None; + result + }) + } + + #[test] + fn usage_limit_reached_error_formats_plus_plan() { + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Plus)), + resets_at: None, + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + assert_eq!( + err.to_string(), + "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again later." + ); + } + + #[test] + fn server_overloaded_maps_to_protocol() { + let err = CodexErr::ServerOverloaded; + assert_eq!( + err.to_codex_protocol_error(), + CodexErrorInfo::ServerOverloaded + ); + } + + #[test] + fn sandbox_denied_uses_aggregated_output_when_stderr_empty() { + let output = ExecToolCallOutput { + exit_code: 77, + stdout: StreamOutput::new(String::new()), + stderr: StreamOutput::new(String::new()), + aggregated_output: StreamOutput::new("aggregate detail".to_string()), + duration: Duration::from_millis(10), + timed_out: false, + }; + let err = CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + }); + assert_eq!(get_error_message_ui(&err), "aggregate detail"); + } + + #[test] + fn sandbox_denied_reports_both_streams_when_available() { + let output = ExecToolCallOutput { + exit_code: 9, + stdout: StreamOutput::new("stdout detail".to_string()), + stderr: StreamOutput::new("stderr detail".to_string()), + aggregated_output: StreamOutput::new(String::new()), + duration: Duration::from_millis(10), + timed_out: false, + }; + let err = CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + }); + assert_eq!(get_error_message_ui(&err), "stderr detail\nstdout detail"); + } + + #[test] + fn sandbox_denied_reports_stdout_when_no_stderr() { + let output = ExecToolCallOutput { + exit_code: 11, + stdout: StreamOutput::new("stdout only".to_string()), + stderr: StreamOutput::new(String::new()), + aggregated_output: StreamOutput::new(String::new()), + duration: Duration::from_millis(8), + timed_out: false, + }; + let err = CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + }); + assert_eq!(get_error_message_ui(&err), "stdout only"); + } + + #[test] + fn to_error_event_handles_response_stream_failed() { + let response = http::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .url(Url::parse("http://example.com").unwrap()) + .body("") + .unwrap(); + let source = Response::from(response).error_for_status_ref().unwrap_err(); + let err = CodexErr::ResponseStreamFailed(ResponseStreamFailed { + source, + request_id: Some("req-123".to_string()), + }); + + let event = err.to_error_event(Some("prefix".to_string())); + + assert_eq!( + event.message, + "prefix: Error while reading the server response: HTTP status client error (429 Too Many Requests) for url (http://example.com/), request id: req-123" + ); + assert_eq!( + event.codex_error_info, + Some(CodexErrorInfo::ResponseStreamConnectionFailed { + http_status_code: Some(429) + }) + ); + } + + #[test] + fn sandbox_denied_reports_exit_code_when_no_output_available() { + let output = ExecToolCallOutput { + exit_code: 13, + stdout: StreamOutput::new(String::new()), + stderr: StreamOutput::new(String::new()), + aggregated_output: StreamOutput::new(String::new()), + duration: Duration::from_millis(5), + timed_out: false, + }; + let err = CodexErr::Sandbox(SandboxErr::Denied { + output: Box::new(output), + network_policy_decision: None, + }); + assert_eq!( + get_error_message_ui(&err), + "command failed inside sandbox with exit code 13" + ); + } + + #[test] + fn usage_limit_reached_error_formats_free_plan() { + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Free)), + resets_at: None, + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + assert_eq!( + err.to_string(), + "You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus), or try again later." + ); + } + + #[test] + fn usage_limit_reached_error_formats_go_plan() { + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Go)), + resets_at: None, + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + assert_eq!( + err.to_string(), + "You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus), or try again later." + ); + } + + #[test] + fn usage_limit_reached_error_formats_default_when_none() { + let err = UsageLimitReachedError { + plan_type: None, + resets_at: None, + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + assert_eq!( + err.to_string(), + "You've hit your usage limit. Try again later." + ); + } + + #[test] + fn usage_limit_reached_error_formats_team_plan() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::hours(1); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Team)), + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + let expected = format!( + "You've hit your usage limit. To get more access now, send a request to your admin or try again at {expected_time}." + ); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn usage_limit_reached_error_formats_business_plan_without_reset() { + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Business)), + resets_at: None, + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + assert_eq!( + err.to_string(), + "You've hit your usage limit. To get more access now, send a request to your admin or try again later." + ); + } + + #[test] + fn usage_limit_reached_error_formats_default_for_other_plans() { + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Enterprise)), + resets_at: None, + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + assert_eq!( + err.to_string(), + "You've hit your usage limit. Try again later." + ); + } + + #[test] + fn usage_limit_reached_error_formats_pro_plan_with_reset() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::hours(1); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Pro)), + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + let expected = format!( + "You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at {expected_time}." + ); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn usage_limit_reached_error_hides_upsell_for_non_codex_limit_name() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::hours(1); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Plus)), + resets_at: Some(resets_at), + rate_limits: Some(Box::new(RateLimitSnapshot { + limit_id: Some("codex_other".to_string()), + limit_name: Some("codex_other".to_string()), + ..rate_limit_snapshot() + })), + promo_message: Some( + "Visit https://chatgpt.com/codex/settings/usage to purchase more credits" + .to_string(), + ), + }; + let expected = format!( + "You've hit your usage limit for codex_other. Switch to another model now, or try again at {expected_time}." + ); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn usage_limit_reached_includes_minutes_when_available() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::minutes(5); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: None, + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + let expected = format!("You've hit your usage limit. Try again at {expected_time}."); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn unexpected_status_cloudflare_html_is_simplified() { + let err = UnexpectedResponseError { + status: StatusCode::FORBIDDEN, + body: "Cloudflare error: Sorry, you have been blocked" + .to_string(), + url: Some("http://example.com/blocked".to_string()), + cf_ray: Some("ray-id".to_string()), + request_id: None, + }; + let status = StatusCode::FORBIDDEN.to_string(); + let url = "http://example.com/blocked"; + assert_eq!( + err.to_string(), + format!("{CLOUDFLARE_BLOCKED_MESSAGE} (status {status}), url: {url}, cf-ray: ray-id") + ); + } + + #[test] + fn unexpected_status_non_html_is_unchanged() { + let err = UnexpectedResponseError { + status: StatusCode::FORBIDDEN, + body: "plain text error".to_string(), + url: Some("http://example.com/plain".to_string()), + cf_ray: None, + request_id: None, + }; + let status = StatusCode::FORBIDDEN.to_string(); + let url = "http://example.com/plain"; + assert_eq!( + err.to_string(), + format!("unexpected status {status}: plain text error, url: {url}") + ); + } + + #[test] + fn unexpected_status_prefers_error_message_when_present() { + let err = UnexpectedResponseError { + status: StatusCode::UNAUTHORIZED, + body: r#"{"error":{"message":"Workspace is not authorized in this region."},"status":401}"# + .to_string(), + url: Some("https://chatgpt.com/backend-api/codex/responses".to_string()), + cf_ray: None, + request_id: Some("req-123".to_string()), + }; + let status = StatusCode::UNAUTHORIZED.to_string(); + assert_eq!( + err.to_string(), + format!( + "unexpected status {status}: Workspace is not authorized in this region., url: https://chatgpt.com/backend-api/codex/responses, request id: req-123" + ) + ); + } + + #[test] + fn unexpected_status_truncates_long_body_with_ellipsis() { + let long_body = "x".repeat(UNEXPECTED_RESPONSE_BODY_MAX_BYTES + 10); + let err = UnexpectedResponseError { + status: StatusCode::BAD_GATEWAY, + body: long_body, + url: Some("http://example.com/long".to_string()), + cf_ray: None, + request_id: Some("req-long".to_string()), + }; + let status = StatusCode::BAD_GATEWAY.to_string(); + let expected_body = format!("{}...", "x".repeat(UNEXPECTED_RESPONSE_BODY_MAX_BYTES)); + assert_eq!( + err.to_string(), + format!( + "unexpected status {status}: {expected_body}, url: http://example.com/long, request id: req-long" + ) + ); + } + + #[test] + fn unexpected_status_includes_cf_ray_and_request_id() { + let err = UnexpectedResponseError { + status: StatusCode::UNAUTHORIZED, + body: "plain text error".to_string(), + url: Some("https://chatgpt.com/backend-api/codex/responses".to_string()), + cf_ray: Some("9c81f9f18f2fa49d-LHR".to_string()), + request_id: Some("req-xyz".to_string()), + }; + let status = StatusCode::UNAUTHORIZED.to_string(); + assert_eq!( + err.to_string(), + format!( + "unexpected status {status}: plain text error, url: https://chatgpt.com/backend-api/codex/responses, cf-ray: 9c81f9f18f2fa49d-LHR, request id: req-xyz" + ) + ); + } + + #[test] + fn usage_limit_reached_includes_hours_and_minutes() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::hours(3) + ChronoDuration::minutes(32); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: Some(PlanType::Known(KnownPlan::Plus)), + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + let expected = format!( + "You've hit your usage limit. Upgrade to Pro (https://chatgpt.com/explore/pro), visit https://chatgpt.com/codex/settings/usage to purchase more credits or try again at {expected_time}." + ); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn usage_limit_reached_includes_days_hours_minutes() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = + base + ChronoDuration::days(2) + ChronoDuration::hours(3) + ChronoDuration::minutes(5); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: None, + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + let expected = format!("You've hit your usage limit. Try again at {expected_time}."); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn usage_limit_reached_less_than_minute() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::seconds(30); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: None, + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: None, + }; + let expected = format!("You've hit your usage limit. Try again at {expected_time}."); + assert_eq!(err.to_string(), expected); + }); + } + + #[test] + fn usage_limit_reached_with_promo_message() { + let base = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap(); + let resets_at = base + ChronoDuration::seconds(30); + with_now_override(base, move || { + let expected_time = format_retry_timestamp(&resets_at); + let err = UsageLimitReachedError { + plan_type: None, + resets_at: Some(resets_at), + rate_limits: Some(Box::new(rate_limit_snapshot())), + promo_message: Some( + "To continue using Codex, start a free trial of today".to_string(), + ), + }; + let expected = format!( + "You've hit your usage limit. To continue using Codex, start a free trial of today, or try again at {expected_time}." + ); + assert_eq!(err.to_string(), expected); + }); + } +} diff --git a/codex-rs/core/src/token_data.rs b/codex-rs/core/src/token_data.rs index 5952d5940d..1ecaf48d38 100644 --- a/codex-rs/core/src/token_data.rs +++ b/codex-rs/core/src/token_data.rs @@ -1,178 +1,4 @@ -use base64::Engine; -use serde::Deserialize; -use serde::Serialize; -use thiserror::Error; - -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)] -pub struct TokenData { - /// Flat info parsed from the JWT in auth.json. - #[serde( - deserialize_with = "deserialize_id_token", - serialize_with = "serialize_id_token" - )] - pub id_token: IdTokenInfo, - - /// This is a JWT. - pub access_token: String, - - pub refresh_token: String, - - pub account_id: Option, -} - -/// Flat subset of useful claims in id_token from auth.json. -#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)] -pub struct IdTokenInfo { - pub email: Option, - /// The ChatGPT subscription plan type - /// (e.g., "free", "plus", "pro", "business", "enterprise", "edu"). - /// (Note: values may vary by backend.) - pub(crate) chatgpt_plan_type: Option, - /// ChatGPT user identifier associated with the token, if present. - pub chatgpt_user_id: Option, - /// Organization/workspace identifier associated with the token, if present. - pub chatgpt_account_id: Option, - pub raw_jwt: String, -} - -impl IdTokenInfo { - pub fn get_chatgpt_plan_type(&self) -> Option { - self.chatgpt_plan_type.as_ref().map(|t| match t { - PlanType::Known(plan) => format!("{plan:?}"), - PlanType::Unknown(s) => s.clone(), - }) - } - - pub fn is_workspace_account(&self) -> bool { - matches!( - self.chatgpt_plan_type, - Some(PlanType::Known( - KnownPlan::Team | KnownPlan::Business | KnownPlan::Enterprise | KnownPlan::Edu - )) - ) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum PlanType { - Known(KnownPlan), - Unknown(String), -} - -impl PlanType { - pub(crate) fn from_raw_value(raw: &str) -> Self { - match raw.to_ascii_lowercase().as_str() { - "free" => Self::Known(KnownPlan::Free), - "go" => Self::Known(KnownPlan::Go), - "plus" => Self::Known(KnownPlan::Plus), - "pro" => Self::Known(KnownPlan::Pro), - "team" => Self::Known(KnownPlan::Team), - "business" => Self::Known(KnownPlan::Business), - "enterprise" => Self::Known(KnownPlan::Enterprise), - "education" | "edu" => Self::Known(KnownPlan::Edu), - _ => Self::Unknown(raw.to_string()), - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub(crate) enum KnownPlan { - Free, - Go, - Plus, - Pro, - Team, - Business, - Enterprise, - Edu, -} - -#[derive(Deserialize)] -struct IdClaims { - #[serde(default)] - email: Option, - #[serde(rename = "https://api.openai.com/profile", default)] - profile: Option, - #[serde(rename = "https://api.openai.com/auth", default)] - auth: Option, -} - -#[derive(Deserialize)] -struct ProfileClaims { - #[serde(default)] - email: Option, -} - -#[derive(Deserialize)] -struct AuthClaims { - #[serde(default)] - chatgpt_plan_type: Option, - #[serde(default)] - chatgpt_user_id: Option, - #[serde(default)] - user_id: Option, - #[serde(default)] - chatgpt_account_id: Option, -} - -#[derive(Debug, Error)] -pub enum IdTokenInfoError { - #[error("invalid ID token format")] - InvalidFormat, - #[error(transparent)] - Base64(#[from] base64::DecodeError), - #[error(transparent)] - Json(#[from] serde_json::Error), -} - -pub fn parse_chatgpt_jwt_claims(jwt: &str) -> Result { - // JWT format: header.payload.signature - let mut parts = jwt.split('.'); - let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) { - (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s), - _ => return Err(IdTokenInfoError::InvalidFormat), - }; - - let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?; - let claims: IdClaims = serde_json::from_slice(&payload_bytes)?; - let email = claims - .email - .or_else(|| claims.profile.and_then(|profile| profile.email)); - - match claims.auth { - Some(auth) => Ok(IdTokenInfo { - email, - raw_jwt: jwt.to_string(), - chatgpt_plan_type: auth.chatgpt_plan_type, - chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id), - chatgpt_account_id: auth.chatgpt_account_id, - }), - None => Ok(IdTokenInfo { - email, - raw_jwt: jwt.to_string(), - chatgpt_plan_type: None, - chatgpt_user_id: None, - chatgpt_account_id: None, - }), - } -} - -fn deserialize_id_token<'de, D>(deserializer: D) -> Result -where - D: serde::Deserializer<'de>, -{ - let s = String::deserialize(deserializer)?; - parse_chatgpt_jwt_claims(&s).map_err(serde::de::Error::custom) -} - -fn serialize_id_token(id_token: &IdTokenInfo, serializer: S) -> Result -where - S: serde::Serializer, -{ - serializer.serialize_str(&id_token.raw_jwt) -} +pub use codex_auth::token_data::*; #[cfg(test)] #[path = "token_data_tests.rs"] diff --git a/codex-rs/core/src/token_data_tests.rs b/codex-rs/core/src/token_data_tests.rs index e599379c18..fcfd4d71ad 100644 --- a/codex-rs/core/src/token_data_tests.rs +++ b/codex-rs/core/src/token_data_tests.rs @@ -1,4 +1,5 @@ use super::*; +use base64::Engine; use pretty_assertions::assert_eq; use serde::Serialize; @@ -95,15 +96,38 @@ fn id_token_info_handles_missing_fields() { #[test] fn workspace_account_detection_matches_workspace_plans() { - let workspace = IdTokenInfo { - chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Business)), - ..IdTokenInfo::default() - }; + #[derive(Serialize)] + struct Header { + alg: &'static str, + typ: &'static str, + } + + fn id_token_info_with_plan(plan: &str) -> IdTokenInfo { + let header = Header { + alg: "none", + typ: "JWT", + }; + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_plan_type": plan + } + }); + + fn b64url_no_pad(bytes: &[u8]) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes) + } + + let header_b64 = b64url_no_pad(&serde_json::to_vec(&header).unwrap()); + let payload_b64 = b64url_no_pad(&serde_json::to_vec(&payload).unwrap()); + let signature_b64 = b64url_no_pad(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + parse_chatgpt_jwt_claims(&fake_jwt).expect("should parse") + } + + let workspace = id_token_info_with_plan("business"); assert_eq!(workspace.is_workspace_account(), true); - let personal = IdTokenInfo { - chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Pro)), - ..IdTokenInfo::default() - }; + let personal = id_token_info_with_plan("pro"); assert_eq!(personal.is_workspace_account(), false); } diff --git a/codex-rs/core/src/util.rs b/codex-rs/core/src/util.rs index 62e872ae6c..084ee3377b 100644 --- a/codex-rs/core/src/util.rs +++ b/codex-rs/core/src/util.rs @@ -4,7 +4,6 @@ use std::time::Duration; use codex_protocol::ThreadId; use rand::Rng; -use tracing::debug; use tracing::error; use crate::parse_command::shlex_join; @@ -52,21 +51,6 @@ pub(crate) fn error_or_panic(message: impl std::string::ToString) { } } -pub(crate) fn try_parse_error_message(text: &str) -> String { - debug!("Parsing server error response: {}", text); - let json = serde_json::from_str::(text).unwrap_or_default(); - if let Some(error) = json.get("error") - && let Some(message) = error.get("message") - && let Some(message_str) = message.as_str() - { - return message_str.to_string(); - } - if text.is_empty() { - return "Unknown error".to_string(); - } - text.to_string() -} - pub fn resolve_path(base: &Path, path: &PathBuf) -> PathBuf { if path.is_absolute() { path.clone() diff --git a/codex-rs/core/src/util_tests.rs b/codex-rs/core/src/util_tests.rs index dd5956bf61..dc377fa4b3 100644 --- a/codex-rs/core/src/util_tests.rs +++ b/codex-rs/core/src/util_tests.rs @@ -1,29 +1,5 @@ use super::*; -#[test] -fn test_try_parse_error_message() { - let text = r#"{ - "error": { - "message": "Your refresh token has already been used to generate a new access token. Please try signing in again.", - "type": "invalid_request_error", - "param": null, - "code": "refresh_token_reused" - } -}"#; - let message = try_parse_error_message(text); - assert_eq!( - message, - "Your refresh token has already been used to generate a new access token. Please try signing in again." - ); -} - -#[test] -fn test_try_parse_error_message_no_error() { - let text = r#"{"message": "test"}"#; - let message = try_parse_error_message(text); - assert_eq!(message, r#"{"message": "test"}"#); -} - #[test] fn feedback_tags_macro_compiles() { #[derive(Debug)]