diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 9592563bece..481e707574b 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -177,12 +177,16 @@ impl AccountRequestProcessor { .set_auth_mode(self.auth_manager.get_api_auth_mode()); } - fn current_account_updated_notification(&self) -> AccountUpdatedNotification { - let auth = self.auth_manager.auth_cached(); + async fn current_account_updated_notification(&self) -> AccountUpdatedNotification { + let provider = create_model_provider( + self.config.model_provider.clone(), + Some(self.auth_manager.clone()), + ); + let auth = provider.auth().await; AccountUpdatedNotification { auth_mode: auth .as_ref() - .map(CodexAuth::api_auth_mode) + .map(CodexAuth::auth_mode) .map(auth_mode_to_api), plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type), } @@ -657,7 +661,7 @@ impl AccountRequestProcessor { self.outgoing .send_server_notification(ServerNotification::AccountUpdated( - self.current_account_updated_notification(), + self.current_account_updated_notification().await, )) .await; } diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index d3f9b7b96b0..bcf03941131 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -10,6 +10,7 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_core::config::Config; use codex_login::AuthKeyringBackendKind; +use codex_login::AuthManager; use codex_login::AuthRouteConfig; use codex_login::CLIENT_ID; use codex_login::CodexAuth; @@ -19,6 +20,7 @@ use codex_login::login_with_api_key; use codex_login::logout_with_revoke; use codex_login::run_device_code_login; use codex_login::run_login_server; +use codex_model_provider::create_model_provider; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ForcedLoginMethod; use codex_utils_cli::CliConfigOverrides; @@ -423,51 +425,50 @@ pub async fn run_login_with_device_code_fallback_to_browser( pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; - let auth_route_config = config.auth_route_config(); + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager)); + let auth = provider.auth().await; + let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); + + if config.model_provider.is_amazon_bedrock() { + if auth_mode == Some(AuthMode::BedrockApiKey) { + eprintln!("Logged in using Amazon Bedrock API key"); + } else { + eprintln!("Using Amazon Bedrock AWS SDK credential chain"); + } + std::process::exit(0); + } - match CodexAuth::from_auth_storage( - &config.codex_home, - config.cli_auth_credentials_store_mode, - Some(&config.chatgpt_base_url), - config.auth_keyring_backend_kind(), - auth_route_config.as_ref(), - ) - .await - { - Ok(Some(auth)) => match auth.auth_mode() { - AuthMode::ApiKey => match auth.get_token() { - Ok(api_key) => { - eprintln!("Logged in using an API key - {}", safe_format_key(&api_key)); - std::process::exit(0); - } - Err(e) => { - eprintln!("Unexpected error retrieving API key: {e}"); - std::process::exit(1); - } - }, - AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens => { - eprintln!("Logged in using ChatGPT"); + match (auth_mode, auth) { + (Some(AuthMode::ApiKey), Some(auth)) => match auth.get_token() { + Ok(api_key) => { + eprintln!("Logged in using an API key - {}", safe_format_key(&api_key)); std::process::exit(0); } - AuthMode::AgentIdentity => { - eprintln!("Logged in using access token"); - std::process::exit(0); - } - AuthMode::PersonalAccessToken => { - eprintln!("Logged in using personal access token"); - std::process::exit(0); - } - AuthMode::BedrockApiKey => { - eprintln!("Logged in using Amazon Bedrock API key"); - std::process::exit(0); + Err(e) => { + eprintln!("Unexpected error retrieving API key: {e}"); + std::process::exit(1); } }, - Ok(None) => { + (Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens), Some(_)) => { + eprintln!("Logged in using ChatGPT"); + std::process::exit(0); + } + (Some(AuthMode::AgentIdentity), Some(_)) => { + eprintln!("Logged in using access token"); + std::process::exit(0); + } + (Some(AuthMode::PersonalAccessToken), Some(_)) => { + eprintln!("Logged in using personal access token"); + std::process::exit(0); + } + (None, None) => { eprintln!("Not logged in"); std::process::exit(1); } - Err(e) => { - eprintln!("Error checking login status: {e}"); + (auth_mode, _) => { + eprintln!("Unexpected provider auth state: mode={auth_mode:?}"); std::process::exit(1); } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 9509cab4566..3936a17cdab 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1596,26 +1596,47 @@ async fn personal_access_token_does_not_offer_unauthorized_recovery() { #[serial(codex_auth_env)] async fn load_auth_keeps_codex_api_key_env_precedence() { let codex_home = tempdir().unwrap(); + let bedrock_auth = BedrockApiKeyAuth { + api_key: "bedrock-api-key-test".to_string(), + region: "us-east-1".to_string(), + }; let record = agent_identity_record(WORKSPACE_ID_ALLOWED); let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity"); let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity); let _api_key_guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env"); - let auth = super::load_auth( - codex_home.path(), + let auth_manager = AuthManager::new( + codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ true, AuthCredentialsStoreMode::File, /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::Direct, - /*agent_identity_authapi_base_url*/ None, /*auth_route_config*/ None, ) - .await - .expect("env auth should load") - .expect("env auth should be present"); + .await; + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); - assert_eq!(auth.api_key(), Some("sk-env")); + crate::auth::login_with_bedrock_api_key( + codex_home.path(), + &bedrock_auth.api_key, + &bedrock_auth.region, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .expect("seed Bedrock auth"); + assert!(auth_manager.reload().await); + + assert_eq!( + auth_manager + .auth_cached() + .and_then(|auth| auth.api_key().map(str::to_string)), + Some("sk-env".to_string()) + ); + assert_eq!( + auth_manager.bedrock_api_key_auth_cached(), + Some(bedrock_auth) + ); } #[tokio::test] diff --git a/codex-rs/login/src/auth/bedrock_api_key.rs b/codex-rs/login/src/auth/bedrock_api_key.rs index 3445878365a..57c02e058eb 100644 --- a/codex-rs/login/src/auth/bedrock_api_key.rs +++ b/codex-rs/login/src/auth/bedrock_api_key.rs @@ -1,21 +1,33 @@ +use std::fmt; use std::path::Path; use codex_config::types::AuthCredentialsStoreMode; use serde::Deserialize; use serde::Serialize; +use super::manager::AuthManager; +use super::manager::load_auth_dot_json; use super::manager::save_auth; use super::storage::AuthDotJson; use super::storage::AuthKeyringBackendKind; use codex_protocol::auth::AuthMode; /// Managed Amazon Bedrock API key persisted in `auth.json`. -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +#[derive(Deserialize, Serialize, Clone, PartialEq, Eq)] pub struct BedrockApiKeyAuth { pub api_key: String, pub region: String, } +impl fmt::Debug for BedrockApiKeyAuth { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BedrockApiKeyAuth") + .field("api_key", &"") + .field("region", &self.region) + .finish() + } +} + /// Writes an `auth.json` that contains only the Amazon Bedrock API key auth. pub fn login_with_bedrock_api_key( codex_home: &Path, @@ -44,6 +56,67 @@ pub fn login_with_bedrock_api_key( ) } +pub(super) fn load_stored_bedrock_api_key( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> std::io::Result> { + load_stored_bedrock_api_key_auth( + codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + .map(|auth| auth.and_then(|auth| auth.bedrock_api_key)) +} + +fn load_stored_bedrock_api_key_auth( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> std::io::Result> { + load_auth_dot_json( + codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + .map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key)) +} + +impl AuthManager { + pub fn has_stored_bedrock_api_key_auth(&self) -> std::io::Result { + load_stored_bedrock_api_key_auth( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + .map(|auth| auth.is_some()) + } + + /// Persists managed Bedrock auth without changing the in-memory auth snapshot. + /// + /// The account coordinator publishes the provider-scoped cache update only + /// after the associated config mutation succeeds. + pub fn persist_bedrock_api_key_auth(&self, api_key: &str, region: &str) -> std::io::Result<()> { + login_with_bedrock_api_key( + &self.codex_home, + api_key, + region, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + } + + /// Reloads only managed Bedrock auth, preserving cached general Codex auth. + pub fn reload_bedrock_api_key_auth(&self) -> std::io::Result { + let bedrock_api_key = load_stored_bedrock_api_key( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + )?; + Ok(self.set_cached_bedrock_api_key_auth(bedrock_api_key)) + } +} + #[cfg(test)] #[path = "bedrock_api_key_tests.rs"] mod tests; diff --git a/codex-rs/login/src/auth/bedrock_api_key_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs index 1f8fab68b2c..a00fbe68e2f 100644 --- a/codex-rs/login/src/auth/bedrock_api_key_tests.rs +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -1,13 +1,15 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_protocol::auth::AuthMode; +use codex_protocol::config_types::ForcedLoginMethod; use pretty_assertions::assert_eq; use serial_test::serial; use tempfile::tempdir; use super::*; +use crate::auth::AuthConfig; use crate::auth::AuthKeyringBackendKind; use crate::auth::AuthManager; -use crate::auth::CodexAuth; +use crate::auth::enforce_login_restrictions; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::FileAuthStorage; @@ -42,20 +44,24 @@ fn bedrock_auth() -> BedrockApiKeyAuth { } } +fn auth_config(codex_home: &std::path::Path, forced_login_method: ForcedLoginMethod) -> AuthConfig { + AuthConfig { + codex_home: codex_home.to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + keyring_backend_kind: AuthKeyringBackendKind::default(), + forced_login_method: Some(forced_login_method), + chatgpt_base_url: None, + forced_chatgpt_workspace_id: None, + auth_route_config: None, + } +} + #[tokio::test] #[serial(codex_auth_env)] -async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> { +async fn managed_bedrock_persistence_preserves_cached_openai_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); storage.save(&api_key_auth())?; - login_with_bedrock_api_key( - codex_home.path(), - "bedrock-api-key-test", - "us-east-1", - AuthCredentialsStoreMode::File, - AuthKeyringBackendKind::default(), - )?; - let auth_manager = AuthManager::new( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, @@ -67,6 +73,8 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> ) .await; + auth_manager.persist_bedrock_api_key_auth("bedrock-api-key-test", "us-east-1")?; + let loaded = storage.load()?.expect("auth should be stored"); let expected = AuthDotJson { auth_mode: Some(AuthMode::BedrockApiKey), @@ -78,24 +86,33 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> bedrock_api_key: Some(bedrock_auth()), }; assert_eq!(loaded, expected); - assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); + assert_eq!(auth_manager.auth_mode(), Some(AuthMode::ApiKey)); + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); + assert!(auth_manager.reload_bedrock_api_key_auth()?); assert_eq!( - auth_manager.auth_cached().and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, - }), + auth_manager + .auth_cached() + .and_then(|auth| auth.api_key().map(str::to_string)), + Some("sk-test-key".to_string()) + ); + assert_eq!( + auth_manager.bedrock_api_key_auth_cached(), Some(bedrock_auth()) ); Ok(()) } +#[test] +fn bedrock_auth_debug_redacts_api_key() { + assert_eq!( + format!("{:?}", bedrock_auth()), + "BedrockApiKeyAuth { api_key: \"\", region: \"us-east-1\" }" + ); +} + #[tokio::test] #[serial(codex_auth_env)] -async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { +async fn auth_manager_logout_removes_bedrock_api_key_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); login_with_bedrock_api_key( @@ -120,12 +137,13 @@ async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { assert_eq!(storage.load()?, None); assert_eq!(auth_manager.auth_cached(), None); + assert_eq!(auth_manager.bedrock_api_key_auth_cached(), None); Ok(()) } #[tokio::test] #[serial(codex_auth_env)] -async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> { +async fn bedrock_only_auth_storage_does_not_create_codex_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); storage.save(&bedrock_only_auth())?; @@ -141,18 +159,12 @@ async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> ) .await; - assert_eq!(auth_manager.auth_mode(), Some(AuthMode::BedrockApiKey)); + assert_eq!(auth_manager.auth_mode(), None); assert_eq!( - auth_manager.auth_cached().and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, - }), + auth_manager.bedrock_api_key_auth_cached(), Some(bedrock_auth()) ); + assert_eq!(auth_manager.auth_cached(), None); Ok(()) } @@ -178,3 +190,35 @@ async fn login_with_api_key_clears_bedrock_api_key() -> anyhow::Result<()> { assert_eq!(storage.load()?, Some(api_key_auth())); Ok(()) } + +#[tokio::test] +async fn forced_chatgpt_login_removes_bedrock_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + storage.save(&bedrock_only_auth())?; + + let error = + enforce_login_restrictions(&auth_config(codex_home.path(), ForcedLoginMethod::Chatgpt)) + .await + .expect_err("Bedrock auth should violate forced ChatGPT login"); + + assert_eq!( + error.to_string(), + "ChatGPT login is required, but an API key is currently being used. Logging out." + ); + assert_eq!(storage.load()?, None); + Ok(()) +} + +#[tokio::test] +async fn forced_api_login_allows_bedrock_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth = bedrock_only_auth(); + storage.save(&auth)?; + + enforce_login_restrictions(&auth_config(codex_home.path(), ForcedLoginMethod::Api)).await?; + + assert_eq!(storage.load()?, Some(auth)); + Ok(()) +} diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 844410874d5..0a9b240c4c4 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -36,6 +36,7 @@ use super::agent_identity::record_needs_task_registration; use super::agent_identity::register_managed_chatgpt_agent_identity; use super::agent_identity::require_agent_identity_authapi_base_url; use super::agent_identity::verified_record_from_jwt; +use super::bedrock_api_key::load_stored_bedrock_api_key; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::agent_identity::AgentIdentityAuth; @@ -1065,7 +1066,7 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( config: &AuthConfig, agent_identity_authapi_base_url: Option<&str>, ) -> std::io::Result<()> { - let Some(auth) = load_auth( + let auth = load_auth( &config.codex_home, /*enable_codex_api_key_env*/ true, config.auth_credentials_store_mode, @@ -1075,13 +1076,23 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( agent_identity_authapi_base_url, config.auth_route_config.as_ref(), ) - .await? - else { + .await?; + let stored_bedrock_api_key = load_stored_bedrock_api_key( + &config.codex_home, + config.auth_credentials_store_mode, + config.keyring_backend_kind, + )?; + let auth_mode = auth.as_ref().map(CodexAuth::auth_mode).or_else(|| { + stored_bedrock_api_key + .as_ref() + .map(|_| AuthMode::BedrockApiKey) + }); + let Some(auth_mode) = auth_mode else { return Ok(()); }; if let Some(required_method) = config.forced_login_method { - let method_violation = match (required_method, auth.auth_mode()) { + let method_violation = match (required_method, auth_mode) { (ForcedLoginMethod::Api, AuthMode::ApiKey) | (ForcedLoginMethod::Api, AuthMode::BedrockApiKey) => None, (ForcedLoginMethod::Chatgpt, AuthMode::Chatgpt) @@ -1113,7 +1124,10 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( } if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { - let chatgpt_account_id = match &auth { + let Some(auth) = auth.as_ref() else { + return Ok(()); + }; + let chatgpt_account_id = match auth { CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => return Ok(()), CodexAuth::AgentIdentity(_) | CodexAuth::PersonalAccessToken(_) => { auth.get_account_id() @@ -1232,7 +1246,10 @@ async fn load_auth( AuthCredentialsStoreMode::Ephemeral, AuthKeyringBackendKind::default(), ); - if let Some(auth_dot_json) = ephemeral_storage.load()? { + if let Some(auth_dot_json) = ephemeral_storage + .load()? + .filter(|auth| !auth.is_bedrock_api_key()) + { let auth = CodexAuth::from_auth_dot_json( codex_home, auth_dot_json, @@ -1281,7 +1298,8 @@ async fn load_auth( keyring_backend_kind, ); let auth_dot_json = match storage.load()? { - Some(auth) => auth, + Some(auth) if !auth.is_bedrock_api_key() => auth, + Some(_) => return Ok(None), None => return Ok(None), }; @@ -1456,6 +1474,11 @@ fn refresh_token_endpoint() -> String { } impl AuthDotJson { + /// Returns whether this stored payload resolves to managed Amazon Bedrock API key auth. + pub fn is_bedrock_api_key(&self) -> bool { + self.resolved_mode() == AuthMode::BedrockApiKey + } + fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { let Some(chatgpt_metadata) = external.chatgpt_metadata() else { return Err(std::io::Error::other( @@ -1534,6 +1557,7 @@ impl AuthDotJson { #[derive(Clone)] struct CachedAuth { auth: Option, + bedrock_api_key: Option, /// Permanent refresh failure cached for the current auth snapshot so /// later refresh attempts for the same credentials fail fast without network. permanent_refresh_failure: Option, @@ -1552,6 +1576,7 @@ impl Debug for CachedAuth { "auth_mode", &self.auth.as_ref().map(CodexAuth::api_auth_mode), ) + .field("has_bedrock_api_key", &self.bedrock_api_key.is_some()) .field( "permanent_refresh_failure", &self @@ -1790,12 +1815,12 @@ impl UnauthorizedRecovery { /// `reload()` is called explicitly. This matches the design goal of avoiding /// different parts of the program seeing inconsistent auth data mid‑run. pub struct AuthManager { - codex_home: PathBuf, + pub(super) codex_home: PathBuf, inner: RwLock, auth_change_tx: watch::Sender, enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + pub(super) auth_credentials_store_mode: AuthCredentialsStoreMode, + pub(super) keyring_backend_kind: AuthKeyringBackendKind, forced_chatgpt_workspace_id: RwLock>>, chatgpt_base_url: Option, agent_identity_authapi_base_url: Option, @@ -1874,6 +1899,13 @@ impl AuthManager { ) -> Self { let agent_identity_authapi_base_url = agent_identity_authapi_base_url(chatgpt_base_url.as_deref()).ok(); + let bedrock_api_key = load_stored_bedrock_api_key( + &codex_home, + auth_credentials_store_mode, + keyring_backend_kind, + ) + .ok() + .flatten(); let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, @@ -1892,6 +1924,7 @@ impl AuthManager { codex_home, inner: RwLock::new(CachedAuth { auth: managed_auth, + bedrock_api_key, permanent_refresh_failure: None, }), auth_change_tx, @@ -1913,6 +1946,7 @@ impl AuthManager { pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { let cached = CachedAuth { auth: Some(auth), + bedrock_api_key: None, permanent_refresh_failure: None, }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); @@ -1939,6 +1973,7 @@ impl AuthManager { pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { let cached = CachedAuth { auth: Some(auth), + bedrock_api_key: None, permanent_refresh_failure: None, }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); @@ -1968,6 +2003,7 @@ impl AuthManager { ) -> Arc { let cached = CachedAuth { auth: Some(auth), + bedrock_api_key: None, permanent_refresh_failure: None, }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); @@ -1999,6 +2035,7 @@ impl AuthManager { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(CachedAuth { auth: None, + bedrock_api_key: None, permanent_refresh_failure: None, }), auth_change_tx, @@ -2023,6 +2060,14 @@ impl AuthManager { self.inner.read().ok().and_then(|c| c.auth.clone()) } + /// Current managed Amazon Bedrock API key without applying Codex auth precedence. + pub fn bedrock_api_key_auth_cached(&self) -> Option { + self.inner + .read() + .ok() + .and_then(|cached| cached.bedrock_api_key.clone()) + } + /// Subscribes to cached auth changes that can affect request recovery. pub fn auth_change_receiver(&self) -> watch::Receiver { self.auth_change_tx.subscribe() @@ -2125,7 +2170,8 @@ impl AuthManager { pub async fn reload(&self) -> bool { tracing::info!("Reloading auth"); let new_auth = self.load_auth_from_storage().await; - self.set_cached_auth(new_auth) + let new_bedrock_api_key = self.load_bedrock_api_key_from_storage(); + self.set_cached_auth_state(new_auth, new_bedrock_api_key) } async fn reload_if_account_id_matches( @@ -2230,6 +2276,60 @@ impl AuthManager { .flatten() } + fn load_bedrock_api_key_from_storage(&self) -> Option { + load_stored_bedrock_api_key( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + .ok() + .flatten() + } + + fn set_cached_auth_state( + &self, + new_auth: Option, + new_bedrock_api_key: Option, + ) -> bool { + if let Ok(mut guard) = self.inner.write() { + let previous = guard.auth.as_ref(); + let auth_changed = !Self::auths_equal(previous, new_auth.as_ref()); + let auth_changed_for_refresh = + !Self::auths_equal_for_refresh(previous, new_auth.as_ref()); + let bedrock_auth_changed = guard.bedrock_api_key != new_bedrock_api_key; + if auth_changed_for_refresh { + guard.permanent_refresh_failure = None; + } + let changed = auth_changed || bedrock_auth_changed; + tracing::info!("Reloaded auth, changed: {changed}"); + guard.auth = new_auth; + guard.bedrock_api_key = new_bedrock_api_key; + if auth_changed_for_refresh || bedrock_auth_changed { + self.auth_change_tx.send_modify(|revision| *revision += 1); + } + changed + } else { + false + } + } + + pub(super) fn set_cached_bedrock_api_key_auth( + &self, + new_bedrock_api_key: Option, + ) -> bool { + if let Ok(mut guard) = self.inner.write() { + let changed = guard.bedrock_api_key != new_bedrock_api_key; + tracing::info!("Reloaded managed Bedrock auth, changed: {changed}"); + guard.bedrock_api_key = new_bedrock_api_key; + if changed { + self.auth_change_tx.send_modify(|revision| *revision += 1); + } + changed + } else { + false + } + } + fn set_cached_auth(&self, new_auth: Option) -> bool { if let Ok(mut guard) = self.inner.write() { let previous = guard.auth.as_ref(); diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 108308c1bfa..af883bffccb 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -38,13 +38,13 @@ use mantle::runtime_base_url; pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, pub(crate) aws: ModelProviderAwsAuthInfo, - auth_manager: Option>, + managed_auth: Option, } impl AmazonBedrockModelProvider { pub(crate) fn new( provider_info: ModelProviderInfo, - auth_manager: Option>, + managed_auth: Option, ) -> Self { let aws = provider_info .aws @@ -56,22 +56,12 @@ impl AmazonBedrockModelProvider { Self { info: provider_info, aws, - auth_manager, + managed_auth, } } fn managed_auth(&self) -> Option { - self.auth_manager - .as_ref() - .and_then(|auth_manager| auth_manager.auth_cached()) - .and_then(|auth| match auth { - CodexAuth::BedrockApiKey(auth) => Some(auth), - CodexAuth::ApiKey(_) - | CodexAuth::Chatgpt(_) - | CodexAuth::ChatgptAuthTokens(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) => None, - }) + self.managed_auth.clone() } async fn auth(&self) -> Option { @@ -125,8 +115,7 @@ impl ModelProvider for AmazonBedrockModelProvider { } fn auth_manager(&self) -> Option> { - self.managed_auth() - .and_then(|_| self.auth_manager.as_ref().cloned()) + None } fn auth(&self) -> ModelProviderFuture<'_, Option> { @@ -206,22 +195,15 @@ mod tests { api_key: "managed-bedrock-api-key".to_string(), region: "us-east-1".to_string(), }; - let auth_manager = - AuthManager::from_auth_for_testing(CodexAuth::BedrockApiKey(managed_auth.clone())); let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: Some("aws-profile-that-should-not-be-loaded".to_string()), region: Some("us-west-2".to_string()), })), - Some(auth_manager.clone()), + Some(managed_auth.clone()), ); - assert!(Arc::ptr_eq( - &provider - .auth_manager() - .expect("managed Bedrock auth manager should be exposed"), - &auth_manager, - )); + assert!(provider.auth_manager().is_none()); assert_eq!( provider.auth().await, Some(CodexAuth::BedrockApiKey(managed_auth)) @@ -254,12 +236,10 @@ mod tests { } #[tokio::test] - async fn openai_auth_is_not_exposed_to_bedrock() { + async fn bedrock_without_managed_auth_uses_aws_account() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( - "openai-api-key", - ))), + /*managed_auth*/ None, ); assert!(provider.auth_manager().is_none()); @@ -279,7 +259,7 @@ mod tests { fn capabilities_disable_unsupported_hosted_tools() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - /*auth_manager*/ None, + /*managed_auth*/ None, ); assert_eq!( @@ -296,7 +276,7 @@ mod tests { fn approval_review_preferred_model_uses_bedrock_gpt_5_4() { let provider = AmazonBedrockModelProvider::new( ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - /*auth_manager*/ None, + /*managed_auth*/ None, ); assert_eq!( diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 5f101e52cd6..ab4addef3fa 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -220,7 +220,10 @@ pub fn create_model_provider( auth_manager: Option>, ) -> SharedModelProvider { if provider_info.is_amazon_bedrock() { - Arc::new(AmazonBedrockModelProvider::new(provider_info, auth_manager)) + let managed_auth = auth_manager + .as_ref() + .and_then(|auth_manager| auth_manager.bedrock_api_key_auth_cached()); + Arc::new(AmazonBedrockModelProvider::new(provider_info, managed_auth)) } else { Arc::new(ConfiguredModelProvider::new(provider_info, auth_manager)) } @@ -337,7 +340,6 @@ mod tests { use std::num::NonZeroU64; use codex_login::auth::AgentIdentityAuthPolicy; - use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; @@ -429,13 +431,6 @@ mod tests { .expect("valid model") } - fn bedrock_api_key_auth() -> CodexAuth { - CodexAuth::BedrockApiKey(BedrockApiKeyAuth { - api_key: "bedrock-api-key-test".to_string(), - region: "us-east-1".to_string(), - }) - } - #[tokio::test] async fn scoped_auth_ignores_scope_for_non_openai_provider() { let provider = create_model_provider( @@ -523,17 +518,6 @@ mod tests { assert!(provider.auth_manager().is_none()); } - #[tokio::test] - async fn create_model_provider_uses_managed_auth_for_amazon_bedrock_provider() { - let auth = bedrock_api_key_auth(); - let provider = create_model_provider( - ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None), - Some(AuthManager::from_auth_for_testing(auth.clone())), - ); - - assert_eq!(provider.auth().await, Some(auth)); - } - #[test] fn openai_provider_returns_unauthenticated_openai_account_state() { let provider = create_model_provider( @@ -589,19 +573,6 @@ mod tests { ); } - #[test] - fn openai_provider_rejects_bedrock_api_key_account_state() { - let provider = create_model_provider( - ModelProviderInfo::create_openai_provider(/*base_url*/ None), - Some(AuthManager::from_auth_for_testing(bedrock_api_key_auth())), - ); - - assert_eq!( - provider.account_state(), - Err(ProviderAccountError::UnsupportedBedrockApiKeyAuth) - ); - } - #[test] fn custom_non_openai_provider_returns_no_account_state() { let provider = create_model_provider(