diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 15d517d398..715225b0a7 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -601,6 +601,9 @@ "background_paginated_rollout_migration": { "type": "boolean" }, + "bedrock_setup_wizard": { + "type": "boolean" + }, "browser_use": { "type": "boolean" }, @@ -5616,6 +5619,9 @@ "background_paginated_rollout_migration": { "type": "boolean" }, + "bedrock_setup_wizard": { + "type": "boolean" + }, "browser_use": { "type": "boolean" }, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 26548db324..80cc9fac25 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -305,6 +305,8 @@ pub enum Feature { ToolCallMcpElicitation, /// Prompt Codex Apps connector auth failures through MCP URL elicitations. AuthElicitation, + /// Offer Amazon Bedrock setup during TUI sign-in onboarding. + BedrockSetupWizard, /// Enable personality selection in the TUI. Personality, /// Enable native artifact tools. @@ -1507,6 +1509,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::BedrockSetupWizard, + key: "bedrock_setup_wizard", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::Personality, key: "personality", diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 878978a1d2..0cd2d7ed24 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -49,6 +49,7 @@ use codex_config::format_config_error_with_source; use codex_config::types::ResumeCwdMode; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; +use codex_features::Feature; use codex_login::AuthConfig; use codex_login::default_client::originator; use codex_login::default_client::set_default_client_residency_requirement; @@ -57,6 +58,7 @@ use codex_login::is_workload_identity_selected; use codex_protocol::ThreadId; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::AltScreenMode; +use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::SandboxMode; #[cfg(target_os = "windows")] use codex_protocol::config_types::WindowsSandboxLevel; @@ -1149,9 +1151,12 @@ async fn run_ratatui_app( // Authentication can change while any interactive onboarding screen is open. startup_account = None; let show_login_screen = should_show_login_screen(login_status, &initial_config); + let bedrock_setup_enabled = + should_show_bedrock_setup_wizard(login_status, &initial_config, &app_server_target); let onboarding_result = run_onboarding_app( OnboardingScreenArgs { show_login_screen, + bedrock_setup_enabled, show_trust_screen: should_show_trust_screen_flag, remote_project_trust, login_status, @@ -1602,6 +1607,14 @@ async fn run_ratatui_app( let use_alt_screen = determine_alt_screen_mode(no_alt_screen, config.tui_alternate_screen); tui.set_alt_screen_enabled(use_alt_screen); + if config.model_provider_id != startup_model_provider { + startup_account = None; + if matches!(&app_server_target, AppServerTarget::Embedded) { + // App-server providers are fixed at startup, so onboarding cannot + // reuse a server initialized before it persisted another provider. + shutdown_app_server_if_present(app_server.take()).await; + } + } let mut app_server = match app_server { Some(app_server) => app_server, None => match startup_draft @@ -1653,9 +1666,6 @@ async fn run_ratatui_app( let bypass_hook_trust_for_startup_review = config.bypass_hook_trust && !is_persistent_resume; let hooks_request_handle = app_server.request_handle(); let hooks_cwd = config.cwd.to_path_buf(); - if config.model_provider_id != startup_model_provider { - startup_account = None; - } let startup_prefetch_started_at = Instant::now(); let startup_prefetch = startup_draft .run_until(&mut tui, async { @@ -1930,6 +1940,25 @@ fn should_show_login_screen(login_status: LoginStatus, config: &Config) -> bool login_status == LoginStatus::NotAuthenticated } +fn should_show_bedrock_setup_wizard( + login_status: LoginStatus, + config: &Config, + app_server_target: &AppServerTarget, +) -> bool { + matches!(app_server_target, AppServerTarget::Embedded) + && should_show_login_screen(login_status, config) + && config.features.enabled(Feature::BedrockSetupWizard) + && config.model_provider_id == "openai" + && config + .config_layer_stack + .effective_config() + .get("model_provider") + .is_none() + && config + .auth_config() + .is_login_method_allowed(ForcedLoginMethod::Api) +} + #[cfg(test)] mod tests { use super::*; @@ -1953,6 +1982,83 @@ mod tests { .await } + #[tokio::test] + async fn bedrock_setup_wizard_requires_eligible_onboarding() -> color_eyre::Result<()> { + let shared_endpoint = RemoteAppServerEndpoint::WebSocket { + websocket_url: "ws://127.0.0.1:4500/".to_string(), + auth_token: None, + }; + let enabled = "[features]\nbedrock_setup_wizard = true\n"; + + for (label, config_toml, login_status, target, expected) in [ + ( + "disabled by default", + "", + LoginStatus::NotAuthenticated, + AppServerTarget::Embedded, + false, + ), + ( + "enabled for the default provider", + enabled, + LoginStatus::NotAuthenticated, + AppServerTarget::Embedded, + true, + ), + ( + "explicit provider", + "model_provider = \"openai\"\n[features]\nbedrock_setup_wizard = true\n", + LoginStatus::NotAuthenticated, + AppServerTarget::Embedded, + false, + ), + ( + "forced ChatGPT login", + "forced_login_method = \"chatgpt\"\n[features]\nbedrock_setup_wizard = true\n", + LoginStatus::NotAuthenticated, + AppServerTarget::Embedded, + false, + ), + ( + "existing authentication", + enabled, + LoginStatus::AuthMode(AuthMode::Chatgpt), + AppServerTarget::Embedded, + false, + ), + ( + "shared local daemon", + enabled, + LoginStatus::NotAuthenticated, + AppServerTarget::LocalDaemon { + endpoint: shared_endpoint.clone(), + }, + false, + ), + ( + "remote app server", + enabled, + LoginStatus::NotAuthenticated, + AppServerTarget::Remote { + endpoint: shared_endpoint, + }, + false, + ), + ] { + let codex_home = TempDir::new()?; + std::fs::write(codex_home.path().join("config.toml"), config_toml)?; + let config = build_config(&codex_home).await?; + + assert_eq!( + should_show_bedrock_setup_wizard(login_status, &config, &target), + expected, + "{label}" + ); + } + + Ok(()) + } + fn write_session_rollout( codex_home: &Path, filename_ts: &str, diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 0bbb899f4a..d242732ca8 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -50,6 +50,7 @@ use crate::key_hint::KeyBinding; use crate::key_hint::KeyBindingListExt; use crate::motion::MotionMode; use crate::motion::shimmer_text; +use crate::onboarding::bedrock::BedrockState; use crate::onboarding::keys; use crate::onboarding::onboarding_screen::KeyboardHandler; use crate::onboarding::onboarding_screen::StepStateProvider; @@ -88,6 +89,8 @@ pub(crate) enum SignInState { ChatGptSuccess, ApiKeyEntry(ApiKeyInputState), ApiKeyConfigured, + Bedrock(BedrockState), + BedrockConfigured, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -95,10 +98,11 @@ pub(crate) enum SignInOption { ChatGpt, DeviceCode, ApiKey, + Bedrock, } const API_KEY_DISABLED_MESSAGE: &str = "API key login is disabled."; -fn onboarding_request_id() -> codex_app_server_protocol::RequestId { +pub(super) fn onboarding_request_id() -> codex_app_server_protocol::RequestId { codex_app_server_protocol::RequestId::String(Uuid::new_v4().to_string()) } @@ -178,6 +182,9 @@ impl ContinueWithDeviceCodeState { impl KeyboardHandler for AuthModeWidget { fn handle_key_event(&mut self, key_event: KeyEvent) { + if self.handle_bedrock_key_event(&key_event) { + return; + } if self.handle_api_key_entry_key_event(&key_event) { return; } @@ -202,6 +209,10 @@ impl KeyboardHandler for AuthModeWidget { self.select_option_by_index(/*index*/ 2); return; } + if keys::SELECT_FOURTH.is_pressed(key_event) { + self.select_option_by_index(/*index*/ 3); + return; + } if keys::CONFIRM.is_pressed(key_event) { let sign_in_state = { (*self.sign_in_state.read().unwrap()).clone() }; match sign_in_state { @@ -222,7 +233,18 @@ impl KeyboardHandler for AuthModeWidget { } fn handle_paste(&mut self, pasted: String) { - let _ = self.handle_api_key_entry_paste(pasted); + let sign_in_state = self.sign_in_state.read().unwrap(); + match &*sign_in_state { + SignInState::Bedrock(_) => { + drop(sign_in_state); + let _ = self.handle_bedrock_paste(&pasted); + } + SignInState::ApiKeyEntry(_) => { + drop(sign_in_state); + let _ = self.handle_api_key_entry_paste(pasted); + } + _ => {} + } } } @@ -236,6 +258,7 @@ pub(crate) struct AuthModeWidget { pub login_status: LoginStatus, pub app_server_request_handle: AppServerRequestHandle, pub auth_config: AuthConfig, + pub bedrock_setup_enabled: bool, pub animations_enabled: bool, pub animations_suppressed: Cell, } @@ -286,18 +309,25 @@ impl AuthModeWidget { self.error.read().unwrap().clone() } - /// Returns whether the auth flow is currently in API-key entry mode. - pub(crate) fn is_api_key_entry_active(&self) -> bool { - self.sign_in_state - .read() - .is_ok_and(|guard| matches!(&*guard, SignInState::ApiKeyEntry(_))) + /// Returns whether the auth flow is currently accepting text input. + pub(crate) fn is_text_entry_active(&self) -> bool { + self.sign_in_state.read().is_ok_and(|guard| match &*guard { + SignInState::ApiKeyEntry(_) => true, + SignInState::Bedrock(state) => state.is_text_entry_active(), + _ => false, + }) } - /// Returns whether the API-key entry field currently contains any text. - pub(crate) fn api_key_entry_has_text(&self) -> bool { - self.sign_in_state.read().is_ok_and( - |guard| matches!(&*guard, SignInState::ApiKeyEntry(state) if !state.value.is_empty()), - ) + /// Returns whether printable quit shortcuts must be treated as text input. + /// + /// OpenAI API-key entry keeps its existing empty-field quit behavior, while + /// Bedrock fields accept printable input from their first character. + pub(crate) fn should_suppress_printable_quit(&self) -> bool { + self.sign_in_state.read().is_ok_and(|guard| match &*guard { + SignInState::ApiKeyEntry(state) => !state.value.is_empty(), + SignInState::Bedrock(state) => state.is_text_entry_active(), + _ => false, + }) } fn confirm_binding(&self) -> KeyBinding { @@ -325,6 +355,9 @@ impl AuthModeWidget { } if self.is_api_login_allowed() { options.push(SignInOption::ApiKey); + if self.bedrock_setup_enabled { + options.push(SignInOption::Bedrock); + } } options } @@ -337,6 +370,9 @@ impl AuthModeWidget { } if self.is_api_login_allowed() { options.push(SignInOption::ApiKey); + if self.bedrock_setup_enabled { + options.push(SignInOption::Bedrock); + } } options } @@ -382,6 +418,13 @@ impl AuthModeWidget { self.disallow_api_login(); } } + SignInOption::Bedrock => { + if self.bedrock_setup_enabled && self.is_api_login_allowed() { + self.start_bedrock_discovery(); + } else if !self.is_api_login_allowed() { + self.disallow_api_login(); + } + } } } @@ -393,17 +436,21 @@ impl AuthModeWidget { } fn render_pick_mode(&self, area: Rect, buf: &mut Buffer) { - let mut lines: Vec = vec![ - Line::from(vec![ - " ".into(), - "Sign in with ChatGPT to use Codex as part of your paid plan".into(), - ]), - Line::from(vec![ - " ".into(), - "or connect an API key for usage-based billing".into(), - ]), - "".into(), - ]; + let mut lines: Vec = if self.bedrock_setup_enabled { + vec![" Choose how you want to use Codex.".into(), "".into()] + } else { + vec![ + Line::from(vec![ + " ".into(), + "Sign in with ChatGPT to use Codex as part of your paid plan".into(), + ]), + Line::from(vec![ + " ".into(), + "or connect an API key for usage-based billing".into(), + ]), + "".into(), + ] + }; let create_mode_item = |idx: usize, selected_mode: SignInOption, @@ -463,10 +510,22 @@ impl AuthModeWidget { lines.extend(create_mode_item( idx, option, - "Provide your own API key", + if self.bedrock_setup_enabled { + "Use an OpenAI API key" + } else { + "Provide your own API key" + }, "Pay for what you use", )); } + SignInOption::Bedrock => { + lines.extend(create_mode_item( + idx, + option, + "Use Amazon Bedrock", + "Connect using your AWS credentials", + )); + } } lines.push("".into()); } @@ -978,8 +1037,11 @@ impl StepStateProvider for AuthModeWidget { | SignInState::ApiKeyEntry(_) | SignInState::ChatGptContinueInBrowser(_) | SignInState::ChatGptDeviceCode(_) - | SignInState::ChatGptSuccessMessage => StepState::InProgress, - SignInState::ChatGptSuccess | SignInState::ApiKeyConfigured => StepState::Complete, + | SignInState::ChatGptSuccessMessage + | SignInState::Bedrock(_) => StepState::InProgress, + SignInState::ChatGptSuccess + | SignInState::ApiKeyConfigured + | SignInState::BedrockConfigured => StepState::Complete, } } } @@ -1009,6 +1071,14 @@ impl WidgetRef for AuthModeWidget { SignInState::ApiKeyConfigured => { self.render_api_key_configured(area, buf); } + SignInState::Bedrock(state) => { + state.render(area, buf, self.error_message()); + } + SignInState::BedrockConfigured => { + Paragraph::new("✓ Amazon Bedrock configured".green()) + .wrap(Wrap { trim: false }) + .render(area, buf); + } } } } @@ -1101,6 +1171,7 @@ mod tests { login_status: LoginStatus::NotAuthenticated, app_server_request_handle: AppServerRequestHandle::InProcess(client.request_handle()), auth_config, + bedrock_setup_enabled: false, animations_enabled: true, animations_suppressed: std::cell::Cell::new(false), }; @@ -1123,6 +1194,70 @@ mod tests { )); } + #[tokio::test] + async fn bedrock_option_requires_feature_and_api_login_permission() { + let (mut widget, _tmp) = widget_forced_chatgpt().await; + widget.auth_config.forced_login_method = None; + assert_eq!( + widget.displayed_sign_in_options(), + vec![ + SignInOption::ChatGpt, + SignInOption::DeviceCode, + SignInOption::ApiKey, + ] + ); + + widget.bedrock_setup_enabled = true; + assert_eq!( + widget.displayed_sign_in_options(), + vec![ + SignInOption::ChatGpt, + SignInOption::DeviceCode, + SignInOption::ApiKey, + SignInOption::Bedrock, + ] + ); + + let area = Rect::new(0, 0, 76, 19); + let mut buffer = Buffer::empty(area); + widget.render_pick_mode(area, &mut buffer); + let mut rows = (area.top()..area.bottom()) + .map(|row| { + (area.left()..area.right()) + .map(|column| buffer[(column, row)].symbol()) + .collect::() + .trim_end() + .to_string() + }) + .collect::>(); + while rows.last().is_some_and(String::is_empty) { + rows.pop(); + } + insta::assert_snapshot!(rows.join("\n"), @r###" + Choose how you want to use Codex. + + > 1. Sign in with ChatGPT + Usage included with Plus, Pro, Business, and Enterprise plans + + 2. Sign in with Device Code + Sign in from another device with a one-time code + + 3. Use an OpenAI API key + Pay for what you use + + 4. Use Amazon Bedrock + Connect using your AWS credentials + + Press enter to continue + "###); + + widget.auth_config.forced_login_method = Some(ForcedLoginMethod::Chatgpt); + assert_eq!( + widget.displayed_sign_in_options(), + vec![SignInOption::ChatGpt, SignInOption::DeviceCode] + ); + } + #[tokio::test] async fn saving_api_key_is_blocked_when_chatgpt_forced() { let (mut widget, _tmp) = widget_forced_chatgpt().await; diff --git a/codex-rs/tui/src/onboarding/bedrock.rs b/codex-rs/tui/src/onboarding/bedrock.rs new file mode 100644 index 0000000000..43a11175e3 --- /dev/null +++ b/codex-rs/tui/src/onboarding/bedrock.rs @@ -0,0 +1,803 @@ +//! Amazon Bedrock credential discovery and setup within the authentication step. + +use super::auth::AuthModeWidget; +use super::auth::SignInState; +use super::auth::onboarding_request_id; +use super::keys; +use crate::key_hint::KeyBindingListExt; +use crate::wrapping::word_wrap_lines; +use codex_app_server_protocol::AwsCredentialType; +use codex_app_server_protocol::BedrockAwsProfile; +use codex_app_server_protocol::BedrockDiscoverParams; +use codex_app_server_protocol::BedrockDiscoverResponse; +use codex_app_server_protocol::BedrockEnvironmentCredential; +use codex_app_server_protocol::BedrockSetupParams; +use codex_app_server_protocol::BedrockSetupResponse; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::LoginAccountParams; +use codex_app_server_protocol::LoginAccountResponse; +use codex_app_server_protocol::RequestId; +use crossterm::event::KeyCode; +use crossterm::event::KeyEvent; +use crossterm::event::KeyEventKind; +use crossterm::event::KeyModifiers; +use ratatui::buffer::Buffer; +use ratatui::layout::Rect; +use ratatui::prelude::Widget; +use ratatui::style::Stylize; +use ratatui::text::Line; +use ratatui::widgets::Paragraph; +use std::sync::PoisonError; + +#[derive(Clone)] +pub(super) struct BedrockState { + view: BedrockView, + highlighted: usize, + profiles: Vec, + environment_credentials: Vec, +} + +#[derive(Clone)] +enum BedrockView { + Discovering(RequestId), + Methods(BedrockMethodList), + ProfileEntry(String), + AccessKeyEntry { + values: [String; 3], + selected_field: usize, + }, + ApiKeyEntry(String), + RegionEntry { + credential: BedrockCredential, + value: String, + }, + EnvironmentInstructions, + Configuring(RequestId), +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum BedrockMethodList { + Detected, + All, +} + +#[derive(Clone)] +enum BedrockCredential { + Profile(String), + Environment, + AccessKeys { + access_key_id: String, + secret_access_key: String, + session_token: Option, + }, + ApiKey(String), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum BedrockMethod { + Profile(usize), + Environment, + OtherMethods, + ManualProfile, + AccessKeys, + EnvironmentInstructions, + ApiKey, +} + +enum BedrockAction { + BackToAuth, + Configure(BedrockCredential, String), +} + +impl BedrockState { + fn discovering(request_id: RequestId) -> Self { + Self { + view: BedrockView::Discovering(request_id), + highlighted: 0, + profiles: Vec::new(), + environment_credentials: Vec::new(), + } + } + + fn discovered(response: BedrockDiscoverResponse) -> Self { + Self { + view: BedrockView::Methods(BedrockMethodList::Detected), + highlighted: 0, + profiles: response.profiles, + environment_credentials: response.environment_credentials, + } + } + + fn methods(&self) -> Vec { + if matches!(self.view, BedrockView::EnvironmentInstructions) { + return vec![BedrockMethod::OtherMethods]; + } + let detected = matches!(self.view, BedrockView::Methods(BedrockMethodList::Detected)); + if detected && !self.profiles.is_empty() { + let mut methods = (0..self.profiles.len()) + .map(BedrockMethod::Profile) + .collect::>(); + if !self.environment_credentials.is_empty() { + methods.push(BedrockMethod::Environment); + } + methods.extend([BedrockMethod::OtherMethods, BedrockMethod::ApiKey]); + return methods; + } + if detected && !self.environment_credentials.is_empty() { + return vec![ + BedrockMethod::Environment, + BedrockMethod::OtherMethods, + BedrockMethod::ApiKey, + ]; + } + let mut methods = vec![ + BedrockMethod::ManualProfile, + BedrockMethod::AccessKeys, + BedrockMethod::EnvironmentInstructions, + ]; + if detected { + methods.push(BedrockMethod::ApiKey); + } + methods + } + + pub(super) fn is_text_entry_active(&self) -> bool { + matches!( + self.view, + BedrockView::ProfileEntry(_) + | BedrockView::AccessKeyEntry { .. } + | BedrockView::ApiKeyEntry(_) + | BedrockView::RegionEntry { .. } + ) + } + + fn active_input_mut(&mut self) -> Option<&mut String> { + match &mut self.view { + BedrockView::ProfileEntry(value) + | BedrockView::ApiKeyEntry(value) + | BedrockView::RegionEntry { value, .. } => Some(value), + BedrockView::AccessKeyEntry { + values, + selected_field, + } => Some(&mut values[*selected_field]), + _ => None, + } + } + + fn enter_region(&mut self, credential: BedrockCredential, value: String) { + self.view = BedrockView::RegionEntry { credential, value }; + } + + fn select_method(&mut self, method: BedrockMethod) -> Option { + self.highlighted = 0; + match method { + BedrockMethod::Profile(index) => { + let profile = self.profiles[index].clone(); + let credential = BedrockCredential::Profile(profile.name); + if let Some(region) = profile.region { + return Some(BedrockAction::Configure(credential, region)); + } + self.enter_region(credential, String::new()); + } + BedrockMethod::Environment => { + let credential = BedrockCredential::Environment; + if let Some(region) = self + .environment_credentials + .iter() + .find_map(|credential| credential.region.clone()) + { + return Some(BedrockAction::Configure(credential, region)); + } + self.enter_region(credential, String::new()); + } + BedrockMethod::OtherMethods => { + self.view = BedrockView::Methods(BedrockMethodList::All); + } + BedrockMethod::ManualProfile => { + self.view = BedrockView::ProfileEntry(String::new()); + } + BedrockMethod::AccessKeys => { + self.view = BedrockView::AccessKeyEntry { + values: Default::default(), + selected_field: 0, + }; + } + BedrockMethod::EnvironmentInstructions => { + self.view = BedrockView::EnvironmentInstructions; + } + BedrockMethod::ApiKey => { + self.view = BedrockView::ApiKeyEntry(String::new()); + } + } + None + } + + fn handle_key_event(&mut self, key_event: &KeyEvent) -> Option { + if keys::CANCEL.is_pressed(*key_event) { + let leave_wizard = matches!( + self.view, + BedrockView::Discovering(_) | BedrockView::Methods(BedrockMethodList::Detected) + ); + if leave_wizard { + return Some(BedrockAction::BackToAuth); + } + self.view = BedrockView::Methods(BedrockMethodList::Detected); + self.highlighted = 0; + return None; + } + if matches!(self.view, BedrockView::Configuring(_)) { + return None; + } + + if matches!( + self.view, + BedrockView::Methods(_) | BedrockView::EnvironmentInstructions + ) { + let methods = self.methods(); + if keys::MOVE_UP.is_pressed(*key_event) { + self.highlighted = (self.highlighted + methods.len() - 1) % methods.len(); + } else if keys::MOVE_DOWN.is_pressed(*key_event) { + self.highlighted = (self.highlighted + 1) % methods.len(); + } else if keys::CONFIRM.is_pressed(*key_event) { + return self.select_method(methods[self.highlighted]); + } else if let KeyCode::Char(digit @ '1'..='9') = key_event.code + && let Some(method) = methods.get(digit as usize - '1' as usize).copied() + { + return self.select_method(method); + } + return None; + } + + if let BedrockView::AccessKeyEntry { + values, + selected_field, + } = &mut self.view + { + if key_event.code == KeyCode::Up || key_event.code == KeyCode::BackTab { + *selected_field = (*selected_field + values.len() - 1) % values.len(); + return None; + } + if key_event.code == KeyCode::Down || key_event.code == KeyCode::Tab { + *selected_field = (*selected_field + 1) % values.len(); + return None; + } + } + + if keys::CONFIRM.is_pressed(*key_event) { + return self.confirm_input(); + } + if let Some(input) = self.active_input_mut() { + match key_event.code { + KeyCode::Backspace => { + input.pop(); + } + KeyCode::Char(character) + if matches!(key_event.kind, KeyEventKind::Press | KeyEventKind::Repeat) + && !key_event.modifiers.intersects( + KeyModifiers::SUPER | KeyModifiers::CONTROL | KeyModifiers::ALT, + ) => + { + input.push(character); + } + _ => {} + } + } + None + } + + fn confirm_input(&mut self) -> Option { + let credential = match &mut self.view { + BedrockView::ProfileEntry(value) if !value.trim().is_empty() => { + BedrockCredential::Profile(value.trim().to_string()) + } + BedrockView::ApiKeyEntry(value) if !value.trim().is_empty() => { + BedrockCredential::ApiKey(value.trim().to_string()) + } + BedrockView::RegionEntry { credential, value } if !value.trim().is_empty() => { + return Some(BedrockAction::Configure( + credential.clone(), + value.trim().to_string(), + )); + } + BedrockView::AccessKeyEntry { + values, + selected_field, + } => { + if *selected_field < 2 { + if !values[*selected_field].trim().is_empty() { + *selected_field += 1; + } + return None; + } + if values[0].trim().is_empty() || values[1].trim().is_empty() { + *selected_field = if values[0].trim().is_empty() { 0 } else { 1 }; + return None; + } + BedrockCredential::AccessKeys { + access_key_id: values[0].trim().to_string(), + secret_access_key: values[1].trim().to_string(), + session_token: (!values[2].trim().is_empty()) + .then(|| values[2].trim().to_string()), + } + } + _ => return None, + }; + if let BedrockCredential::Profile(profile) = &credential + && let Some(region) = self.profiles.iter().find_map(|discovered| { + (discovered.name == *profile) + .then(|| discovered.region.clone()) + .flatten() + }) + { + return Some(BedrockAction::Configure(credential, region)); + } + self.enter_region(credential, String::new()); + None + } + + pub(super) fn render(&self, area: Rect, buf: &mut Buffer, error: Option) { + if area.width == 0 || area.height == 0 { + return; + } + let mut lines: Vec = vec![ + Line::from(vec!["> ".into(), "Set up Amazon Bedrock".bold()]), + "".into(), + ]; + match &self.view { + BedrockView::Discovering(_) => { + lines.push(" Checking for existing AWS credentials...".dim().into()); + } + BedrockView::Configuring(_) => { + lines.push(" Setting up Amazon Bedrock...".dim().into()); + } + BedrockView::Methods(list) => { + if *list == BedrockMethodList::Detected && self.profiles.len() == 1 { + let profile = &self.profiles[0]; + lines.push(format!(" AWS profile detected: {}", profile.name).into()); + if let Some(region) = &profile.region { + lines.push(format!(" Region: {region}").dim().into()); + } + } else if *list == BedrockMethodList::Detected && self.profiles.len() > 1 { + lines.push(" Choose an AWS profile.".into()); + } else if *list == BedrockMethodList::Detected + && !self.environment_credentials.is_empty() + { + lines.push(" AWS credentials detected in your environment.".into()); + } else { + if *list == BedrockMethodList::Detected { + lines.push(" No AWS credentials found.".into()); + } + lines.push(" Choose how you authenticate with AWS.".into()); + } + lines.push("".into()); + self.render_methods(&mut lines); + } + BedrockView::ProfileEntry(value) => { + lines.push(" Enter the name of your AWS profile.".into()); + lines.push("".into()); + lines.push(Line::from(vec![ + " AWS profile: ".into(), + value.clone().cyan(), + ])); + } + BedrockView::ApiKeyEntry(value) => { + lines.push(" Enter your Amazon Bedrock API key.".into()); + lines.push("".into()); + let mut masked_value = "•".repeat(value.chars().count().saturating_sub(1)); + if let Some(character) = value.chars().last() { + masked_value.push(character); + } + lines.push(Line::from(vec![ + " Bedrock API key: ".into(), + masked_value.cyan(), + ])); + } + BedrockView::RegionEntry { value, .. } => { + lines.push(" Enter the AWS Region to use with Amazon Bedrock.".into()); + lines.push("".into()); + lines.push(Line::from(vec![ + " AWS Region: ".into(), + value.clone().cyan(), + ])); + } + BedrockView::AccessKeyEntry { + values, + selected_field, + } => { + lines.push(" Enter your AWS access keys.".into()); + lines.push("".into()); + for (index, label) in [ + "AWS access key ID", + "AWS secret access key", + "AWS session token (optional)", + ] + .into_iter() + .enumerate() + { + let marker = if index == *selected_field { ">" } else { " " }; + let value = if index == 0 { + values[index].clone() + } else if index == *selected_field { + let mut masked_value = + "•".repeat(values[index].chars().count().saturating_sub(1)); + if let Some(character) = values[index].chars().last() { + masked_value.push(character); + } + masked_value + } else { + "•".repeat(values[index].chars().count()) + }; + let line = format!("{marker} {label}: {value}"); + lines.push(if index == *selected_field { + line.cyan().into() + } else { + line.into() + }); + } + } + BedrockView::EnvironmentInstructions => { + lines.push( + " Configure AWS credentials in your environment, then restart Codex.".into(), + ); + lines.push("".into()); + lines.push(Line::from(vec![ + " Setup guide: ".into(), + "https://learn.chatgpt.com/docs/amazon-bedrock" + .cyan() + .underlined(), + ])); + lines.push("".into()); + self.render_methods(&mut lines); + } + } + let mut footer = Vec::new(); + if !matches!(self.view, BedrockView::Discovering(_)) { + footer.push("".into()); + if !matches!(self.view, BedrockView::Configuring(_)) { + footer.push(Line::from(vec![ + " Press ".dim(), + keys::CONFIRM[0].into(), + " to continue".dim(), + ])); + } + footer.push(Line::from(vec![ + " Press ".dim(), + keys::CANCEL[0].into(), + " to go back".dim(), + ])); + } + if let Some(error) = error { + footer.push("".into()); + footer.push(error.red().into()); + } + let mut lines = word_wrap_lines(lines, usize::from(area.width)); + let footer = word_wrap_lines(footer, usize::from(area.width)); + if lines.len() + footer.len() <= usize::from(area.height) || footer.is_empty() { + lines.extend(footer); + Paragraph::new(lines).render(area, buf); + return; + } + + let footer_height = u16::try_from(footer.len()) + .unwrap_or(u16::MAX) + .min(area.height.saturating_sub(1)); + let content_height = area.height.saturating_sub(footer_height); + let highlighted_row = lines + .iter() + .enumerate() + .skip(1) + .find_map(|(index, line)| { + line.spans + .first() + .is_some_and(|span| span.content.starts_with("> ")) + .then_some(index) + }) + .unwrap_or_default(); + let scroll = highlighted_row + .saturating_add(2) + .saturating_sub(usize::from(content_height)) + .min(lines.len().saturating_sub(usize::from(content_height))); + let content_area = Rect { + height: content_height, + ..area + }; + let footer_area = Rect { + y: area.y.saturating_add(content_height), + height: footer_height, + ..area + }; + Paragraph::new(lines) + .scroll((u16::try_from(scroll).unwrap_or(u16::MAX), 0)) + .render(content_area, buf); + Paragraph::new(footer).render(footer_area, buf); + } + + fn render_methods(&self, lines: &mut Vec>) { + for (index, method) in self.methods().into_iter().enumerate() { + let (title, description) = match method { + BedrockMethod::Profile(profile_index) => { + let profile = &self.profiles[profile_index]; + let title = if self.profiles.len() == 1 { + format!("Continue with {}", profile.name) + } else { + profile.name.clone() + }; + let description = if self.profiles.len() == 1 { + "Use your existing AWS credentials".to_string() + } else { + profile.region.clone().unwrap_or_default() + }; + (title, description) + } + BedrockMethod::Environment => { + let description = if self.environment_credentials.iter().any(|credential| { + credential.credential_type == AwsCredentialType::BedrockApiKey + }) { + "Use your existing Amazon Bedrock API key" + } else { + "Use your existing AWS credentials" + }; + ( + "Continue with detected credentials".to_string(), + description.to_string(), + ) + } + BedrockMethod::OtherMethods => ( + if matches!(self.view, BedrockView::EnvironmentInstructions) { + "Choose another sign-in method" + } else { + "Other AWS sign-in methods" + } + .to_string(), + if matches!(self.view, BedrockView::EnvironmentInstructions) { + "" + } else { + "Use another profile, access keys, or environment variables" + } + .to_string(), + ), + BedrockMethod::ManualProfile => ( + "AWS profile".to_string(), + "Use AWS SSO or a named profile".to_string(), + ), + BedrockMethod::AccessKeys => ( + "AWS access keys".to_string(), + "Enter an access key ID and secret access key".to_string(), + ), + BedrockMethod::EnvironmentInstructions => ( + "Environment variables".to_string(), + "Configure AWS credentials in your environment, then return here.".to_string(), + ), + BedrockMethod::ApiKey => ( + "Bedrock API key".to_string(), + "Enter a Bedrock API key".to_string(), + ), + }; + let selected = index == self.highlighted; + let marker = if selected { ">" } else { " " }; + let title_line = format!("{marker} {}. {title}", index + 1); + lines.push(if selected { + title_line.cyan().into() + } else { + title_line.into() + }); + if !description.is_empty() + && (!matches!(self.view, BedrockView::Methods(BedrockMethodList::Detected)) + || !self.profiles.is_empty() + || self.environment_credentials.is_empty()) + { + lines.push(format!(" {description}").dim().into()); + } + lines.push("".into()); + } + } +} + +impl AuthModeWidget { + pub(super) fn start_bedrock_discovery(&mut self) { + let request_id = onboarding_request_id(); + *self + .sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner) = + SignInState::Bedrock(BedrockState::discovering(request_id.clone())); + *self.error.write().unwrap_or_else(PoisonError::into_inner) = None; + self.request_frame.schedule_frame(); + + let request_handle = self.app_server_request_handle.clone(); + let sign_in_state = self.sign_in_state.clone(); + let error = self.error.clone(); + let request_frame = self.request_frame.clone(); + tokio::spawn(async move { + let result = request_handle + .request_typed::(ClientRequest::BedrockDiscover { + request_id: request_id.clone(), + params: BedrockDiscoverParams {}, + }) + .await; + let mut guard = sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner); + if !matches!( + &*guard, + SignInState::Bedrock(state) + if matches!(&state.view, BedrockView::Discovering(active_request_id) + if *active_request_id == request_id) + ) { + return; + } + match result { + Ok(response) => { + *guard = SignInState::Bedrock(BedrockState::discovered(response)); + } + Err(err) => { + *guard = + SignInState::Bedrock(BedrockState::discovered(BedrockDiscoverResponse { + profiles: Vec::new(), + environment_credentials: Vec::new(), + })); + *error.write().unwrap_or_else(PoisonError::into_inner) = + Some(format!("Unable to check AWS credentials: {err}")); + } + } + drop(guard); + request_frame.schedule_frame(); + }); + } + + pub(super) fn handle_bedrock_key_event(&mut self, key_event: &KeyEvent) -> bool { + let (action, fallback) = { + let mut guard = self + .sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner); + let SignInState::Bedrock(state) = &mut *guard else { + return false; + }; + *self.error.write().unwrap_or_else(PoisonError::into_inner) = None; + let action = state.handle_key_event(key_event); + (action, state.clone()) + }; + + match action { + Some(BedrockAction::BackToAuth) => { + *self + .sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner) = SignInState::PickMode; + } + Some(BedrockAction::Configure(credential, region)) => { + self.start_bedrock_setup(credential, region, fallback); + } + None => {} + } + self.request_frame.schedule_frame(); + true + } + + pub(super) fn handle_bedrock_paste(&mut self, pasted: &str) -> bool { + let pasted = pasted.trim(); + if pasted.is_empty() { + return false; + } + let mut guard = self + .sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner); + let SignInState::Bedrock(state) = &mut *guard else { + return false; + }; + let Some(input) = state.active_input_mut() else { + return false; + }; + input.push_str(pasted); + *self.error.write().unwrap_or_else(PoisonError::into_inner) = None; + drop(guard); + self.request_frame.schedule_frame(); + true + } + + fn start_bedrock_setup( + &mut self, + credential: BedrockCredential, + region: String, + mut fallback: BedrockState, + ) { + fallback.enter_region(credential.clone(), region.clone()); + let mut configuring = fallback.clone(); + let request_id = onboarding_request_id(); + configuring.view = BedrockView::Configuring(request_id.clone()); + *self + .sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner) = SignInState::Bedrock(configuring); + + let request_handle = self.app_server_request_handle.clone(); + let sign_in_state = self.sign_in_state.clone(); + let error = self.error.clone(); + let request_frame = self.request_frame.clone(); + tokio::spawn(async move { + let result = match credential { + BedrockCredential::ApiKey(api_key) => request_handle + .request_typed::(ClientRequest::LoginAccount { + request_id: request_id.clone(), + params: LoginAccountParams::AmazonBedrock { api_key, region }, + }) + .await + .map_err(|error| error.to_string()) + .and_then(|response| match response { + LoginAccountResponse::AmazonBedrock {} => Ok(()), + response => Err(format!( + "Unexpected account/login/start response: {response:?}" + )), + }), + BedrockCredential::Profile(profile) => request_handle + .request_typed::(ClientRequest::BedrockSetup { + request_id: request_id.clone(), + params: BedrockSetupParams::Profile { profile, region }, + }) + .await + .map(|_| ()) + .map_err(|error| error.to_string()), + BedrockCredential::Environment => request_handle + .request_typed::(ClientRequest::BedrockSetup { + request_id: request_id.clone(), + params: BedrockSetupParams::Environment { region }, + }) + .await + .map(|_| ()) + .map_err(|error| error.to_string()), + BedrockCredential::AccessKeys { + access_key_id, + secret_access_key, + session_token, + } => request_handle + .request_typed::(ClientRequest::LoginAccount { + request_id: request_id.clone(), + params: LoginAccountParams::AmazonBedrockAccessKeys { + access_key_id, + secret_access_key, + session_token, + region, + }, + }) + .await + .map_err(|error| error.to_string()) + .and_then(|response| match response { + LoginAccountResponse::AmazonBedrock {} => Ok(()), + response => Err(format!( + "Unexpected account/login/start response: {response:?}" + )), + }), + }; + let mut guard = sign_in_state + .write() + .unwrap_or_else(PoisonError::into_inner); + if !matches!( + &*guard, + SignInState::Bedrock(state) + if matches!(&state.view, BedrockView::Configuring(active_request_id) + if *active_request_id == request_id) + ) { + return; + } + match result { + Ok(()) => { + *error.write().unwrap_or_else(PoisonError::into_inner) = None; + *guard = SignInState::BedrockConfigured; + } + Err(err) => { + *error.write().unwrap_or_else(PoisonError::into_inner) = + Some(format!("Unable to set up Amazon Bedrock: {err}")); + *guard = SignInState::Bedrock(fallback); + } + } + drop(guard); + request_frame.schedule_frame(); + }); + } +} + +#[cfg(test)] +#[path = "bedrock_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/onboarding/bedrock_tests.rs b/codex-rs/tui/src/onboarding/bedrock_tests.rs new file mode 100644 index 0000000000..d64baa6f93 --- /dev/null +++ b/codex-rs/tui/src/onboarding/bedrock_tests.rs @@ -0,0 +1,498 @@ +use super::*; +use pretty_assertions::assert_eq; + +fn render_visible(state: &BedrockState) -> String { + let area = Rect::new(0, 0, 72, 24); + let mut buffer = Buffer::empty(area); + state.render(area, &mut buffer, /*error*/ None); + let mut rows = (area.top()..area.bottom()) + .map(|row| { + (area.left()..area.right()) + .map(|column| buffer[(column, row)].symbol()) + .collect::() + .trim_end() + .to_string() + }) + .collect::>(); + while rows.last().is_some_and(String::is_empty) { + rows.pop(); + } + rows.join("\n") +} + +#[test] +fn discovery_prioritizes_profiles_and_keeps_bedrock_api_key_last() { + let discovering = BedrockState::discovering(RequestId::Integer(1)); + insta::assert_snapshot!(render_visible(&discovering), @r###" + > Set up Amazon Bedrock + + Checking for existing AWS credentials... + "###); + + let mut configuring = discovering; + configuring.view = BedrockView::Configuring(RequestId::Integer(2)); + insta::assert_snapshot!(render_visible(&configuring), @r###" + > Set up Amazon Bedrock + + Setting up Amazon Bedrock... + + Press esc to go back + "###); + assert!( + configuring + .handle_key_event(&KeyEvent::new(KeyCode::Esc, KeyModifiers::NONE)) + .is_none() + ); + assert!(matches!( + configuring.view, + BedrockView::Methods(BedrockMethodList::Detected) + )); + + let profile = BedrockAwsProfile { + name: "engineering".to_string(), + region: Some("us-east-2".to_string()), + }; + let environment = BedrockEnvironmentCredential { + credential_type: AwsCredentialType::AccessKeys, + region: Some("us-west-2".to_string()), + }; + for (profiles, environment_credentials, expected) in [ + ( + vec![profile.clone()], + vec![environment.clone()], + vec![ + BedrockMethod::Profile(0), + BedrockMethod::Environment, + BedrockMethod::OtherMethods, + BedrockMethod::ApiKey, + ], + ), + ( + vec![ + profile.clone(), + BedrockAwsProfile { + name: "production".to_string(), + region: Some("us-west-2".to_string()), + }, + ], + vec![environment.clone()], + vec![ + BedrockMethod::Profile(0), + BedrockMethod::Profile(1), + BedrockMethod::Environment, + BedrockMethod::OtherMethods, + BedrockMethod::ApiKey, + ], + ), + ( + Vec::new(), + vec![environment], + vec![ + BedrockMethod::Environment, + BedrockMethod::OtherMethods, + BedrockMethod::ApiKey, + ], + ), + ( + Vec::new(), + Vec::new(), + vec![ + BedrockMethod::ManualProfile, + BedrockMethod::AccessKeys, + BedrockMethod::EnvironmentInstructions, + BedrockMethod::ApiKey, + ], + ), + ] { + let state = BedrockState::discovered(BedrockDiscoverResponse { + profiles, + environment_credentials, + }); + assert_eq!(state.methods(), expected); + } + + let mut profile_and_bearer = BedrockState::discovered(BedrockDiscoverResponse { + profiles: vec![profile.clone()], + environment_credentials: vec![BedrockEnvironmentCredential { + credential_type: AwsCredentialType::BedrockApiKey, + region: Some("us-east-2".to_string()), + }], + }); + assert_eq!( + profile_and_bearer.methods(), + vec![ + BedrockMethod::Profile(0), + BedrockMethod::Environment, + BedrockMethod::OtherMethods, + BedrockMethod::ApiKey, + ] + ); + insta::assert_snapshot!(render_visible(&profile_and_bearer), @r###" + > Set up Amazon Bedrock + + AWS profile detected: engineering + Region: us-east-2 + + > 1. Continue with engineering + Use your existing AWS credentials + + 2. Continue with detected credentials + Use your existing Amazon Bedrock API key + + 3. Other AWS sign-in methods + Use another profile, access keys, or environment variables + + 4. Bedrock API key + Enter a Bedrock API key + + + Press enter to continue + Press esc to go back + "###); + profile_and_bearer.select_method(BedrockMethod::OtherMethods); + assert_eq!( + profile_and_bearer.methods(), + vec![ + BedrockMethod::ManualProfile, + BedrockMethod::AccessKeys, + BedrockMethod::EnvironmentInstructions, + ] + ); + insta::assert_snapshot!(render_visible(&profile_and_bearer), @r###" + > Set up Amazon Bedrock + + Choose how you authenticate with AWS. + + > 1. AWS profile + Use AWS SSO or a named profile + + 2. AWS access keys + Enter an access key ID and secret access key + + 3. Environment variables + Configure AWS credentials in your environment, then return here. + + + Press enter to continue + Press esc to go back + "###); + + let profile_state = BedrockState::discovered(BedrockDiscoverResponse { + profiles: vec![profile], + environment_credentials: Vec::new(), + }); + insta::assert_snapshot!(render_visible(&profile_state), @r###" + > Set up Amazon Bedrock + + AWS profile detected: engineering + Region: us-east-2 + + > 1. Continue with engineering + Use your existing AWS credentials + + 2. Other AWS sign-in methods + Use another profile, access keys, or environment variables + + 3. Bedrock API key + Enter a Bedrock API key + + + Press enter to continue + Press esc to go back + "###); + + let multiple_profiles = BedrockState::discovered(BedrockDiscoverResponse { + profiles: vec![ + BedrockAwsProfile { + name: "engineering".to_string(), + region: Some("us-east-2".to_string()), + }, + BedrockAwsProfile { + name: "production".to_string(), + region: Some("us-west-2".to_string()), + }, + ], + environment_credentials: Vec::new(), + }); + insta::assert_snapshot!(render_visible(&multiple_profiles), @r###" + > Set up Amazon Bedrock + + Choose an AWS profile. + + > 1. engineering + us-east-2 + + 2. production + us-west-2 + + 3. Other AWS sign-in methods + Use another profile, access keys, or environment variables + + 4. Bedrock API key + Enter a Bedrock API key + + + Press enter to continue + Press esc to go back + "###); + + let environment_state = BedrockState::discovered(BedrockDiscoverResponse { + profiles: Vec::new(), + environment_credentials: vec![BedrockEnvironmentCredential { + credential_type: AwsCredentialType::AccessKeys, + region: Some("us-west-2".to_string()), + }], + }); + insta::assert_snapshot!(render_visible(&environment_state), @r###" + > Set up Amazon Bedrock + + AWS credentials detected in your environment. + + > 1. Continue with detected credentials + + 2. Other AWS sign-in methods + + 3. Bedrock API key + + + Press enter to continue + Press esc to go back + "###); + + let empty_state = BedrockState::discovered(BedrockDiscoverResponse { + profiles: Vec::new(), + environment_credentials: Vec::new(), + }); + insta::assert_snapshot!(render_visible(&empty_state), @r###" + > Set up Amazon Bedrock + + No AWS credentials found. + Choose how you authenticate with AWS. + + > 1. AWS profile + Use AWS SSO or a named profile + + 2. AWS access keys + Enter an access key ID and secret access key + + 3. Environment variables + Configure AWS credentials in your environment, then return here. + + 4. Bedrock API key + Enter a Bedrock API key + + + Press enter to continue + Press esc to go back + "###); + + let mut many_profiles = BedrockState::discovered(BedrockDiscoverResponse { + profiles: (0..12) + .map(|index| BedrockAwsProfile { + name: format!("engineering-{index}"), + region: Some("us-east-2".to_string()), + }) + .collect(), + environment_credentials: Vec::new(), + }); + many_profiles.highlighted = many_profiles.methods().len() - 1; + let area = Rect::new(0, 0, 72, 10); + let mut buffer = Buffer::empty(area); + many_profiles.render(area, &mut buffer, /*error*/ None); + let visible = (area.top()..area.bottom()) + .flat_map(|row| (area.left()..area.right()).map(move |column| (column, row))) + .map(|position| buffer[position].symbol()) + .collect::(); + assert!(visible.contains("Bedrock API key")); + assert!(visible.contains("Press enter to continue")); + assert!(visible.contains("Press esc to go back")); + + let narrow_area = Rect::new(0, 0, 24, 12); + let mut narrow_buffer = Buffer::empty(narrow_area); + empty_state.render( + narrow_area, + &mut narrow_buffer, + Some("AWS credentials unavailable in this terminal".to_string()), + ); + let narrow_visible = (narrow_area.top()..narrow_area.bottom()) + .map(|row| { + (narrow_area.left()..narrow_area.right()) + .map(|column| narrow_buffer[(column, row)].symbol()) + .collect::() + }) + .collect::>() + .join("\n"); + assert!(narrow_visible.contains("Press enter to")); + assert!(narrow_visible.contains("continue")); + assert!(narrow_visible.contains("Press esc to go back")); + assert!(narrow_visible.contains("AWS credentials")); + assert!(narrow_visible.contains("unavailable")); +} + +#[test] +fn credential_entries_share_region_and_keep_aws_secrets_hidden() { + for (profiles, environment_credentials, method) in [ + ( + vec![BedrockAwsProfile { + name: "engineering".to_string(), + region: None, + }], + Vec::new(), + BedrockMethod::Profile(0), + ), + ( + Vec::new(), + vec![BedrockEnvironmentCredential { + credential_type: AwsCredentialType::AccessKeys, + region: None, + }], + BedrockMethod::Environment, + ), + ] { + let mut state = BedrockState::discovered(BedrockDiscoverResponse { + profiles, + environment_credentials, + }); + state.select_method(method); + assert!(matches!(&state.view, BedrockView::RegionEntry { .. })); + assert!(state.is_text_entry_active()); + } + + let mut access_keys = BedrockState::discovered(BedrockDiscoverResponse { + profiles: Vec::new(), + environment_credentials: Vec::new(), + }); + access_keys.view = BedrockView::AccessKeyEntry { + values: [ + "AKIAEXAMPLE".to_string(), + "secret-value".to_string(), + "session-value".to_string(), + ], + selected_field: 1, + }; + let visible = render_visible(&access_keys); + insta::assert_snapshot!(visible, @r###" + > Set up Amazon Bedrock + + Enter your AWS access keys. + + AWS access key ID: AKIAEXAMPLE + > AWS secret access key: •••••••••••e + AWS session token (optional): ••••••••••••• + + Press enter to continue + Press esc to go back + "###); + assert!(visible.contains("AKIAEXAMPLE")); + assert!(visible.contains("AWS secret access key: •••••••••••e")); + assert!(!visible.contains("secret-value")); + assert!(!visible.contains("session-value")); + + access_keys.handle_key_event(&KeyEvent::new(KeyCode::Char('x'), KeyModifiers::NONE)); + assert!(render_visible(&access_keys).contains("AWS secret access key: ••••••••••••x")); + + access_keys.handle_key_event(&KeyEvent::new(KeyCode::Down, KeyModifiers::NONE)); + let visible = render_visible(&access_keys); + insta::assert_snapshot!(visible, @r###" + > Set up Amazon Bedrock + + Enter your AWS access keys. + + AWS access key ID: AKIAEXAMPLE + AWS secret access key: ••••••••••••• + > AWS session token (optional): ••••••••••••e + + Press enter to continue + Press esc to go back + "###); + assert!(!visible.contains("secret-value")); + assert!(!visible.contains("session-value")); + + access_keys.view = BedrockView::ApiKeyEntry("bedrock-api-key".to_string()); + let visible = render_visible(&access_keys); + insta::assert_snapshot!(visible, @r###" + > Set up Amazon Bedrock + + Enter your Amazon Bedrock API key. + + Bedrock API key: ••••••••••••••y + + Press enter to continue + Press esc to go back + "###); + assert!(!visible.contains("bedrock-api-key")); + access_keys.handle_key_event(&KeyEvent::new(KeyCode::Char('z'), KeyModifiers::NONE)); + assert!(render_visible(&access_keys).contains("Bedrock API key: •••••••••••••••z")); + + access_keys.view = BedrockView::ProfileEntry("engineering".to_string()); + insta::assert_snapshot!(render_visible(&access_keys), @r###" + > Set up Amazon Bedrock + + Enter the name of your AWS profile. + + AWS profile: engineering + + Press enter to continue + Press esc to go back + "###); + + access_keys.enter_region(BedrockCredential::Environment, "us-west-1".to_string()); + insta::assert_snapshot!(render_visible(&access_keys), @r###" + > Set up Amazon Bedrock + + Enter the AWS Region to use with Amazon Bedrock. + + AWS Region: us-west-1 + + Press enter to continue + Press esc to go back + "###); + + access_keys.view = BedrockView::EnvironmentInstructions; + insta::assert_snapshot!(render_visible(&access_keys), @r###" + > Set up Amazon Bedrock + + Configure AWS credentials in your environment, then restart Codex. + + Setup guide: https://learn.chatgpt.com/docs/amazon-bedrock + + > 1. Choose another sign-in method + + + Press enter to continue + Press esc to go back + "###); + + for view in [ + BedrockView::ProfileEntry(String::new()), + BedrockView::ApiKeyEntry(String::new()), + BedrockView::RegionEntry { + credential: BedrockCredential::Environment, + value: String::new(), + }, + BedrockView::AccessKeyEntry { + values: Default::default(), + selected_field: 0, + }, + ] { + access_keys.view = view; + assert!(access_keys.is_text_entry_active()); + } + + for credential in [ + BedrockCredential::Profile("engineering".to_string()), + BedrockCredential::Environment, + BedrockCredential::AccessKeys { + access_key_id: "AKIAEXAMPLE".to_string(), + secret_access_key: "secret-value".to_string(), + session_token: None, + }, + BedrockCredential::ApiKey("bedrock-api-key".to_string()), + ] { + access_keys.enter_region(credential, "us-west-1".to_string()); + assert!(matches!( + &access_keys.view, + BedrockView::RegionEntry { value, .. } if value == "us-west-1" + )); + } +} diff --git a/codex-rs/tui/src/onboarding/keys.rs b/codex-rs/tui/src/onboarding/keys.rs index 5d782a553f..d573f47d33 100644 --- a/codex-rs/tui/src/onboarding/keys.rs +++ b/codex-rs/tui/src/onboarding/keys.rs @@ -23,6 +23,7 @@ pub(crate) const SELECT_SECOND: [KeyBinding; 2] = [ key_hint::plain(KeyCode::Char('n')), ]; pub(crate) const SELECT_THIRD: [KeyBinding; 1] = [key_hint::plain(KeyCode::Char('3'))]; +pub(crate) const SELECT_FOURTH: [KeyBinding; 1] = [key_hint::plain(KeyCode::Char('4'))]; pub(crate) const CONFIRM: [KeyBinding; 1] = [key_hint::plain(KeyCode::Enter)]; pub(crate) const CANCEL: [KeyBinding; 1] = [key_hint::plain(KeyCode::Esc)]; pub(crate) const QUIT: [KeyBinding; 3] = [ diff --git a/codex-rs/tui/src/onboarding/mod.rs b/codex-rs/tui/src/onboarding/mod.rs index 016d086c54..c5438a1e51 100644 --- a/codex-rs/tui/src/onboarding/mod.rs +++ b/codex-rs/tui/src/onboarding/mod.rs @@ -1,4 +1,5 @@ mod auth; +mod bedrock; mod keys; pub(crate) mod onboarding_screen; mod trust_directory; diff --git a/codex-rs/tui/src/onboarding/onboarding_screen.rs b/codex-rs/tui/src/onboarding/onboarding_screen.rs index b0417ff95f..9ef2a74dca 100644 --- a/codex-rs/tui/src/onboarding/onboarding_screen.rs +++ b/codex-rs/tui/src/onboarding/onboarding_screen.rs @@ -92,6 +92,7 @@ pub(crate) struct OnboardingScreenArgs { pub show_trust_screen: bool, pub remote_project_trust: Option, pub show_login_screen: bool, + pub bedrock_setup_enabled: bool, pub login_status: LoginStatus, pub app_server_request_handle: Option, pub config: Config, @@ -103,10 +104,10 @@ pub(crate) struct OnboardingResult { } #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -struct ApiKeyEntryContext { - /// True when onboarding is currently rendering the API-key entry state. +struct TextEntryContext { + /// Whether onboarding is currently rendering a text-entry state. active: bool, - /// True when the API-key input field currently contains user text. + /// Whether printable quit shortcuts should be consumed as text input. has_text: bool, } @@ -116,6 +117,7 @@ impl OnboardingScreen { show_trust_screen, remote_project_trust, show_login_screen, + bedrock_setup_enabled, login_status, app_server_request_handle, config, @@ -147,6 +149,7 @@ impl OnboardingScreen { login_status, app_server_request_handle, auth_config, + bedrock_setup_enabled, animations_enabled: config.animations, animations_suppressed: std::cell::Cell::new(false), })); @@ -287,14 +290,14 @@ impl OnboardingScreen { } } - fn api_key_entry_context(&self) -> ApiKeyEntryContext { + fn text_entry_context(&self) -> TextEntryContext { self.steps .iter() .find_map(|step| { if let Step::Auth(widget) = step { - Some(ApiKeyEntryContext { - active: widget.is_api_key_entry_active(), - has_text: widget.api_key_entry_has_text(), + Some(TextEntryContext { + active: widget.is_text_entry_active(), + has_text: widget.should_suppress_printable_quit(), }) } else { None @@ -307,19 +310,17 @@ impl OnboardingScreen { impl KeyboardHandler for OnboardingScreen { /// Route key events to onboarding steps while preserving text-entry safety. /// - /// In API-key entry mode, printable quit bindings are suppressed only after - /// the user has started typing in the API-key field. This keeps the - /// printable `q` quit key usable on an empty field while protecting in-progress - /// text entry from accidental exits. Control/alt quit chords still work as - /// emergency exits. + /// OpenAI API-key entry suppresses printable quit bindings after typing + /// starts, while Bedrock fields suppress them from the first character. + /// Control/alt quit chords still work as emergency exits. fn handle_key_event(&mut self, key_event: KeyEvent) { if !matches!(key_event.kind, KeyEventKind::Press | KeyEventKind::Repeat) { return; } - let api_key_entry_context = self.api_key_entry_context(); + let text_entry_context = self.text_entry_context(); let should_quit = key_event.kind == KeyEventKind::Press && keys::QUIT.is_pressed(key_event) - && !suppress_quit_while_typing_api_key(key_event, api_key_entry_context); + && !suppress_quit_while_typing(key_event, text_entry_context); if should_quit { if self.is_auth_in_progress() { self.cancel_auth_if_active(); @@ -369,16 +370,11 @@ impl KeyboardHandler for OnboardingScreen { /// Returns `true` when a quit shortcut should be ignored as text input. /// -/// This only applies while API-key entry is active and the key is a printable -/// character without control/alt modifiers and there is already text in the -/// input field. Empty input intentionally does not trigger suppression so -/// the printable `q` quit key can still exit onboarding. -fn suppress_quit_while_typing_api_key( - key_event: KeyEvent, - api_key_entry_context: ApiKeyEntryContext, -) -> bool { - api_key_entry_context.active - && api_key_entry_context.has_text +/// This only applies when the active text-entry state suppresses printable quit +/// shortcuts and the key has no control/alt modifiers. +fn suppress_quit_while_typing(key_event: KeyEvent, text_entry_context: TextEntryContext) -> bool { + text_entry_context.active + && text_entry_context.has_text && matches!(key_event.code, KeyCode::Char(_)) && !key_event .modifiers @@ -712,13 +708,13 @@ async fn persist_selected_trust( #[cfg(test)] mod tests { - use super::ApiKeyEntryContext; use super::OnboardingScreen; use super::Step; use super::StepStateProvider; + use super::TextEntryContext; use super::discard_pending_input_on_trust_step_transition; use super::persist_selected_trust; - use super::suppress_quit_while_typing_api_key; + use super::suppress_quit_while_typing; use crate::onboarding::onboarding_screen::KeyboardHandler; use crate::onboarding::trust_directory::TrustDirectorySelection; use crate::onboarding::trust_directory::TrustDirectoryWidget; @@ -731,9 +727,9 @@ mod tests { #[test] fn suppresses_printable_quit_key_during_api_key_entry() { - let suppressed = suppress_quit_while_typing_api_key( + let suppressed = suppress_quit_while_typing( KeyEvent::new(KeyCode::Char('q'), KeyModifiers::NONE), - ApiKeyEntryContext { + TextEntryContext { active: true, has_text: true, }, @@ -743,9 +739,9 @@ mod tests { #[test] fn does_not_suppress_printable_quit_key_when_api_key_input_is_empty() { - let suppressed = suppress_quit_while_typing_api_key( + let suppressed = suppress_quit_while_typing( KeyEvent::new(KeyCode::Char('q'), KeyModifiers::NONE), - ApiKeyEntryContext { + TextEntryContext { active: true, has_text: false, }, @@ -755,9 +751,9 @@ mod tests { #[test] fn does_not_suppress_control_quit_key_during_api_key_entry() { - let suppressed = suppress_quit_while_typing_api_key( + let suppressed = suppress_quit_while_typing( KeyEvent::new(KeyCode::Char('x'), KeyModifiers::CONTROL), - ApiKeyEntryContext { + TextEntryContext { active: true, has_text: true, }, @@ -767,9 +763,9 @@ mod tests { #[test] fn does_not_suppress_when_not_in_api_key_entry() { - let suppressed = suppress_quit_while_typing_api_key( + let suppressed = suppress_quit_while_typing( KeyEvent::new(KeyCode::Char('x'), KeyModifiers::NONE), - ApiKeyEntryContext { + TextEntryContext { active: false, has_text: true, },