app-server-test-client: add managed 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-amazon-bedrock-login` initializes the experimental app-server API, logs in with an Amazon
+Bedrock API key, waits for the login completion notification, and verifies that `account/read`
+reports a Codex-managed Amazon Bedrock account. 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"
+export AWS_BEARER_TOKEN_BEDROCK="<BEDROCK_API_KEY>"
+
+cargo build -p codex-cli --bin codex
+cargo run -p codex-app-server-test-client -- \
+  --codex-bin ./target/debug/codex \
+  test-amazon-bedrock-login \
+  --region us-west-2
+```
+
+The key can instead be passed with `--api-key`, but the environment variable avoids exposing it in
+shell history. The test client redacts `apiKey` from its outbound request log. After login, unset
+`AWS_BEARER_TOKEN_BEDROCK` and 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

@@ -25,6 +25,7 @@ use anyhow::bail;
 use clap::ArgAction;
 use clap::Parser;
 use clap::Subcommand;
+use codex_app_server_protocol::Account;
 use codex_app_server_protocol::AccountLoginCompletedNotification;
 use codex_app_server_protocol::AskForApproval;
 use codex_app_server_protocol::ClientInfo;
@@ -37,7 +38,9 @@ use codex_app_server_protocol::DynamicToolSpec;
 use codex_app_server_protocol::FileChangeApprovalDecision;
 use codex_app_server_protocol::FileChangeRequestApprovalParams;
 use codex_app_server_protocol::FileChangeRequestApprovalResponse;
+use codex_app_server_protocol::GetAccountParams;
 use codex_app_server_protocol::GetAccountRateLimitsResponse;
+use codex_app_server_protocol::GetAccountResponse;
 use codex_app_server_protocol::InitializeCapabilities;
 use codex_app_server_protocol::InitializeParams;
 use codex_app_server_protocol::InitializeResponse;
@@ -70,6 +73,7 @@ use codex_app_server_protocol::UserInput as V2UserInput;
 use codex_core::config::Config;
 use codex_otel::OtelProvider;
 use codex_otel::current_span_w3c_trace_context;
+use codex_protocol::account::AmazonBedrockCredentialSource;
 use codex_protocol::dynamic_tools::normalize_dynamic_tool_specs;
 use codex_protocol::openai_models::ReasoningEffort;
 use codex_protocol::protocol::W3cTraceContext;
@@ -236,6 +240,15 @@ enum CliCommand {
         #[arg(long, default_value_t = false)]
         device_code: bool,
     },
+    /// Log in with a Codex-managed Amazon Bedrock API key and verify the account state.
+    TestAmazonBedrockLogin {
+        /// Amazon Bedrock API key. Defaults to AWS_BEARER_TOKEN_BEDROCK.
+        #[arg(long, env = "AWS_BEARER_TOKEN_BEDROCK", hide_env_values = true)]
+        api_key: String,
+        /// AWS Region for the Amazon Bedrock Mantle endpoint.
+        #[arg(long, env = "AWS_REGION")]
+        region: String,
+    },
     /// Fetch the current account rate limits from the Codex app-server.
     GetAccountRateLimits,
     /// List the available models from the Codex app-server.
@@ -420,6 +433,11 @@ pub async fn run() -> Result<()> {
             let endpoint = resolve_endpoint(codex_bin, url)?;
             test_login(&endpoint, &config_overrides, device_code).await
         }
+        CliCommand::TestAmazonBedrockLogin { api_key, region } => {
+            ensure_dynamic_tools_unused(&dynamic_tools, "test-amazon-bedrock-login")?;
+            let endpoint = resolve_endpoint(codex_bin, url)?;
+            test_amazon_bedrock_login(&endpoint, &config_overrides, api_key, region).await
+        }
         CliCommand::GetAccountRateLimits => {
             ensure_dynamic_tools_unused(&dynamic_tools, "get-account-rate-limits")?;
             let endpoint = resolve_endpoint(codex_bin, url)?;
@@ -1158,7 +1176,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 {
@@ -1177,6 +1195,80 @@ async fn test_login(
     .await
 }
 
+async fn test_amazon_bedrock_login(
+    endpoint: &Endpoint,
+    config_overrides: &[String],
+    api_key: String,
+    region: String,
+) -> Result<()> {
+    with_client(
+        "test-amazon-bedrock-login",
+        endpoint,
+        config_overrides,
+        |client| {
+            let initialize =
+                client.initialize_with_experimental_api(/*experimental_api*/ true)?;
+            println!("< initialize response: {initialize:?}");
+
+            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:?}");
+            if login_response != (LoginAccountResponse::AmazonBedrock {}) {
+                bail!("expected Amazon Bedrock login response, got {login_response:?}");
+            }
+
+            let completion =
+                client.wait_for_account_login_completion(/*expected_login_id*/ None)?;
+            println!("< account/login/completed notification: {completion:?}");
+            if !completion.success {
+                bail!(
+                    "Amazon Bedrock login failed: {}",
+                    completion
+                        .error
+                        .as_deref()
+                        .unwrap_or("unknown error from account/login/completed")
+                );
+            }
+
+            let request_id = client.request_id();
+            let account: GetAccountResponse = client.send_request(
+                ClientRequest::GetAccount {
+                    request_id: request_id.clone(),
+                    params: GetAccountParams {
+                        refresh_token: false,
+                    },
+                },
+                request_id,
+                "account/read",
+            )?;
+            println!("< account/read response: {account:?}");
+            let expected_account = GetAccountResponse {
+                account: Some(Account::AmazonBedrock {
+                    credential_source: AmazonBedrockCredentialSource::CodexManaged,
+                }),
+                requires_openai_auth: false,
+            };
+            if account != expected_account {
+                bail!("expected managed Amazon Bedrock account, got {account:?}");
+            }
+
+            println!("Amazon Bedrock login succeeded.");
+            Ok(())
+        },
+    )
+    .await
+}
+
 async fn get_account_rate_limits(endpoint: &Endpoint, config_overrides: &[String]) -> Result<()> {
     with_client(
         "get-account-rate-limits",
@@ -1799,7 +1891,7 @@ 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()?;
@@ -1807,7 +1899,7 @@ impl CodexClient {
             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 +2048,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)
     }