From 8414dc291e58e6a5c8db0e1489bc1d6ac1a5dacc Mon Sep 17 00:00:00 2001 From: celia-oai Date: Sat, 13 Jun 2026 19:43:43 -0700 Subject: [PATCH 1/3] pp-server: add managed Bedrock login and logout --- .../schema/json/ClientRequest.json | 25 + .../codex_app_server_protocol.schemas.json | 41 ++ .../codex_app_server_protocol.v2.schemas.json | 41 ++ .../schema/json/v2/LoginAccountParams.json | 25 + .../schema/json/v2/LoginAccountResponse.json | 16 + .../typescript/v2/LoginAccountParams.ts | 2 +- .../typescript/v2/LoginAccountResponse.ts | 2 +- .../src/protocol/common.rs | 28 + .../src/protocol/v2/account.rs | 8 + codex-rs/app-server/README.md | 30 +- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/account_processor.rs | 123 ++++- .../src/request_processors/bedrock_auth.rs | 129 +++++ .../request_processors/config_processor.rs | 2 +- .../tests/common/test_app_server.rs | 14 + codex-rs/app-server/tests/suite/v2/account.rs | 506 ++++++++++++++++++ codex-rs/login/src/auth/manager.rs | 5 + .../src/amazon_bedrock/mantle.rs | 7 +- .../model-provider/src/amazon_bedrock/mod.rs | 1 + codex-rs/model-provider/src/lib.rs | 2 + 20 files changed, 979 insertions(+), 29 deletions(-) create mode 100644 codex-rs/app-server/src/request_processors/bedrock_auth.rs diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 2f39f821b6a8..dc4ad86b7c57 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1529,6 +1529,31 @@ ], "title": "ChatgptAuthTokensLoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock login is experimental.", + "properties": { + "apiKey": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockLoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "apiKey", + "region", + "type" + ], + "title": "AmazonBedrockLoginAccountParams", + "type": "object" } ] }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 0eedef903a09..75b5c15a64b1 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -11569,6 +11569,31 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock login is experimental.", + "properties": { + "apiKey": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "apiKey", + "region", + "type" + ], + "title": "AmazonBedrockv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" @@ -11662,6 +11687,22 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountResponse", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountResponseType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockv2::LoginAccountResponse", + "type": "object" } ], "title": "LoginAccountResponse" diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index b428f252586b..7fa21d15c128 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -7973,6 +7973,31 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock login is experimental.", + "properties": { + "apiKey": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "apiKey", + "region", + "type" + ], + "title": "AmazonBedrockv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" @@ -8066,6 +8091,22 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountResponse", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountResponseType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockv2::LoginAccountResponse", + "type": "object" } ], "title": "LoginAccountResponse" diff --git a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json index ab7b852c9185..474b7666cf82 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountParams.json @@ -89,6 +89,31 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountParams", "type": "object" + }, + { + "description": "[UNSTABLE] Managed Amazon Bedrock login is experimental.", + "properties": { + "apiKey": { + "type": "string" + }, + "region": { + "type": "string" + }, + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountParamsType", + "type": "string" + } + }, + "required": [ + "apiKey", + "region", + "type" + ], + "title": "AmazonBedrockv2::LoginAccountParams", + "type": "object" } ], "title": "LoginAccountParams" diff --git a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json index a800bffccd94..802440d637bf 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/LoginAccountResponse.json @@ -87,6 +87,22 @@ ], "title": "ChatgptAuthTokensv2::LoginAccountResponse", "type": "object" + }, + { + "properties": { + "type": { + "enum": [ + "amazonBedrock" + ], + "title": "AmazonBedrockv2::LoginAccountResponseType", + "type": "string" + } + }, + "required": [ + "type" + ], + "title": "AmazonBedrockv2::LoginAccountResponse", + "type": "object" } ], "title": "LoginAccountResponse" diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts index e6f1e2ed4369..88ee4434cfe8 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountParams.ts @@ -18,4 +18,4 @@ chatgptAccountId: string, * When `null`, Codex attempts to derive the plan type from access-token * claims. If unavailable, the plan defaults to `unknown`. */ -chatgptPlanType?: string | null, }; +chatgptPlanType?: string | null, } | { "type": "amazonBedrock", apiKey: string, region: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts index 34bccd6578e1..5a9f34ead3c1 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/LoginAccountResponse.ts @@ -14,4 +14,4 @@ verificationUrl: string, /** * One-time code the user must enter after signing in. */ -userCode: string, } | { "type": "chatgptAuthTokens", }; +userCode: string, } | { "type": "chatgptAuthTokens", } | { "type": "amazonBedrock", }; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index d8db99d4ffeb..a78be7ef41f8 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -2716,6 +2716,34 @@ mod tests { Ok(()) } + #[test] + fn serialize_account_login_amazon_bedrock() -> Result<()> { + let request = ClientRequest::LoginAccount { + request_id: RequestId::Integer(2), + params: v2::LoginAccountParams::AmazonBedrock { + api_key: "secret".to_string(), + region: "us-west-2".to_string(), + }, + }; + assert_eq!( + json!({ + "method": "account/login/start", + "id": 2, + "params": { + "type": "amazonBedrock", + "apiKey": "secret", + "region": "us-west-2" + } + }), + serde_json::to_value(&request)?, + ); + assert_eq!( + json!({"type": "amazonBedrock"}), + serde_json::to_value(v2::LoginAccountResponse::AmazonBedrock {})?, + ); + Ok(()) + } + #[test] fn serialize_account_login_chatgpt() -> Result<()> { let request = ClientRequest::LoginAccount { diff --git a/codex-rs/app-server-protocol/src/protocol/v2/account.rs b/codex-rs/app-server-protocol/src/protocol/v2/account.rs index 788c8f7aafda..5268c318a6e5 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/account.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/account.rs @@ -100,6 +100,11 @@ pub enum LoginAccountParams { #[ts(optional = nullable)] chatgpt_plan_type: Option, }, + /// [UNSTABLE] Managed Amazon Bedrock login is experimental. + #[experimental("account/login/start.amazonBedrock")] + #[serde(rename = "amazonBedrock", rename_all = "camelCase")] + #[ts(rename = "amazonBedrock", rename_all = "camelCase")] + AmazonBedrock { api_key: String, region: String }, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] @@ -133,6 +138,9 @@ pub enum LoginAccountResponse { #[serde(rename = "chatgptAuthTokens", rename_all = "camelCase")] #[ts(rename = "chatgptAuthTokens", rename_all = "camelCase")] ChatgptAuthTokens {}, + #[serde(rename = "amazonBedrock", rename_all = "camelCase")] + #[ts(rename = "amazonBedrock", rename_all = "camelCase")] + AmazonBedrock {}, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index a7986753ee27..e267167eee95 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1919,16 +1919,17 @@ Codex supports these authentication modes. The current mode is surfaced in `acco - **API key (`apiKey`)**: Caller supplies an OpenAI API key via `account/login/start` with `type: "apiKey"`. The API key is saved and used for API requests. - **ChatGPT managed (`chatgpt`)** (recommended): Codex owns the ChatGPT OAuth flow and refresh tokens. Start via `account/login/start` with `type: "chatgpt"` for the browser flow or `type: "chatgptDeviceCode"` for device code; Codex persists tokens to disk and refreshes them automatically. +- **Codex managed Amazon Bedrock auth (`amazonBedrock`, experimental)**: Caller supplies an Amazon Bedrock API key and region via `account/login/start` with `type: "amazonBedrock"`. The client must enable the `experimentalApi` initialization capability for Codex-managed Amazon Bedrock login. Codex saves the credential and writes `model_provider = "amazon-bedrock"` to the user config. The parameterless `account/logout` method clears managed Bedrock credentials on a best-effort basis. - **Personal access token (`personalAccessToken`)**: Codex uses a ChatGPT-backed personal access token loaded outside the app-server login RPCs, such as with `codex login --with-access-token` or `CODEX_ACCESS_TOKEN`. ### API Overview - `account/read` — fetch current account info; optionally refresh tokens. -- `account/login/start` — begin login (`apiKey`, `chatgpt`, `chatgptDeviceCode`). +- `account/login/start` — begin login (`apiKey`, `chatgpt`, `chatgptDeviceCode`, `amazonBedrock`). - `account/login/completed` (notify) — emitted when a login attempt finishes (success or error). - `account/login/cancel` — cancel a pending managed ChatGPT login by `loginId`. -- `account/logout` — sign out; triggers `account/updated`. -- `account/updated` (notify) — emitted whenever auth mode changes (`authMode`: `apikey`, `chatgpt`, `personalAccessToken`, or `null`) and includes the current ChatGPT `planType` when available. +- `account/logout` — sign out without params; also clears Codex-managed Amazon Bedrock credentials on a best-effort basis; triggers `account/updated` when the OpenAI auth mode changes. +- `account/updated` (notify) — emitted whenever auth mode changes (`authMode`: `apikey`, `bedrockApiKey`, `chatgpt`, `personalAccessToken`, or `null`) and includes the current ChatGPT `planType` when available. - `account/rateLimits/read` — fetch ChatGPT rate limits, an optional effective monthly credit limit, and the earned rate-limit resets currently available, including expiry details when provided by the backend. Rate-limit updates arrive via `account/rateLimits/updated` (notify); reset-credit data is snapshot-only. - `account/rateLimitResetCredit/consume` — consume one earned reset using a caller-provided idempotency key, optionally selecting a reset-credit ID returned by `account/rateLimits/read`. - `account/usage/read` — fetch ChatGPT account token-activity summary and daily buckets. @@ -1999,6 +2000,29 @@ Field notes: { "method": "account/updated", "params": { "authMode": "chatgpt", "planType": "plus" } } ``` +### 3) Log in with an Amazon Bedrock API key + +This experimental flow requires the client to initialize with `experimentalApi: true`. + +1. Send: + ```json + { + "method": "account/login/start", + "id": 3, + "params": { "type": "amazonBedrock", "apiKey": "…", "region": "us-west-2" } + } + ``` +2. Expect: + ```json + { "id": 3, "result": { "type": "amazonBedrock" } } + ``` +3. Notifications: + ```json + { "method": "account/login/completed", "params": { "loginId": null, "success": true, "error": null } } + ``` + +Codex stores the key and region in its configured auth store and writes `model_provider = "amazon-bedrock"` to the active user config. + ### 4) Log in with ChatGPT (device code flow) 1. Start: diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index f03351468faa..119ea6179061 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -495,6 +495,7 @@ use codex_app_server_protocol::ServerRequest; mod account_processor; mod apps_processor; +mod bedrock_auth; mod catalog_processor; mod command_exec_processor; mod config_processor; 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 9592563bece0..f317cc45a007 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,6 +1,12 @@ +use super::bedrock_auth::has_managed_login as has_managed_bedrock_login; +use super::bedrock_auth::login as login_managed_bedrock; +use super::bedrock_auth::logout as logout_managed_bedrock; use super::*; use crate::auth_mode::auth_mode_to_api; use chrono::DateTime; +use codex_login::auth::read_codex_api_key_from_env; +use codex_login::read_codex_access_token_from_env; +use codex_model_provider::is_supported_amazon_bedrock_region; mod rate_limit_resets; @@ -274,6 +280,10 @@ impl AccountRequestProcessor { ) .await; } + LoginAccountParams::AmazonBedrock { api_key, region } => { + self.login_amazon_bedrock_v2(request_id, api_key, region) + .await; + } } Ok(()) } @@ -337,6 +347,64 @@ impl AccountRequestProcessor { } } + async fn login_amazon_bedrock_v2( + &self, + request_id: ConnectionRequestId, + api_key: String, + region: String, + ) { + let result = async { + if self.auth_manager.is_external_chatgpt_auth_active() { + return Err(self.external_auth_active_error()); + } + if read_codex_access_token_from_env().is_some() + || (self.auth_manager.codex_api_key_env_enabled() + && read_codex_api_key_from_env().is_some()) + { + return Err(invalid_request( + "Amazon Bedrock login is unavailable while Codex auth is supplied through the environment.", + )); + } + if matches!( + self.config.forced_login_method, + Some(ForcedLoginMethod::Chatgpt) + ) { + return Err(invalid_request( + "Amazon Bedrock login is disabled. Use ChatGPT login instead.", + )); + } + + let api_key = api_key.trim(); + if api_key.is_empty() { + return Err(invalid_request("Amazon Bedrock API key must not be empty.")); + } + let region = region.trim(); + if !is_supported_amazon_bedrock_region(region) { + return Err(invalid_request(format!( + "Amazon Bedrock Mantle does not support region `{region}`" + ))); + } + + { + let mut guard = self.active_login.lock().await; + if let Some(active) = guard.take() { + drop(active); + } + } + + login_managed_bedrock(&self.config, &self.config_manager, api_key, region).await?; + Ok(LoginAccountResponse::AmazonBedrock {}) + } + .await; + let logged_in = result.is_ok(); + self.outgoing.send_result(request_id, result).await; + + if logged_in { + self.emit_login_completed_notification(/*login_id*/ None) + .await; + } + } + // Build options for a ChatGPT login attempt; performs validation. async fn login_chatgpt_common( &self, @@ -644,6 +712,20 @@ impl AccountRequestProcessor { ) .await; + self.emit_login_success_notifications(login_id).await; + } + + async fn emit_login_success_notifications(&self, login_id: Option) { + self.emit_login_completed_notification(login_id).await; + + self.outgoing + .send_server_notification(ServerNotification::AccountUpdated( + self.current_account_updated_notification(), + )) + .await; + } + + async fn emit_login_completed_notification(&self, login_id: Option) { let payload_login_completed = AccountLoginCompletedNotification { login_id: login_id.map(|id| id.to_string()), success: true, @@ -654,12 +736,6 @@ impl AccountRequestProcessor { payload_login_completed, )) .await; - - self.outgoing - .send_server_notification(ServerNotification::AccountUpdated( - self.current_account_updated_notification(), - )) - .await; } async fn send_chatgpt_login_completion_notifications( @@ -709,7 +785,7 @@ impl AccountRequestProcessor { } } - async fn logout_common(&self) -> std::result::Result, JSONRPCErrorError> { + async fn logout_common(&self) -> Result, JSONRPCErrorError> { // Cancel any active login attempt. { let mut guard = self.active_login.lock().await; @@ -718,6 +794,14 @@ impl AccountRequestProcessor { } } + if has_managed_bedrock_login(&self.config)? { + logout_managed_bedrock(&self.config, &self.config_manager).await?; + return Ok(None); + } + if self.config.model_provider.is_amazon_bedrock() { + return Ok(None); + } + match self.auth_manager.logout_with_revoke().await { Ok(_) => {} Err(err) => { @@ -733,25 +817,20 @@ impl AccountRequestProcessor { .await; // Reflect the current auth method after logout (likely None). - Ok(self - .auth_manager - .auth_cached() - .as_ref() - .map(CodexAuth::api_auth_mode) - .map(auth_mode_to_api)) + Ok(Some(AccountUpdatedNotification { + auth_mode: self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode) + .map(auth_mode_to_api), + plan_type: None, + })) } async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> { let result = self.logout_common().await; - let account_updated = - result - .as_ref() - .ok() - .cloned() - .map(|auth_mode| AccountUpdatedNotification { - auth_mode, - plan_type: None, - }); + let account_updated = result.as_ref().ok().cloned().flatten(); self.outgoing .send_result(request_id, result.map(|_| LogoutAccountResponse {})) .await; diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs new file mode 100644 index 000000000000..72e1866c2d51 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -0,0 +1,129 @@ +use super::config_processor::map_error as map_config_error; +use crate::config_manager::ConfigManager; +use crate::error_code::internal_error; +use codex_app_server_protocol::ConfigLayerSource; +use codex_app_server_protocol::ConfigReadParams; +use codex_app_server_protocol::ConfigValueWriteParams; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::MergeStrategy; +use codex_core::config::Config; +use codex_login::AuthDotJson; +use codex_login::load_auth_dot_json; +use codex_login::login_with_bedrock_api_key; +use codex_login::logout as logout_stored_auth; +use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID; + +struct UserModelProviderState { + model_provider: Option, + version: Option, +} + +pub(super) fn has_managed_login(config: &Config) -> Result { + load_stored_auth(config).map(|auth| auth.is_some()) +} + +pub(super) async fn login( + config: &Config, + config_manager: &ConfigManager, + api_key: &str, + region: &str, +) -> Result<(), JSONRPCErrorError> { + login_with_bedrock_api_key( + &config.codex_home, + api_key, + region, + config.cli_auth_credentials_store_mode, + ) + .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; + + write_user_model_provider( + config_manager, + serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), + /*expected_version*/ None, + ) + .await +} + +pub(super) async fn logout( + config: &Config, + config_manager: &ConfigManager, +) -> Result<(), JSONRPCErrorError> { + let user_model_provider = user_model_provider_state(config_manager).await?; + logout_stored_auth(&config.codex_home, config.cli_auth_credentials_store_mode) + .map_err(|err| internal_error(format!("logout failed: {err}")))?; + + if user_model_provider.model_provider.as_deref() == Some(AMAZON_BEDROCK_PROVIDER_ID) { + write_user_model_provider( + config_manager, + serde_json::Value::Null, + user_model_provider.version, + ) + .await?; + } + + Ok(()) +} + +fn load_stored_auth(config: &Config) -> Result, JSONRPCErrorError> { + load_auth_dot_json(&config.codex_home, config.cli_auth_credentials_store_mode) + .map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key)) + .map_err(|err| internal_error(format!("failed to read stored auth: {err}"))) +} + +async fn user_model_provider_state( + config_manager: &ConfigManager, +) -> Result { + let user_config_path = config_manager + .user_config_path() + .map_err(|err| internal_error(format!("failed to resolve user config path: {err}")))?; + let response = config_manager + .read(ConfigReadParams { + include_layers: true, + cwd: None, + }) + .await + .map_err(|err| internal_error(format!("failed to read user config: {err}")))?; + let layer = response + .layers + .unwrap_or_default() + .into_iter() + .find(|layer| { + matches!( + &layer.name, + ConfigLayerSource::User { file, .. } if file == &user_config_path + ) + }); + let Some(layer) = layer else { + return Ok(UserModelProviderState { + model_provider: None, + version: None, + }); + }; + let model_provider = layer + .config + .get("model_provider") + .and_then(serde_json::Value::as_str) + .map(str::to_string); + Ok(UserModelProviderState { + model_provider, + version: Some(layer.version), + }) +} + +async fn write_user_model_provider( + config_manager: &ConfigManager, + value: serde_json::Value, + expected_version: Option, +) -> Result<(), JSONRPCErrorError> { + config_manager + .write_value(ConfigValueWriteParams { + key_path: "model_provider".to_string(), + value, + merge_strategy: MergeStrategy::Replace, + file_path: None, + expected_version, + }) + .await + .map(|_| ()) + .map_err(map_config_error) +} diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index b6dbb30b0834..0eeb1aaca98a 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -555,7 +555,7 @@ fn map_network_unix_socket_permission_to_api( } } -fn map_error(err: ConfigManagerError) -> JSONRPCErrorError { +pub(super) fn map_error(err: ConfigManagerError) -> JSONRPCErrorError { if let Some(code) = err.write_error_code() { return config_write_error(code, err.to_string()); } diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index cbb8ff161e0f..3cedbee2b143 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -1392,6 +1392,20 @@ impl TestAppServer { self.send_request("account/login/start", Some(params)).await } + /// Send an `account/login/start` JSON-RPC request for managed Amazon Bedrock login. + pub async fn send_login_account_amazon_bedrock_request( + &mut self, + api_key: &str, + region: &str, + ) -> anyhow::Result { + let params = serde_json::json!({ + "type": "amazonBedrock", + "apiKey": api_key, + "region": region, + }); + self.send_request("account/login/start", Some(params)).await + } + /// Send an `account/login/start` JSON-RPC request for ChatGPT login. pub async fn send_login_account_chatgpt_request(&mut self) -> anyhow::Result { let params = serde_json::json!({ diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 3cd98e2e52d6..a11bcf2345db 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -5,24 +5,29 @@ use app_test_support::to_response; use app_test_support::ChatGptAuthFixture; use app_test_support::ChatGptIdTokenClaims; +use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::encode_id_token; use app_test_support::write_chatgpt_auth; use app_test_support::write_models_cache; use chrono::Duration as ChronoDuration; use chrono::Utc; use codex_app_server_protocol::Account; +use codex_app_server_protocol::AccountLoginCompletedNotification; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse; +use codex_app_server_protocol::ClientInfo; use codex_app_server_protocol::GetAccountParams; use codex_app_server_protocol::GetAccountResponse; use codex_app_server_protocol::GetAuthStatusParams; use codex_app_server_protocol::GetAuthStatusResponse; +use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::LoginAccountResponse; @@ -33,9 +38,12 @@ use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::TurnCompletedNotification; use codex_app_server_protocol::TurnStatus; use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::CLIENT_ID_OVERRIDE_ENV_VAR; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; +use codex_login::auth::BedrockApiKeyAuth; +use codex_login::load_auth_dot_json; use codex_login::login_with_api_key; use codex_login::login_with_bedrock_api_key; use codex_protocol::account::AmazonBedrockCredentialSource; @@ -141,6 +149,51 @@ shell_snapshot = false std::fs::write(config_toml, contents) } +fn read_config_toml(codex_home: &Path) -> Result { + Ok(toml::from_str(&std::fs::read_to_string( + codex_home.join("config.toml"), + )?)?) +} + +fn aws_managed_bedrock_config() -> CreateConfigTomlParams { + CreateConfigTomlParams { + model_provider_id: Some("amazon-bedrock".to_string()), + extra_provider_config: Some( + r#"[model_providers.amazon-bedrock.aws] +profile = "codex-bedrock" +region = "us-west-2" +"# + .to_string(), + ), + ..Default::default() + } +} + +async fn read_account(mcp: &mut TestAppServer) -> Result { + let request_id = mcp + .send_get_account_request(GetAccountParams { + refresh_token: false, + }) + .await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + to_response(response) +} + +async fn assert_no_account_updated(mcp: &mut TestAppServer) { + assert!( + timeout( + Duration::from_millis(500), + mcp.read_stream_until_notification_message("account/updated"), + ) + .await + .is_err() + ); +} + async fn mock_device_code_usercode(server: &MockServer, interval_seconds: u64) { Mock::given(method("POST")) .and(path("/api/accounts/deviceauth/usercode")) @@ -961,6 +1014,459 @@ async fn login_account_api_key_succeeds_and_notifies() -> Result<()> { Ok(()) } +#[tokio::test] +async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + requires_openai_auth: Some(true), + ..Default::default() + }, + )?; + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + )?; + let mut expected_config = read_config_toml(codex_home.path())?; + expected_config + .as_table_mut() + .expect("config should be a table") + .insert( + "model_provider".to_string(), + toml::Value::String("amazon-bedrock".to_string()), + ); + let mut expected_config_after_logout = expected_config.clone(); + expected_config_after_logout + .as_table_mut() + .expect("config should be a table") + .remove("model_provider"); + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_login_account_amazon_bedrock_request(" managed-bedrock-api-key ", " us-west-2 ") + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + Some(AuthDotJson { + auth_mode: Some(AuthMode::BedrockApiKey), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(BedrockApiKeyAuth { + api_key: "managed-bedrock-api-key".to_string(), + region: "us-west-2".to_string(), + }), + }) + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + let ServerNotification::AccountLoginCompleted(payload) = notification.try_into()? else { + bail!("unexpected notification") + }; + assert_eq!( + payload, + AccountLoginCompletedNotification { + login_id: None, + success: true, + error: None, + } + ); + assert_no_account_updated(&mut mcp).await; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::ApiKey {}), + requires_openai_auth: true, + } + ); + + let request_id = mcp.send_logout_account_request().await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + assert_eq!( + read_config_toml(codex_home.path())?, + expected_config_after_logout + ); + assert_no_account_updated(&mut mcp).await; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::ApiKey {}), + requires_openai_auth: true, + } + ); + Ok(()) +} + +#[tokio::test] +async fn managed_bedrock_login_requires_experimental_api_but_logout_is_best_effort() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let mut mcp = TestAppServer::new(codex_home.path()).await?; + let initialized = mcp + .initialize_with_capabilities( + ClientInfo { + name: DEFAULT_CLIENT_NAME.to_string(), + title: None, + version: "0.1.0".to_string(), + }, + Some(InitializeCapabilities { + experimental_api: false, + ..Default::default() + }), + ) + .await?; + assert!(matches!(initialized, JSONRPCMessage::Response(_))); + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "account/login/start.amazonBedrock requires experimentalApi capability" + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + login_with_bedrock_api_key( + codex_home.path(), + "managed-bedrock-api-key", + "us-west-2", + AuthCredentialsStoreMode::File, + )?; + let request_id = mcp.send_logout_account_request().await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + Ok(()) +} + +#[tokio::test] +async fn login_and_logout_managed_bedrock_require_restart_when_provider_is_active() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + let mut expected_config = read_config_toml(codex_home.path())?; + expected_config + .as_table_mut() + .expect("config should be a table") + .remove("model_provider"); + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + assert_no_account_updated(&mut mcp).await; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + credential_source: AmazonBedrockCredentialSource::AwsManaged, + }), + requires_openai_auth: false, + } + ); + + assert!(codex_home.path().join("auth.json").exists()); + let request_id = mcp.send_logout_account_request().await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + + assert_no_account_updated(&mut mcp).await; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + credential_source: AmazonBedrockCredentialSource::AwsManaged, + }), + requires_openai_auth: false, + } + ); + Ok(()) +} + +#[tokio::test] +async fn logout_account_aws_managed_bedrock_preserves_openai_auth_and_config() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + login_with_api_key( + codex_home.path(), + "sk-test-key", + AuthCredentialsStoreMode::File, + )?; + let expected_auth = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?; + let expected_config = read_config_toml(codex_home.path())?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp.send_logout_account_request().await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + expected_auth + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert_no_account_updated(&mut mcp).await; + Ok(()) +} + +#[tokio::test] +async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_changes() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let expected_config = read_config_toml(codex_home.path())?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_login_account_amazon_bedrock_request(" ", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "Amazon Bedrock API key must not be empty." + ); + + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-1") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "Amazon Bedrock Mantle does not support region `us-west-1`" + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + + drop(mcp); + let mut mcp = TestAppServer::new_with_env( + codex_home.path(), + &[ + ("OPENAI_API_KEY", None), + ("CODEX_ACCESS_TOKEN", Some("env-access-token")), + ], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + error.error.message, + "Amazon Bedrock login is unavailable while Codex auth is supplied through the environment." + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + + Ok(()) +} + +#[tokio::test] +async fn login_account_amazon_bedrock_persists_under_provider_override() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let mut expected_config = read_config_toml(codex_home.path())?; + expected_config + .as_table_mut() + .expect("config should be a table") + .insert( + "model_provider".to_string(), + toml::Value::String("amazon-bedrock".to_string()), + ); + + let mut mcp = TestAppServer::new_with_args( + codex_home.path(), + &["-c", "model_provider=\"mock_provider\""], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let request_id = mcp + .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + Some(AuthDotJson { + auth_mode: Some(AuthMode::BedrockApiKey), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(BedrockApiKeyAuth { + api_key: "managed-bedrock-api-key".to_string(), + region: "us-west-2".to_string(), + }), + }) + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + assert_no_account_updated(&mut mcp).await; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: None, + requires_openai_auth: false, + } + ); + Ok(()) +} + +#[tokio::test] +async fn logout_account_managed_bedrock_preserves_changed_provider() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + login_with_bedrock_api_key( + codex_home.path(), + "managed-bedrock-api-key", + "us-west-2", + AuthCredentialsStoreMode::File, + )?; + let expected_config = read_config_toml(codex_home.path())?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp.send_logout_account_request().await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + to_response::(response)?, + LogoutAccountResponse {} + ); + assert_eq!( + load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + None + ); + assert_eq!(read_config_toml(codex_home.path())?, expected_config); + assert_no_account_updated(&mut mcp).await; + Ok(()) +} + #[tokio::test] async fn login_account_api_key_rejected_when_forced_chatgpt() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 844410874d58..f8a5fee38b46 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1456,6 +1456,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() == ApiAuthMode::BedrockApiKey + } + fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { let Some(chatgpt_metadata) = external.chatgpt_metadata() else { return Err(std::io::Error::other( diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 9727802d1196..e31ce96edc01 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -39,8 +39,13 @@ pub(super) fn region_from_config(aws: &ModelProviderAwsAuthInfo) -> Option bool { + BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) +} + pub(super) fn base_url(region: &str) -> Result { - if BEDROCK_MANTLE_SUPPORTED_REGIONS.contains(®ion) { + if is_supported_amazon_bedrock_region(region) { Ok(format!("https://bedrock-mantle.{region}.api.aws/openai/v1")) } else { Err(CodexErr::Fatal(format!( diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 108308c1bfae..914ddd4ca75e 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -31,6 +31,7 @@ use crate::provider::ProviderCapabilities; use auth::resolve_provider_auth; pub(crate) use catalog::static_model_catalog; use catalog::with_default_only_service_tier; +pub use mantle::is_supported_amazon_bedrock_region; use mantle::runtime_base_url; /// Runtime provider for Amazon Bedrock's OpenAI-compatible Mantle endpoint. diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 15967c4fe033..f8d1dd40475a 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -7,11 +7,13 @@ mod provider; pub use auth::AgentIdentitySessionFallback; pub use auth::ProviderAuthScope; pub use auth::ResolvedProviderAuth; +pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use auth::auth_provider_from_auth; pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; +pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; pub use codex_protocol::account::ProviderAccount; pub use provider::ModelProvider; pub use provider::ModelProviderFuture; From bcd4adfb9017a02f29741cbba1c702e64d8dc055 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Sat, 13 Jun 2026 23:16:31 -0700 Subject: [PATCH 2/3] changes --- codex-rs/app-server/src/request_processors.rs | 1 - .../request_processors/account_processor.rs | 66 ++++---- .../src/request_processors/bedrock_auth.rs | 19 ++- .../src/request_processors/plugins.rs | 1 + codex-rs/app-server/tests/suite/v2/account.rs | 159 ++++++++++++------ codex-rs/cli/src/doctor.rs | 8 +- codex-rs/cli/src/login.rs | 35 ++-- codex-rs/core/src/client.rs | 13 +- codex-rs/login/src/auth/auth_tests.rs | 35 +++- .../login/src/auth/bedrock_api_key_tests.rs | 74 +++++--- codex-rs/login/src/auth/manager.rs | 139 +++++++++++---- codex-rs/login/src/lib.rs | 1 + .../model-provider/src/amazon_bedrock/mod.rs | 59 +++---- codex-rs/model-provider/src/auth.rs | 30 ---- codex-rs/model-provider/src/provider.rs | 56 ++---- justfile | 11 +- 16 files changed, 425 insertions(+), 282 deletions(-) diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 119ea6179061..5b205f07d15d 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -36,7 +36,6 @@ use codex_app_server_protocol::AppTemplateUnavailableReason; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::AskForApproval; -use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; 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 f317cc45a007..eeb80bdab382 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -4,8 +4,6 @@ use super::bedrock_auth::logout as logout_managed_bedrock; use super::*; use crate::auth_mode::auth_mode_to_api; use chrono::DateTime; -use codex_login::auth::read_codex_api_key_from_env; -use codex_login::read_codex_access_token_from_env; use codex_model_provider::is_supported_amazon_bedrock_region; mod rate_limit_resets; @@ -357,14 +355,6 @@ impl AccountRequestProcessor { if self.auth_manager.is_external_chatgpt_auth_active() { return Err(self.external_auth_active_error()); } - if read_codex_access_token_from_env().is_some() - || (self.auth_manager.codex_api_key_env_enabled() - && read_codex_api_key_from_env().is_some()) - { - return Err(invalid_request( - "Amazon Bedrock login is unavailable while Codex auth is supplied through the environment.", - )); - } if matches!( self.config.forced_login_method, Some(ForcedLoginMethod::Chatgpt) @@ -393,6 +383,7 @@ impl AccountRequestProcessor { } login_managed_bedrock(&self.config, &self.config_manager, api_key, region).await?; + self.auth_manager.reload().await; Ok(LoginAccountResponse::AmazonBedrock {}) } .await; @@ -400,7 +391,7 @@ impl AccountRequestProcessor { self.outgoing.send_result(request_id, result).await; if logged_in { - self.emit_login_completed_notification(/*login_id*/ None) + self.send_login_success_notifications(/*login_id*/ None) .await; } } @@ -712,20 +703,6 @@ impl AccountRequestProcessor { ) .await; - self.emit_login_success_notifications(login_id).await; - } - - async fn emit_login_success_notifications(&self, login_id: Option) { - self.emit_login_completed_notification(login_id).await; - - self.outgoing - .send_server_notification(ServerNotification::AccountUpdated( - self.current_account_updated_notification(), - )) - .await; - } - - async fn emit_login_completed_notification(&self, login_id: Option) { let payload_login_completed = AccountLoginCompletedNotification { login_id: login_id.map(|id| id.to_string()), success: true, @@ -736,6 +713,12 @@ impl AccountRequestProcessor { payload_login_completed, )) .await; + + self.outgoing + .send_server_notification(ServerNotification::AccountUpdated( + self.current_account_updated_notification(), + )) + .await; } async fn send_chatgpt_login_completion_notifications( @@ -785,7 +768,7 @@ impl AccountRequestProcessor { } } - async fn logout_common(&self) -> Result, JSONRPCErrorError> { + async fn logout_common(&self) -> Result, JSONRPCErrorError> { // Cancel any active login attempt. { let mut guard = self.active_login.lock().await; @@ -796,6 +779,7 @@ impl AccountRequestProcessor { if has_managed_bedrock_login(&self.config)? { logout_managed_bedrock(&self.config, &self.config_manager).await?; + self.auth_manager.reload().await; return Ok(None); } if self.config.model_provider.is_amazon_bedrock() { @@ -817,20 +801,30 @@ impl AccountRequestProcessor { .await; // Reflect the current auth method after logout (likely None). - Ok(Some(AccountUpdatedNotification { - auth_mode: self - .auth_manager - .auth_cached() - .as_ref() - .map(CodexAuth::api_auth_mode) - .map(auth_mode_to_api), - plan_type: None, - })) + Ok(self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode) + .map(auth_mode_to_api)) } async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> { + let skip_account_updated = matches!(has_managed_bedrock_login(&self.config), Ok(true)) + || self.config.model_provider.is_amazon_bedrock(); let result = self.logout_common().await; - let account_updated = result.as_ref().ok().cloned().flatten(); + let account_updated = if skip_account_updated { + None + } else { + result + .as_ref() + .ok() + .cloned() + .map(|auth_mode| AccountUpdatedNotification { + auth_mode, + plan_type: None, + }) + }; self.outgoing .send_result(request_id, result.map(|_| LogoutAccountResponse {})) .await; diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs index 72e1866c2d51..eb2cf5122386 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -33,6 +33,7 @@ pub(super) async fn login( api_key, region, config.cli_auth_credentials_store_mode, + config.auth_keyring_backend_kind(), ) .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; @@ -49,8 +50,12 @@ pub(super) async fn logout( config_manager: &ConfigManager, ) -> Result<(), JSONRPCErrorError> { let user_model_provider = user_model_provider_state(config_manager).await?; - logout_stored_auth(&config.codex_home, config.cli_auth_credentials_store_mode) - .map_err(|err| internal_error(format!("logout failed: {err}")))?; + logout_stored_auth( + &config.codex_home, + config.cli_auth_credentials_store_mode, + config.auth_keyring_backend_kind(), + ) + .map_err(|err| internal_error(format!("logout failed: {err}")))?; if user_model_provider.model_provider.as_deref() == Some(AMAZON_BEDROCK_PROVIDER_ID) { write_user_model_provider( @@ -65,9 +70,13 @@ pub(super) async fn logout( } fn load_stored_auth(config: &Config) -> Result, JSONRPCErrorError> { - load_auth_dot_json(&config.codex_home, config.cli_auth_credentials_store_mode) - .map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key)) - .map_err(|err| internal_error(format!("failed to read stored auth: {err}"))) + load_auth_dot_json( + &config.codex_home, + config.cli_auth_credentials_store_mode, + config.auth_keyring_backend_kind(), + ) + .map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key)) + .map_err(|err| internal_error(format!("failed to read stored auth: {err}"))) } async fn user_model_provider_state( diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 35c12d49c1aa..c291f5055f78 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -19,6 +19,7 @@ use codex_core_plugins::remote::RemoteAppTemplateUnavailableReason; use codex_core_plugins::remote::is_valid_remote_plugin_id; use codex_core_plugins::remote::validate_remote_plugin_id; use codex_core_plugins::remote_bundle::RemotePluginBundleInstallError; +use codex_login::AuthMode; use codex_mcp::McpOAuthLoginSupport; use codex_mcp::oauth_login_support; use codex_mcp::should_retry_without_scopes; diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index a11bcf2345db..0a2e34410b4c 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -13,6 +13,7 @@ use chrono::Duration as ChronoDuration; use chrono::Utc; use codex_app_server_protocol::Account; use codex_app_server_protocol::AccountLoginCompletedNotification; +use codex_app_server_protocol::AccountUpdatedNotification; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; @@ -155,6 +156,14 @@ fn read_config_toml(codex_home: &Path) -> Result { )?)?) } +fn load_file_auth(codex_home: &Path) -> Result> { + Ok(load_auth_dot_json( + codex_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?) +} + fn aws_managed_bedrock_config() -> CreateConfigTomlParams { CreateConfigTomlParams { model_provider_id: Some("amazon-bedrock".to_string()), @@ -1028,6 +1037,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu codex_home.path(), "sk-test-key", AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), )?; let mut expected_config = read_config_toml(codex_home.path())?; expected_config @@ -1061,7 +1071,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu ); assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + load_file_auth(codex_home.path())?, Some(AuthDotJson { auth_mode: Some(AuthMode::BedrockApiKey), openai_api_key: None, @@ -1093,11 +1103,25 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu error: None, } ); - assert_no_account_updated(&mut mcp).await; + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + let ServerNotification::AccountUpdated(payload) = notification.try_into()? else { + bail!("unexpected notification") + }; + assert_eq!( + payload, + AccountUpdatedNotification { + auth_mode: Some(AuthMode::BedrockApiKey), + plan_type: None, + } + ); assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { - account: Some(Account::ApiKey {}), + account: None, requires_openai_auth: true, } ); @@ -1112,10 +1136,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu to_response::(response)?, LogoutAccountResponse {} ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None - ); + assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!( read_config_toml(codex_home.path())?, expected_config_after_logout @@ -1124,7 +1145,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { - account: Some(Account::ApiKey {}), + account: None, requires_openai_auth: true, } ); @@ -1163,15 +1184,13 @@ async fn managed_bedrock_login_requires_experimental_api_but_logout_is_best_effo error.error.message, "account/login/start.amazonBedrock requires experimentalApi capability" ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None - ); + assert_eq!(load_file_auth(codex_home.path())?, None); login_with_bedrock_api_key( codex_home.path(), "managed-bedrock-api-key", "us-west-2", AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), )?; let request_id = mcp.send_logout_account_request().await?; let response: JSONRPCResponse = timeout( @@ -1183,27 +1202,24 @@ async fn managed_bedrock_login_requires_experimental_api_but_logout_is_best_effo to_response::(response)?, LogoutAccountResponse {} ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None - ); + assert_eq!(load_file_auth(codex_home.path())?, None); Ok(()) } #[tokio::test] -async fn login_and_logout_managed_bedrock_require_restart_when_provider_is_active() -> Result<()> { +async fn login_and_logout_managed_bedrock_update_active_provider() -> Result<()> { let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), aws_managed_bedrock_config())?; + + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let mut expected_config = read_config_toml(codex_home.path())?; expected_config .as_table_mut() .expect("config should be a table") .remove("model_provider"); - let mut mcp = - TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") .await?; @@ -1221,12 +1237,16 @@ async fn login_and_logout_managed_bedrock_require_restart_when_provider_is_activ mcp.read_stream_until_notification_message("account/login/completed"), ) .await??; - assert_no_account_updated(&mut mcp).await; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { account: Some(Account::AmazonBedrock { - credential_source: AmazonBedrockCredentialSource::AwsManaged, + credential_source: AmazonBedrockCredentialSource::CodexManaged, }), requires_openai_auth: false, } @@ -1243,10 +1263,7 @@ async fn login_and_logout_managed_bedrock_require_restart_when_provider_is_activ to_response::(response)?, LogoutAccountResponse {} ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None - ); + assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); assert_no_account_updated(&mut mcp).await; @@ -1270,8 +1287,9 @@ async fn logout_account_aws_managed_bedrock_preserves_openai_auth_and_config() - codex_home.path(), "sk-test-key", AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), )?; - let expected_auth = load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?; + let expected_auth = load_file_auth(codex_home.path())?; let expected_config = read_config_toml(codex_home.path())?; let mut mcp = @@ -1288,10 +1306,7 @@ async fn logout_account_aws_managed_bedrock_preserves_openai_auth_and_config() - to_response::(response)?, LogoutAccountResponse {} ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - expected_auth - ); + assert_eq!(load_file_auth(codex_home.path())?, expected_auth); assert_eq!(read_config_toml(codex_home.path())?, expected_config); assert_no_account_updated(&mut mcp).await; Ok(()) @@ -1301,11 +1316,11 @@ async fn logout_account_aws_managed_bedrock_preserves_openai_auth_and_config() - async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_changes() -> Result<()> { let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; - let expected_config = read_config_toml(codex_home.path())?; let mut mcp = TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let expected_config = read_config_toml(codex_home.path())?; let request_id = mcp .send_login_account_amazon_bedrock_request(" ", "us-west-2") @@ -1332,13 +1347,17 @@ async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_change error.error.message, "Amazon Bedrock Mantle does not support region `us-west-1`" ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None - ); + assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); - drop(mcp); + Ok(()) +} + +#[tokio::test] +async fn login_account_amazon_bedrock_allows_codex_environment_auth() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + let mut mcp = TestAppServer::new_with_env( codex_home.path(), &[ @@ -1348,23 +1367,67 @@ async fn login_account_amazon_bedrock_rejects_invalid_credentials_without_change ) .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let mut expected_config = read_config_toml(codex_home.path())?; + expected_config + .as_table_mut() + .expect("config should be a table") + .insert( + "model_provider".to_string(), + toml::Value::String("amazon-bedrock".to_string()), + ); let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") .await?; - let error = timeout( + let response: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), ) .await??; assert_eq!( - error.error.message, - "Amazon Bedrock login is unavailable while Codex auth is supplied through the environment." + to_response::(response)?, + LoginAccountResponse::AmazonBedrock {} ); assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None + load_file_auth(codex_home.path())?, + Some(AuthDotJson { + auth_mode: Some(AuthMode::BedrockApiKey), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + bedrock_api_key: Some(BedrockApiKeyAuth { + api_key: "managed-bedrock-api-key".to_string(), + region: "us-west-2".to_string(), + }), + }) ); assert_eq!(read_config_toml(codex_home.path())?, expected_config); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/login/completed"), + ) + .await??; + + drop(mcp); + let mut mcp = TestAppServer::new_with_env( + codex_home.path(), + &[ + ("OPENAI_API_KEY", None), + ("CODEX_ACCESS_TOKEN", Some("env-access-token")), + ], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + assert_eq!( + read_account(&mut mcp).await?, + GetAccountResponse { + account: Some(Account::AmazonBedrock { + credential_source: AmazonBedrockCredentialSource::CodexManaged, + }), + requires_openai_auth: false, + } + ); Ok(()) } @@ -1401,7 +1464,7 @@ async fn login_account_amazon_bedrock_persists_under_provider_override() -> Resu LoginAccountResponse::AmazonBedrock {} ); assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, + load_file_auth(codex_home.path())?, Some(AuthDotJson { auth_mode: Some(AuthMode::BedrockApiKey), openai_api_key: None, @@ -1441,6 +1504,7 @@ async fn logout_account_managed_bedrock_preserves_changed_provider() -> Result<( "managed-bedrock-api-key", "us-west-2", AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), )?; let expected_config = read_config_toml(codex_home.path())?; @@ -1458,10 +1522,7 @@ async fn logout_account_managed_bedrock_preserves_changed_provider() -> Result<( to_response::(response)?, LogoutAccountResponse {} ); - assert_eq!( - load_auth_dot_json(codex_home.path(), AuthCredentialsStoreMode::File)?, - None - ); + assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); assert_no_account_updated(&mut mcp).await; Ok(()) diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index ef587cd1d4eb..95cb63f26c03 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -47,7 +47,6 @@ use codex_login::AuthDotJson; use codex_login::AuthManager; use codex_login::CODEX_ACCESS_TOKEN_ENV_VAR; use codex_login::CODEX_API_KEY_ENV_VAR; -use codex_login::CodexAuth; use codex_login::OPENAI_API_KEY_ENV_VAR; use codex_login::default_client::build_reqwest_client; use codex_login::default_client::default_headers; @@ -2321,9 +2320,10 @@ async fn websocket_reachability_check( let runtime_provider = create_model_provider(provider.clone(), auth_manager); let auth = runtime_provider.auth().await; + let auth_mode = runtime_provider.auth_mode(auth.as_ref()); details.push(format!( "auth mode: {}", - auth.as_ref().map(auth_mode_name).unwrap_or("none") + auth_mode.map(auth_mode_name).unwrap_or("none") )); let api_provider = match runtime_provider.api_provider().await { @@ -2461,8 +2461,8 @@ fn websocket_error_detail(err: &ApiError) -> String { } } -fn auth_mode_name(auth: &CodexAuth) -> &'static str { - match auth.auth_mode() { +fn auth_mode_name(auth_mode: AuthMode) -> &'static str { + match auth_mode { AuthMode::ApiKey => "api_key", AuthMode::Chatgpt => "chatgpt", AuthMode::ChatgptAuthTokens => "chatgpt_auth_tokens", diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index d3f9b7b96b02..f96d15de4f25 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -14,6 +14,7 @@ use codex_login::AuthRouteConfig; use codex_login::CLIENT_ID; use codex_login::CodexAuth; use codex_login::ServerOptions; +use codex_login::load_auth_dot_json; use codex_login::login_with_access_token; use codex_login::login_with_api_key; use codex_login::logout_with_revoke; @@ -434,8 +435,8 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { ) .await { - Ok(Some(auth)) => match auth.auth_mode() { - AuthMode::ApiKey => match auth.get_token() { + Ok(Some(auth)) => match auth { + CodexAuth::ApiKey(_) => match auth.get_token() { Ok(api_key) => { eprintln!("Logged in using an API key - {}", safe_format_key(&api_key)); std::process::exit(0); @@ -445,26 +446,38 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { std::process::exit(1); } }, - AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens => { + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { eprintln!("Logged in using ChatGPT"); std::process::exit(0); } - AuthMode::AgentIdentity => { + CodexAuth::AgentIdentity(_) => { eprintln!("Logged in using access token"); std::process::exit(0); } - AuthMode::PersonalAccessToken => { + CodexAuth::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); - } }, Ok(None) => { - eprintln!("Not logged in"); - std::process::exit(1); + match load_auth_dot_json( + &config.codex_home, + config.cli_auth_credentials_store_mode, + config.auth_keyring_backend_kind(), + ) { + Ok(Some(auth)) if auth.is_bedrock_api_key() => { + eprintln!("Logged in using Amazon Bedrock API key"); + std::process::exit(0); + } + Ok(_) => { + eprintln!("Not logged in"); + std::process::exit(1); + } + Err(e) => { + eprintln!("Error checking login status: {e}"); + std::process::exit(1); + } + } } Err(e) => { eprintln!("Error checking login status: {e}"); diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 43a758b0e851..9156061ee41a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -216,6 +216,7 @@ struct ModelClientState { /// share the same auth/provider setup flow. struct CurrentClientSetup { auth: Option, + auth_mode: Option, api_provider: ApiProvider, api_auth: SharedAuthProvider, agent_identity_telemetry: Option, @@ -536,7 +537,7 @@ impl ModelClient { let request_telemetry = Self::build_request_telemetry( session_telemetry, AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), + client_setup.auth_mode, client_setup.api_auth.as_ref(), client_setup.agent_identity_telemetry.clone(), PendingUnauthorizedRetry::default(), @@ -669,7 +670,7 @@ impl ModelClient { let request_telemetry = Self::build_request_telemetry( session_telemetry, AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), + client_setup.auth_mode, client_setup.api_auth.as_ref(), client_setup.agent_identity_telemetry.clone(), PendingUnauthorizedRetry::default(), @@ -920,6 +921,7 @@ impl ModelClient { /// lockstep when auth/provider resolution changes. async fn current_client_setup(&self) -> Result { let auth = self.state.provider.auth().await; + let auth_mode = self.state.provider.auth_mode(auth.as_ref()); let api_provider = self.state.provider.api_provider().await?; let resolved_auth = self .state @@ -932,6 +934,7 @@ impl ModelClient { .await?; Ok(CurrentClientSetup { auth, + auth_mode, api_provider, api_auth: resolved_auth.auth, agent_identity_telemetry: resolved_auth.agent_identity_telemetry, @@ -1235,7 +1238,7 @@ impl ModelClientSession { )) })?; let auth_context = AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), + client_setup.auth_mode, client_setup.api_auth.as_ref(), client_setup.agent_identity_telemetry.clone(), PendingUnauthorizedRetry::default(), @@ -1374,7 +1377,7 @@ impl ModelClientSession { let client_setup = self.client.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_auth_context = AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), + client_setup.auth_mode, client_setup.api_auth.as_ref(), client_setup.agent_identity_telemetry.clone(), pending_retry, @@ -1502,7 +1505,7 @@ impl ModelClientSession { loop { let client_setup = self.client.current_client_setup().await?; let request_auth_context = AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), + client_setup.auth_mode, client_setup.api_auth.as_ref(), client_setup.agent_identity_telemetry.clone(), pending_retry, diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 9509cab45662..158630f12ec2 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_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs index 1f8fab68b2c9..282957fc3ad5 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,6 +44,17 @@ 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, + } +} + #[tokio::test] #[serial(codex_auth_env)] async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> { @@ -78,18 +91,12 @@ 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(), 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(()) } @@ -120,12 +127,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 +149,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 +180,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 f8a5fee38b46..80f5d3d4992f 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -73,7 +73,6 @@ pub enum CodexAuth { ChatgptAuthTokens(ChatgptAuthTokens), AgentIdentity(AgentIdentityAuth), PersonalAccessToken(PersonalAccessTokenAuth), - BedrockApiKey(BedrockApiKeyAuth), } /// Policy for resolving Agent Identity auth from a broader Codex auth snapshot. @@ -146,7 +145,6 @@ impl PartialEq for CodexAuth { fn eq(&self, other: &Self) -> bool { match (self, other) { (Self::PersonalAccessToken(a), Self::PersonalAccessToken(b)) => a == b, - (Self::BedrockApiKey(a), Self::BedrockApiKey(b)) => a == b, _ => self.api_auth_mode() == other.api_auth_mode(), } } @@ -350,12 +348,9 @@ impl CodexAuth { .await; } if auth_mode == AuthMode::BedrockApiKey { - let Some(auth) = auth_dot_json.bedrock_api_key else { - return Err(std::io::Error::other( - "Bedrock API key auth is missing a Bedrock API key.", - )); - }; - return Ok(Self::BedrockApiKey(auth)); + return Err(std::io::Error::other( + "Amazon Bedrock API key auth is provider-specific", + )); } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); @@ -380,7 +375,9 @@ impl CodexAuth { AuthMode::PersonalAccessToken => { unreachable!("personal access token mode is handled above") } - AuthMode::BedrockApiKey => unreachable!("bedrock api key mode is handled above"), + AuthMode::BedrockApiKey => { + unreachable!("bedrock api key mode is rejected above") + } } } @@ -460,7 +457,6 @@ impl CodexAuth { Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => AuthMode::Chatgpt, Self::AgentIdentity(_) => AuthMode::AgentIdentity, Self::PersonalAccessToken(_) => AuthMode::PersonalAccessToken, - Self::BedrockApiKey(_) => AuthMode::BedrockApiKey, } } @@ -472,7 +468,6 @@ impl CodexAuth { Self::ChatgptAuthTokens(_) => AuthMode::ChatgptAuthTokens, Self::AgentIdentity(_) => AuthMode::AgentIdentity, Self::PersonalAccessToken(_) => AuthMode::PersonalAccessToken, - Self::BedrockApiKey(_) => AuthMode::BedrockApiKey, } } @@ -507,8 +502,7 @@ impl CodexAuth { Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) | Self::AgentIdentity(_) - | Self::PersonalAccessToken(_) - | Self::BedrockApiKey(_) => None, + | Self::PersonalAccessToken(_) => None, } } @@ -537,9 +531,6 @@ impl CodexAuth { "agent identity auth does not expose a bearer token", )), Self::PersonalAccessToken(auth) => Ok(auth.access_token().to_string()), - Self::BedrockApiKey(_) => Err(std::io::Error::other( - "Bedrock API key auth does not expose a Codex bearer token", - )), } } @@ -612,10 +603,7 @@ impl CodexAuth { let state = match self { Self::Chatgpt(auth) => &auth.state, Self::ChatgptAuthTokens(auth) => &auth.state, - Self::ApiKey(_) - | Self::AgentIdentity(_) - | Self::PersonalAccessToken(_) - | Self::BedrockApiKey(_) => return None, + Self::ApiKey(_) | Self::AgentIdentity(_) | Self::PersonalAccessToken(_) => return None, }; #[expect(clippy::unwrap_used)] state.auth_dot_json.lock().unwrap().clone() @@ -1039,6 +1027,19 @@ pub fn load_auth_dot_json( storage.load() } +fn load_stored_bedrock_api_key( + 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.and_then(|auth| auth.bedrock_api_key)) +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct AuthConfig { pub codex_home: PathBuf, @@ -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,8 +1124,11 @@ 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 { - CodexAuth::ApiKey(_) | CodexAuth::BedrockApiKey(_) => return Ok(()), + let Some(auth) = auth.as_ref() else { + return Ok(()); + }; + let chatgpt_account_id = match auth { + CodexAuth::ApiKey(_) => 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), }; @@ -1539,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, @@ -1557,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 @@ -1879,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, @@ -1897,6 +1924,7 @@ impl AuthManager { codex_home, inner: RwLock::new(CachedAuth { auth: managed_auth, + bedrock_api_key, permanent_refresh_failure: None, }), auth_change_tx, @@ -1918,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); @@ -1944,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); @@ -2004,6 +2034,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, @@ -2028,6 +2059,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() @@ -2130,7 +2169,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( @@ -2184,7 +2224,6 @@ impl AuthManager { _ => false, }, (AuthMode::PersonalAccessToken, AuthMode::PersonalAccessToken) => a == b, - (AuthMode::BedrockApiKey, AuthMode::BedrockApiKey) => a == b, _ => false, }, _ => false, @@ -2235,6 +2274,43 @@ 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 + } + } + fn set_cached_auth(&self, new_auth: Option) -> bool { if let Ok(mut guard) = self.inner.write() { let previous = guard.auth.as_ref(); @@ -2458,8 +2534,7 @@ impl AuthManager { } CodexAuth::ApiKey(_) | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) - | CodexAuth::BedrockApiKey(_) => Ok(()), + | CodexAuth::PersonalAccessToken(_) => Ok(()), }; if let Err(RefreshTokenError::Permanent(error)) = &result { self.record_permanent_refresh_failure_if_unchanged(&attempted_auth, error); diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 44c2f1c4d343..54c6da3f778b 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -24,6 +24,7 @@ pub use auth::AuthDotJson; pub use auth::AuthKeyringBackendKind; pub use auth::AuthManager; pub use auth::AuthManagerConfig; +pub use auth::AuthMode; pub use auth::CLIENT_ID; pub use auth::CLIENT_ID_OVERRIDE_ENV_VAR; pub use auth::CODEX_ACCESS_TOKEN_ENV_VAR; diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 914ddd4ca75e..ba7efd0845a8 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -10,6 +10,7 @@ use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; use codex_login::AuthManager; +use codex_login::AuthMode; use codex_login::CodexAuth; use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_4_MODEL_ID; @@ -39,13 +40,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 @@ -57,26 +58,16 @@ 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 { - self.managed_auth().map(CodexAuth::BedrockApiKey) + None } async fn api_provider(&self) -> Result { @@ -126,14 +117,17 @@ 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> { Box::pin(AmazonBedrockModelProvider::auth(self)) } + fn auth_mode(&self, _auth: Option<&CodexAuth>) -> Option { + self.managed_auth().map(|_| AuthMode::BedrockApiKey) + } + fn account_state(&self) -> ProviderAccountResult { let credential_source = if self.managed_auth().is_some() { AmazonBedrockCredentialSource::CodexManaged @@ -207,25 +201,20 @@ 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()); + let auth = provider.auth().await; + assert_eq!(auth, None); assert_eq!( - provider.auth().await, - Some(CodexAuth::BedrockApiKey(managed_auth)) + provider.auth_mode(auth.as_ref()), + Some(AuthMode::BedrockApiKey) ); assert_eq!( provider.account_state(), @@ -255,16 +244,16 @@ 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()); - assert_eq!(provider.auth().await, None); + let auth = provider.auth().await; + assert_eq!(auth, None); + assert_eq!(provider.auth_mode(auth.as_ref()), None); assert_eq!( provider.account_state(), Ok(ProviderAccountState { @@ -280,7 +269,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!( @@ -297,7 +286,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/auth.rs b/codex-rs/model-provider/src/auth.rs index 5b667e454973..8e6d63edd99b 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -13,16 +13,12 @@ use codex_login::auth::AgentIdentityAuth; use codex_login::auth::AgentIdentityAuthError; use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider_info::ModelProviderInfo; -use codex_protocol::error::CodexErr; use codex_protocol::protocol::SessionSource; use http::HeaderMap; use http::HeaderValue; use crate::bearer_auth_provider::BearerAuthProvider; -const BEDROCK_API_KEY_UNSUPPORTED_MESSAGE: &str = - "Bedrock API key auth is only supported by the Amazon Bedrock model provider"; - #[derive(Clone, Debug)] pub struct ProviderAuthScope { pub agent_identity_policy: AgentIdentityAuthPolicy, @@ -139,12 +135,6 @@ pub(crate) fn resolve_provider_auth( auth: Option<&CodexAuth>, provider: &ModelProviderInfo, ) -> codex_protocol::error::Result { - if matches!(auth, Some(CodexAuth::BedrockApiKey(_))) { - return Err(CodexErr::UnsupportedOperation( - BEDROCK_API_KEY_UNSUPPORTED_MESSAGE.to_string(), - )); - } - if let Some(auth) = bearer_auth_for_provider(provider)? { return Ok(Arc::new(auth)); } @@ -243,7 +233,6 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { CodexAuth::AgentIdentity(auth) => { Arc::new(AgentIdentityAuthProvider { auth: auth.clone() }) } - CodexAuth::BedrockApiKey(_) => unreachable!("{BEDROCK_API_KEY_UNSUPPORTED_MESSAGE}"), CodexAuth::ApiKey(_) | CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) @@ -261,7 +250,6 @@ mod tests { use codex_login::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; use codex_login::auth::AgentIdentityAuthRecord; - use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_protocol::account::PlanType; @@ -387,24 +375,6 @@ mod tests { assert!(auth.to_auth_headers().is_empty()); } - - #[test] - fn openai_provider_rejects_bedrock_api_key_auth() { - let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); - let auth = CodexAuth::BedrockApiKey(BedrockApiKeyAuth { - api_key: "bedrock-api-key-test".to_string(), - region: "us-east-1".to_string(), - }); - - match resolve_provider_auth(Some(&auth), &provider) { - Err(CodexErr::UnsupportedOperation(message)) => { - assert_eq!(message, BEDROCK_API_KEY_UNSUPPORTED_MESSAGE); - } - Err(err) => panic!("unexpected auth error: {err:?}"), - Ok(_) => panic!("Bedrock API key auth should be rejected"), - } - } - #[tokio::test] async fn first_party_run_scope_uses_agent_assertion_and_exposes_telemetry() { let auth = CodexAuth::AgentIdentity( diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 5f101e52cd6b..d8c1f1e2f84c 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -8,6 +8,7 @@ use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; use codex_login::AuthManager; +use codex_login::AuthMode; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::OpenAiModelsManager; @@ -58,7 +59,6 @@ pub struct ProviderAccountState { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ProviderAccountError { MissingChatgptAccountDetails, - UnsupportedBedrockApiKeyAuth, } impl fmt::Display for ProviderAccountError { @@ -67,12 +67,6 @@ impl fmt::Display for ProviderAccountError { Self::MissingChatgptAccountDetails => { write!(f, "plan type is required for chatgpt authentication") } - Self::UnsupportedBedrockApiKeyAuth => { - write!( - f, - "Bedrock API key auth is only supported by the Amazon Bedrock model provider" - ) - } } } } @@ -144,6 +138,11 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { /// Returns the current provider-scoped auth value, if one is configured. fn auth(&self) -> ModelProviderFuture<'_, Option>; + /// Returns the reporting mode for the resolved provider auth. + fn auth_mode(&self, auth: Option<&CodexAuth>) -> Option { + auth.map(CodexAuth::auth_mode) + } + /// Returns the current app-visible account state for this provider. fn account_state(&self) -> ProviderAccountResult; @@ -156,8 +155,7 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { fn api_provider(&self) -> ModelProviderFuture<'_, codex_protocol::error::Result> { Box::pin(async move { let auth = self.auth().await; - self.info() - .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode)) + self.info().to_api_provider(self.auth_mode(auth.as_ref())) }) } @@ -220,7 +218,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)) } @@ -281,9 +282,6 @@ impl ModelProvider for ConfiguredModelProvider { }) .map(|auth| match &auth { CodexAuth::ApiKey(_) => Ok(ProviderAccount::ApiKey), - CodexAuth::BedrockApiKey(_) => { - Err(ProviderAccountError::UnsupportedBedrockApiKeyAuth) - } CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) | CodexAuth::AgentIdentity(_) @@ -337,7 +335,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 +426,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 +513,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 +568,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( diff --git a/justfile b/justfile index d4a0e6566c67..ecad2a5e6e69 100644 --- a/justfile +++ b/justfile @@ -1,10 +1,12 @@ set working-directory := "codex-rs" -set positional-arguments +set positional-arguments := true + export JUST_SHELL := justfile_directory() / "scripts/just-shell.py" + set shell := ["python3", "-c", 'import os, runpy; runpy.run_path(os.environ["JUST_SHELL"], run_name="__main__")'] set windows-shell := ["python", "-c", 'import os, runpy; runpy.run_path(os.environ["JUST_SHELL"], run_name="__main__")'] -rust_min_stack := "8388608" # 8 MiB +rust_min_stack := "8388608" python := if os_family() == "windows" { "python" } else { "python3" } # Display help @@ -12,7 +14,9 @@ help: just -l # `codex` + alias c := codex + codex *args: cargo run --bin codex -- {args} @@ -76,6 +80,7 @@ install: # # Run `cargo install --locked cargo-nextest` if you don't have it installed. # Prefer this for routine local runs. Workspace crate features are banned, so + # there should be no need to add `--all-features`. [unix] test *args: @@ -86,6 +91,7 @@ test *args: $env:RUST_MIN_STACK = "{{ rust_min_stack }}"; $env:NEXTEST_PROFILE = "local"; cargo nextest run --no-fail-fast @($args | Select-Object -Skip 1) # Run from the repository root so scripts that resolve paths from `cwd` see + # the same layout they use in GitHub Actions. [no-cd] test-github-scripts: @@ -101,6 +107,7 @@ bench-smoke: # Build and run Codex from source using Bazel. # On Unix, use `[no-cd]` and `--run_under="cd $PWD &&"` to ensure Bazel runs + # the command in the current working directory. [no-cd] [unix] From f83e13db553c0f4ab0d9f1b83434db8a21df8b6a Mon Sep 17 00:00:00 2001 From: celia-oai Date: Tue, 16 Jun 2026 17:27:23 -0700 Subject: [PATCH 3/3] changes --- codex-rs/app-server/README.md | 3 +- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/account_processor.rs | 44 +++++---- .../src/request_processors/bedrock_auth.rs | 49 +--------- .../src/request_processors/plugins.rs | 1 - codex-rs/app-server/tests/suite/v2/account.rs | 92 +++++++++---------- codex-rs/login/src/auth/auth_tests.rs | 2 +- codex-rs/login/src/auth/bedrock_api_key.rs | 63 +++++++++++++ .../login/src/auth/bedrock_api_key_tests.rs | 26 ++++-- codex-rs/login/src/auth/manager.rs | 28 ++---- codex-rs/login/src/lib.rs | 2 +- codex-rs/model-provider/src/lib.rs | 4 +- 12 files changed, 167 insertions(+), 148 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index e267167eee95..4fdef79ccda6 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1928,7 +1928,7 @@ Codex supports these authentication modes. The current mode is surfaced in `acco - `account/login/start` — begin login (`apiKey`, `chatgpt`, `chatgptDeviceCode`, `amazonBedrock`). - `account/login/completed` (notify) — emitted when a login attempt finishes (success or error). - `account/login/cancel` — cancel a pending managed ChatGPT login by `loginId`. -- `account/logout` — sign out without params; also clears Codex-managed Amazon Bedrock credentials on a best-effort basis; triggers `account/updated` when the OpenAI auth mode changes. +- `account/logout` — sign out without params; also clears Codex-managed Amazon Bedrock credentials on a best-effort basis; triggers `account/updated` after successful logout. - `account/updated` (notify) — emitted whenever auth mode changes (`authMode`: `apikey`, `bedrockApiKey`, `chatgpt`, `personalAccessToken`, or `null`) and includes the current ChatGPT `planType` when available. - `account/rateLimits/read` — fetch ChatGPT rate limits, an optional effective monthly credit limit, and the earned rate-limit resets currently available, including expiry details when provided by the backend. Rate-limit updates arrive via `account/rateLimits/updated` (notify); reset-credit data is snapshot-only. - `account/rateLimitResetCredit/consume` — consume one earned reset using a caller-provided idempotency key, optionally selecting a reset-credit ID returned by `account/rateLimits/read`. @@ -2019,6 +2019,7 @@ This experimental flow requires the client to initialize with `experimentalApi: 3. Notifications: ```json { "method": "account/login/completed", "params": { "loginId": null, "success": true, "error": null } } + { "method": "account/updated", "params": { "authMode": "bedrockApiKey", "planType": null } } ``` Codex stores the key and region in its configured auth store and writes `model_provider = "amazon-bedrock"` to the active user config. diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 5b205f07d15d..119ea6179061 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -36,6 +36,7 @@ use codex_app_server_protocol::AppTemplateUnavailableReason; use codex_app_server_protocol::AppsListParams; use codex_app_server_protocol::AppsListResponse; use codex_app_server_protocol::AskForApproval; +use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::CancelLoginAccountParams; use codex_app_server_protocol::CancelLoginAccountResponse; use codex_app_server_protocol::CancelLoginAccountStatus; 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 eeb80bdab382..02999e29a59c 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1,6 +1,6 @@ -use super::bedrock_auth::has_managed_login as has_managed_bedrock_login; -use super::bedrock_auth::login as login_managed_bedrock; -use super::bedrock_auth::logout as logout_managed_bedrock; +use super::bedrock_auth::clear_user_model_provider_if_bedrock; +use super::bedrock_auth::set_user_model_provider_to_bedrock; +use super::bedrock_auth::user_model_provider_state; use super::*; use crate::auth_mode::auth_mode_to_api; use chrono::DateTime; @@ -184,9 +184,9 @@ impl AccountRequestProcessor { fn current_account_updated_notification(&self) -> AccountUpdatedNotification { let auth = self.auth_manager.auth_cached(); AccountUpdatedNotification { - auth_mode: auth - .as_ref() - .map(CodexAuth::api_auth_mode) + auth_mode: self + .auth_manager + .account_auth_mode_cached() .map(auth_mode_to_api), plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type), } @@ -382,8 +382,13 @@ impl AccountRequestProcessor { } } - login_managed_bedrock(&self.config, &self.config_manager, api_key, region).await?; - self.auth_manager.reload().await; + self.auth_manager + .login_with_bedrock_api_key(api_key, region) + .await + .map_err(|err| { + internal_error(format!("failed to save Amazon Bedrock auth: {err}")) + })?; + set_user_model_provider_to_bedrock(&self.config_manager).await?; Ok(LoginAccountResponse::AmazonBedrock {}) } .await; @@ -777,9 +782,17 @@ impl AccountRequestProcessor { } } - if has_managed_bedrock_login(&self.config)? { - logout_managed_bedrock(&self.config, &self.config_manager).await?; - self.auth_manager.reload().await; + if self + .auth_manager + .has_stored_bedrock_api_key_auth() + .map_err(|err| internal_error(format!("failed to read stored auth: {err}")))? + { + let user_model_provider = user_model_provider_state(&self.config_manager).await?; + self.auth_manager + .logout() + .await + .map_err(|err| internal_error(format!("logout failed: {err}")))?; + clear_user_model_provider_if_bedrock(&self.config_manager, user_model_provider).await?; return Ok(None); } if self.config.model_provider.is_amazon_bedrock() { @@ -810,12 +823,8 @@ impl AccountRequestProcessor { } async fn logout_v2(&self, request_id: ConnectionRequestId) -> Result<(), JSONRPCErrorError> { - let skip_account_updated = matches!(has_managed_bedrock_login(&self.config), Ok(true)) - || self.config.model_provider.is_amazon_bedrock(); let result = self.logout_common().await; - let account_updated = if skip_account_updated { - None - } else { + let account_updated = result .as_ref() .ok() @@ -823,8 +832,7 @@ impl AccountRequestProcessor { .map(|auth_mode| AccountUpdatedNotification { auth_mode, plan_type: None, - }) - }; + }); self.outgoing .send_result(request_id, result.map(|_| LogoutAccountResponse {})) .await; diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs index eb2cf5122386..333d905d95f6 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -6,37 +6,16 @@ use codex_app_server_protocol::ConfigReadParams; use codex_app_server_protocol::ConfigValueWriteParams; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::MergeStrategy; -use codex_core::config::Config; -use codex_login::AuthDotJson; -use codex_login::load_auth_dot_json; -use codex_login::login_with_bedrock_api_key; -use codex_login::logout as logout_stored_auth; use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID; -struct UserModelProviderState { +pub(super) struct UserModelProviderState { model_provider: Option, version: Option, } -pub(super) fn has_managed_login(config: &Config) -> Result { - load_stored_auth(config).map(|auth| auth.is_some()) -} - -pub(super) async fn login( - config: &Config, +pub(super) async fn set_user_model_provider_to_bedrock( config_manager: &ConfigManager, - api_key: &str, - region: &str, ) -> Result<(), JSONRPCErrorError> { - login_with_bedrock_api_key( - &config.codex_home, - api_key, - region, - config.cli_auth_credentials_store_mode, - config.auth_keyring_backend_kind(), - ) - .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?; - write_user_model_provider( config_manager, serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID), @@ -45,18 +24,10 @@ pub(super) async fn login( .await } -pub(super) async fn logout( - config: &Config, +pub(super) async fn clear_user_model_provider_if_bedrock( config_manager: &ConfigManager, + user_model_provider: UserModelProviderState, ) -> Result<(), JSONRPCErrorError> { - let user_model_provider = user_model_provider_state(config_manager).await?; - logout_stored_auth( - &config.codex_home, - config.cli_auth_credentials_store_mode, - config.auth_keyring_backend_kind(), - ) - .map_err(|err| internal_error(format!("logout failed: {err}")))?; - if user_model_provider.model_provider.as_deref() == Some(AMAZON_BEDROCK_PROVIDER_ID) { write_user_model_provider( config_manager, @@ -69,17 +40,7 @@ pub(super) async fn logout( Ok(()) } -fn load_stored_auth(config: &Config) -> Result, JSONRPCErrorError> { - load_auth_dot_json( - &config.codex_home, - config.cli_auth_credentials_store_mode, - config.auth_keyring_backend_kind(), - ) - .map(|auth| auth.filter(AuthDotJson::is_bedrock_api_key)) - .map_err(|err| internal_error(format!("failed to read stored auth: {err}"))) -} - -async fn user_model_provider_state( +pub(super) async fn user_model_provider_state( config_manager: &ConfigManager, ) -> Result { let user_config_path = config_manager diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index c291f5055f78..35c12d49c1aa 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -19,7 +19,6 @@ use codex_core_plugins::remote::RemoteAppTemplateUnavailableReason; use codex_core_plugins::remote::is_valid_remote_plugin_id; use codex_core_plugins::remote::validate_remote_plugin_id; use codex_core_plugins::remote_bundle::RemotePluginBundleInstallError; -use codex_login::AuthMode; use codex_mcp::McpOAuthLoginSupport; use codex_mcp::oauth_login_support; use codex_mcp::should_retry_without_scopes; diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 0a2e34410b4c..cb0e570a29ba 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -49,6 +49,7 @@ use codex_login::login_with_api_key; use codex_login::login_with_bedrock_api_key; use codex_protocol::account::AmazonBedrockCredentialSource; use codex_protocol::account::PlanType as AccountPlanType; +use codex_protocol::auth::AuthMode as DomainAuthMode; use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::json; @@ -192,15 +193,26 @@ async fn read_account(mcp: &mut TestAppServer) -> Result { to_response(response) } -async fn assert_no_account_updated(mcp: &mut TestAppServer) { - assert!( - timeout( - Duration::from_millis(500), - mcp.read_stream_until_notification_message("account/updated"), - ) - .await - .is_err() +async fn assert_account_updated( + mcp: &mut TestAppServer, + auth_mode: Option, +) -> Result<()> { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("account/updated"), + ) + .await??; + let ServerNotification::AccountUpdated(payload) = notification.try_into()? else { + bail!("unexpected notification") + }; + assert_eq!( + payload, + AccountUpdatedNotification { + auth_mode, + plan_type: None, + } ); + Ok(()) } async fn mock_device_code_usercode(server: &MockServer, interval_seconds: u64) { @@ -1039,6 +1051,9 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), )?; + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let mut expected_config = read_config_toml(codex_home.path())?; expected_config .as_table_mut() @@ -1053,10 +1068,6 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu .expect("config should be a table") .remove("model_provider"); - let mut mcp = - TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let request_id = mcp .send_login_account_amazon_bedrock_request(" managed-bedrock-api-key ", " us-west-2 ") .await?; @@ -1073,7 +1084,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu assert_eq!( load_file_auth(codex_home.path())?, Some(AuthDotJson { - auth_mode: Some(AuthMode::BedrockApiKey), + auth_mode: Some(DomainAuthMode::BedrockApiKey), openai_api_key: None, tokens: None, last_refresh: None, @@ -1103,21 +1114,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu error: None, } ); - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("account/updated"), - ) - .await??; - let ServerNotification::AccountUpdated(payload) = notification.try_into()? else { - bail!("unexpected notification") - }; - assert_eq!( - payload, - AccountUpdatedNotification { - auth_mode: Some(AuthMode::BedrockApiKey), - plan_type: None, - } - ); + assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { @@ -1141,7 +1138,7 @@ async fn login_and_logout_amazon_bedrock_persist_changes_until_restart() -> Resu read_config_toml(codex_home.path())?, expected_config_after_logout ); - assert_no_account_updated(&mut mcp).await; + assert_account_updated(&mut mcp, None).await?; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { @@ -1203,6 +1200,7 @@ async fn managed_bedrock_login_requires_experimental_api_but_logout_is_best_effo LogoutAccountResponse {} ); assert_eq!(load_file_auth(codex_home.path())?, None); + assert_account_updated(&mut mcp, None).await?; Ok(()) } @@ -1237,11 +1235,7 @@ async fn login_and_logout_managed_bedrock_update_active_provider() -> Result<()> mcp.read_stream_until_notification_message("account/login/completed"), ) .await??; - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("account/updated"), - ) - .await??; + assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { @@ -1266,7 +1260,7 @@ async fn login_and_logout_managed_bedrock_update_active_provider() -> Result<()> assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); - assert_no_account_updated(&mut mcp).await; + assert_account_updated(&mut mcp, None).await?; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { @@ -1290,11 +1284,11 @@ async fn logout_account_aws_managed_bedrock_preserves_openai_auth_and_config() - AuthKeyringBackendKind::default(), )?; let expected_auth = load_file_auth(codex_home.path())?; - let expected_config = read_config_toml(codex_home.path())?; let mut mcp = TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let expected_config = read_config_toml(codex_home.path())?; let request_id = mcp.send_logout_account_request().await?; let response: JSONRPCResponse = timeout( @@ -1308,7 +1302,7 @@ async fn logout_account_aws_managed_bedrock_preserves_openai_auth_and_config() - ); assert_eq!(load_file_auth(codex_home.path())?, expected_auth); assert_eq!(read_config_toml(codex_home.path())?, expected_config); - assert_no_account_updated(&mut mcp).await; + assert_account_updated(&mut mcp, None).await?; Ok(()) } @@ -1390,7 +1384,7 @@ async fn login_account_amazon_bedrock_allows_codex_environment_auth() -> Result< assert_eq!( load_file_auth(codex_home.path())?, Some(AuthDotJson { - auth_mode: Some(AuthMode::BedrockApiKey), + auth_mode: Some(DomainAuthMode::BedrockApiKey), openai_api_key: None, tokens: None, last_refresh: None, @@ -1436,6 +1430,13 @@ async fn login_account_amazon_bedrock_allows_codex_environment_auth() -> Result< async fn login_account_amazon_bedrock_persists_under_provider_override() -> Result<()> { let codex_home = TempDir::new()?; create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + + let mut mcp = TestAppServer::new_with_args( + codex_home.path(), + &["-c", "model_provider=\"mock_provider\""], + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let mut expected_config = read_config_toml(codex_home.path())?; expected_config .as_table_mut() @@ -1444,13 +1445,6 @@ async fn login_account_amazon_bedrock_persists_under_provider_override() -> Resu "model_provider".to_string(), toml::Value::String("amazon-bedrock".to_string()), ); - - let mut mcp = TestAppServer::new_with_args( - codex_home.path(), - &["-c", "model_provider=\"mock_provider\""], - ) - .await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; let request_id = mcp .send_login_account_amazon_bedrock_request("managed-bedrock-api-key", "us-west-2") .await?; @@ -1466,7 +1460,7 @@ async fn login_account_amazon_bedrock_persists_under_provider_override() -> Resu assert_eq!( load_file_auth(codex_home.path())?, Some(AuthDotJson { - auth_mode: Some(AuthMode::BedrockApiKey), + auth_mode: Some(DomainAuthMode::BedrockApiKey), openai_api_key: None, tokens: None, last_refresh: None, @@ -1484,7 +1478,7 @@ async fn login_account_amazon_bedrock_persists_under_provider_override() -> Resu mcp.read_stream_until_notification_message("account/login/completed"), ) .await??; - assert_no_account_updated(&mut mcp).await; + assert_account_updated(&mut mcp, Some(AuthMode::BedrockApiKey)).await?; assert_eq!( read_account(&mut mcp).await?, GetAccountResponse { @@ -1506,11 +1500,11 @@ async fn logout_account_managed_bedrock_preserves_changed_provider() -> Result<( AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), )?; - let expected_config = read_config_toml(codex_home.path())?; let mut mcp = TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let expected_config = read_config_toml(codex_home.path())?; let request_id = mcp.send_logout_account_request().await?; let response: JSONRPCResponse = timeout( @@ -1524,7 +1518,7 @@ async fn logout_account_managed_bedrock_preserves_changed_provider() -> Result<( ); assert_eq!(load_file_auth(codex_home.path())?, None); assert_eq!(read_config_toml(codex_home.path())?, expected_config); - assert_no_account_updated(&mut mcp).await; + assert_account_updated(&mut mcp, None).await?; Ok(()) } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 158630f12ec2..3936a17cdabf 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1609,9 +1609,9 @@ async fn load_auth_keeps_codex_api_key_env_precedence() { 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; diff --git a/codex-rs/login/src/auth/bedrock_api_key.rs b/codex-rs/login/src/auth/bedrock_api_key.rs index 3445878365a1..066ee2aadc5e 100644 --- a/codex-rs/login/src/auth/bedrock_api_key.rs +++ b/codex-rs/login/src/auth/bedrock_api_key.rs @@ -4,6 +4,8 @@ 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; @@ -44,6 +46,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 { + /// Current cached account auth mode, including provider-specific managed auth. + pub fn account_auth_mode_cached(&self) -> Option { + self.get_api_auth_mode().or_else(|| { + self.bedrock_api_key_auth_cached() + .map(|_| AuthMode::BedrockApiKey) + }) + } + + 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()) + } + + pub async fn login_with_bedrock_api_key( + &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, + )?; + self.reload().await; + Ok(()) + } +} + #[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 282957fc3ad5..0aa25a8866a5 100644 --- a/codex-rs/login/src/auth/bedrock_api_key_tests.rs +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -52,23 +52,16 @@ fn auth_config(codex_home: &std::path::Path, forced_login_method: ForcedLoginMet 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 auth_manager_login_with_bedrock_api_key_replaces_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, @@ -80,6 +73,10 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> ) .await; + auth_manager + .login_with_bedrock_api_key("bedrock-api-key-test", "us-east-1") + .await?; + let loaded = storage.load()?.expect("auth should be stored"); let expected = AuthDotJson { auth_mode: Some(AuthMode::BedrockApiKey), @@ -92,6 +89,10 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> }; assert_eq!(loaded, expected); assert_eq!(auth_manager.auth_mode(), None); + assert_eq!( + auth_manager.account_auth_mode_cached(), + Some(AuthMode::BedrockApiKey) + ); assert_eq!( auth_manager.bedrock_api_key_auth_cached(), Some(bedrock_auth()) @@ -102,7 +103,7 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> #[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( @@ -128,6 +129,7 @@ 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); + assert_eq!(auth_manager.account_auth_mode_cached(), None); Ok(()) } @@ -150,6 +152,10 @@ async fn bedrock_only_auth_storage_does_not_create_codex_auth() -> anyhow::Resul .await; assert_eq!(auth_manager.auth_mode(), None); + assert_eq!( + auth_manager.account_auth_mode_cached(), + Some(AuthMode::BedrockApiKey) + ); assert_eq!( auth_manager.bedrock_api_key_auth_cached(), Some(bedrock_auth()) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 80f5d3d4992f..4fddb0044253 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; @@ -644,10 +645,7 @@ impl CodexAuth { ) -> std::io::Result> { match self { Self::AgentIdentity(auth) => Ok(Some(auth.clone())), - Self::ApiKey(_) - | Self::ChatgptAuthTokens(_) - | Self::PersonalAccessToken(_) - | Self::BedrockApiKey(_) => Ok(None), + Self::ApiKey(_) | Self::ChatgptAuthTokens(_) | Self::PersonalAccessToken(_) => Ok(None), Self::Chatgpt(_) => { if policy == AgentIdentityAuthPolicy::JwtOnly { return Ok(None); @@ -1027,19 +1025,6 @@ pub fn load_auth_dot_json( storage.load() } -fn load_stored_bedrock_api_key( - 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.and_then(|auth| auth.bedrock_api_key)) -} - #[derive(Debug, Clone, PartialEq, Eq)] pub struct AuthConfig { pub codex_home: PathBuf, @@ -1476,7 +1461,7 @@ 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() == ApiAuthMode::BedrockApiKey + self.resolved_mode() == AuthMode::BedrockApiKey } fn from_external_tokens(external: &ExternalAuthTokens) -> std::io::Result { @@ -1815,12 +1800,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, @@ -2003,6 +1988,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); diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 54c6da3f778b..ad0b08fcd780 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -24,7 +24,6 @@ pub use auth::AuthDotJson; pub use auth::AuthKeyringBackendKind; pub use auth::AuthManager; pub use auth::AuthManagerConfig; -pub use auth::AuthMode; pub use auth::CLIENT_ID; pub use auth::CLIENT_ID_OVERRIDE_ENV_VAR; pub use auth::CODEX_ACCESS_TOKEN_ENV_VAR; @@ -55,5 +54,6 @@ pub use auth::read_openai_api_key_from_env; pub use auth::save_auth; pub use auth_env_telemetry::AuthEnvTelemetry; pub use auth_env_telemetry::collect_auth_env_telemetry; +pub use codex_protocol::auth::AuthMode; pub use outbound_proxy::AuthRouteConfig; pub use token_data::TokenData; diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index f8d1dd40475a..2a23343dfc36 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -4,16 +4,16 @@ mod bearer_auth_provider; mod models_endpoint; mod provider; +pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use auth::AgentIdentitySessionFallback; pub use auth::ProviderAuthScope; pub use auth::ResolvedProviderAuth; -pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use auth::auth_provider_from_auth; pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; -pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; pub use codex_model_provider_info::AMAZON_BEDROCK_PROVIDER_ID; +pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; pub use codex_protocol::account::ProviderAccount; pub use provider::ModelProvider; pub use provider::ModelProviderFuture;