diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5f7569dfddda..c1bc93f9dd28 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1761,6 +1761,7 @@ dependencies = [ "codex-app-server-protocol", "codex-git-utils", "codex-login", + "codex-model-provider", "codex-plugin", "codex-protocol", "codex-utils-absolute-path", @@ -1844,6 +1845,7 @@ dependencies = [ "codex-git-utils", "codex-login", "codex-mcp", + "codex-model-provider", "codex-model-provider-info", "codex-models-manager", "codex-otel", @@ -2033,9 +2035,11 @@ name = "codex-backend-client" version = "0.0.0" dependencies = [ "anyhow", + "codex-api", "codex-backend-openapi-models", "codex-client", "codex-login", + "codex-model-provider", "codex-protocol", "pretty_assertions", "reqwest", @@ -2059,11 +2063,11 @@ dependencies = [ "anyhow", "clap", "codex-app-server-protocol", - "codex-config", "codex-connectors", "codex-core", "codex-git-utils", "codex-login", + "codex-model-provider", "codex-utils-cargo-bin", "codex-utils-cli", "pretty_assertions", @@ -2190,7 +2194,6 @@ version = "0.0.0" dependencies = [ "anyhow", "async-trait", - "base64 0.22.1", "chrono", "clap", "codex-client", @@ -2199,6 +2202,7 @@ dependencies = [ "codex-core", "codex-git-utils", "codex-login", + "codex-model-provider", "codex-tui", "codex-utils-cli", "crossterm", @@ -2223,6 +2227,7 @@ dependencies = [ "anyhow", "async-trait", "chrono", + "codex-api", "codex-backend-client", "codex-git-utils", "serde", @@ -2441,6 +2446,7 @@ dependencies = [ "codex-exec-server", "codex-git-utils", "codex-login", + "codex-model-provider", "codex-plugin", "codex-protocol", "codex-utils-absolute-path", @@ -2468,6 +2474,7 @@ dependencies = [ "codex-config", "codex-exec-server", "codex-login", + "codex-model-provider", "codex-otel", "codex-protocol", "codex-skills", @@ -2826,6 +2833,7 @@ dependencies = [ "codex-config", "codex-exec-server", "codex-login", + "codex-model-provider", "codex-otel", "codex-plugin", "codex-protocol", @@ -2885,6 +2893,7 @@ name = "codex-model-provider" version = "0.0.0" dependencies = [ "async-trait", + "codex-agent-identity", "codex-api", "codex-aws-auth", "codex-client", diff --git a/codex-rs/analytics/Cargo.toml b/codex-rs/analytics/Cargo.toml index f706814d4193..918e7edc720e 100644 --- a/codex-rs/analytics/Cargo.toml +++ b/codex-rs/analytics/Cargo.toml @@ -16,6 +16,7 @@ workspace = true codex-app-server-protocol = { workspace = true } codex-git-utils = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-plugin = { workspace = true } codex-protocol = { workspace = true } os_info = { workspace = true } diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 1a4b5defe9cd..f842a4c126bd 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -307,16 +307,9 @@ async fn send_track_events( let Some(auth) = auth_manager.auth().await else { return; }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { return; } - let access_token = match auth.get_token() { - Ok(token) => token, - Err(_) => return, - }; - let Some(account_id) = auth.get_account_id() else { - return; - }; let base_url = base_url.trim_end_matches('/'); let url = format!("{base_url}/codex/analytics-events/events"); @@ -325,8 +318,7 @@ async fn send_track_events( let response = create_client() .post(&url) .timeout(ANALYTICS_EVENTS_TIMEOUT) - .bearer_auth(&access_token) - .header("chatgpt-account-id", &account_id) + .headers(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()) .header("Content-Type", "application/json") .json(&payload) .send() diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 339bc20f10f0..b38b1e28120c 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -49,6 +49,7 @@ codex-file-search = { workspace = true } codex-chatgpt = { workspace = true } codex-login = { workspace = true } codex-mcp = { workspace = true } +codex-model-provider = { workspace = true } codex-models-manager = { workspace = true } codex-protocol = { workspace = true } codex-app-server-protocol = { workspace = true } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 61ba231b6a51..9881814d0980 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1920,7 +1920,7 @@ impl CodexMessageProcessor { }); }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "chatgpt authentication required to notify workspace owner".to_string(), @@ -1975,7 +1975,7 @@ impl CodexMessageProcessor { }); }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "chatgpt authentication required to read rate limits".to_string(), @@ -6198,7 +6198,7 @@ impl CodexMessageProcessor { let auth = self.auth_manager.auth().await; if !config .features - .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth)) + .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) { self.outgoing .send_response( diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 57fa6e21e0c7..725b72269089 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -1078,7 +1078,7 @@ impl MessageProcessor { let auth = self.auth_manager.auth().await; if !config.features.apps_enabled_for_auth( auth.as_ref() - .is_some_and(codex_login::CodexAuth::is_chatgpt_auth), + .is_some_and(codex_login::CodexAuth::uses_codex_backend), ) { return; } diff --git a/codex-rs/app-server/src/transport/remote_control/enroll.rs b/codex-rs/app-server/src/transport/remote_control/enroll.rs index dbe18c8355db..6737dc9445bf 100644 --- a/codex-rs/app-server/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server/src/transport/remote_control/enroll.rs @@ -29,7 +29,7 @@ pub(super) struct RemoteControlEnrollment { #[derive(Debug, Clone, PartialEq, Eq)] pub(super) struct RemoteControlConnectionAuth { - pub(super) bearer_token: String, + pub(super) auth_headers: HeaderMap, pub(super) account_id: String, } @@ -202,7 +202,7 @@ pub(super) async fn enroll_remote_control_server( let http_request = client .post(enroll_url) .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT) - .bearer_auth(&auth.bearer_token) + .headers(auth.auth_headers.clone()) .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id) .json(&request); @@ -445,7 +445,7 @@ mod tests { let err = enroll_remote_control_server( &remote_control_target, &RemoteControlConnectionAuth { - bearer_token: "Access Token".to_string(), + auth_headers: HeaderMap::new(), account_id: "account_id".to_string(), }, ) diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 4eb58a87f2c9..1457eecf341c 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -680,11 +680,7 @@ fn build_remote_control_websocket_request( "x-codex-protocol-version", REMOTE_CONTROL_PROTOCOL_VERSION, )?; - set_remote_control_header( - headers, - "authorization", - &format!("Bearer {}", auth.bearer_token), - )?; + headers.extend(auth.auth_headers.clone()); set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?; if let Some(subscribe_cursor) = subscribe_cursor { set_remote_control_header( @@ -712,7 +708,7 @@ pub(crate) async fn load_remote_control_auth( reloaded = true; continue; }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { break auth; } if auth.get_account_id().is_none() && !reloaded { @@ -723,7 +719,7 @@ pub(crate) async fn load_remote_control_auth( break auth; }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { return Err(io::Error::new( ErrorKind::PermissionDenied, "remote control requires ChatGPT authentication; API key auth is not supported", @@ -731,7 +727,7 @@ pub(crate) async fn load_remote_control_auth( } Ok(RemoteControlConnectionAuth { - bearer_token: auth.get_token().map_err(io::Error::other)?, + auth_headers: codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers(), account_id: auth.get_account_id().ok_or_else(|| { io::Error::new( ErrorKind::WouldBlock, diff --git a/codex-rs/backend-client/Cargo.toml b/codex-rs/backend-client/Cargo.toml index 1707d45b1b67..d2e374ae2a0e 100644 --- a/codex-rs/backend-client/Cargo.toml +++ b/codex-rs/backend-client/Cargo.toml @@ -17,8 +17,10 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } codex-backend-openapi-models = { path = "../codex-backend-openapi-models" } +codex-api = { workspace = true } codex-client = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-protocol = { workspace = true } [dev-dependencies] diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index b96395b01515..8f7fe924a082 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -5,6 +5,7 @@ use crate::types::RateLimitReachedKind as BackendRateLimitReachedKind; use crate::types::RateLimitStatusPayload; use crate::types::TurnAttemptsSiblingTurnsResponse; use anyhow::Result; +use codex_api::AuthProvider; use codex_client::build_reqwest_client_with_custom_ca; use codex_client::with_chatgpt_cloudflare_cookie_store; use codex_login::CodexAuth; @@ -118,6 +119,7 @@ pub struct Client { base_url: String, http: reqwest::Client, bearer_token: Option, + auth_headers: HeaderMap, user_agent: Option, chatgpt_account_id: Option, chatgpt_account_is_fedramp: bool, @@ -146,6 +148,7 @@ impl Client { base_url, http, bearer_token: None, + auth_headers: HeaderMap::new(), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -154,17 +157,14 @@ impl Client { } pub fn from_auth(base_url: impl Into, auth: &CodexAuth) -> Result { - let token = auth.get_token().map_err(anyhow::Error::from)?; - let mut client = Self::new(base_url)? + Ok(Self::new(base_url)? .with_user_agent(get_codex_user_agent()) - .with_bearer_token(token); - if let Some(account_id) = auth.get_account_id() { - client = client.with_chatgpt_account_id(account_id); - } - if auth.is_fedramp_account() { - client = client.with_fedramp_routing_header(); - } - Ok(client) + .with_auth_provider(codex_model_provider::auth_provider_from_auth(auth).as_ref())) + } + + pub fn with_auth_provider(mut self, auth: &dyn AuthProvider) -> Self { + auth.add_auth_headers(&mut self.auth_headers); + self } pub fn with_bearer_token(mut self, token: impl Into) -> Self { @@ -207,6 +207,7 @@ impl Client { h.insert(AUTHORIZATION, hv); } } + h.extend(self.auth_headers.clone()); if let Some(acc) = &self.chatgpt_account_id && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") && let Ok(hv) = HeaderValue::from_str(acc) @@ -820,6 +821,7 @@ mod tests { base_url: "https://example.test".to_string(), http: reqwest::Client::new(), bearer_token: None, + auth_headers: HeaderMap::new(), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -834,6 +836,7 @@ mod tests { base_url: "https://chatgpt.com/backend-api".to_string(), http: reqwest::Client::new(), bearer_token: None, + auth_headers: HeaderMap::new(), user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, diff --git a/codex-rs/chatgpt/Cargo.toml b/codex-rs/chatgpt/Cargo.toml index 354449934ade..ce9aa627d435 100644 --- a/codex-rs/chatgpt/Cargo.toml +++ b/codex-rs/chatgpt/Cargo.toml @@ -12,10 +12,10 @@ anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } codex-app-server-protocol = { workspace = true } codex-connectors = { workspace = true } -codex-config = { workspace = true } codex-core = { workspace = true } codex-git-utils = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-utils-cli = { workspace = true } serde = { workspace = true, features = ["derive"] } tokio = { workspace = true, features = ["full"] } diff --git a/codex-rs/chatgpt/src/apply_command.rs b/codex-rs/chatgpt/src/apply_command.rs index 1a9553955d24..70fe4481db70 100644 --- a/codex-rs/chatgpt/src/apply_command.rs +++ b/codex-rs/chatgpt/src/apply_command.rs @@ -6,7 +6,6 @@ use codex_git_utils::ApplyGitRequest; use codex_git_utils::apply_git_patch; use codex_utils_cli::CliConfigOverrides; -use crate::chatgpt_token::init_chatgpt_token_from_auth; use crate::get_task::GetTaskResponse; use crate::get_task::OutputItem; use crate::get_task::PrOutputItem; @@ -32,9 +31,6 @@ pub async fn run_apply_command( ) .await?; - init_chatgpt_token_from_auth(&config.codex_home, config.cli_auth_credentials_store_mode) - .await?; - let task_response = get_task(&config, apply_cli.task_id).await?; apply_diff_from_task(task_response, cwd).await } diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index fa3a63dadbb4..0f9bef956f1f 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -1,9 +1,7 @@ use codex_core::config::Config; +use codex_login::AuthManager; use codex_login::default_client::create_client; -use crate::chatgpt_token::get_chatgpt_token_data; -use crate::chatgpt_token::init_chatgpt_token_from_auth; - use anyhow::Context; use serde::de::DeserializeOwned; use std::time::Duration; @@ -22,24 +20,28 @@ pub(crate) async fn chatgpt_get_request_with_timeout( timeout: Option, ) -> anyhow::Result { let chatgpt_base_url = &config.chatgpt_base_url; - init_chatgpt_token_from_auth(&config.codex_home, config.cli_auth_credentials_store_mode) - .await?; + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + let auth = auth_manager + .auth() + .await + .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + anyhow::ensure!( + auth.uses_codex_backend(), + "ChatGPT backend requests require Codex backend auth" + ); + anyhow::ensure!( + auth.get_account_id().is_some(), + "ChatGPT account ID not available, please re-run `codex login`" + ); // Make direct HTTP request to ChatGPT backend API with the token let client = create_client(); let url = format!("{chatgpt_base_url}{path}"); - let token = - get_chatgpt_token_data().ok_or_else(|| anyhow::anyhow!("ChatGPT token not available"))?; - - let account_id = token.account_id.ok_or_else(|| { - anyhow::anyhow!("ChatGPT account ID not available, please re-run `codex login`") - }); - let mut request = client .get(&url) - .bearer_auth(&token.access_token) - .header("chatgpt-account-id", account_id?) + .headers(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()) .header("Content-Type", "application/json"); if let Some(timeout) = timeout { diff --git a/codex-rs/chatgpt/src/chatgpt_token.rs b/codex-rs/chatgpt/src/chatgpt_token.rs deleted file mode 100644 index fe19c3015e86..000000000000 --- a/codex-rs/chatgpt/src/chatgpt_token.rs +++ /dev/null @@ -1,36 +0,0 @@ -use codex_config::types::AuthCredentialsStoreMode; -use codex_login::AuthManager; -use codex_login::token_data::TokenData; -use std::path::Path; -use std::sync::LazyLock; -use std::sync::RwLock; - -static CHATGPT_TOKEN: LazyLock>> = LazyLock::new(|| RwLock::new(None)); - -pub fn get_chatgpt_token_data() -> Option { - CHATGPT_TOKEN.read().ok()?.clone() -} - -pub fn set_chatgpt_token_data(value: TokenData) { - if let Ok(mut guard) = CHATGPT_TOKEN.write() { - *guard = Some(value); - } -} - -/// Initialize the ChatGPT token from auth.json file -pub async fn init_chatgpt_token_from_auth( - codex_home: &Path, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result<()> { - let auth_manager = AuthManager::new( - codex_home.to_path_buf(), - /*enable_codex_api_key_env*/ false, - auth_credentials_store_mode, - /*chatgpt_base_url*/ None, - ); - if let Some(auth) = auth_manager.auth().await { - let token_data = auth.get_token_data()?; - set_chatgpt_token_data(token_data); - } - Ok(()) -} diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index 4c6f05a68163..62e804094071 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -2,8 +2,6 @@ use std::collections::HashSet; use std::time::Duration; use crate::chatgpt_client::chatgpt_get_request_with_timeout; -use crate::chatgpt_token::get_chatgpt_token_data; -use crate::chatgpt_token::init_chatgpt_token_from_auth; use codex_app_server_protocol::AppInfo; use codex_connectors::AllConnectorsCacheKey; @@ -23,22 +21,32 @@ use codex_core::plugins::PluginsManager; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::originator; -use codex_login::token_data::TokenData; const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); async fn apps_enabled(config: &Config) -> bool { - let auth_manager = AuthManager::shared( - config.codex_home.to_path_buf(), - /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, - Some(config.chatgpt_base_url.clone()), - ); + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); let auth = auth_manager.auth().await; config .features - .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth)) + .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) +} + +async fn connector_auth(config: &Config) -> anyhow::Result { + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + let auth = auth_manager + .auth() + .await + .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + anyhow::ensure!( + auth.uses_codex_backend(), + "ChatGPT connectors require Codex backend auth" + ); + Ok(auth) } + pub async fn list_connectors(config: &Config) -> anyhow::Result> { if !apps_enabled(config).await { return Ok(Vec::new()); @@ -66,14 +74,8 @@ pub async fn list_cached_all_connectors(config: &Config) -> Option> return Some(Vec::new()); } - if init_chatgpt_token_from_auth(&config.codex_home, config.cli_auth_credentials_store_mode) - .await - .is_err() - { - return None; - } - let token_data = get_chatgpt_token_data()?; - let cache_key = all_connectors_cache_key(config, &token_data); + let auth = connector_auth(config).await.ok()?; + let cache_key = all_connectors_cache_key(config, &auth); let connectors = codex_connectors::cached_all_connectors(&cache_key)?; let connectors = merge_plugin_connectors( connectors, @@ -95,15 +97,11 @@ pub async fn list_all_connectors_with_options( if !apps_enabled(config).await { return Ok(Vec::new()); } - init_chatgpt_token_from_auth(&config.codex_home, config.cli_auth_credentials_store_mode) - .await?; - - let token_data = - get_chatgpt_token_data().ok_or_else(|| anyhow::anyhow!("ChatGPT token not available"))?; - let cache_key = all_connectors_cache_key(config, &token_data); + let auth = connector_auth(config).await?; + let cache_key = all_connectors_cache_key(config, &auth); let connectors = codex_connectors::list_all_connectors_with_options( cache_key, - token_data.id_token.is_workspace_account(), + auth.is_workspace_account(), force_refetch, |path| async move { chatgpt_get_request_with_timeout::( @@ -128,12 +126,12 @@ pub async fn list_all_connectors_with_options( )) } -fn all_connectors_cache_key(config: &Config, token_data: &TokenData) -> AllConnectorsCacheKey { +fn all_connectors_cache_key(config: &Config, auth: &CodexAuth) -> AllConnectorsCacheKey { AllConnectorsCacheKey::new( config.chatgpt_base_url.clone(), - token_data.account_id.clone(), - token_data.id_token.chatgpt_user_id.clone(), - token_data.id_token.is_workspace_account(), + auth.get_account_id(), + auth.get_chatgpt_user_id(), + auth.is_workspace_account(), ) } diff --git a/codex-rs/chatgpt/src/lib.rs b/codex-rs/chatgpt/src/lib.rs index 0d39bb932db5..057478db1856 100644 --- a/codex-rs/chatgpt/src/lib.rs +++ b/codex-rs/chatgpt/src/lib.rs @@ -1,5 +1,4 @@ pub mod apply_command; mod chatgpt_client; -mod chatgpt_token; pub mod connectors; pub mod get_task; diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 0b50aa834fa7..7390ea0ab4c9 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -169,13 +169,7 @@ fn verify_cache_signature(payload_bytes: &[u8], signature: &str) -> bool { } fn auth_identity(auth: &CodexAuth) -> (Option, Option) { - let token_data = auth.get_token_data().ok(); - let chatgpt_user_id = token_data - .as_ref() - .and_then(|token_data| token_data.id_token.chatgpt_user_id.as_deref()) - .map(str::to_owned); - let account_id = auth.get_account_id(); - (chatgpt_user_id, account_id) + (auth.get_chatgpt_user_id(), auth.get_account_id()) } fn cache_payload_bytes(payload: &CloudRequirementsCacheSignedPayload) -> Option> { @@ -331,7 +325,7 @@ impl CloudRequirementsService { let Some(plan_type) = auth.account_plan_type() else { return Ok(None); }; - if !auth.is_chatgpt_auth() + if !auth.uses_codex_backend() || !(plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) { return Ok(None); @@ -551,7 +545,7 @@ impl CloudRequirementsService { let Some(plan_type) = auth.account_plan_type() else { return false; }; - if !auth.is_chatgpt_auth() + if !auth.uses_codex_backend() || !(plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) { return false; diff --git a/codex-rs/cloud-tasks-client/Cargo.toml b/codex-rs/cloud-tasks-client/Cargo.toml index cdfcba47b828..929c3e313629 100644 --- a/codex-rs/cloud-tasks-client/Cargo.toml +++ b/codex-rs/cloud-tasks-client/Cargo.toml @@ -15,6 +15,7 @@ workspace = true anyhow = { workspace = true } async-trait = { workspace = true } chrono = { workspace = true, features = ["serde"] } +codex-api = { workspace = true } codex-backend-client = { workspace = true } codex-git-utils = { workspace = true } serde = { version = "1", features = ["derive"] } diff --git a/codex-rs/cloud-tasks-client/src/http.rs b/codex-rs/cloud-tasks-client/src/http.rs index 4ea098022737..874209243550 100644 --- a/codex-rs/cloud-tasks-client/src/http.rs +++ b/codex-rs/cloud-tasks-client/src/http.rs @@ -14,6 +14,7 @@ use crate::api::TaskText; use chrono::DateTime; use chrono::Utc; +use codex_api::AuthProvider; use codex_backend_client as backend; use codex_backend_client::CodeTaskDetailsResponseExt; use codex_git_utils::ApplyGitRequest; @@ -42,6 +43,11 @@ impl HttpClient { self } + pub fn with_auth_provider(mut self, auth: &dyn AuthProvider) -> Self { + self.backend = self.backend.clone().with_auth_provider(auth); + self + } + pub fn with_chatgpt_account_id(mut self, account_id: impl Into) -> Self { self.backend = self.backend.clone().with_chatgpt_account_id(account_id); self diff --git a/codex-rs/cloud-tasks/Cargo.toml b/codex-rs/cloud-tasks/Cargo.toml index 30e8b73a8fe2..6429c1edcd4b 100644 --- a/codex-rs/cloud-tasks/Cargo.toml +++ b/codex-rs/cloud-tasks/Cargo.toml @@ -13,7 +13,6 @@ workspace = true [dependencies] anyhow = { workspace = true } -base64 = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-client = { workspace = true } @@ -23,6 +22,7 @@ codex-cloud-tasks-mock-client = { workspace = true } codex-core = { workspace = true } codex-git-utils = { workspace = true } codex-login = { path = "../login" } +codex-model-provider = { workspace = true } codex-tui = { workspace = true } codex-utils-cli = { workspace = true } crossterm = { workspace = true, features = ["event-stream"] } diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index 7006d52b921d..9cc45f36c661 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -87,23 +87,17 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result append_error_log(format!("auth: mode=ChatGPT account_id={acc}")); } - let token = match auth.get_token() { - Ok(t) if !t.is_empty() => t, - _ => { - eprintln!( - "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." - ); - std::process::exit(1); - } - }; + if !auth.uses_codex_backend() { + eprintln!( + "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." + ); + std::process::exit(1); + } - http = http.with_bearer_token(token.clone()); - if let Some(acc) = auth - .get_account_id() - .or_else(|| util::extract_chatgpt_account_id(&token)) - { + let auth_provider = codex_model_provider::auth_provider_from_auth(&auth); + http = http.with_auth_provider(auth_provider.as_ref()); + if let Some(acc) = auth.get_account_id() { append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}")); - http = http.with_chatgpt_account_id(acc); } Ok(BackendContext { diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index 525ea3b5945a..fc94b268ded5 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -1,4 +1,3 @@ -use base64::Engine as _; use chrono::DateTime; use chrono::Local; use chrono::Utc; @@ -42,23 +41,6 @@ pub fn normalize_base_url(input: &str) -> String { base_url } -/// Extract the ChatGPT account id from a JWT token, when present. -pub fn extract_chatgpt_account_id(token: &str) -> Option { - let mut parts = token.split('.'); - let (_h, payload_b64, _s) = 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 None, - }; - let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(payload_b64) - .ok()?; - let v: serde_json::Value = serde_json::from_slice(&payload_bytes).ok()?; - v.get("https://api.openai.com/auth") - .and_then(|auth| auth.get("chatgpt_account_id")) - .and_then(|id| id.as_str()) - .map(str::to_string) -} - pub async fn load_auth_manager() -> Option { // TODO: pass in cli overrides once cloud tasks properly support them. let config = Config::load_with_cli_overrides(Vec::new()).await.ok()?; @@ -73,8 +55,6 @@ pub async fn load_auth_manager() -> Option { /// Build headers for ChatGPT-backed requests: `User-Agent`, optional `Authorization`, /// and optional `ChatGPT-Account-Id`. pub async fn build_chatgpt_headers() -> HeaderMap { - use reqwest::header::AUTHORIZATION; - use reqwest::header::HeaderName; use reqwest::header::HeaderValue; use reqwest::header::USER_AGENT; @@ -87,21 +67,9 @@ pub async fn build_chatgpt_headers() -> HeaderMap { ); if let Some(am) = load_auth_manager().await && let Some(auth) = am.auth().await - && let Ok(tok) = auth.get_token() - && !tok.is_empty() + && auth.uses_codex_backend() { - let v = format!("Bearer {tok}"); - if let Ok(hv) = HeaderValue::from_str(&v) { - headers.insert(AUTHORIZATION, hv); - } - if let Some(acc) = auth - .get_account_id() - .or_else(|| extract_chatgpt_account_id(&tok)) - && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") - && let Ok(hv) = HeaderValue::from_str(&acc) - { - headers.insert(name, hv); - } + headers.extend(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()); } headers } diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index e1130c770740..41394a22584b 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -34,6 +34,13 @@ pub trait AuthProvider: Send + Sync { /// used by telemetry and non-HTTP request paths. fn add_auth_headers(&self, headers: &mut HeaderMap); + /// Returns any auth headers that are available without request body access. + fn to_auth_headers(&self) -> HeaderMap { + let mut headers = HeaderMap::new(); + self.add_auth_headers(&mut headers); + headers + } + /// Applies auth to a complete outbound request and returns the request to send. /// /// The input `request` is moved into this method. Implementations may mutate diff --git a/codex-rs/codex-mcp/Cargo.toml b/codex-rs/codex-mcp/Cargo.toml index 0aec1f3aaf66..f99f4b79026f 100644 --- a/codex-rs/codex-mcp/Cargo.toml +++ b/codex-rs/codex-mcp/Cargo.toml @@ -18,6 +18,7 @@ codex-async-utils = { workspace = true } codex-config = { workspace = true } codex-exec-server = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-otel = { workspace = true } codex-plugin = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 97053cbe53e8..8e6aea91cc36 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -205,24 +205,15 @@ fn codex_apps_mcp_bearer_token_env_var() -> Option { } } -fn codex_apps_mcp_bearer_token(auth: Option<&CodexAuth>) -> Option { - let token = auth.and_then(|auth| auth.get_token().ok())?; - let token = token.trim(); - if token.is_empty() { - None - } else { - Some(token.to_string()) - } -} - fn codex_apps_mcp_http_headers(auth: Option<&CodexAuth>) -> Option> { - let mut headers = HashMap::new(); - if let Some(token) = codex_apps_mcp_bearer_token(auth) { - headers.insert("Authorization".to_string(), format!("Bearer {token}")); - } - if let Some(account_id) = auth.and_then(CodexAuth::get_account_id) { - headers.insert("ChatGPT-Account-ID".to_string(), account_id); - } + let auth = auth.filter(|auth| auth.uses_codex_backend())?; + let headers = codex_model_provider::auth_provider_from_auth(auth).to_auth_headers(); + let headers: HashMap<_, _> = headers + .iter() + .filter_map(|(name, value)| { + Some((name.as_str().to_string(), value.to_str().ok()?.to_string())) + }) + .collect(); if headers.is_empty() { None } else { @@ -293,7 +284,7 @@ pub fn with_codex_apps_mcp( auth: Option<&CodexAuth>, config: &McpConfig, ) -> HashMap { - if config.apps_enabled && auth.is_some_and(CodexAuth::is_chatgpt_auth) { + if config.apps_enabled && auth.is_some_and(CodexAuth::uses_codex_backend) { servers.insert( CODEX_APPS_MCP_SERVER_NAME.to_string(), codex_apps_mcp_server_config(config, auth), diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager.rs b/codex-rs/codex-mcp/src/mcp_connection_manager.rs index 10c48e040e85..8b8d3fe35cad 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager.rs @@ -120,21 +120,10 @@ fn sha1_hex(s: &str) -> String { } pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCacheKey { - let token_data = auth.and_then(|auth| auth.get_token_data().ok()); - let account_id = token_data - .as_ref() - .and_then(|token_data| token_data.account_id.clone()); - let chatgpt_user_id = token_data - .as_ref() - .and_then(|token_data| token_data.id_token.chatgpt_user_id.clone()); - let is_workspace_account = token_data - .as_ref() - .is_some_and(|token_data| token_data.id_token.is_workspace_account()); - CodexAppsToolsCacheKey { - account_id, - chatgpt_user_id, - is_workspace_account, + account_id: auth.and_then(CodexAuth::get_account_id), + chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), + is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), } } diff --git a/codex-rs/core-plugins/Cargo.toml b/codex-rs/core-plugins/Cargo.toml index 0372d9a14acb..6b214d394e43 100644 --- a/codex-rs/core-plugins/Cargo.toml +++ b/codex-rs/core-plugins/Cargo.toml @@ -19,6 +19,7 @@ codex-core-skills = { workspace = true } codex-exec-server = { workspace = true } codex-git-utils = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-plugin = { workspace = true } codex-protocol = { workspace = true } codex-utils-absolute-path = { workspace = true } diff --git a/codex-rs/core-plugins/src/remote.rs b/codex-rs/core-plugins/src/remote.rs index 89dabc841b2c..a4850fea028a 100644 --- a/codex-rs/core-plugins/src/remote.rs +++ b/codex-rs/core-plugins/src/remote.rs @@ -606,7 +606,7 @@ fn ensure_chatgpt_auth(auth: Option<&CodexAuth>) -> Result<&CodexAuth, RemotePlu let Some(auth) = auth else { return Err(RemotePluginCatalogError::AuthRequired); }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { return Err(RemotePluginCatalogError::UnsupportedAuthMode); } Ok(auth) @@ -616,16 +616,9 @@ fn authenticated_request( request: RequestBuilder, auth: &CodexAuth, ) -> Result { - let token = auth - .get_token() - .map_err(RemotePluginCatalogError::AuthToken)?; - let mut request = request + Ok(request .timeout(REMOTE_PLUGIN_CATALOG_TIMEOUT) - .bearer_auth(token); - if let Some(account_id) = auth.get_account_id() { - request = request.header("chatgpt-account-id", account_id); - } - Ok(request) + .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers())) } async fn send_and_decode Deserialize<'de>>( diff --git a/codex-rs/core-skills/Cargo.toml b/codex-rs/core-skills/Cargo.toml index 355374114a14..4324d29dee94 100644 --- a/codex-rs/core-skills/Cargo.toml +++ b/codex-rs/core-skills/Cargo.toml @@ -19,6 +19,7 @@ codex-app-server-protocol = { workspace = true } codex-config = { workspace = true } codex-exec-server = { workspace = true } codex-login = { workspace = true } +codex-model-provider = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-skills = { workspace = true } diff --git a/codex-rs/core-skills/src/remote.rs b/codex-rs/core-skills/src/remote.rs index 2dc620b864d8..1ca7cd0cb768 100644 --- a/codex-rs/core-skills/src/remote.rs +++ b/codex-rs/core-skills/src/remote.rs @@ -48,11 +48,11 @@ fn as_query_product_surface(product_surface: RemoteSkillProductSurface) -> &'sta } } -fn ensure_chatgpt_auth(auth: Option<&CodexAuth>) -> Result<&CodexAuth> { +fn ensure_codex_backend_auth(auth: Option<&CodexAuth>) -> Result<&CodexAuth> { let Some(auth) = auth else { anyhow::bail!("chatgpt authentication required for remote skill scopes"); }; - if !auth.is_chatgpt_auth() { + if !auth.uses_codex_backend() { anyhow::bail!( "chatgpt authentication required for remote skill scopes; api key auth is not supported" ); @@ -94,7 +94,7 @@ pub async fn list_remote_skills( enabled: Option, ) -> Result> { let base_url = chatgpt_base_url.trim_end_matches('/'); - let auth = ensure_chatgpt_auth(auth)?; + let auth = ensure_codex_backend_auth(auth)?; let url = format!("{base_url}/hazelnuts"); let product_surface = as_query_product_surface(product_surface); @@ -108,17 +108,11 @@ pub async fn list_remote_skills( } let client = build_reqwest_client(); - let mut request = client + let request = client .get(&url) .timeout(REMOTE_SKILLS_API_TIMEOUT) - .query(&query_params); - let token = auth - .get_token() - .context("Failed to read auth token for remote skills")?; - request = request.bearer_auth(token); - if let Some(account_id) = auth.get_account_id() { - request = request.header("chatgpt-account-id", account_id); - } + .query(&query_params) + .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()); let response = request .send() .await @@ -150,20 +144,15 @@ pub async fn export_remote_skill( auth: Option<&CodexAuth>, skill_id: &str, ) -> Result { - let auth = ensure_chatgpt_auth(auth)?; + let auth = ensure_codex_backend_auth(auth)?; let client = build_reqwest_client(); let base_url = chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/hazelnuts/{skill_id}/export"); - let mut request = client.get(&url).timeout(REMOTE_SKILLS_API_TIMEOUT); - - let token = auth - .get_token() - .context("Failed to read auth token for remote skills")?; - request = request.bearer_auth(token); - if let Some(account_id) = auth.get_account_id() { - request = request.header("chatgpt-account-id", account_id); - } + let request = client + .get(&url) + .timeout(REMOTE_SKILLS_API_TIMEOUT) + .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()); let response = request .send() diff --git a/codex-rs/core/src/arc_monitor.rs b/codex-rs/core/src/arc_monitor.rs index ecd7f3966628..08b7465178f3 100644 --- a/codex-rs/core/src/arc_monitor.rs +++ b/codex-rs/core/src/arc_monitor.rs @@ -9,7 +9,6 @@ use crate::compact::content_items_to_text; use crate::event_mapping::is_contextual_user_message_content; use crate::session::session::Session; use crate::session::turn_context::TurnContext; -use codex_login::CodexAuth; use codex_login::default_client::build_reqwest_client; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; @@ -104,28 +103,15 @@ pub(crate) async fn monitor_action( ) -> ArcMonitorOutcome { let auth = match turn_context.auth_manager.as_ref() { Some(auth_manager) => match auth_manager.auth().await { - Some(auth) if auth.is_chatgpt_auth() => Some(auth), + Some(auth) if auth.uses_codex_backend() => Some(auth), _ => None, }, None => None, }; - let token = if let Some(token) = read_non_empty_env_var(CODEX_ARC_MONITOR_TOKEN) { - token - } else { - let Some(auth) = auth.as_ref() else { - return ArcMonitorOutcome::Ok; - }; - match auth.get_token() { - Ok(token) => token, - Err(err) => { - warn!( - error = %err, - "skipping safety monitor because auth token is unavailable" - ); - return ArcMonitorOutcome::Ok; - } - } - }; + let env_token = read_non_empty_env_var(CODEX_ARC_MONITOR_TOKEN); + if env_token.is_none() && auth.is_none() { + return ArcMonitorOutcome::Ok; + } let url = read_non_empty_env_var(CODEX_ARC_MONITOR_ENDPOINT_OVERRIDE).unwrap_or_else(|| { format!( @@ -143,13 +129,12 @@ pub(crate) async fn monitor_action( let body = build_arc_monitor_request(sess, turn_context, action, protection_client_callsite).await; let client = build_reqwest_client(); - let mut request = client - .post(&url) - .timeout(ARC_MONITOR_TIMEOUT) - .json(&body) - .bearer_auth(token); - if let Some(account_id) = auth.as_ref().and_then(CodexAuth::get_account_id) { - request = request.header("chatgpt-account-id", account_id); + let mut request = client.post(&url).timeout(ARC_MONITOR_TIMEOUT).json(&body); + if let Some(token) = env_token { + request = request.bearer_auth(token); + } else if let Some(auth) = auth.as_ref() { + request = + request.headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()); } let response = match request.send().await { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 77022029f1d3..ed75f5c783a9 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1120,7 +1120,7 @@ impl ModelClientSession { fn responses_request_compression(&self, auth: Option<&CodexAuth>) -> Compression { if self.client.state.enable_request_compression - && auth.is_some_and(CodexAuth::is_chatgpt_auth) + && auth.is_some_and(CodexAuth::uses_codex_backend) && self.client.state.provider.info().is_openai() { Compression::Zstd diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index d6e6d1f9c072..0e0d4a600839 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -14,7 +14,6 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; use codex_api::upload_local_file; use codex_login::CodexAuth; -use codex_model_provider::BearerAuthProvider; use serde_json::Value as JsonValue; pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( @@ -109,17 +108,15 @@ async fn build_uploaded_local_argument_value( "ChatGPT auth is required to upload local files for Codex Apps tools".to_string(), ); }; - let token_data = auth - .get_token_data() - .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?; - let upload_auth = BearerAuthProvider { - token: Some(token_data.access_token), - account_id: token_data.account_id, - is_fedramp_account: auth.is_fedramp_account(), - }; + if !auth.uses_codex_backend() { + return Err( + "ChatGPT auth is required to upload local files for Codex Apps tools".to_string(), + ); + } + let upload_auth = codex_model_provider::auth_provider_from_auth(auth); let uploaded = upload_local_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), - &upload_auth, + upload_auth.as_ref(), &resolved_path, ) .await diff --git a/codex-rs/core/src/plugins/manager.rs b/codex-rs/core/src/plugins/manager.rs index 7d9b426b2608..0a9e0c17c691 100644 --- a/codex-rs/core/src/plugins/manager.rs +++ b/codex-rs/core/src/plugins/manager.rs @@ -125,21 +125,11 @@ fn featured_plugin_ids_cache_key( config: &Config, auth: Option<&CodexAuth>, ) -> FeaturedPluginIdsCacheKey { - let token_data = auth.and_then(|auth| auth.get_token_data().ok()); - let account_id = token_data - .as_ref() - .and_then(|token_data| token_data.account_id.clone()); - let chatgpt_user_id = token_data - .as_ref() - .and_then(|token_data| token_data.id_token.chatgpt_user_id.clone()); - let is_workspace_account = token_data - .as_ref() - .is_some_and(|token_data| token_data.id_token.is_workspace_account()); FeaturedPluginIdsCacheKey { chatgpt_base_url: config.chatgpt_base_url.clone(), - account_id, - chatgpt_user_id, - is_workspace_account, + account_id: auth.and_then(CodexAuth::get_account_id), + chatgpt_user_id: auth.and_then(CodexAuth::get_chatgpt_user_id), + is_workspace_account: auth.is_some_and(CodexAuth::is_workspace_account), } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d45357e3e58c..e384521c73a4 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -45,7 +45,6 @@ use chrono::Local; use chrono::Utc; use codex_analytics::AnalyticsEventsClient; use codex_analytics::SubAgentThreadStartedInput; -use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::McpServerElicitationRequest; use codex_app_server_protocol::McpServerElicitationRequestParams; use codex_config::types::OAuthCredentialsStoreMode; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 998f016c9671..aca7ddbb452e 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -4,10 +4,7 @@ use codex_model_provider::create_model_provider; use codex_protocol::protocol::TurnEnvironmentSelection; pub(super) fn image_generation_tool_auth_allowed(auth_manager: Option<&AuthManager>) -> bool { - matches!( - auth_manager.and_then(AuthManager::auth_mode), - Some(AuthMode::Chatgpt) - ) + auth_manager.is_some_and(AuthManager::current_auth_uses_codex_backend) } #[derive(Clone, Debug)] @@ -93,13 +90,11 @@ impl TurnContext { } pub(crate) fn apps_enabled(&self) -> bool { - let is_chatgpt_auth = self + let uses_codex_backend = self .auth_manager .as_deref() - .and_then(AuthManager::auth_cached) - .as_ref() - .is_some_and(CodexAuth::is_chatgpt_auth); - self.features.apps_enabled_for_auth(is_chatgpt_auth) + .is_some_and(AuthManager::current_auth_uses_codex_backend); + self.features.apps_enabled_for_auth(uses_codex_backend) } pub(crate) async fn with_model(&self, model: String, models_manager: &ModelsManager) -> Self { diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index e8f81f39fac0..5f2dc9cfc8bc 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -39,6 +39,10 @@ impl AgentIdentityAuth { &self.record } + pub fn process_task_id(&self) -> Option<&str> { + self.process_task_id.get().map(String::as_str) + } + pub async fn ensure_runtime(&self, chatgpt_base_url: Option) -> std::io::Result<()> { self.process_task_id .get_or_try_init(|| async { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 6cc87386f55d..419c6a4bac41 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -397,6 +397,11 @@ impl CodexAuth { }) } + pub fn is_workspace_account(&self) -> bool { + self.account_plan_type() + .is_some_and(AccountPlanType::is_workspace_account) + } + /// Returns `None` if token-backed ChatGPT auth is unavailable. fn get_current_auth_json(&self) -> Option { let state = match self { @@ -1709,6 +1714,13 @@ impl AuthManager { self.auth_cached().as_ref().map(CodexAuth::auth_mode) } + pub fn current_auth_uses_codex_backend(&self) -> bool { + matches!( + self.auth_mode(), + Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::AgentIdentity) + ) + } + fn is_stale_for_proactive_refresh(auth: &CodexAuth) -> bool { let chatgpt_auth = match auth { CodexAuth::Chatgpt(chatgpt_auth) => chatgpt_auth, diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index ad8cad4e81f5..d60c00716007 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -15,6 +15,7 @@ workspace = true [dependencies] async-trait = { workspace = true } codex-api = { workspace = true } +codex-agent-identity = { workspace = true } codex-aws-auth = { workspace = true } codex-client = { workspace = true } codex-login = { workspace = true } diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 64640dcc960e..a4f0f3025860 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -1,12 +1,72 @@ use std::sync::Arc; +use codex_agent_identity::AgentIdentityKey; +use codex_agent_identity::AgentTaskAuthorizationTarget; +use codex_agent_identity::authorization_header_for_agent_task; +use codex_api::AuthProvider; use codex_api::SharedAuthProvider; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; +use http::HeaderMap; +use http::HeaderValue; use crate::bearer_auth_provider::BearerAuthProvider; +#[derive(Clone, Debug)] +struct CodexAuthProvider { + auth: CodexAuth, +} + +impl AuthProvider for CodexAuthProvider { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + let header_value = match &self.auth { + CodexAuth::AgentIdentity(auth) => { + let record = auth.record(); + let process_task_id = auth.process_task_id().ok_or_else(|| { + std::io::Error::other("agent identity process task is not initialized") + }); + process_task_id.and_then(|task_id| { + authorization_header_for_agent_task( + AgentIdentityKey { + agent_runtime_id: &record.agent_runtime_id, + private_key_pkcs8_base64: &record.agent_private_key, + }, + AgentTaskAuthorizationTarget { + agent_runtime_id: &record.agent_runtime_id, + task_id, + }, + ) + .map_err(std::io::Error::other) + }) + } + CodexAuth::ApiKey(_) => self.auth.api_key().map_or_else( + || Err(std::io::Error::other("API key auth missing API key")), + |api_key| Ok(format!("Bearer {api_key}")), + ), + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { + self.auth.get_token().map(|token| format!("Bearer {token}")) + } + }; + + if let Ok(header_value) = header_value + && let Ok(header) = HeaderValue::from_str(&header_value) + { + let _ = headers.insert(http::header::AUTHORIZATION, header); + } + + if let Some(account_id) = self.auth.get_account_id() + && let Ok(header) = HeaderValue::from_str(&account_id) + { + let _ = headers.insert("ChatGPT-Account-ID", header); + } + + if self.auth.is_fedramp_account() { + let _ = headers.insert("X-OpenAI-Fedramp", HeaderValue::from_static("true")); + } + } +} + /// Returns the provider-scoped auth manager when this provider uses command-backed auth. /// /// Providers without custom auth continue using the caller-supplied base manager, when present. @@ -20,45 +80,35 @@ pub(crate) fn auth_manager_for_provider( } } -fn bearer_auth_provider_from_auth( +pub(crate) fn resolve_provider_auth( auth: Option<&CodexAuth>, provider: &ModelProviderInfo, -) -> codex_protocol::error::Result { +) -> codex_protocol::error::Result { + if let Some(auth) = bearer_auth_for_provider(provider)? { + return Ok(Arc::new(auth)); + } + + Ok(match auth { + Some(auth) => auth_provider_from_auth(auth), + None => Arc::new(BearerAuthProvider::empty()), + }) +} + +fn bearer_auth_for_provider( + provider: &ModelProviderInfo, +) -> codex_protocol::error::Result> { if let Some(api_key) = provider.api_key()? { - return Ok(BearerAuthProvider { - token: Some(api_key), - account_id: None, - is_fedramp_account: false, - }); + return Ok(Some(BearerAuthProvider::new(api_key))); } if let Some(token) = provider.experimental_bearer_token.clone() { - return Ok(BearerAuthProvider { - token: Some(token), - account_id: None, - is_fedramp_account: false, - }); + return Ok(Some(BearerAuthProvider::new(token))); } - if let Some(auth) = auth { - let token = auth.get_token()?; - Ok(BearerAuthProvider { - token: Some(token), - account_id: auth.get_account_id(), - is_fedramp_account: auth.is_fedramp_account(), - }) - } else { - Ok(BearerAuthProvider { - token: None, - account_id: None, - is_fedramp_account: false, - }) - } + Ok(None) } -pub(crate) fn resolve_provider_auth( - auth: Option<&CodexAuth>, - provider: &ModelProviderInfo, -) -> codex_protocol::error::Result { - Ok(Arc::new(bearer_auth_provider_from_auth(auth, provider)?)) +/// Builds request-header auth for a first-party Codex auth snapshot. +pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { + Arc::new(CodexAuthProvider { auth: auth.clone() }) } diff --git a/codex-rs/model-provider/src/bearer_auth_provider.rs b/codex-rs/model-provider/src/bearer_auth_provider.rs index 5a24ca6f78da..7685e4139a3f 100644 --- a/codex-rs/model-provider/src/bearer_auth_provider.rs +++ b/codex-rs/model-provider/src/bearer_auth_provider.rs @@ -11,6 +11,18 @@ pub struct BearerAuthProvider { } impl BearerAuthProvider { + pub fn new(token: String) -> Self { + Self { + token: Some(token), + account_id: None, + is_fedramp_account: false, + } + } + + pub fn empty() -> Self { + Self::default() + } + pub fn for_test(token: Option<&str>, account_id: Option<&str>) -> Self { Self { token: token.map(str::to_string), diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index f12c6a914a92..0db7b9c8d6a1 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -3,6 +3,7 @@ mod auth; mod bearer_auth_provider; mod provider; +pub use auth::auth_provider_from_auth; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use provider::ModelProvider; diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index c029960a7039..34f9f7a781fe 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -9,7 +9,6 @@ use codex_api::ReqwestTransport; use codex_api::TransportError; use codex_api::auth_header_telemetry; use codex_api::map_api_error; -use codex_app_server_protocol::AuthMode; use codex_feedback::FeedbackRequestTags; use codex_feedback::emit_feedback_request_tags_with_auth_env; use codex_login::AuthEnvTelemetry; @@ -407,11 +406,13 @@ impl ModelsManager { return Ok(()); } - let auth_mode = self + let uses_codex_backend = self .provider - .auth_manager() - .and_then(|auth_manager| auth_manager.auth_mode()); - if auth_mode != Some(AuthMode::Chatgpt) && !self.provider.info().has_command_auth() { + .auth() + .await + .as_ref() + .is_some_and(CodexAuth::uses_codex_backend); + if !uses_codex_backend && !self.provider.info().has_command_auth() { if matches!( refresh_strategy, RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached @@ -536,12 +537,12 @@ impl ModelsManager { remote_models.sort_by(|a, b| a.priority.cmp(&b.priority)); let mut presets: Vec = remote_models.into_iter().map(Into::into).collect(); - let auth_mode = self + let uses_codex_backend = self .provider .auth_manager() - .and_then(|auth_manager| auth_manager.auth_mode()); - let chatgpt_mode = matches!(auth_mode, Some(AuthMode::Chatgpt)); - presets = ModelPreset::filter_by_auth(presets, chatgpt_mode); + .as_deref() + .is_some_and(AuthManager::current_auth_uses_codex_backend); + presets = ModelPreset::filter_by_auth(presets, uses_codex_backend); ModelPreset::mark_default_by_picker_visibility(&mut presets); diff --git a/codex-rs/protocol/src/account.rs b/codex-rs/protocol/src/account.rs index bb46329a51d9..cc6632a5091a 100644 --- a/codex-rs/protocol/src/account.rs +++ b/codex-rs/protocol/src/account.rs @@ -35,6 +35,12 @@ impl PlanType { pub fn is_business_like(self) -> bool { matches!(self, Self::Business | Self::EnterpriseCbpUsageBased) } + + pub fn is_workspace_account(self) -> bool { + self.is_team_like() + || self.is_business_like() + || matches!(self, Self::Enterprise | Self::Edu) + } } #[cfg(test)] @@ -83,5 +89,19 @@ mod tests { assert_eq!(PlanType::Business.is_business_like(), true); assert_eq!(PlanType::EnterpriseCbpUsageBased.is_business_like(), true); assert_eq!(PlanType::Team.is_business_like(), false); + + assert_eq!(PlanType::Team.is_workspace_account(), true); + assert_eq!( + PlanType::SelfServeBusinessUsageBased.is_workspace_account(), + true + ); + assert_eq!(PlanType::Business.is_workspace_account(), true); + assert_eq!( + PlanType::EnterpriseCbpUsageBased.is_workspace_account(), + true + ); + assert_eq!(PlanType::Enterprise.is_workspace_account(), true); + assert_eq!(PlanType::Edu.is_workspace_account(), true); + assert_eq!(PlanType::Pro.is_workspace_account(), false); } } diff --git a/codex-rs/protocol/src/auth.rs b/codex-rs/protocol/src/auth.rs index 99e067bf2408..604f47855bd2 100644 --- a/codex-rs/protocol/src/auth.rs +++ b/codex-rs/protocol/src/auth.rs @@ -50,6 +50,14 @@ pub enum KnownPlan { } impl KnownPlan { + pub fn is_team_like(self) -> bool { + matches!(self, Self::Team | Self::SelfServeBusinessUsageBased) + } + + pub fn is_business_like(self) -> bool { + matches!(self, Self::Business | Self::EnterpriseCbpUsageBased) + } + pub fn display_name(self) -> &'static str { match self { Self::Free => "Free", @@ -83,15 +91,9 @@ impl KnownPlan { } pub fn is_workspace_account(self) -> bool { - matches!( - self, - Self::Team - | Self::SelfServeBusinessUsageBased - | Self::Business - | Self::EnterpriseCbpUsageBased - | Self::Enterprise - | Self::Edu - ) + self.is_team_like() + || self.is_business_like() + || matches!(self, Self::Enterprise | Self::Edu) } } diff --git a/codex-rs/protocol/src/error.rs b/codex-rs/protocol/src/error.rs index f421db8af5ed..d2601f2424cc 100644 --- a/codex-rs/protocol/src/error.rs +++ b/codex-rs/protocol/src/error.rs @@ -472,40 +472,44 @@ impl std::fmt::Display for UsageLimitReachedError { } let message = match self.plan_type.as_ref() { - Some(PlanType::Known(KnownPlan::Plus)) => 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{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ), - Some(PlanType::Known( + Some(PlanType::Known(plan)) => match plan { + KnownPlan::Plus => 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{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ), + plan if plan.is_team_like() || plan.is_business_like() => { + format!( + "You've hit your usage limit. To get more access now, send a request to your admin{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ) + } + KnownPlan::Free | KnownPlan::Go => { + format!( + "You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus),{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ) + } + KnownPlan::Pro | KnownPlan::ProLite => format!( + "You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits{}", + retry_suffix_after_or(self.resets_at.as_ref()) + ), + KnownPlan::Enterprise | KnownPlan::Edu => format!( + "You've hit your usage limit.{}", + retry_suffix(self.resets_at.as_ref()) + ), KnownPlan::Team | KnownPlan::SelfServeBusinessUsageBased | KnownPlan::Business - | KnownPlan::EnterpriseCbpUsageBased, - )) => { - format!( - "You've hit your usage limit. To get more access now, send a request to your admin{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ) - } - Some(PlanType::Known(KnownPlan::Free)) | Some(PlanType::Known(KnownPlan::Go)) => { + | KnownPlan::EnterpriseCbpUsageBased => { + unreachable!("team-like and business-like plans are handled above") + } + }, + Some(PlanType::Unknown(_)) | None => { format!( - "You've hit your usage limit. Upgrade to Plus to continue using Codex (https://chatgpt.com/explore/plus),{}", - retry_suffix_after_or(self.resets_at.as_ref()) + "You've hit your usage limit.{}", + retry_suffix(self.resets_at.as_ref()) ) } - Some(PlanType::Known(KnownPlan::Pro | KnownPlan::ProLite)) => format!( - "You've hit your usage limit. Visit https://chatgpt.com/codex/settings/usage to purchase more credits{}", - retry_suffix_after_or(self.resets_at.as_ref()) - ), - Some(PlanType::Known(KnownPlan::Enterprise)) - | Some(PlanType::Known(KnownPlan::Edu)) => format!( - "You've hit your usage limit.{}", - retry_suffix(self.resets_at.as_ref()) - ), - Some(PlanType::Unknown(_)) | None => format!( - "You've hit your usage limit.{}", - retry_suffix(self.resets_at.as_ref()) - ), }; write!(f, "{message}")