feat: add managed Amazon Bedrock login

codex-rs/app-server-test-client/README.md

@@ -22,6 +22,30 @@ cargo run -p codex-app-server-test-client -- model-list
 When Codex asks a question, choose a numbered option (or `o` for a free-form answer when offered)
 and the client will send the response and continue streaming the same turn.
 
+## Testing Codex-managed Amazon Bedrock login
+
+`test-login --amazon-bedrock` initializes the experimental app-server API, sends an
+`account/login/start` request with an Amazon Bedrock API key, and waits for the
+`account/login/completed` and `account/updated` notifications. Login replaces the current primary
+credential and sets `model_provider = "amazon-bedrock"`, so use an isolated `CODEX_HOME` when
+testing.
+
+```bash
+export CODEX_HOME="$(mktemp -d)"
+printf 'cli_auth_credentials_store = "file"\n' > "$CODEX_HOME/config.toml"
+
+cargo build -p codex-cli --bin codex
+cargo run -p codex-app-server-test-client -- \
+  --codex-bin ./target/debug/codex \
+  test-login \
+  --amazon-bedrock \
+  --api-key "<BEDROCK_API_KEY>" \
+  --region us-west-2
+```
+
+The test client redacts `apiKey` from its outbound request log. After login, start a fresh Codex
+process with the same `CODEX_HOME` to verify that it uses the persisted managed credential.
+
 ## Testing Plugin Analytics
 
 The `plugin-analytics-smoke` command exercises `plugin/installed`, plugin

codex-rs/app-server-test-client/src/lib.rs

@@ -230,11 +230,20 @@ enum CliCommand {
         #[arg(long)]
         abort_on: Option<usize>,
     },
-    /// Trigger the ChatGPT login flow and wait for completion.
+    /// Trigger a ChatGPT or Amazon Bedrock login flow.
     TestLogin {
         /// Use the device-code login flow instead of the browser callback flow.
-        #[arg(long, default_value_t = false)]
+        #[arg(long, default_value_t = false, conflicts_with = "amazon_bedrock")]
         device_code: bool,
+        /// Use a Codex-managed Amazon Bedrock API key.
+        #[arg(long, default_value_t = false, conflicts_with = "device_code")]
+        amazon_bedrock: bool,
+        /// Amazon Bedrock API key.
+        #[arg(long, value_name = "API_KEY")]
+        api_key: Option<String>,
+        /// AWS Region for the Amazon Bedrock Mantle endpoint.
+        #[arg(long, value_name = "REGION")]
+        region: Option<String>,
     },
     /// Fetch the current account rate limits from the Codex app-server.
     GetAccountRateLimits,
@@ -313,6 +322,12 @@ enum CliCommand {
     },
 }
 
+enum TestLoginMode {
+    ChatgptBrowser,
+    ChatgptDeviceCode,
+    AmazonBedrock { api_key: String, region: String },
+}
+
 pub async fn run() -> Result<()> {
     let Cli {
         codex_bin,
@@ -415,10 +430,24 @@ pub async fn run() -> Result<()> {
             )
             .await
         }
-        CliCommand::TestLogin { device_code } => {
+        CliCommand::TestLogin {
+            device_code,
+            amazon_bedrock,
+            api_key,
+            region,
+        } => {
             ensure_dynamic_tools_unused(&dynamic_tools, "test-login")?;
             let endpoint = resolve_endpoint(codex_bin, url)?;
-            test_login(&endpoint, &config_overrides, device_code).await
+            let mode = if amazon_bedrock {
+                let api_key = api_key.context("--api-key is required with --amazon-bedrock")?;
+                let region = region.context("--region is required with --amazon-bedrock")?;
+                TestLoginMode::AmazonBedrock { api_key, region }
+            } else if device_code {
+                TestLoginMode::ChatgptDeviceCode
+            } else {
+                TestLoginMode::ChatgptBrowser
+            };
+            test_login(&endpoint, &config_overrides, mode).await
         }
         CliCommand::GetAccountRateLimits => {
             ensure_dynamic_tools_unused(&dynamic_tools, "get-account-rate-limits")?;
@@ -1128,16 +1157,45 @@ async fn send_follow_up_v2(
 async fn test_login(
     endpoint: &Endpoint,
     config_overrides: &[String],
-    device_code: bool,
+    mode: TestLoginMode,
 ) -> Result<()> {
     with_client("test-login", endpoint, config_overrides, |client| {
         let initialize = client.initialize()?;
         println!("< initialize response: {initialize:?}");
 
-        let login_response = if device_code {
-            client.login_account_chatgpt_device_code()?
-        } else {
-            client.login_account_chatgpt()?
+        let login_response = match mode {
+            TestLoginMode::ChatgptBrowser => client.login_account_chatgpt()?,
+            TestLoginMode::ChatgptDeviceCode => client.login_account_chatgpt_device_code()?,
+            TestLoginMode::AmazonBedrock { api_key, region } => {
+                let request_id = client.request_id();
+                let login_response: LoginAccountResponse = client.send_request(
+                    ClientRequest::LoginAccount {
+                        request_id: request_id.clone(),
+                        params: codex_app_server_protocol::LoginAccountParams::AmazonBedrock {
+                            api_key,
+                            region,
+                        },
+                    },
+                    request_id,
+                    "account/login/start",
+                )?;
+                println!("< account/login/start response: {login_response:?}");
+
+                let completion =
+                    client.wait_for_account_login_completion(/*expected_login_id*/ None)?;
+                println!("< account/login/completed notification: {completion:?}");
+
+                loop {
+                    let notification = client.next_notification()?;
+                    if let Ok(ServerNotification::AccountUpdated(account_updated)) =
+                        ServerNotification::try_from(notification)
+                    {
+                        println!("< account/updated notification: {account_updated:?}");
+                        break;
+                    }
+                }
+                return Ok(());
+            }
         };
         println!("< account/login/start response: {login_response:?}");
         let login_id = match login_response {
@@ -1158,7 +1216,7 @@ async fn test_login(
             _ => bail!("expected chatgpt login response"),
         };
 
-        let completion = client.wait_for_account_login_completion(&login_id)?;
+        let completion = client.wait_for_account_login_completion(Some(&login_id))?;
         println!("< account/login/completed notification: {completion:?}");
 
         if completion.success {
@@ -1799,15 +1857,15 @@ impl CodexClient {
 
     fn wait_for_account_login_completion(
         &mut self,
-        expected_login_id: &str,
+        expected_login_id: Option<&str>,
     ) -> Result<AccountLoginCompletedNotification> {
         loop {
             let notification = self.next_notification()?;
 
             if let Ok(server_notification) = ServerNotification::try_from(notification) {
                 match server_notification {
                     ServerNotification::AccountLoginCompleted(completion) => {
-                        if completion.login_id.as_deref() == Some(expected_login_id) {
+                        if completion.login_id.as_deref() == expected_login_id {
                             return Ok(completion);
                         }
 
@@ -1956,7 +2014,13 @@ impl CodexClient {
             .context("client request was not a valid JSON-RPC request")?;
         request.trace = current_span_w3c_trace_context();
         let request_json = serde_json::to_string(&request)?;
-        let request_pretty = serde_json::to_string_pretty(&request)?;
+        let mut request_for_logging = serde_json::to_value(&request)?;
+        if request.method == "account/login/start"
+            && let Some(api_key) = request_for_logging.pointer_mut("/params/apiKey")
+        {
+            *api_key = Value::String("<redacted>".to_string());
+        }
+        let request_pretty = serde_json::to_string_pretty(&request_for_logging)?;
         print_multiline_with_prefix("> ", &request_pretty);
         self.write_payload(&request_json)
     }

codex-rs/app-server/src/request_processors.rs

@@ -387,6 +387,7 @@ use codex_login::ServerOptions as LoginServerOptions;
 use codex_login::ShutdownHandle;
 use codex_login::complete_device_code_login;
 use codex_login::login_with_api_key;
+use codex_login::login_with_bedrock_api_key;
 use codex_login::oauth_client_id;
 use codex_login::request_device_code;
 use codex_login::run_login_server;
@@ -497,6 +498,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;

codex-rs/app-server/src/request_processors/account_processor.rs

@@ -1,7 +1,9 @@
+use super::bedrock_auth::set_user_model_provider_to_bedrock;
 use super::*;
 use crate::auth_mode::auth_mode_to_api;
 use crate::external_auth::ExternalAuthBridge;
 use chrono::DateTime;
+use codex_model_provider::is_supported_amazon_bedrock_region;
 
 mod rate_limit_resets;
 
@@ -191,6 +193,13 @@ impl AccountRequestProcessor {
         }
     }
 
+    async fn load_latest_config(&self) -> Result<Config, JSONRPCErrorError> {
+        self.config_manager
+            .load_latest_config(/*fallback_cwd*/ None)
+            .await
+            .map_err(|err| internal_error(format!("failed to reload config: {err}")))
+    }
+
     async fn maybe_refresh_plugin_caches_for_current_config(
         config_manager: &ConfigManager,
         thread_manager: &Arc<ThreadManager>,
@@ -293,8 +302,9 @@ impl AccountRequestProcessor {
                 )
                 .await;
             }
-            LoginAccountParams::AmazonBedrock { .. } => {
-                return Err(invalid_request("Amazon Bedrock login is not implemented"));
+            LoginAccountParams::AmazonBedrock { api_key, region } => {
+                self.login_amazon_bedrock_v2(request_id, api_key, region)
+                    .await;
             }
         }
         Ok(())
@@ -359,6 +369,65 @@ 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 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);
+                }
+            }
+
+            set_user_model_provider_to_bedrock(&self.config_manager).await?;
+            login_with_bedrock_api_key(
+                &self.config.codex_home,
+                api_key,
+                region,
+                self.config.cli_auth_credentials_store_mode,
+                self.config.auth_keyring_backend_kind(),
+            )
+            .map_err(|err| internal_error(format!("failed to save Amazon Bedrock auth: {err}")))?;
+            self.auth_manager.reload().await;
+            Ok(LoginAccountResponse::AmazonBedrock {})
+        }
+        .await;
+        let logged_in = result.is_ok();
+        self.outgoing.send_result(request_id, result).await;
+
+        if logged_in {
+            self.send_login_success_notifications(/*login_id*/ None)
+                .await;
+        }
+    }
+
     // Build options for a ChatGPT login attempt; performs validation.
     async fn login_chatgpt_common(
         &self,
@@ -837,7 +906,8 @@ impl AccountRequestProcessor {
         // Determine whether auth is required based on the active model provider.
         // If a custom provider is configured with `requires_openai_auth == false`,
         // then no auth step is required; otherwise, default to requiring auth.
-        let requires_openai_auth = self.config.model_provider.requires_openai_auth;
+        let config = self.load_latest_config().await?;
+        let requires_openai_auth = config.model_provider.requires_openai_auth;
 
         let response = if !requires_openai_auth {
             GetAuthStatusResponse {
@@ -905,10 +975,9 @@ impl AccountRequestProcessor {
 
         self.refresh_token_if_requested(do_refresh).await;
 
-        let provider = create_model_provider(
-            self.config.model_provider.clone(),
-            Some(self.auth_manager.clone()),
-        );
+        let config = self.load_latest_config().await?;
+        let provider =
+            create_model_provider(config.model_provider, Some(self.auth_manager.clone()));
         let account_state = match provider.account_state() {
             Ok(account_state) => account_state,
             Err(err) => return Err(invalid_request(err.to_string())),

codex-rs/app-server/src/request_processors/bedrock_auth.rs

@@ -0,0 +1,35 @@
+use super::config_processor::map_error as map_config_error;
+use crate::config_manager::ConfigManager;
+use codex_app_server_protocol::ConfigValueWriteParams;
+use codex_app_server_protocol::JSONRPCErrorError;
+use codex_app_server_protocol::MergeStrategy;
+use codex_model_provider::AMAZON_BEDROCK_PROVIDER_ID;
+
+pub(super) async fn set_user_model_provider_to_bedrock(
+    config_manager: &ConfigManager,
+) -> Result<(), JSONRPCErrorError> {
+    write_user_model_provider(
+        config_manager,
+        serde_json::json!(AMAZON_BEDROCK_PROVIDER_ID),
+        /*expected_version*/ None,
+    )
+    .await
+}
+
+async fn write_user_model_provider(
+    config_manager: &ConfigManager,
+    value: serde_json::Value,
+    expected_version: Option<String>,
+) -> 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)
+}

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());
     }

codex-rs/app-server/tests/common/test_app_server.rs

@@ -1325,6 +1325,20 @@ impl TestAppServer {
         self.send_login_account_request(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<i64> {
+        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<i64> {
         let params = serde_json::json!({