Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions codex-rs/app-server-test-client/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
90 changes: 77 additions & 13 deletions codex-rs/app-server-test-client/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")?;
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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)
}
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/app-server/src/request_processors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
90 changes: 83 additions & 7 deletions codex-rs/app-server/src/request_processors/account_processor.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -191,6 +193,20 @@ impl AccountRequestProcessor {
}
}

async fn load_latest_config(&self) -> Config {
match self
.config_manager
.load_latest_config(/*fallback_cwd*/ None)
.await
{
Ok(config) => config,
Err(err) => {
tracing::warn!("failed to reload config, using startup config: {err}");
self.config.as_ref().clone()
}
}
}

async fn maybe_refresh_plugin_caches_for_current_config(
config_manager: &ConfigManager,
thread_manager: &Arc<ThreadManager>,
Expand Down Expand Up @@ -293,8 +309,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(())
Expand Down Expand Up @@ -359,6 +376,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(
Comment thread
celia-oai marked this conversation as resolved.
&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;
Comment thread
celia-oai marked this conversation as resolved.

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,
Expand Down Expand Up @@ -837,7 +913,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 {
Expand Down Expand Up @@ -905,10 +982,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())),
Expand Down
Loading
Loading