diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index dbd1e24a33..31daabdd6a 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1033,6 +1033,9 @@ "fast_mode": { "type": "boolean" }, + "gateway_setup_wizard": { + "type": "boolean" + }, "gcp_setup_wizard": { "type": "boolean" }, @@ -8056,6 +8059,9 @@ "fast_mode": { "type": "boolean" }, + "gateway_setup_wizard": { + "type": "boolean" + }, "gcp_setup_wizard": { "type": "boolean" }, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 6164574f65..6101738b54 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -477,6 +477,8 @@ pub enum Feature { BedrockSetupWizard, /// Offer Google Cloud Vertex AI setup during TUI sign-in onboarding. GcpSetupWizard, + /// Offer custom gateway setup during TUI sign-in onboarding. + GatewaySetupWizard, /// Removed compatibility flag retained as a no-op. Personality, /// Enable native artifact tools. @@ -2262,6 +2264,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: false, }, + FeatureSpec { + id: Feature::GatewaySetupWizard, + key: "gateway_setup_wizard", + stage: Stage::Stable, + default_enabled: false, + }, FeatureSpec { id: Feature::Personality, key: "personality", diff --git a/codex-rs/tui/src/app/provider.rs b/codex-rs/tui/src/app/provider.rs index 7c742fecb6..2281789bbe 100644 --- a/codex-rs/tui/src/app/provider.rs +++ b/codex-rs/tui/src/app/provider.rs @@ -14,22 +14,18 @@ use crate::pager_overlay::Overlay; use crate::pager_overlay::PagerAction; use crate::pager_overlay::StaticOverlay; use crate::tui::Tui; -use codex_app_server_client::TypedRequestError; -use codex_app_server_protocol::ClientRequest; -use codex_app_server_protocol::ModelProviderListParams; use codex_app_server_protocol::ModelProviderListResponse; use codex_app_server_protocol::ModelProviderSummary; use codex_app_server_protocol::ProviderSettings; -use codex_app_server_protocol::RequestId; use crossterm::event::KeyCode; use ratatui::style::Stylize; use ratatui::text::Line; use std::sync::Arc; use tokio::sync::watch; -use uuid::Uuid; mod editor; mod fields; +pub(crate) mod rpc; mod write; use fields::Field; use fields::provider_fields; @@ -81,6 +77,16 @@ fn back() -> SelectionItem { } } +pub(crate) fn load_error(title: &str, error: String) -> SelectionViewParams { + SelectionViewParams { + title: Some(title.into()), + subtitle: Some(error), + items: vec![back()], + allow_input_when_disconnected: true, + ..Default::default() + } +} + pub(crate) fn loading() -> (SelectionViewParams, Event) { loading_with_notice(/*notice*/ None) } @@ -117,16 +123,12 @@ pub(super) fn handle(app: &mut App, tui: &mut Tui, server: &AppServerSession, ev let handle = server.request_handle(); let offline = app.reconnect.offline; tokio::spawn(async move { - let request = ClientRequest::ModelProviderList { - request_id: RequestId::String(Uuid::new_v4().to_string()), - params: ModelProviderListParams {}, - }; let result = if offline { Err("Reconnect to the app server to browse model providers.".into()) } else { tokio::select! { _ = dismissed.changed() => return, - result = handle.request_typed(request) => result.map_err(request_error), + result = rpc::list_providers(&handle) => result, } }; tx.send(AppEvent::Provider(Event::Loaded { @@ -156,13 +158,7 @@ pub(super) fn handle(app: &mut App, tui: &mut Tui, server: &AppServerSession, ev &app.keymap.list, notice, ), - Err(error) => SelectionViewParams { - title: Some("Model providers".into()), - subtitle: Some(error), - items: vec![back()], - allow_input_when_disconnected: true, - ..Default::default() - }, + Err(error) => load_error("Model providers", error), }; chat.replace_selection_view_if_present(LOADING_VIEW, view); } @@ -242,17 +238,18 @@ pub(crate) fn refresh_list_if_present(chat: &mut crate::chatwidget::ChatWidget) .then_some(event) } -fn request_error(err: TypedRequestError) -> String { - if let TypedRequestError::Server { method, source } = &err - && (source.code == -32601 - || (source.code == -32600 - && source - .message - .starts_with(&format!("Invalid request: unknown variant `{method}`")))) - { - return "Update the connected app server to browse model providers.".into(); +pub(crate) fn provider_item(provider: &ModelProviderSummary, default_id: &str) -> SelectionItem { + let is_default = provider.id == default_id; + SelectionItem { + name: provider.name.clone(), + name_prefix_spans: vec![if is_default { + "● ".green() + } else { + "○ ".dim() + }], + description: Some(provider.id.clone()), + ..Default::default() } - format!("Could not load model providers: {err}") } fn list( @@ -303,14 +300,7 @@ fn list( let catalog = Arc::clone(&catalog); let dismissed = dismissed.clone(); SelectionItem { - name: provider.name.clone(), - name_prefix_spans: vec![if is_default { - "● ".green() - } else { - "○ ".dim() - }], is_current: current == Some(provider.id.as_str()), - description: Some(provider.id.clone()), secondary_action, actions: vec![Box::new(move |tx| { tx.send(AppEvent::Provider(Event::Details { @@ -319,7 +309,7 @@ fn list( dismissed: dismissed.clone(), })); })], - ..Default::default() + ..provider_item(provider, &page.default_provider_id) } }) .collect(); diff --git a/codex-rs/tui/src/app/provider/rpc.rs b/codex-rs/tui/src/app/provider/rpc.rs new file mode 100644 index 0000000000..b807a6a571 --- /dev/null +++ b/codex-rs/tui/src/app/provider/rpc.rs @@ -0,0 +1,62 @@ +//! List providers and select a default from either settings or sign-in. + +use codex_app_server_client::AppServerRequestHandle; +use codex_app_server_client::TypedRequestError; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::ConfigValueWriteParams; +use codex_app_server_protocol::ConfigWriteResponse; +use codex_app_server_protocol::MergeStrategy; +use codex_app_server_protocol::ModelProviderListParams; +use codex_app_server_protocol::ModelProviderListResponse; +use codex_app_server_protocol::RequestId; +use uuid::Uuid; + +pub(crate) async fn list_providers( + handle: &AppServerRequestHandle, +) -> Result { + handle + .request_typed(ClientRequest::ModelProviderList { + request_id: RequestId::String(Uuid::new_v4().to_string()), + params: ModelProviderListParams {}, + }) + .await + .map_err(request_error) +} + +pub(crate) async fn set_default_provider( + handle: &AppServerRequestHandle, + id: &str, +) -> Result { + handle + .request_typed(ClientRequest::ConfigValueWrite { + request_id: RequestId::String(Uuid::new_v4().to_string()), + params: ConfigValueWriteParams { + key_path: "model_provider".into(), + value: id.into(), + merge_strategy: MergeStrategy::Replace, + file_path: None, + expected_version: None, + }, + }) + .await +} + +pub(crate) fn override_reason(response: ConfigWriteResponse) -> String { + response + .overridden_metadata + .map(|metadata| metadata.message) + .unwrap_or_else(|| "A higher-priority setting overrides this choice.".into()) +} + +pub(super) fn request_error(err: TypedRequestError) -> String { + if let TypedRequestError::Server { method, source } = &err + && (source.code == -32601 + || (source.code == -32600 + && source + .message + .starts_with(&format!("Invalid request: unknown variant `{method}`")))) + { + return "Update the connected app server to browse model providers.".into(); + } + format!("Could not load model providers: {err}") +} diff --git a/codex-rs/tui/src/app/provider/write.rs b/codex-rs/tui/src/app/provider/write.rs index 51d724e164..38f381a779 100644 --- a/codex-rs/tui/src/app/provider/write.rs +++ b/codex-rs/tui/src/app/provider/write.rs @@ -11,12 +11,10 @@ use super::edit_view; use super::editor; use super::loading_with_notice; use super::progress; +use super::rpc; use codex_app_server_client::AppServerRequestHandle; use codex_app_server_client::TypedRequestError; use codex_app_server_protocol::ClientRequest; -use codex_app_server_protocol::ConfigValueWriteParams; -use codex_app_server_protocol::ConfigWriteResponse; -use codex_app_server_protocol::MergeStrategy; use codex_app_server_protocol::ModelProviderConfigWriteParams; use codex_app_server_protocol::ModelProviderConfigWriteResponse; use codex_app_server_protocol::ProviderSettingsPatch; @@ -120,21 +118,14 @@ impl Write { } }), Self::Default(id) => { - let response: ConfigWriteResponse = handle - .request_typed(ClientRequest::ConfigValueWrite { - request_id: RequestId::String(Uuid::new_v4().to_string()), - params: ConfigValueWriteParams { - key_path: "model_provider".into(), - value: id.clone().into(), - merge_strategy: MergeStrategy::Replace, - file_path: None, - expected_version: None, - }, - }) + let response = rpc::set_default_provider(handle, id) .await .map_err(|error| match error { - TypedRequestError::Server { .. } => format!("Could not save provider: {error}"), - TypedRequestError::Transport { .. } | TypedRequestError::Deserialize { .. } => { + TypedRequestError::Server { .. } => { + format!("Could not save provider: {error}") + } + TypedRequestError::Transport { .. } + | TypedRequestError::Deserialize { .. } => { // The server may have committed the write before its reply was lost. *provider_changed = true; format!("Could not confirm whether {id} was saved: {error}") @@ -143,11 +134,12 @@ impl Write { match response.status { WriteStatus::Ok => { *provider_changed = true; - Ok(SaveNotice::Saved(format!("/new now uses {id}. Restart the app server and TUI to refresh /model."))) + Ok(SaveNotice::Saved(format!( + "/new now uses {id}. Restart the app server and TUI to refresh /model." + ))) } WriteStatus::OkOverridden => { - let reason = response.overridden_metadata.map(|metadata| metadata.message) - .unwrap_or_else(|| "A higher-priority setting overrides this choice.".into()); + let reason = rpc::override_reason(response); Ok(SaveNotice::Saved(format!("Saved {id}. {reason}"))) } } diff --git a/codex-rs/tui/src/app/provider_tests.rs b/codex-rs/tui/src/app/provider_tests.rs index 74531ad297..c1f7799df7 100644 --- a/codex-rs/tui/src/app/provider_tests.rs +++ b/codex-rs/tui/src/app/provider_tests.rs @@ -11,6 +11,7 @@ use crate::config_panel::ResumeCurrentAvailability; use crate::model_catalog::ModelCatalog; use crate::tui::TuiEvent; use crate::tui::test_support::make_test_tui; +use codex_app_server_client::TypedRequestError; use codex_protocol::ThreadId; use color_eyre::eyre::Result; use crossterm::event::KeyCode; @@ -349,7 +350,7 @@ fn unsupported_list_method_explains_server_compatibility() { }, }; assert_eq!( - request_error(error), + rpc::request_error(error), "Update the connected app server to browse model providers." ); } diff --git a/codex-rs/tui/src/bottom_pane/list_selection_view.rs b/codex-rs/tui/src/bottom_pane/list_selection_view.rs index 6dca72b967..5b10995f14 100644 --- a/codex-rs/tui/src/bottom_pane/list_selection_view.rs +++ b/codex-rs/tui/src/bottom_pane/list_selection_view.rs @@ -661,7 +661,7 @@ impl ListSelectionView { limit.min(len.max(/*other*/ 1)) } - fn selected_actual_idx(&self) -> Option { + pub(crate) fn selected_actual_idx(&self) -> Option { self.state .selected_idx .and_then(|visible_idx| self.filtered_indices.get(visible_idx).copied()) diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 93a3c64941..115f0835c2 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -1310,22 +1310,34 @@ async fn run_ratatui_app( // Initialize high-fidelity session event logging if enabled. session_log::maybe_init(&initial_config); + let connect = async |target: &mut AppServerTarget, + config: &Config, + bundle: &CloudConfigBundleLoader, + state_db: &mut Option| { + start_app_server( + target, + arg0_paths.clone(), + config.clone(), + cli_kv_overrides.clone(), + loader_overrides.clone(), + strict_config, + bundle.clone(), + feedback.clone(), + log_db.clone(), + state_db, + environment_manager.clone(), + embedded_network_policy.clone(), + ) + .await + }; let startup_app_server = startup_draft .run_until( &mut tui, - start_app_server( + connect( &mut app_server_target, - arg0_paths.clone(), - initial_config.clone(), - cli_kv_overrides.clone(), - loader_overrides.clone(), - strict_config, - cloud_config_bundle.clone(), - feedback.clone(), - log_db.clone(), + &initial_config, + &cloud_config_bundle, &mut state_db, - environment_manager.clone(), - embedded_network_policy.clone(), ), ) .await; @@ -1388,7 +1400,7 @@ async fn run_ratatui_app( let should_show_trust_screen_flag = false; #[cfg(target_os = "windows")] let mut trust_decision_was_made = false; - let startup_model_provider = initial_config.model_provider_id.clone(); + let mut startup_model_provider = initial_config.model_provider_id.clone(); let Some(active_app_server) = app_server.as_mut() else { unreachable!("app server should exist when auth is required"); }; @@ -1423,25 +1435,26 @@ 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, requires_openai_auth); - let bedrock_setup_enabled = should_show_provider_setup_wizard( - login_status, - requires_openai_auth, - &initial_config, - &app_server_target, - Feature::BedrockSetupWizard, - ); - let gcp_setup_enabled = should_show_provider_setup_wizard( - login_status, - requires_openai_auth, - &initial_config, - &app_server_target, - Feature::GcpSetupWizard, - ); + let show_wizard = |feature| { + should_show_provider_setup_wizard( + login_status, + requires_openai_auth, + &initial_config, + &app_server_target, + feature, + ) + }; let onboarding_result = run_onboarding_app( OnboardingScreenArgs { show_login_screen, - bedrock_setup_enabled, - gcp_setup_enabled, + bedrock_setup_enabled: show_wizard(Feature::BedrockSetupWizard), + gcp_setup_enabled: show_wizard(Feature::GcpSetupWizard), + gateway_setup_enabled: show_wizard(Feature::GatewaySetupWizard) + && app_server.as_ref().is_some_and(|server| { + server + .client_provider_override_for_new_chats(&initial_config) + .is_none() + }), show_trust_screen: should_show_trust_screen_flag, remote_project_trust: None, login_status, @@ -1483,6 +1496,57 @@ async fn run_ratatui_app( trust_decision_was_made = !uses_remote_workspace && onboarding_result.directory_trust_persisted; } + if let Some(selected_id) = onboarding_result.selected_gateway { + // Account state belongs to the provider captured by the embedded server at startup. + let mut selected = startup_draft + .run_until( + &mut tui, + load_config_or_exit_with_fallback_cwd( + cli_kv_overrides.clone(), + overrides.clone(), + loader_overrides.clone(), + cloud_config_bundle.clone(), + strict_config, + /*fallback_cwd*/ None, + managed_worktree.as_ref(), + ), + ) + .await?; + color_eyre::eyre::ensure!( + selected.model_provider_id == selected_id, + "The current configuration selects {} instead of {selected_id}.", + selected.model_provider_id + ); + embedded_network_policy.bind_config(&mut selected); + shutdown_app_server_if_present(app_server.take()).await; + let client = startup_draft + .run_until( + &mut tui, + connect( + &mut app_server_target, + &selected, + &cloud_config_bundle, + &mut state_db, + ), + ) + .await??; + let mut server = + AppServerSession::new(backend_id, client, app_server_target.thread_params_mode()) + .with_local_codex_home(&selected.codex_home) + .with_remote_cwd_override(remote_cwd_override.clone()); + if let Err(err) = boxed_future!(onboarding::authenticate_gateway( + &mut tui, + &mut server, + &selected + )) + .await + { + shutdown_startup_session(Some(server), &mut terminal_restore_guard).await; + return Err(err); + } + startup_model_provider = selected_id; + app_server = Some(server); + } let reloaded_config = startup_draft .run_until(&mut tui, async { // If this onboarding run included the login step, always refresh the cloud config @@ -1915,19 +1979,11 @@ async fn run_ratatui_app( None => match startup_draft .run_until( &mut tui, - start_app_server( + connect( &mut app_server_target, - arg0_paths.clone(), - config.clone(), - cli_kv_overrides.clone(), - loader_overrides.clone(), - strict_config, - cloud_config_bundle.clone(), - feedback.clone(), - log_db.clone(), + &config, + &cloud_config_bundle, &mut state_db, - environment_manager.clone(), - embedded_network_policy.clone(), ), ) .await @@ -2562,6 +2618,11 @@ fn should_show_provider_setup_wizard( supported_target && should_show_login_screen(login_status, requires_openai_auth) && config.features.enabled(feature) + && (feature != Feature::GatewaySetupWizard + || config + .config_layer_stack + .required_model_provider() + .is_none()) && config.model_provider_id == "openai" && config .auth_config() diff --git a/codex-rs/tui/src/lib_tests.rs b/codex-rs/tui/src/lib_tests.rs index 9516e002e1..40e360aa3f 100644 --- a/codex-rs/tui/src/lib_tests.rs +++ b/codex-rs/tui/src/lib_tests.rs @@ -81,7 +81,7 @@ async fn provider_setup_wizards_require_eligible_onboarding() -> color_eyre::Res websocket_url: "ws://127.0.0.1:4500/".to_string(), auth_token: None, }; - let enabled = "[features]\nbedrock_setup_wizard = true\ngcp_setup_wizard = true\n"; + let enabled = "[features]\nbedrock_setup_wizard = true\ngcp_setup_wizard = true\ngateway_setup_wizard = true\n"; for (label, config_toml, login_status, target, expected) in [ ( @@ -89,42 +89,49 @@ async fn provider_setup_wizards_require_eligible_onboarding() -> color_eyre::Res "", LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [true, false], + [true, false, false], ), ( "explicitly disabled", - "[features]\nbedrock_setup_wizard = false\ngcp_setup_wizard = false\n", + "[features]\nbedrock_setup_wizard = false\ngcp_setup_wizard = false\ngateway_setup_wizard = false\n", LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [false, false], + [false, false, false], ), ( "only GCP enabled", "[features]\nbedrock_setup_wizard = false\ngcp_setup_wizard = true\n", LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [false, true], + [false, true, false], + ), + ( + "only custom gateway enabled", + "[features]\nbedrock_setup_wizard = false\ngateway_setup_wizard = true\n", + LoginStatus::NotAuthenticated, + AppServerTarget::Embedded, + [false, false, true], ), ( "enabled for the default provider", enabled, LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [true, true], + [true, true, true], ), ( "explicit OpenAI provider with default flags", "model_provider = \"openai\"\n", LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [true, false], + [true, false, false], ), ( - "explicit OpenAI provider with both wizards enabled", - "model_provider = \"openai\"\n[features]\nbedrock_setup_wizard = true\ngcp_setup_wizard = true\n", + "explicit OpenAI provider with wizards enabled", + "model_provider = \"openai\"\n[features]\nbedrock_setup_wizard = true\ngcp_setup_wizard = true\ngateway_setup_wizard = true\n", LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [true, true], + [true, true, true], ), ( "custom provider requiring OpenAI auth", @@ -138,24 +145,25 @@ requires_openai_auth = true [features] bedrock_setup_wizard = true gcp_setup_wizard = true +gateway_setup_wizard = true "#, LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [false, false], + [false, false, false], ), ( "forced ChatGPT login", - "forced_login_method = \"chatgpt\"\n[features]\nbedrock_setup_wizard = true\ngcp_setup_wizard = true\n", + "forced_login_method = \"chatgpt\"\n[features]\nbedrock_setup_wizard = true\ngcp_setup_wizard = true\ngateway_setup_wizard = true\n", LoginStatus::NotAuthenticated, AppServerTarget::Embedded, - [false, false], + [false, false, false], ), ( "existing authentication", enabled, LoginStatus::AuthMode(AuthMode::Chatgpt), AppServerTarget::Embedded, - [false, false], + [false, false, false], ), ( "shared local daemon", @@ -165,7 +173,7 @@ gcp_setup_wizard = true allow_embedded_fallback: true, endpoint: shared_endpoint.clone(), }, - [true, false], + [true, false, false], ), ( "remote app server", @@ -174,15 +182,19 @@ gcp_setup_wizard = true AppServerTarget::Remote { endpoint: shared_endpoint, }, - [false, false], + [false, false, false], ), ] { let codex_home = TempDir::new()?; std::fs::write(codex_home.path().join("config.toml"), config_toml)?; let config = build_config(&codex_home).await?; - for (feature, expected) in [Feature::BedrockSetupWizard, Feature::GcpSetupWizard] - .into_iter() - .zip(expected) + for (feature, expected) in [ + Feature::BedrockSetupWizard, + Feature::GcpSetupWizard, + Feature::GatewaySetupWizard, + ] + .into_iter() + .zip(expected) { assert_eq!( should_show_provider_setup_wizard( @@ -205,6 +217,26 @@ gcp_setup_wizard = true } } + let home = TempDir::new()?; + std::fs::write(home.path().join("config.toml"), enabled)?; + let requirements = home.path().join("requirements.toml"); + std::fs::write(&requirements, "model_provider = 'openai'")?; + let config = ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .loader_overrides(LoaderOverrides { + system_requirements_path: Some(requirements), + ..LoaderOverrides::without_managed_config_for_tests() + }) + .build() + .await?; + assert!(!should_show_provider_setup_wizard( + LoginStatus::NotAuthenticated, + /*requires_openai_auth*/ true, + &config, + &AppServerTarget::Embedded, + Feature::GatewaySetupWizard, + )); + Ok(()) } diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index bb87405cc7..14ed755b21 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -53,6 +53,7 @@ use crate::key_hint::KeyBindingListExt; use crate::motion::MotionMode; use crate::motion::shimmer_text; use crate::onboarding::bedrock::BedrockState; +use crate::onboarding::gateway::GatewayStep; use crate::onboarding::gcp::GcpState; use crate::onboarding::keys; use crate::onboarding::onboarding_screen::KeyboardHandler; @@ -91,6 +92,16 @@ pub(crate) enum SignInState { BedrockConfigured, Gcp(GcpState), GcpConfigured, + GatewaySelected(String), + GatewayOAuth { + provider_id: String, + url: Option, + }, + GatewayOAuthFailed { + provider_id: String, + message: String, + }, + GatewayConfigured, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -101,6 +112,7 @@ pub(crate) enum SignInOption { ApiKey, Bedrock, Gcp, + Gateway, } impl SignInOption { @@ -198,7 +210,10 @@ impl KeyboardHandler for AuthModeWidget { { return; } - if self.handle_gcp_key_event(&key_event) || self.handle_bedrock_key_event(&key_event) { + if self.handle_gcp_key_event(&key_event) + || self.handle_bedrock_key_event(&key_event) + || self.handle_gateway_key_event(key_event) + { return; } if self.handle_api_key_entry_key_event(&key_event) { @@ -228,9 +243,16 @@ impl KeyboardHandler for AuthModeWidget { self.select_option_by_index(/*index*/ 2, key_event); return; } + if keys::SELECT_FOURTH.is_pressed(key_event) { + self.select_option_by_index(/*index*/ 3, key_event); + return; + } if keys::CONFIRM.is_pressed(key_event) { let sign_in_state = { (*self.sign_in_state.read().unwrap()).clone() }; match sign_in_state { + SignInState::GatewayOAuthFailed { provider_id, .. } => { + self.start_gateway_oauth(provider_id) + } SignInState::PickMode => { self.activate_highlighted_mode(key_event); } @@ -274,7 +296,6 @@ impl KeyboardHandler for AuthModeWidget { } } -#[derive(Clone)] #[allow(dead_code)] pub(crate) struct AuthModeWidget { pub request_frame: FrameRequester, @@ -288,6 +309,8 @@ pub(crate) struct AuthModeWidget { pub auth_config: AuthConfig, pub bedrock_setup_enabled: bool, pub gcp_setup_enabled: bool, + pub gateway_setup_enabled: bool, + pub(super) gateway: Option, pub animations_enabled: bool, pub animations_suppressed: Cell, } @@ -330,7 +353,7 @@ impl AuthModeWidget { self.request_frame.schedule_frame(); } - fn set_error(&self, message: Option) { + pub(super) fn set_error(&self, message: Option) { *self.error.write().unwrap() = message; } @@ -388,6 +411,9 @@ impl AuthModeWidget { if self.gcp_setup_enabled { options.push(SignInOption::Gcp); } + if self.gateway_setup_enabled { + options.push(SignInOption::Gateway); + } return options; } let mut options = vec![SignInOption::ChatGpt]; @@ -437,16 +463,20 @@ impl AuthModeWidget { fn activate_highlighted_mode(&mut self, key_event: KeyEvent) { let option = self.highlighted_mode; if option == SignInOption::OtherMethods || option.is_alternative() { - // Enter and printable shortcuts can repeat as plain presses without a release. - self.activation_key_guard = Some(ActivationKeyGuard::new_at( - KeyBinding::from_event(key_event), - Instant::now(), - /*reported_repeats_only*/ false, - )); + self.guard_activation(key_event); } self.handle_sign_in_option(option); } + pub(super) fn guard_activation(&mut self, key_event: KeyEvent) { + // Enter and printable shortcuts can repeat as plain presses without a release. + self.activation_key_guard = Some(ActivationKeyGuard::new_at( + KeyBinding::from_event(key_event), + Instant::now(), + /*reported_repeats_only*/ false, + )); + } + fn handle_sign_in_option(&mut self, option: SignInOption) { match option { SignInOption::OtherMethods => { @@ -486,6 +516,13 @@ impl AuthModeWidget { self.disallow_api_login(); } } + SignInOption::Gateway => { + if self.gateway_setup_enabled && self.is_api_login_allowed() { + self.start_gateway_picker(); + } else if !self.is_api_login_allowed() { + self.disallow_api_login(); + } + } } } @@ -499,7 +536,8 @@ impl AuthModeWidget { fn render_pick_mode(&self, area: Rect, buf: &mut Buffer) { let mut lines: Vec = if self.highlighted_mode.is_alternative() { vec![" Sign in another way".into(), "".into()] - } else if self.bedrock_setup_enabled || self.gcp_setup_enabled { + } else if self.bedrock_setup_enabled || self.gcp_setup_enabled || self.gateway_setup_enabled + { vec![" Choose how you want to use Codex.".into(), "".into()] } else { vec![ @@ -561,7 +599,10 @@ impl AuthModeWidget { idx, option, "Sign in another way", - if self.bedrock_setup_enabled || self.gcp_setup_enabled { + if self.bedrock_setup_enabled + || self.gcp_setup_enabled + || self.gateway_setup_enabled + { "Use an API key or another provider" } else { "Use an OpenAI API key" @@ -608,6 +649,14 @@ impl AuthModeWidget { "Connect using your Google Cloud credentials", )); } + SignInOption::Gateway => { + lines.extend(create_mode_item( + idx, + option, + "Use a custom gateway", + "Choose a saved gateway", + )); + } } lines.push("".into()); } @@ -942,7 +991,7 @@ impl AuthModeWidget { true } - fn start_api_key_entry(&mut self) { + pub(super) fn start_api_key_entry(&mut self) { if !self.is_api_login_allowed() { self.disallow_api_login(); return; @@ -1032,7 +1081,7 @@ impl AuthModeWidget { } /// Kicks off the ChatGPT auth flow and keeps the UI state consistent with the attempt. - fn start_chatgpt_login(&mut self) { + pub(super) fn start_chatgpt_login(&mut self) { // If we're already authenticated with ChatGPT, don't start a new login – // just proceed to the success message flow. if self.handle_existing_chatgpt_login() { @@ -1150,17 +1199,25 @@ impl StepStateProvider for AuthModeWidget { | SignInState::ChatGptDeviceCode(_) | SignInState::ChatGptSuccessMessage | SignInState::Bedrock(_) - | SignInState::Gcp(_) => StepState::InProgress, + | SignInState::Gcp(_) + | SignInState::GatewayOAuth { .. } + | SignInState::GatewayOAuthFailed { .. } => StepState::InProgress, SignInState::ChatGptSuccess | SignInState::ApiKeyConfigured | SignInState::BedrockConfigured - | SignInState::GcpConfigured => StepState::Complete, + | SignInState::GcpConfigured + | SignInState::GatewaySelected(_) + | SignInState::GatewayConfigured => StepState::Complete, } } } impl WidgetRef for AuthModeWidget { fn render_ref(&self, area: Rect, buf: &mut Buffer) { + if let Some(wizard) = &self.gateway { + wizard.render(area, buf); + return; + } let sign_in_state = self.sign_in_state.read().unwrap(); match &*sign_in_state { SignInState::PickMode => { @@ -1190,6 +1247,34 @@ impl WidgetRef for AuthModeWidget { SignInState::Gcp(state) => { state.render(area, buf, self.error_message()); } + SignInState::GatewayOAuth { url, .. } => { + let mut lines = vec![Line::from( + "Finish signing in to your gateway in the browser.".bold(), + )]; + if let Some(url) = url { + lines.push(Line::from(url.clone().cyan().underlined())); + } + lines.push(Line::from("Press Ctrl+C to quit.".dim())); + Paragraph::new(lines) + .wrap(Wrap { trim: false }) + .render(area, buf); + if let Some(url) = url { + mark_url_hyperlink(buf, area, url); + } + } + SignInState::GatewayOAuthFailed { message, .. } => { + Paragraph::new(vec![ + Line::from(message.clone().red()), + Line::from("Press Enter to retry or Ctrl+C to quit.".dim()), + ]) + .wrap(Wrap { trim: false }) + .render(area, buf); + } + SignInState::GatewaySelected(_) | SignInState::GatewayConfigured => { + Paragraph::new("✓ Custom gateway configured".green()) + .wrap(Wrap { trim: false }) + .render(area, buf); + } SignInState::GcpConfigured => { Paragraph::new("✓ Google Cloud Vertex AI configured".green()) .wrap(Wrap { trim: false }) @@ -1235,7 +1320,8 @@ mod tests { use codex_app_server_client::InProcessClientStartArgs; use codex_app_server_client::RemoteAppServerEndpoint; use codex_arg0::Arg0DispatchPaths; - use codex_cloud_config::chatgpt_config_bundle_loader_for_storage; + use codex_config::CloudConfigBundleLoader; + use codex_config::LoaderOverrides; use codex_features::Feature; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; @@ -1332,7 +1418,13 @@ mod tests { pub(super) async fn widget_forced_chatgpt() -> (AuthModeWidget, TempDir) { let codex_home = TempDir::new().unwrap(); let codex_home_path = codex_home.path().to_path_buf(); + // Embedded startup can refresh cloud config using the host's workload identity. + let loader_overrides = LoaderOverrides { + ignore_managed_requirements: true, + ..LoaderOverrides::without_managed_config_for_tests() + }; let config = ConfigBuilder::default() + .loader_overrides(loader_overrides.clone()) .codex_home(codex_home_path.clone()) .build() .await @@ -1343,14 +1435,9 @@ mod tests { arg0_paths: Arg0DispatchPaths::default(), config: Arc::new(config), cli_overrides: Vec::new(), - loader_overrides: Default::default(), + loader_overrides, strict_config: false, - cloud_config_bundle: chatgpt_config_bundle_loader_for_storage( - auth_config.clone(), - /*enable_codex_api_key_env*/ false, - ) - .await - .expect("test cloud config loader"), + cloud_config_bundle: CloudConfigBundleLoader::default(), embedded_network_policy: Default::default(), feedback: codex_feedback::CodexFeedback::new(), log_db: None, @@ -1384,6 +1471,8 @@ mod tests { auth_config, bedrock_setup_enabled: false, gcp_setup_enabled: false, + gateway_setup_enabled: false, + gateway: None, animations_enabled: true, animations_suppressed: std::cell::Cell::new(false), }; @@ -1809,3 +1898,7 @@ mod activation_key_tests; #[cfg(test)] #[path = "gcp_tests.rs"] mod gcp_tests; + +#[cfg(test)] +#[path = "gateway_tests.rs"] +mod gateway_tests; diff --git a/codex-rs/tui/src/onboarding/gateway.rs b/codex-rs/tui/src/onboarding/gateway.rs new file mode 100644 index 0000000000..6ed70532d1 --- /dev/null +++ b/codex-rs/tui/src/onboarding/gateway.rs @@ -0,0 +1,203 @@ +//! Sign in by selecting a saved custom gateway as the default provider. + +use super::auth::AuthModeWidget; +use super::auth::SignInState; +use super::keys; +use crate::app::provider::provider_item; +use crate::app::provider::rpc; +use crate::app_event_sender::AppEventSender; +use crate::bottom_pane::BottomPaneView; +use crate::bottom_pane::ListSelectionView; +use crate::bottom_pane::SelectionViewParams; +use crate::bottom_pane::popup_consts::accept_cancel_hint_line; +use crate::key_hint::KeyBindingListExt; +use crate::keymap::ListAction; +use crate::keymap::RuntimeKeymap; +use crate::render::renderable::Renderable; +use codex_app_server_client::AppServerRequestHandle; +use codex_app_server_protocol::ModelProviderListResponse; +use codex_app_server_protocol::ProviderSettings; +use codex_app_server_protocol::WriteStatus; +use crossterm::event::KeyEvent; +use crossterm::event::KeyEventKind; +use futures::future::LocalBoxFuture; +use ratatui::buffer::Buffer; +use ratatui::layout::Rect; +use ratatui::style::Stylize; +use ratatui::text::Line; +use ratatui::widgets::Widget; +use tokio::sync::mpsc::unbounded_channel; + +// Requests stay on the onboarding thread so a redraw can stop polling without losing the reply. +pub(super) enum GatewayStep { + Loading(LocalBoxFuture<'static, Self>), + LoadError(Box), + Picking(Box), + Selecting(LocalBoxFuture<'static, Result>), +} + +impl GatewayStep { + fn load(handle: AppServerRequestHandle) -> Self { + Self::Loading(Box::pin(async move { + match rpc::list_providers(&handle).await { + Ok(mut page) if page.can_change_default => { + page.data + .retain(|p| matches!(p.settings, ProviderSettings::CustomGateway { .. })); + if page.data.is_empty() { + Self::load_error("No custom gateways are configured.".into()) + } else { + Self::Picking(Box::new(Picker::new( + page, /*selected*/ None, /*notice*/ None, + ))) + } + } + Ok(_) => Self::load_error( + "The default provider is set by configuration policy or a launch override." + .into(), + ), + Err(message) => Self::load_error(message), + } + })) + } + + fn load_error(message: String) -> Self { + let (tx, _) = unbounded_channel(); + Self::LoadError(Box::new(ListSelectionView::new( + crate::app::provider::load_error("Use a custom gateway", message), + AppEventSender::new(tx), + RuntimeKeymap::defaults().list, + ))) + } + + pub(super) fn render(&self, area: Rect, buf: &mut Buffer) { + match self { + Self::Loading(_) => Widget::render(Line::from("Loading gateways…".dim()), area, buf), + Self::Selecting(_) => { + Widget::render(Line::from("Saving gateway settings…".dim()), area, buf) + } + Self::LoadError(list) => list.render(area, buf), + Self::Picking(picker) => picker.list.render(area, buf), + } + } +} + +pub(super) struct Picker { + page: ModelProviderListResponse, + list: ListSelectionView, +} + +impl Picker { + fn new( + page: ModelProviderListResponse, + selected: Option, + notice: Option, + ) -> Self { + let items = page + .data + .iter() + .map(|p| provider_item(p, &page.default_provider_id)) + .collect(); + let keymap = RuntimeKeymap::defaults().list; + // Like the other startup pickers, consume the selection directly before the app event loop exists. + let (tx, _) = unbounded_channel(); + let list = ListSelectionView::new( + SelectionViewParams { + title: Some("Use a custom gateway".into()), + footer_hint: Some(accept_cancel_hint_line( + keymap.primary_hint(ListAction::Accept), + "to use as default", + keymap.primary_hint(ListAction::Cancel), + "to go back", + )), + footer_note: notice.map(Into::into), + initial_selected_idx: selected, + auto_number_shortcuts: false, + items, + ..Default::default() + }, + AppEventSender::new(tx), + keymap, + ); + Self { page, list } + } +} + +impl AuthModeWidget { + pub(super) fn start_gateway_picker(&mut self) { + self.set_error(/*message*/ None); + self.gateway = Some(GatewayStep::load(self.app_server_request_handle.clone())); + self.request_frame.schedule_frame(); + } + + pub(super) async fn advance_gateway_request(&mut self) { + self.gateway = match &mut self.gateway { + Some(GatewayStep::Loading(request)) => Some(request.await), + Some(GatewayStep::Selecting(request)) => match request.await { + Ok(id) => { + *self + .sign_in_state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + SignInState::GatewaySelected(id); + None + } + Err(picker) => Some(picker), + }, + None | Some(GatewayStep::Picking(_) | GatewayStep::LoadError(_)) => { + std::future::pending().await + } + }; + self.request_frame.schedule_frame(); + } + + pub(super) fn handle_gateway_key_event(&mut self, key: KeyEvent) -> bool { + let Some(step) = self.gateway.take() else { + return false; + }; + self.gateway = match step { + GatewayStep::Loading(_) | GatewayStep::Picking(_) | GatewayStep::LoadError(_) + if keys::CANCEL.is_pressed(key) => + { + None + } + GatewayStep::LoadError(mut view) => { + view.handle_key_event(key); + if view.take_last_selected_index().is_some() { + self.guard_activation(key); + None + } else { + Some(GatewayStep::LoadError(view)) + } + } + GatewayStep::Picking(mut picker) + if key.kind == KeyEventKind::Press || !keys::CONFIRM.is_pressed(key) => + { + picker.list.handle_key_event(key); + if let Some(index) = picker.list.take_last_selected_index() { + self.guard_activation(key); + let id = picker.page.data[index].id.clone(); + let handle = self.app_server_request_handle.clone(); + Some(GatewayStep::Selecting(Box::pin(async move { + let message = match rpc::set_default_provider(&handle, &id).await { + Ok(response) => match response.status { + WriteStatus::Ok => return Ok(id), + WriteStatus::OkOverridden => rpc::override_reason(response), + }, + Err(err) => err.to_string(), + }; + Err(GatewayStep::Picking(Box::new(Picker::new( + picker.page, + Some(index), + Some(message), + )))) + }))) + } else { + Some(GatewayStep::Picking(picker)) + } + } + step => Some(step), + }; + self.request_frame.schedule_frame(); + true + } +} diff --git a/codex-rs/tui/src/onboarding/gateway_login.rs b/codex-rs/tui/src/onboarding/gateway_login.rs new file mode 100644 index 0000000000..842c1da479 --- /dev/null +++ b/codex-rs/tui/src/onboarding/gateway_login.rs @@ -0,0 +1,172 @@ +//! Finish authentication against the embedded server running the selected gateway. + +use super::auth::AuthModeWidget; +use super::auth::SignInState; +use super::auth::onboarding_request_id; +use super::onboarding_screen::OnboardingScreen; +use super::onboarding_screen::OnboardingScreenArgs; +use super::onboarding_screen::run_onboarding_screen; +use crate::AppServerTarget; +use crate::app_server_session::AppServerSession; +use crate::get_login_status; +use crate::legacy_core::config::Config; +use crate::should_show_login_screen; +use crate::startup_draft::StartupCancelled; +use crate::tui::Tui; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::GatewayOAuthCancelResponse; +use codex_app_server_protocol::GatewayOAuthChangedNotification; +use codex_app_server_protocol::GatewayOAuthLoginResponse; +use codex_app_server_protocol::GatewayOAuthReadResponse; +use codex_app_server_protocol::GatewayOAuthStatus; +use codex_protocol::config_types::ForcedLoginMethod; +use color_eyre::Result; +use color_eyre::eyre::OptionExt; + +pub(crate) async fn authenticate_gateway( + tui: &mut Tui, + server: &mut AppServerSession, + config: &Config, +) -> Result<()> { + let (login_status, account) = get_login_status(server).await?; + let primary_login = should_show_login_screen(login_status, account.requires_openai_auth); + // Some gateways require both credentials. The primary login must finish first. + for login in [Login::Primary, Login::GatewayOAuth] { + match login { + Login::Primary if !primary_login => continue, + Login::GatewayOAuth if config.model_provider.gateway_oauth.is_none() => continue, + Login::GatewayOAuth | Login::Primary => {} + } + let handle = server.request_handle(); + let mut screen = OnboardingScreen::new( + tui, + OnboardingScreenArgs { + show_login_screen: true, + gateway_setup_enabled: false, + gcp_setup_enabled: false, + bedrock_setup_enabled: false, + show_trust_screen: false, + remote_project_trust: None, + login_status, + app_server_target: AppServerTarget::Embedded, + app_server_request_handle: Some(handle.clone()), + config: config.clone(), + }, + ) + .await; + let widget = screen + .auth_widget_mut() + .ok_or_eyre("The gateway login screen is unavailable.")?; + match login { + Login::Primary + if widget + .auth_config + .is_login_method_allowed(ForcedLoginMethod::Chatgpt) => + { + widget.start_chatgpt_login(); + } + Login::Primary => widget.start_api_key_entry(), + Login::GatewayOAuth => widget.start_gateway_oauth(config.model_provider_id.clone()), + } + if run_onboarding_screen(screen, Some(handle), Some(server), tui) + .await? + .should_exit + { + if matches!(login, Login::GatewayOAuth) { + let _ = server + .request_handle() + .request_typed::( + ClientRequest::GatewayOAuthCancel { + request_id: onboarding_request_id(), + params: None, + }, + ) + .await; + } + return Err( + std::io::Error::new(std::io::ErrorKind::Interrupted, StartupCancelled).into(), + ); + } + } + Ok(()) +} + +enum Login { + Primary, + GatewayOAuth, +} + +impl AuthModeWidget { + pub(super) fn start_gateway_oauth(&mut self, provider_id: String) { + *self + .sign_in_state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = SignInState::GatewayOAuth { + provider_id: provider_id.clone(), + url: None, + }; + let state = self.sign_in_state.clone(); + let frame = self.request_frame.clone(); + let handle = self.app_server_request_handle.clone(); + tokio::spawn(async move { + let result: Result<()> = async { + let status = handle + .request_typed::(ClientRequest::GatewayOAuthRead { + request_id: onboarding_request_id(), + params: None, + }) + .await?; + color_eyre::eyre::ensure!( + status.provider_id == provider_id, + "The selected gateway changed during sign-in." + ); + if status.required && status.status != Some(GatewayOAuthStatus::Succeeded) { + handle + .request_typed::( + ClientRequest::GatewayOAuthLogin { + request_id: onboarding_request_id(), + params: None, + }, + ) + .await?; + } + Ok(()) + } + .await; + *state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = match result { + Ok(()) => SignInState::GatewayConfigured, + Err(err) => SignInState::GatewayOAuthFailed { + provider_id, + message: err.to_string(), + }, + }; + frame.schedule_frame(); + }); + } + + pub(super) fn on_gateway_oauth_changed( + &mut self, + notification: GatewayOAuthChangedNotification, + ) { + let Some(auth_url) = notification.auth_url else { + return; + }; + let mut state = self + .sign_in_state + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let SignInState::GatewayOAuth { provider_id, url } = &mut *state + && notification.provider_id == *provider_id + && url.as_ref() != Some(&auth_url) + { + *url = Some(auth_url.clone()); + drop(state); + self.request_frame.schedule_frame(); + tokio::task::spawn_blocking(move || { + let _ = webbrowser::open(&auth_url); + }); + } + } +} diff --git a/codex-rs/tui/src/onboarding/gateway_tests.rs b/codex-rs/tui/src/onboarding/gateway_tests.rs new file mode 100644 index 0000000000..52a7eae606 --- /dev/null +++ b/codex-rs/tui/src/onboarding/gateway_tests.rs @@ -0,0 +1,79 @@ +#![cfg(test)] + +use super::*; +use crossterm::event::KeyCode; +use pretty_assertions::assert_eq; + +fn screen(widget: &AuthModeWidget, height: u16) -> String { + let area = Rect::new(0, 0, 72, height); + let mut terminal = crate::custom_terminal::Terminal::with_options( + crate::test_backend::VT100Backend::new(area.width, area.height), + ) + .unwrap(); + terminal.set_viewport_area(area); + terminal + .draw(|frame| widget.render_ref(area, frame.buffer_mut())) + .unwrap(); + terminal + .backend() + .to_string() + .lines() + .map(str::trim_end) + .collect::>() + .join("\n") +} + +#[tokio::test] +async fn select_a_saved_gateway_during_sign_in() -> color_eyre::Result<()> { + let (mut widget, home) = super::tests::widget_forced_chatgpt().await; + widget.auth_config.forced_login_method = None; + widget.bedrock_setup_enabled = true; + widget.gcp_setup_enabled = true; + widget.gateway_setup_enabled = true; + + widget.handle_key_event(KeyCode::Char('3').into()); + insta::assert_snapshot!("gateway_sign_in_menu", screen(&widget, /*height*/ 8)); + widget.handle_key_event(KeyCode::Char('4').into()); + widget.advance_gateway_request().await; + insta::assert_snapshot!("gateway_load_error", screen(&widget, /*height*/ 8)); + widget.handle_key_event(KeyCode::Esc.into()); + + widget.handle_key_event(KeyEvent { + kind: KeyEventKind::Release, + ..KeyCode::Char('4').into() + }); + std::fs::write( + home.path().join("config.toml"), + r#" +[model_providers.office] +name = "Office" +base_url = "https://office.example.test/v1" +wire_api = "responses" +"#, + )?; + widget.handle_key_event(KeyCode::Char('4').into()); + widget.advance_gateway_request().await; + insta::assert_snapshot!("gateway_saved_picker", screen(&widget, /*height*/ 8)); + widget.handle_key_event(KeyCode::Enter.into()); + widget.advance_gateway_request().await; + assert!(matches!(widget.get_step_state(), StepState::Complete)); + let saved: toml::Value = + toml::from_str(&std::fs::read_to_string(home.path().join("config.toml"))?)?; + assert_eq!(saved["model_provider"].as_str(), Some("office")); + Ok(()) +} + +#[tokio::test] +async fn gateway_oauth_displays_the_browser_link_and_retry_action() { + let (widget, _home) = super::tests::widget_forced_chatgpt().await; + *widget.sign_in_state.write().unwrap() = SignInState::GatewayOAuth { + provider_id: "office".into(), + url: Some("https://office.example.test/authorize".into()), + }; + insta::assert_snapshot!("gateway_oauth_browser", screen(&widget, /*height*/ 4)); + *widget.sign_in_state.write().unwrap() = SignInState::GatewayOAuthFailed { + provider_id: "office".into(), + message: "Gateway sign-in expired.".into(), + }; + insta::assert_snapshot!("gateway_oauth_failed", screen(&widget, /*height*/ 4)); +} diff --git a/codex-rs/tui/src/onboarding/keys.rs b/codex-rs/tui/src/onboarding/keys.rs index 6037a370b0..ff7eacac44 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 COPY_LINK: [KeyBinding; 1] = [key_hint::plain(KeyCode::Char('c'))]; diff --git a/codex-rs/tui/src/onboarding/mod.rs b/codex-rs/tui/src/onboarding/mod.rs index ef93716976..9bc1552279 100644 --- a/codex-rs/tui/src/onboarding/mod.rs +++ b/codex-rs/tui/src/onboarding/mod.rs @@ -1,5 +1,8 @@ mod auth; mod bedrock; +mod gateway; +mod gateway_login; +pub(crate) use gateway_login::authenticate_gateway; mod gcp; pub(crate) use bedrock::check_gov_cloud; mod keys; diff --git a/codex-rs/tui/src/onboarding/onboarding_screen.rs b/codex-rs/tui/src/onboarding/onboarding_screen.rs index db40dd0d71..83c426a11f 100644 --- a/codex-rs/tui/src/onboarding/onboarding_screen.rs +++ b/codex-rs/tui/src/onboarding/onboarding_screen.rs @@ -107,6 +107,7 @@ pub(crate) struct OnboardingScreenArgs { pub show_login_screen: bool, pub bedrock_setup_enabled: bool, pub gcp_setup_enabled: bool, + pub gateway_setup_enabled: bool, pub login_status: LoginStatus, pub app_server_target: crate::AppServerTarget, pub app_server_request_handle: Option, @@ -116,6 +117,7 @@ pub(crate) struct OnboardingScreenArgs { #[derive(Default)] pub(crate) struct OnboardingResult { pub directory_trust_persisted: bool, + pub selected_gateway: Option, pub should_exit: bool, } @@ -135,6 +137,7 @@ impl OnboardingScreen { show_login_screen, bedrock_setup_enabled, gcp_setup_enabled, + gateway_setup_enabled, login_status, app_server_target, app_server_request_handle, @@ -188,6 +191,8 @@ impl OnboardingScreen { auth_config, bedrock_setup_enabled, gcp_setup_enabled, + gateway_setup_enabled, + gateway: None, animations_enabled: local_settings.tui.animations && local_settings.tui.effects.shimmer, animations_suppressed: std::cell::Cell::new(false), @@ -306,7 +311,7 @@ impl OnboardingScreen { } } - fn auth_widget_mut(&mut self) -> Option<&mut AuthModeWidget> { + pub(super) fn auth_widget_mut(&mut self) -> Option<&mut AuthModeWidget> { self.steps.iter_mut().find_map(|step| match step { Step::Auth(widget) => Some(widget), Step::Welcome(_) | Step::TrustDirectory(_) | Step::GovCloudGuidance { .. } => None, @@ -320,6 +325,11 @@ impl OnboardingScreen { widget.on_account_login_completed(notification); } } + ServerNotification::GatewayOAuthChanged(notification) => { + if let Some(widget) = self.auth_widget_mut() { + widget.on_gateway_oauth_changed(notification); + } + } ServerNotification::AccountUpdated(notification) => { if let Some(widget) = self.auth_widget_mut() { widget.on_account_updated(notification); @@ -599,7 +609,7 @@ pub(crate) async fn run_onboarding_app( run_onboarding_screen(screen, request_handle, app_server, tui).await } -async fn run_onboarding_screen( +pub(super) async fn run_onboarding_screen( onboarding_screen: OnboardingScreen, app_server_request_handle: Option, app_server: Option<&mut AppServerSession>, @@ -646,6 +656,12 @@ async fn run_onboarding_screen_inner( while !onboarding_screen.is_done() { tokio::select! { + () = async { + match onboarding_screen.auth_widget_mut() { + Some(auth) => auth.advance_gateway_request().await, + None => std::future::pending().await, + } + } => {} event = tui_events.next() => { if let Some(event) = event { tui.screen_size_for_event(&event)?; @@ -799,6 +815,16 @@ async fn run_onboarding_screen_inner( } Ok(OnboardingResult { directory_trust_persisted, + selected_gateway: onboarding_screen.auth_widget_mut().and_then(|widget| { + match &*widget + .sign_in_state + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + { + SignInState::GatewaySelected(id) => Some(id.clone()), + _ => None, + } + }), should_exit: onboarding_screen.should_exit(), }) } diff --git a/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_load_error.snap b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_load_error.snap new file mode 100644 index 0000000000..a47f391cb6 --- /dev/null +++ b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_load_error.snap @@ -0,0 +1,9 @@ +--- +source: tui/src/onboarding/gateway_tests.rs +expression: "screen(&widget, 8)" +--- + Use a custom gateway + No custom gateways are configured. + + +› 1. Back diff --git a/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_oauth_browser.snap b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_oauth_browser.snap new file mode 100644 index 0000000000..913b4b405e --- /dev/null +++ b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_oauth_browser.snap @@ -0,0 +1,7 @@ +--- +source: tui/src/onboarding/gateway_tests.rs +expression: "screen(&widget, 4)" +--- +Finish signing in to your gateway in the browser. +https://office.example.test/authorize +Press Ctrl+C to quit. diff --git a/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_oauth_failed.snap b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_oauth_failed.snap new file mode 100644 index 0000000000..5af304135e --- /dev/null +++ b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_oauth_failed.snap @@ -0,0 +1,6 @@ +--- +source: tui/src/onboarding/gateway_tests.rs +expression: "screen(&widget, 4)" +--- +Gateway sign-in expired. +Press Enter to retry or Ctrl+C to quit. diff --git a/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_saved_picker.snap b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_saved_picker.snap new file mode 100644 index 0000000000..26cdd9aa58 --- /dev/null +++ b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_saved_picker.snap @@ -0,0 +1,12 @@ +--- +source: tui/src/onboarding/gateway_tests.rs +expression: screen(&widget) +--- + Use a custom gateway + + +› ○ Office office + + + + Press enter to use as default or esc to go back diff --git a/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_sign_in_menu.snap b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_sign_in_menu.snap new file mode 100644 index 0000000000..8fc332b81b --- /dev/null +++ b/codex-rs/tui/src/onboarding/snapshots/codex_tui__onboarding__auth__gateway_tests__gateway_sign_in_menu.snap @@ -0,0 +1,12 @@ +--- +source: tui/src/onboarding/gateway_tests.rs +expression: terminal.backend().to_string() +--- + Sign in another way + +> 1. Use an OpenAI API key + 2. Use Amazon Bedrock + 3. Use Google Cloud Vertex AI + 4. Use a custom gateway + Press enter to continue + Press esc to go back diff --git a/codex-rs/tui/src/startup_orchestration.rs b/codex-rs/tui/src/startup_orchestration.rs index 7397feb331..c247a8ea61 100644 --- a/codex-rs/tui/src/startup_orchestration.rs +++ b/codex-rs/tui/src/startup_orchestration.rs @@ -575,13 +575,17 @@ pub(super) async fn run_main_inner( && !app_server_target.uses_remote_workspace(); if auto_start_daemon && daemon_exclusion.is_none() - && should_show_provider_setup_wizard( - LoginStatus::NotAuthenticated, - config.model_provider.requires_openai_auth, - &config, - &AppServerTarget::Embedded, - Feature::GcpSetupWizard, - ) + && [Feature::GcpSetupWizard, Feature::GatewaySetupWizard] + .into_iter() + .any(|feature| { + should_show_provider_setup_wizard( + LoginStatus::NotAuthenticated, + config.model_provider.requires_openai_auth, + &config, + &AppServerTarget::Embedded, + feature, + ) + }) && startup_draft .run_until( config @@ -593,8 +597,13 @@ pub(super) async fn run_main_inner( .flatten() .is_none() { - // GCP setup resolves credential paths and runs gcloud in the launch environment. - daemon_exclusion = Some("GCP sign-in"); + // These wizards use credentials from the launch environment and rebuild the + // embedded server with the selected provider after onboarding. + daemon_exclusion = Some(if config.features.enabled(Feature::GcpSetupWizard) { + "GCP sign-in" + } else { + "custom gateway sign-in" + }); app_server_target = AppServerTarget::Embedded; } if auto_start_daemon @@ -813,6 +822,7 @@ pub(super) async fn run_main_inner( } } Some("GCP sign-in") => "gcp_sign_in", + Some("custom gateway sign-in") => "gateway_sign_in", Some(daemon_startup::WSL_DRVFS_EXCLUSION) => "wsl_drvfs", Some("this Windows launcher") => "windows_detach_restricted", Some("daemon feature settings") => "daemon_compatibility", diff --git a/codex-rs/tui/tests/suite/focus_palette.rs b/codex-rs/tui/tests/suite/focus_palette.rs index 084599ce1c..62bd74fd72 100644 --- a/codex-rs/tui/tests/suite/focus_palette.rs +++ b/codex-rs/tui/tests/suite/focus_palette.rs @@ -385,6 +385,9 @@ impl PtyCodex { .env_remove("STY") .env("TERM_PROGRAM", "kitty") .env_remove("TERM_PROGRAM_VERSION") + // Exercise the fixture's login state even on hosts configured with workload identity. + .env_remove(codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR) + .env_remove(codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR) .env("OPENAI_API_KEY", "focus-palette-test") .env("CODEX_HOME", codex_home.path()) .stdin(stdin) diff --git a/codex-rs/tui/tests/suite/provider_defaults.rs b/codex-rs/tui/tests/suite/provider_defaults.rs index 16e8f344f3..5c6cfb555b 100644 --- a/codex-rs/tui/tests/suite/provider_defaults.rs +++ b/codex-rs/tui/tests/suite/provider_defaults.rs @@ -132,3 +132,101 @@ async fn resume_lists_all_providers_and_fork_keeps_provider_filter() -> Result<( } Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_gateway_authenticates_before_its_first_chat() -> Result<()> { + use core_test_support::responses; + + let server = responses::start_mock_server().await; + let reply = responses::mount_sse_once_match( + &server, + wiremock::matchers::body_string_contains(r#"\"thread_source\":\"user\""#), + responses::sse(vec![ + responses::ev_assistant_message("msg", "gateway handoff verified"), + responses::ev_completed("resp"), + ]), + ) + .await; + let workspace = tempfile::tempdir()?; + let home = tempfile::tempdir()?; + let cwd = workspace.path().canonicalize()?; + let profile = home.path().join("work.config.toml"); + let auth = home.path().join("auth.json"); + let url = server.uri(); + let trusted = toml::Value::String(cwd.display().to_string()); + std::fs::write( + &profile, + format!( + r#" +model_provider = "openai" +model = "gpt-5.1" +forced_login_method = "api" +cli_auth_credentials_store = "file" +suppress_unstable_features_warning = true +[features] +gateway_setup_wizard = true +bedrock_setup_wizard = false +gcp_setup_wizard = false +api_key_model_discovery = false +[projects.{trusted}] +trust_level = "trusted" +[model_providers.office] +name = "Office gateway" +base_url = "{url}/v1" +wire_api = "responses" +supports_websockets = false +requires_openai_auth = true +"# + ), + )?; + let mut terminal = PtyCodex::start( + &cwd, + home, + &["--profile", "work", "--no-daemon", "verify the gateway"], + )?; + terminal + .wait_for_screen("Sign in another way") + .with_context(|| terminal.screen_contents())?; + // Startup drains input after its first draw. Retry the menu shortcut until it is accepted. + let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 30); + while !terminal.screen_contains("Use a custom gateway") { + ensure!( + std::time::Instant::now() < deadline, + "gateway option missing: {}", + terminal.screen_contents() + ); + terminal.write_input(b"2")?; + terminal.read_output(Duration::from_millis(/*millis*/ 50))?; + } + // Arrow navigation avoids carrying the menu shortcut across two screens. + terminal.write_input(b"\x1b[B\r")?; + terminal.wait_for_screen("Office gateway")?; + terminal.write_input(b"\x1b[H\r")?; + terminal.wait_for_screen("Paste or type your API key")?; + let deadline = std::time::Instant::now() + Duration::from_secs(/*secs*/ 30); + while !auth.is_file() && terminal.screen_contains("Paste or type your API key") { + ensure!( + std::time::Instant::now() < deadline, + "API key was not saved: {}", + terminal.screen_contents() + ); + terminal.write_input(b"\r")?; + terminal.read_output(Duration::from_millis(/*millis*/ 50))?; + } + terminal.wait_for_screen("gateway handoff verified")?; + let first_chat = reply + .requests() + .into_iter() + .find(|request| { + request + .message_input_texts("user") + .iter() + .any(|text| text == "verify the gateway") + }) + .context("the selected gateway received the first chat")?; + assert_eq!( + first_chat.header("authorization").as_deref(), + Some("Bearer focus-palette-test") + ); + Ok(()) +}