diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index b1afc77872..c49c49f9e4 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/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="" + +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 diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index 0e32365d12..b6b485fa10 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/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 { 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("".to_string()); + } + let request_pretty = serde_json::to_string_pretty(&request_for_logging)?; print_multiline_with_prefix("> ", &request_pretty); self.write_payload(&request_json) }