From 2a5cfbf358ba72787bc01fb26a696badfaaacf04 Mon Sep 17 00:00:00 2001 From: celia-oai Date: Wed, 8 Jul 2026 16:20:49 -0700 Subject: [PATCH] app-server-test-client: fold Bedrock into test-login --- codex-rs/app-server-test-client/README.md | 11 +- codex-rs/app-server-test-client/src/lib.rs | 155 ++++++++------------- 2 files changed, 61 insertions(+), 105 deletions(-) diff --git a/codex-rs/app-server-test-client/README.md b/codex-rs/app-server-test-client/README.md index c49c49f9e4..7c5d0a3e07 100644 --- a/codex-rs/app-server-test-client/README.md +++ b/codex-rs/app-server-test-client/README.md @@ -24,10 +24,10 @@ 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. +`test-login --amazon-bedrock` initializes the experimental app-server API and sends an +`account/login/start` request with an Amazon Bedrock API key. 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)" @@ -37,7 +37,8 @@ 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 \ + test-login \ + --amazon-bedrock \ --region us-west-2 ``` diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index b6b485fa10..a3fd631eed 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -25,7 +25,6 @@ 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; @@ -38,9 +37,7 @@ 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; @@ -73,7 +70,6 @@ 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; @@ -234,20 +230,20 @@ enum CliCommand { #[arg(long)] abort_on: Option, }, - /// 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, - }, - /// Log in with a Codex-managed Amazon Bedrock API key and verify the account state. - TestAmazonBedrockLogin { + /// 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. Defaults to AWS_BEARER_TOKEN_BEDROCK. #[arg(long, env = "AWS_BEARER_TOKEN_BEDROCK", hide_env_values = true)] - api_key: String, + api_key: Option, /// AWS Region for the Amazon Bedrock Mantle endpoint. #[arg(long, env = "AWS_REGION")] - region: String, + region: Option, }, /// Fetch the current account rate limits from the Codex app-server. GetAccountRateLimits, @@ -326,6 +322,12 @@ enum CliCommand { }, } +enum TestLoginMode { + ChatgptBrowser, + ChatgptDeviceCode, + AmazonBedrock { api_key: String, region: String }, +} + pub async fn run() -> Result<()> { let Cli { codex_bin, @@ -428,15 +430,27 @@ 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 - } - 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 + let mode = if amazon_bedrock { + let api_key = api_key.context( + "--api-key or AWS_BEARER_TOKEN_BEDROCK is required with --amazon-bedrock", + )?; + let region = + region.context("--region or AWS_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")?; @@ -1146,16 +1160,31 @@ 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:?}"); + return Ok(()); + } }; println!("< account/login/start response: {login_response:?}"); let login_id = match login_response { @@ -1176,7 +1205,7 @@ async fn test_login( _ => bail!("expected chatgpt login response"), }; - let completion = client.wait_for_account_login_completion(Some(&login_id))?; + let completion = client.wait_for_account_login_completion(&login_id)?; println!("< account/login/completed notification: {completion:?}"); if completion.success { @@ -1195,80 +1224,6 @@ 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", @@ -1891,7 +1846,7 @@ impl CodexClient { fn wait_for_account_login_completion( &mut self, - expected_login_id: Option<&str>, + expected_login_id: &str, ) -> Result { loop { let notification = self.next_notification()?; @@ -1899,7 +1854,7 @@ impl CodexClient { if let Ok(server_notification) = ServerNotification::try_from(notification) { match server_notification { ServerNotification::AccountLoginCompleted(completion) => { - if completion.login_id.as_deref() == expected_login_id { + if completion.login_id.as_deref() == Some(expected_login_id) { return Ok(completion); }