valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplify auth flow and reconcile differences between ChatGPT and API Key auth (#3189)

Browse Source 
This PR does the following:
* Adds the ability to paste or type an API key.
* Removes the `preferred_auth_method` config option. The last login
method is always persisted in auth.json, so this isn't needed.
* If OPENAI_API_KEY env variable is defined, the value is used to
prepopulate the new UI. The env variable is otherwise ignored by the
CLI.
* Adds a new MCP server entry point "login_api_key" so we can implement
this same API key behavior for the VS Code extension.
<img width="473" height="140" alt="Screenshot 2025-09-04 at 3 51 04 PM"
src="https://github.com/user-attachments/assets/c11bbd5b-8a4d-4d71-90fd-34130460f9d9"
/>
<img width="726" height="254" alt="Screenshot 2025-09-04 at 3 51 32 PM"
src="https://github.com/user-attachments/assets/6cc76b34-309a-4387-acbc-15ee5c756db9"
/>
This commit is contained in:
Eric Traut
2025-09-11 09:16:34 -07:00
committed by GitHub
parent 377af75730
commit e13b35ecb0
24 changed files with 412 additions and 494 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
codex-rs/tui/src/lib.rs
View File

@@ -308,7 +308,7 @@ async fn run_ratatui_app(
..
} = cli;
let auth_manager = AuthManager::shared(config.codex_home.clone(), config.preferred_auth_method);
let auth_manager = AuthManager::shared(config.codex_home.clone());
let login_status = get_login_status(&config);
let should_show_onboarding =
should_show_onboarding(login_status, &config, should_show_trust_screen);
@@ -392,7 +392,7 @@ fn get_login_status(config: &Config) -> LoginStatus {
// Reading the OpenAI API key is an async operation because it may need
// to refresh the token. Block on it.
let codex_home = config.codex_home.clone();
match CodexAuth::from_codex_home(&codex_home, config.preferred_auth_method) {
match CodexAuth::from_codex_home(&codex_home) {
Ok(Some(auth)) => LoginStatus::AuthMode(auth.mode),
Ok(None) => LoginStatus::NotAuthenticated,
Err(err) => {
@@ -460,60 +460,28 @@ fn should_show_login_screen(login_status: LoginStatus, config: &Config) -> bool
return false;
}
match login_status {
LoginStatus::NotAuthenticated => true,
LoginStatus::AuthMode(method) => method != config.preferred_auth_method,
}
login_status == LoginStatus::NotAuthenticated
}
#[cfg(test)]
mod tests {
use super::*;
fn make_config(preferred: AuthMode) -> Config {
let mut cfg = Config::load_from_base_config_with_overrides(
fn make_config() -> Config {
Config::load_from_base_config_with_overrides(
ConfigToml::default(),
ConfigOverrides::default(),
std::env::temp_dir(),
)
.expect("load default config");
cfg.preferred_auth_method = preferred;
cfg
.expect("load default config")
}
#[test]
fn shows_login_when_not_authenticated() {
let cfg = make_config(AuthMode::ChatGPT);
let cfg = make_config();
assert!(should_show_login_screen(
LoginStatus::NotAuthenticated,
&cfg
));
}
#[test]
fn shows_login_when_api_key_but_prefers_chatgpt() {
let cfg = make_config(AuthMode::ChatGPT);
assert!(should_show_login_screen(
LoginStatus::AuthMode(AuthMode::ApiKey),
&cfg
))
}
#[test]
fn hides_login_when_api_key_and_prefers_api_key() {
let cfg = make_config(AuthMode::ApiKey);
assert!(!should_show_login_screen(
LoginStatus::AuthMode(AuthMode::ApiKey),
&cfg
))
}
#[test]
fn hides_login_when_chatgpt_and_prefers_chatgpt() {
let cfg = make_config(AuthMode::ChatGPT);
assert!(!should_show_login_screen(
LoginStatus::AuthMode(AuthMode::ChatGPT),
&cfg
))
}
}

303
codex-rs/tui/src/onboarding/auth.rs
View File

@@ -2,12 +2,17 @@
use codex_core::AuthManager;
use codex_core::auth::CLIENT_ID;
use codex_core::auth::login_with_api_key;
use codex_core::auth::read_openai_api_key_from_env;
use codex_login::ServerOptions;
use codex_login::ShutdownHandle;
use codex_login::run_login_server;
use crossterm::event::KeyCode;
use crossterm::event::KeyEvent;
use crossterm::event::KeyModifiers;
use ratatui::buffer::Buffer;
use ratatui::layout::Constraint;
use ratatui::layout::Layout;
use ratatui::layout::Rect;
use ratatui::prelude::Widget;
use ratatui::style::Color;
@@ -15,6 +20,9 @@ use ratatui::style::Modifier;
use ratatui::style::Style;
use ratatui::style::Stylize;
use ratatui::text::Line;
use ratatui::widgets::Block;
use ratatui::widgets::BorderType;
use ratatui::widgets::Borders;
use ratatui::widgets::Paragraph;
use ratatui::widgets::WidgetRef;
use ratatui::widgets::Wrap;
@@ -38,8 +46,14 @@ pub(crate) enum SignInState {
ChatGptContinueInBrowser(ContinueInBrowserState),
ChatGptSuccessMessage,
ChatGptSuccess,
EnvVarMissing,
EnvVarFound,
ApiKeyEntry(ApiKeyInputState),
ApiKeyConfigured,
}
#[derive(Clone, Default)]
pub(crate) struct ApiKeyInputState {
value: String,
prepopulated_from_env: bool,
}
#[derive(Clone)]
@@ -59,6 +73,10 @@ impl Drop for ContinueInBrowserState {
impl KeyboardHandler for AuthModeWidget {
fn handle_key_event(&mut self, key_event: KeyEvent) {
if self.handle_api_key_entry_key_event(&key_event) {
return;
}
match key_event.code {
KeyCode::Up | KeyCode::Char('k') => {
self.highlighted_mode = AuthMode::ChatGPT;
@@ -69,7 +87,7 @@ impl KeyboardHandler for AuthModeWidget {
KeyCode::Char('1') => {
self.start_chatgpt_login();
}
KeyCode::Char('2') => self.verify_api_key(),
KeyCode::Char('2') => self.start_api_key_entry(),
KeyCode::Enter => {
let sign_in_state = { (*self.sign_in_state.read().unwrap()).clone() };
match sign_in_state {
@@ -78,12 +96,9 @@ impl KeyboardHandler for AuthModeWidget {
self.start_chatgpt_login();
}
AuthMode::ApiKey => {
self.verify_api_key();
self.start_api_key_entry();
}
},
SignInState::EnvVarMissing => {
*self.sign_in_state.write().unwrap() = SignInState::PickMode;
}
SignInState::ChatGptSuccessMessage => {
*self.sign_in_state.write().unwrap() = SignInState::ChatGptSuccess;
}
@@ -101,6 +116,10 @@ impl KeyboardHandler for AuthModeWidget {
_ => {}
}
}
fn handle_paste(&mut self, pasted: String) {
let _ = self.handle_api_key_entry_paste(pasted);
}
}
#[derive(Clone)]
@@ -111,7 +130,6 @@ pub(crate) struct AuthModeWidget {
pub sign_in_state: Arc<RwLock<SignInState>>,
pub codex_home: PathBuf,
pub login_status: LoginStatus,
pub preferred_auth_method: AuthMode,
pub auth_manager: Arc<AuthManager>,
}
@@ -129,24 +147,6 @@ impl AuthModeWidget {
"".into(),
];
// If the user is already authenticated but the method differs from their
// preferred auth method, show a brief explanation.
if let LoginStatus::AuthMode(current) = self.login_status
&& current != self.preferred_auth_method
{
let to_label = |mode: AuthMode| match mode {
AuthMode::ApiKey => "API key",
AuthMode::ChatGPT => "ChatGPT",
};
let msg = format!(
" You’re currently using {} while your preferred method is {}.",
to_label(current),
to_label(self.preferred_auth_method)
);
lines.push(msg.into());
lines.push("".into());
}
let create_mode_item = |idx: usize,
selected_mode: AuthMode,
text: &str,
@@ -175,29 +175,17 @@ impl AuthModeWidget {
vec![line1, line2]
};
let chatgpt_label = if matches!(self.login_status, LoginStatus::AuthMode(AuthMode::ChatGPT))
{
"Continue using ChatGPT"
} else {
"Sign in with ChatGPT"
};
lines.extend(create_mode_item(
0,
AuthMode::ChatGPT,
chatgpt_label,
"Sign in with ChatGPT",
"Usage included with Plus, Pro, and Team plans",
));
let api_key_label = if matches!(self.login_status, LoginStatus::AuthMode(AuthMode::ApiKey))
{
"Continue using API key"
} else {
"Provide your own API key"
};
lines.extend(create_mode_item(
1,
AuthMode::ApiKey,
api_key_label,
"Provide your own API key",
"Pay for what you use",
));
lines.push("".into());
@@ -282,26 +270,213 @@ impl AuthModeWidget {
.render(area, buf);
}
fn render_env_var_found(&self, area: Rect, buf: &mut Buffer) {
let lines = vec!["✓ Using OPENAI_API_KEY".fg(Color::Green).into()];
fn render_api_key_configured(&self, area: Rect, buf: &mut Buffer) {
let lines = vec![
"✓ API key configured".fg(Color::Green).into(),
"".into(),
" Codex will use usage-based billing with your API key.".into(),
];
Paragraph::new(lines)
.wrap(Wrap { trim: false })
.render(area, buf);
}
fn render_env_var_missing(&self, area: Rect, buf: &mut Buffer) {
let lines = vec![
" To use Codex with the OpenAI API, set OPENAI_API_KEY in your environment"
.fg(Color::Cyan)
.into(),
"".into(),
" Press Enter to return".dim().into(),
];
fn render_api_key_entry(&self, area: Rect, buf: &mut Buffer, state: &ApiKeyInputState) {
let [intro_area, input_area, footer_area] = Layout::vertical([
Constraint::Min(4),
Constraint::Length(3),
Constraint::Min(2),
])
.areas(area);
Paragraph::new(lines)
let mut intro_lines: Vec<Line> = vec![
Line::from(vec![
"> ".into(),
"Use your own OpenAI API key for usage-based billing".bold(),
]),
"".into(),
" Paste or type your API key below. It will be stored locally in auth.json.".into(),
"".into(),
];
if state.prepopulated_from_env {
intro_lines.push(" Detected OPENAI_API_KEY environment variable.".into());
intro_lines.push(
" Paste a different key if you prefer to use another account."
.dim()
.into(),
);
intro_lines.push("".into());
}
Paragraph::new(intro_lines)
.wrap(Wrap { trim: false })
.render(area, buf);
.render(intro_area, buf);
let content_line: Line = if state.value.is_empty() {
vec!["Paste or type your API key".dim()].into()
} else {
Line::from(state.value.clone())
};
Paragraph::new(content_line)
.wrap(Wrap { trim: false })
.block(
Block::default()
.title("API key")
.borders(Borders::ALL)
.border_type(BorderType::Rounded)
.border_style(Style::default().fg(Color::Cyan)),
)
.render(input_area, buf);
let mut footer_lines: Vec<Line> = vec![
" Press Enter to save".dim().into(),
" Press Esc to go back".dim().into(),
];
if let Some(error) = &self.error {
footer_lines.push("".into());
footer_lines.push(error.as_str().red().into());
}
Paragraph::new(footer_lines)
.wrap(Wrap { trim: false })
.render(footer_area, buf);
}
fn handle_api_key_entry_key_event(&mut self, key_event: &KeyEvent) -> bool {
let mut should_save: Option<String> = None;
let mut should_request_frame = false;
{
let mut guard = self.sign_in_state.write().unwrap();
if let SignInState::ApiKeyEntry(state) = &mut *guard {
match key_event.code {
KeyCode::Esc => {
*guard = SignInState::PickMode;
self.error = None;
should_request_frame = true;
}
KeyCode::Enter => {
let trimmed = state.value.trim().to_string();
if trimmed.is_empty() {
self.error = Some("API key cannot be empty".to_string());
should_request_frame = true;
} else {
should_save = Some(trimmed);
}
}
KeyCode::Backspace => {
if state.prepopulated_from_env {
state.value.clear();
state.prepopulated_from_env = false;
} else {
state.value.pop();
}
self.error = None;
should_request_frame = true;
}
KeyCode::Char(c)
if !key_event.modifiers.contains(KeyModifiers::CONTROL)
&& !key_event.modifiers.contains(KeyModifiers::ALT) =>
{
if state.prepopulated_from_env {
state.value.clear();
state.prepopulated_from_env = false;
}
state.value.push(c);
self.error = None;
should_request_frame = true;
}
_ => {}
}
// handled; let guard drop before potential save
} else {
return false;
}
}
if let Some(api_key) = should_save {
self.save_api_key(api_key);
} else if should_request_frame {
self.request_frame.schedule_frame();
}
true
}
fn handle_api_key_entry_paste(&mut self, pasted: String) -> bool {
let trimmed = pasted.trim();
if trimmed.is_empty() {
return false;
}
let mut guard = self.sign_in_state.write().unwrap();
if let SignInState::ApiKeyEntry(state) = &mut *guard {
if state.prepopulated_from_env {
state.value = trimmed.to_string();
state.prepopulated_from_env = false;
} else {
state.value.push_str(trimmed);
}
self.error = None;
} else {
return false;
}
drop(guard);
self.request_frame.schedule_frame();
true
}
fn start_api_key_entry(&mut self) {
self.error = None;
let prefill_from_env = read_openai_api_key_from_env();
let mut guard = self.sign_in_state.write().unwrap();
match &mut *guard {
SignInState::ApiKeyEntry(state) => {
if state.value.is_empty() {
if let Some(prefill) = prefill_from_env.clone() {
state.value = prefill;
state.prepopulated_from_env = true;
} else {
state.prepopulated_from_env = false;
}
}
}
_ => {
*guard = SignInState::ApiKeyEntry(ApiKeyInputState {
value: prefill_from_env.clone().unwrap_or_default(),
prepopulated_from_env: prefill_from_env.is_some(),
});
}
}
drop(guard);
self.request_frame.schedule_frame();
}
fn save_api_key(&mut self, api_key: String) {
match login_with_api_key(&self.codex_home, &api_key) {
Ok(()) => {
self.error = None;
self.login_status = LoginStatus::AuthMode(AuthMode::ApiKey);
self.auth_manager.reload();
*self.sign_in_state.write().unwrap() = SignInState::ApiKeyConfigured;
}
Err(err) => {
self.error = Some(format!("Failed to save API key: {err}"));
let mut guard = self.sign_in_state.write().unwrap();
if let SignInState::ApiKeyEntry(existing) = &mut *guard {
if existing.value.is_empty() {
existing.value.push_str(&api_key);
}
existing.prepopulated_from_env = false;
} else {
*guard = SignInState::ApiKeyEntry(ApiKeyInputState {
value: api_key,
prepopulated_from_env: false,
});
}
}
}
self.request_frame.schedule_frame();
}
fn start_chatgpt_login(&mut self) {
@@ -354,18 +529,6 @@ impl AuthModeWidget {
}
}
}
/// TODO: Read/write from the correct hierarchy config overrides + auth json + OPENAI_API_KEY.
fn verify_api_key(&mut self) {
if matches!(self.login_status, LoginStatus::AuthMode(AuthMode::ApiKey)) {
// We already have an API key configured (e.g., from auth.json or env),
// so mark this step complete immediately.
*self.sign_in_state.write().unwrap() = SignInState::EnvVarFound;
} else {
*self.sign_in_state.write().unwrap() = SignInState::EnvVarMissing;
}
self.request_frame.schedule_frame();
}
}
impl StepStateProvider for AuthModeWidget {
@@ -373,10 +536,10 @@ impl StepStateProvider for AuthModeWidget {
let sign_in_state = self.sign_in_state.read().unwrap();
match &*sign_in_state {
SignInState::PickMode
| SignInState::EnvVarMissing
| SignInState::ApiKeyEntry(_)
| SignInState::ChatGptContinueInBrowser(_)
| SignInState::ChatGptSuccessMessage => StepState::InProgress,
SignInState::ChatGptSuccess | SignInState::EnvVarFound => StepState::Complete,
SignInState::ChatGptSuccess | SignInState::ApiKeyConfigured => StepState::Complete,
}
}
}
@@ -397,11 +560,11 @@ impl WidgetRef for AuthModeWidget {
SignInState::ChatGptSuccess => {
self.render_chatgpt_success(area, buf);
}
SignInState::EnvVarMissing => {
self.render_env_var_missing(area, buf);
SignInState::ApiKeyEntry(state) => {
self.render_api_key_entry(area, buf, state);
}
SignInState::EnvVarFound => {
self.render_env_var_found(area, buf);
SignInState::ApiKeyConfigured => {
self.render_api_key_configured(area, buf);
}
}
}

26
codex-rs/tui/src/onboarding/onboarding_screen.rs
View File

@@ -34,6 +34,7 @@ enum Step {
pub(crate) trait KeyboardHandler {
fn handle_key_event(&mut self, key_event: KeyEvent);
fn handle_paste(&mut self, _pasted: String) {}
}
pub(crate) enum StepState {
@@ -69,7 +70,6 @@ impl OnboardingScreen {
auth_manager,
config,
} = args;
let preferred_auth_method = config.preferred_auth_method;
let cwd = config.cwd.clone();
let codex_home = config.codex_home.clone();
let mut steps: Vec<Step> = vec![Step::Welcome(WelcomeWidget {
@@ -84,7 +84,6 @@ impl OnboardingScreen {
codex_home: codex_home.clone(),
login_status,
auth_manager,
preferred_auth_method,
}))
}
let is_git_repo = get_git_repo_root(&cwd).is_some();
@@ -194,6 +193,17 @@ impl KeyboardHandler for OnboardingScreen {
};
self.request_frame.schedule_frame();
}
fn handle_paste(&mut self, pasted: String) {
if pasted.is_empty() {
return;
}
if let Some(active_step) = self.current_steps_mut().into_iter().last() {
active_step.handle_paste(pasted);
}
self.request_frame.schedule_frame();
}
}
impl WidgetRef for &OnboardingScreen {
@@ -263,6 +273,14 @@ impl KeyboardHandler for Step {
Step::TrustDirectory(widget) => widget.handle_key_event(key_event),
}
}
fn handle_paste(&mut self, pasted: String) {
match self {
Step::Welcome(_) => {}
Step::Auth(widget) => widget.handle_paste(pasted),
Step::TrustDirectory(widget) => widget.handle_paste(pasted),
}
}
}
impl StepStateProvider for Step {
@@ -312,12 +330,14 @@ pub(crate) async fn run_onboarding_app(
TuiEvent::Key(key_event) => {
onboarding_screen.handle_key_event(key_event);
}
TuiEvent::Paste(text) => {
onboarding_screen.handle_paste(text);
}
TuiEvent::Draw => {
let _ = tui.draw(u16::MAX, |frame| {
frame.render_widget_ref(&onboarding_screen, frame.area());
});
}
_ => {}
}
}
}