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

Add forced_chatgpt_workspace_id and forced_login_method configuration options (#5303)

Browse Source 
This PR adds support for configs to specify a forced login method
(chatgpt or api) as well as a forced chatgpt account id. This lets
enterprises uses [managed
configs](https://developers.openai.com/codex/security#managed-configuration)
to force all employees to use their company's workspace instead of their
own or any other.

When a workspace id is set, a query param is sent to the login flow
which auto-selects the given workspace or errors if the user isn't a
member of it.

This PR is large but a large % of it is tests, wiring, and required
formatting changes.

API login with chatgpt forced
<img width="1592" height="116" alt="CleanShot 2025-10-19 at 22 40 04"
src="https://github.com/user-attachments/assets/560c6bb4-a20a-4a37-95af-93df39d057dd"
/>

ChatGPT login with api forced
<img width="1018" height="100" alt="CleanShot 2025-10-19 at 22 40 29"
src="https://github.com/user-attachments/assets/d010bbbb-9c8d-4227-9eda-e55bf043b4af"
/>

Onboarding with api forced
<img width="892" height="460" alt="CleanShot 2025-10-19 at 22 41 02"
src="https://github.com/user-attachments/assets/cc0ed45c-b257-4d62-a32e-6ca7514b5edd"
/>

Onboarding with ChatGPT forced
<img width="1154" height="426" alt="CleanShot 2025-10-19 at 22 41 27"
src="https://github.com/user-attachments/assets/41c41417-dc68-4bb4-b3e7-3b7769f7e6a1"
/>

Logging in with the wrong workspace
<img width="2222" height="84" alt="CleanShot 2025-10-19 at 22 42 31"
src="https://github.com/user-attachments/assets/0ff4222c-f626-4dd3-b035-0b7fe998a046"
/>
This commit is contained in:
Gabriel Peal
2025-10-20 08:50:54 -07:00
committed by GitHub
parent d01f91ecec
commit d87f87e25b
19 changed files with 920 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -11,6 +11,7 @@ use codex_core::BUILT_IN_OSS_MODEL_PROVIDER_ID;
use codex_core::CodexAuth;
use codex_core::INTERACTIVE_SESSION_SOURCES;
use codex_core::RolloutRecorder;
use codex_core::auth::enforce_login_restrictions;
use codex_core::config::Config;
use codex_core::config::ConfigOverrides;
use codex_core::find_conversation_path_by_id_str;
@@ -193,6 +194,12 @@ pub async fn run_main(
let config = load_config_or_exit(cli_kv_overrides.clone(), overrides.clone()).await;
#[allow(clippy::print_stderr)]
if let Err(err) = enforce_login_restrictions(&config).await {
eprintln!("{err}");
std::process::exit(1);
}
let active_profile = config.active_profile.clone();
let log_dir = codex_core::config::log_dir(&config)?;
std::fs::create_dir_all(&log_dir)?;

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

@@ -28,6 +28,7 @@ use ratatui::widgets::WidgetRef;
use ratatui::widgets::Wrap;
use codex_app_server_protocol::AuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
use std::sync::RwLock;
use crate::LoginStatus;
@@ -50,6 +51,8 @@ pub(crate) enum SignInState {
ApiKeyConfigured,
}
const API_KEY_DISABLED_MESSAGE: &str = "API key login is disabled.";
#[derive(Clone, Default)]
pub(crate) struct ApiKeyInputState {
value: String,
@@ -79,25 +82,41 @@ impl KeyboardHandler for AuthModeWidget {
match key_event.code {
KeyCode::Up | KeyCode::Char('k') => {
self.highlighted_mode = AuthMode::ChatGPT;
if self.is_chatgpt_login_allowed() {
self.highlighted_mode = AuthMode::ChatGPT;
}
}
KeyCode::Down | KeyCode::Char('j') => {
self.highlighted_mode = AuthMode::ApiKey;
if self.is_api_login_allowed() {
self.highlighted_mode = AuthMode::ApiKey;
}
}
KeyCode::Char('1') => {
self.start_chatgpt_login();
if self.is_chatgpt_login_allowed() {
self.start_chatgpt_login();
}
}
KeyCode::Char('2') => {
if self.is_api_login_allowed() {
self.start_api_key_entry();
} else {
self.disallow_api_login();
}
}
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 {
SignInState::PickMode => match self.highlighted_mode {
AuthMode::ChatGPT => {
AuthMode::ChatGPT if self.is_chatgpt_login_allowed() => {
self.start_chatgpt_login();
}
AuthMode::ApiKey => {
AuthMode::ApiKey if self.is_api_login_allowed() => {
self.start_api_key_entry();
}
AuthMode::ChatGPT => {}
AuthMode::ApiKey => {
self.disallow_api_login();
}
},
SignInState::ChatGptSuccessMessage => {
*self.sign_in_state.write().unwrap() = SignInState::ChatGptSuccess;
@@ -131,9 +150,26 @@ pub(crate) struct AuthModeWidget {
pub codex_home: PathBuf,
pub login_status: LoginStatus,
pub auth_manager: Arc<AuthManager>,
pub forced_chatgpt_workspace_id: Option<String>,
pub forced_login_method: Option<ForcedLoginMethod>,
}
impl AuthModeWidget {
fn is_api_login_allowed(&self) -> bool {
!matches!(self.forced_login_method, Some(ForcedLoginMethod::Chatgpt))
}
fn is_chatgpt_login_allowed(&self) -> bool {
!matches!(self.forced_login_method, Some(ForcedLoginMethod::Api))
}
fn disallow_api_login(&mut self) {
self.highlighted_mode = AuthMode::ChatGPT;
self.error = Some(API_KEY_DISABLED_MESSAGE.to_string());
*self.sign_in_state.write().unwrap() = SignInState::PickMode;
self.request_frame.schedule_frame();
}
fn render_pick_mode(&self, area: Rect, buf: &mut Buffer) {
let mut lines: Vec<Line> = vec![
Line::from(vec![
@@ -176,20 +212,34 @@ impl AuthModeWidget {
vec![line1, line2]
};
let chatgpt_description = if self.is_chatgpt_login_allowed() {
"Usage included with Plus, Pro, and Team plans"
} else {
"ChatGPT login is disabled"
};
lines.extend(create_mode_item(
0,
AuthMode::ChatGPT,
"Sign in with ChatGPT",
"Usage included with Plus, Pro, and Team plans",
));
lines.push("".into());
lines.extend(create_mode_item(
1,
AuthMode::ApiKey,
"Provide your own API key",
"Pay for what you use",
chatgpt_description,
));
lines.push("".into());
if self.is_api_login_allowed() {
lines.extend(create_mode_item(
1,
AuthMode::ApiKey,
"Provide your own API key",
"Pay for what you use",
));
lines.push("".into());
} else {
lines.push(
" API key login is disabled by this workspace. Sign in with ChatGPT to continue."
.dim()
.into(),
);
lines.push("".into());
}
lines.push(
// AE: Following styles.md, this should probably be Cyan because it's a user input tip.
// But leaving this for a future cleanup.
@@ -428,6 +478,10 @@ impl AuthModeWidget {
}
fn start_api_key_entry(&mut self) {
if !self.is_api_login_allowed() {
self.disallow_api_login();
return;
}
self.error = None;
let prefill_from_env = read_openai_api_key_from_env();
let mut guard = self.sign_in_state.write().unwrap();
@@ -454,6 +508,10 @@ impl AuthModeWidget {
}
fn save_api_key(&mut self, api_key: String) {
if !self.is_api_login_allowed() {
self.disallow_api_login();
return;
}
match login_with_api_key(&self.codex_home, &api_key) {
Ok(()) => {
self.error = None;
@@ -491,7 +549,11 @@ impl AuthModeWidget {
}
self.error = None;
let opts = ServerOptions::new(self.codex_home.clone(), CLIENT_ID.to_string());
let opts = ServerOptions::new(
self.codex_home.clone(),
CLIENT_ID.to_string(),
self.forced_chatgpt_workspace_id.clone(),
);
match run_login_server(opts) {
Ok(child) => {
let sign_in_state = self.sign_in_state.clone();
@@ -571,3 +633,54 @@ impl WidgetRef for AuthModeWidget {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
fn widget_forced_chatgpt() -> (AuthModeWidget, TempDir) {
let codex_home = TempDir::new().unwrap();
let codex_home_path = codex_home.path().to_path_buf();
let widget = AuthModeWidget {
request_frame: FrameRequester::test_dummy(),
highlighted_mode: AuthMode::ChatGPT,
error: None,
sign_in_state: Arc::new(RwLock::new(SignInState::PickMode)),
codex_home: codex_home_path.clone(),
login_status: LoginStatus::NotAuthenticated,
auth_manager: AuthManager::shared(codex_home_path, false),
forced_chatgpt_workspace_id: None,
forced_login_method: Some(ForcedLoginMethod::Chatgpt),
};
(widget, codex_home)
}
#[test]
fn api_key_flow_disabled_when_chatgpt_forced() {
let (mut widget, _tmp) = widget_forced_chatgpt();
widget.start_api_key_entry();
assert_eq!(widget.error.as_deref(), Some(API_KEY_DISABLED_MESSAGE));
assert!(matches!(
&*widget.sign_in_state.read().unwrap(),
SignInState::PickMode
));
}
#[test]
fn saving_api_key_is_blocked_when_chatgpt_forced() {
let (mut widget, _tmp) = widget_forced_chatgpt();
widget.save_api_key("sk-test".to_string());
assert_eq!(widget.error.as_deref(), Some(API_KEY_DISABLED_MESSAGE));
assert!(matches!(
&*widget.sign_in_state.read().unwrap(),
SignInState::PickMode
));
assert_eq!(widget.login_status, LoginStatus::NotAuthenticated);
}
}

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

@@ -12,6 +12,7 @@ use ratatui::widgets::Clear;
use ratatui::widgets::WidgetRef;
use codex_app_server_protocol::AuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
use crate::LoginStatus;
use crate::onboarding::auth::AuthModeWidget;
@@ -83,6 +84,8 @@ impl OnboardingScreen {
config,
} = args;
let cwd = config.cwd.clone();
let forced_chatgpt_workspace_id = config.forced_chatgpt_workspace_id.clone();
let forced_login_method = config.forced_login_method;
let codex_home = config.codex_home;
let mut steps: Vec<Step> = Vec::new();
if show_windows_wsl_screen {
@@ -93,14 +96,20 @@ impl OnboardingScreen {
tui.frame_requester(),
)));
if show_login_screen {
let highlighted_mode = match forced_login_method {
Some(ForcedLoginMethod::Api) => AuthMode::ApiKey,
_ => AuthMode::ChatGPT,
};
steps.push(Step::Auth(AuthModeWidget {
request_frame: tui.frame_requester(),
highlighted_mode: AuthMode::ChatGPT,
highlighted_mode,
error: None,
sign_in_state: Arc::new(RwLock::new(SignInState::PickMode)),
codex_home: codex_home.clone(),
login_status,
auth_manager,
forced_chatgpt_workspace_id,
forced_login_method,
}))
}
let is_git_repo = get_git_repo_root(&cwd).is_some();