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/login/src/device_code_auth.rs
View File

@@ -186,6 +186,13 @@ pub async fn run_device_code_login(opts: ServerOptions) -> std::io::Result<()> {
.await
.map_err(|err| std::io::Error::other(format!("device code exchange failed: {err}")))?;
if let Err(message) = crate::server::ensure_workspace_allowed(
opts.forced_chatgpt_workspace_id.as_deref(),
&tokens.id_token,
) {
return Err(io::Error::new(io::ErrorKind::PermissionDenied, message));
}
crate::server::persist_tokens_async(
&opts.codex_home,
None,

98
codex-rs/login/src/server.rs
View File

@@ -38,10 +38,15 @@ pub struct ServerOptions {
pub port: u16,
pub open_browser: bool,
pub force_state: Option<String>,
pub forced_chatgpt_workspace_id: Option<String>,
}
impl ServerOptions {
pub fn new(codex_home: PathBuf, client_id: String) -> Self {
pub fn new(
codex_home: PathBuf,
client_id: String,
forced_chatgpt_workspace_id: Option<String>,
) -> Self {
Self {
codex_home,
client_id,
@@ -49,6 +54,7 @@ impl ServerOptions {
port: DEFAULT_PORT,
open_browser: true,
force_state: None,
forced_chatgpt_workspace_id,
}
}
}
@@ -104,7 +110,14 @@ pub fn run_login_server(opts: ServerOptions) -> io::Result<LoginServer> {
let server = Arc::new(server);
let redirect_uri = format!("http://localhost:{actual_port}/auth/callback");
let auth_url = build_authorize_url(&opts.issuer, &opts.client_id, &redirect_uri, &pkce, &state);
let auth_url = build_authorize_url(
&opts.issuer,
&opts.client_id,
&redirect_uri,
&pkce,
&state,
opts.forced_chatgpt_workspace_id.as_deref(),
);
if opts.open_browser {
let _ = webbrowser::open(&auth_url);
@@ -240,6 +253,13 @@ async fn process_request(
.await
{
Ok(tokens) => {
if let Err(message) = ensure_workspace_allowed(
opts.forced_chatgpt_workspace_id.as_deref(),
&tokens.id_token,
) {
eprintln!("Workspace restriction error: {message}");
return login_error_response(&message);
}
// Obtain API key via token-exchange and persist
let api_key = obtain_api_key(&opts.issuer, &opts.client_id, &tokens.id_token)
.await
@@ -358,22 +378,35 @@ fn build_authorize_url(
redirect_uri: &str,
pkce: &PkceCodes,
state: &str,
forced_chatgpt_workspace_id: Option<&str>,
) -> String {
let query = vec![
("response_type", "code"),
("client_id", client_id),
("redirect_uri", redirect_uri),
("scope", "openid profile email offline_access"),
("code_challenge", &pkce.code_challenge),
("code_challenge_method", "S256"),
("id_token_add_organizations", "true"),
("codex_cli_simplified_flow", "true"),
("state", state),
("originator", originator().value.as_str()),
let mut query = vec![
("response_type".to_string(), "code".to_string()),
("client_id".to_string(), client_id.to_string()),
("redirect_uri".to_string(), redirect_uri.to_string()),
(
"scope".to_string(),
"openid profile email offline_access".to_string(),
),
(
"code_challenge".to_string(),
pkce.code_challenge.to_string(),
),
("code_challenge_method".to_string(), "S256".to_string()),
("id_token_add_organizations".to_string(), "true".to_string()),
("codex_cli_simplified_flow".to_string(), "true".to_string()),
("state".to_string(), state.to_string()),
(
"originator".to_string(),
originator().value.as_str().to_string(),
),
];
if let Some(workspace_id) = forced_chatgpt_workspace_id {
query.push(("allowed_workspace_id".to_string(), workspace_id.to_string()));
}
let qs = query
.into_iter()
.map(|(k, v)| format!("{}={}", k, urlencoding::encode(v)))
.map(|(k, v)| format!("{k}={}", urlencoding::encode(&v)))
.collect::<Vec<_>>()
.join("&");
format!("{issuer}/oauth/authorize?{qs}")
@@ -616,6 +649,43 @@ fn jwt_auth_claims(jwt: &str) -> serde_json::Map<String, serde_json::Value> {
serde_json::Map::new()
}
pub(crate) fn ensure_workspace_allowed(
expected: Option<&str>,
id_token: &str,
) -> Result<(), String> {
let Some(expected) = expected else {
return Ok(());
};
let claims = jwt_auth_claims(id_token);
let Some(actual) = claims.get("chatgpt_account_id").and_then(JsonValue::as_str) else {
return Err("Login is restricted to a specific workspace, but the token did not include an chatgpt_account_id claim.".to_string());
};
if actual == expected {
Ok(())
} else {
Err(format!("Login is restricted to workspace id {expected}."))
}
}
// Respond to the oauth server with an error so the code becomes unusable by anybody else.
fn login_error_response(message: &str) -> HandledRequest {
let mut headers = Vec::new();
if let Ok(header) = Header::from_bytes(&b"Content-Type"[..], &b"text/plain; charset=utf-8"[..])
{
headers.push(header);
}
HandledRequest::ResponseAndExit {
headers,
body: message.as_bytes().to_vec(),
result: Err(io::Error::new(
io::ErrorKind::PermissionDenied,
message.to_string(),
)),
}
}
pub(crate) async fn obtain_api_key(
issuer: &str,
client_id: &str,

54
codex-rs/login/tests/suite/device_code_login.rs
View File

@@ -97,7 +97,11 @@ async fn mock_oauth_token_single(server: &MockServer, jwt: String) {
}
fn server_opts(codex_home: &tempfile::TempDir, issuer: String) -> ServerOptions {
let mut opts = ServerOptions::new(codex_home.path().to_path_buf(), "client-id".to_string());
let mut opts = ServerOptions::new(
codex_home.path().to_path_buf(),
"client-id".to_string(),
None,
);
opts.issuer = issuer;
opts.open_browser = false;
opts
@@ -139,6 +143,42 @@ async fn device_code_login_integration_succeeds() {
assert_eq!(tokens.account_id.as_deref(), Some("acct_321"));
}
#[tokio::test]
async fn device_code_login_rejects_workspace_mismatch() {
skip_if_no_network!();
let codex_home = tempdir().unwrap();
let mock_server = MockServer::start().await;
mock_usercode_success(&mock_server).await;
mock_poll_token_two_step(&mock_server, Arc::new(AtomicUsize::new(0)), 404).await;
let jwt = make_jwt(json!({
"https://api.openai.com/auth": {
"chatgpt_account_id": "acct_321",
"organization_id": "org-actual"
}
}));
mock_oauth_token_single(&mock_server, jwt).await;
let issuer = mock_server.uri();
let mut opts = server_opts(&codex_home, issuer);
opts.forced_chatgpt_workspace_id = Some("org-required".to_string());
let err = run_device_code_login(opts)
.await
.expect_err("device code login should fail when workspace mismatches");
assert_eq!(err.kind(), std::io::ErrorKind::PermissionDenied);
let auth_path = get_auth_file(codex_home.path());
assert!(
!auth_path.exists(),
"auth.json should not be created when workspace validation fails"
);
}
#[tokio::test]
async fn device_code_login_integration_handles_usercode_http_failure() {
skip_if_no_network!();
@@ -183,7 +223,11 @@ async fn device_code_login_integration_persists_without_api_key_on_exchange_fail
let issuer = mock_server.uri();
let mut opts = ServerOptions::new(codex_home.path().to_path_buf(), "client-id".to_string());
let mut opts = ServerOptions::new(
codex_home.path().to_path_buf(),
"client-id".to_string(),
None,
);
opts.issuer = issuer;
opts.open_browser = false;
@@ -226,7 +270,11 @@ async fn device_code_login_integration_handles_error_payload() {
let issuer = mock_server.uri();
let mut opts = ServerOptions::new(codex_home.path().to_path_buf(), "client-id".to_string());
let mut opts = ServerOptions::new(
codex_home.path().to_path_buf(),
"client-id".to_string(),
None,
);
opts.issuer = issuer;
opts.open_browser = false;

80
codex-rs/login/tests/suite/login_server_e2e.rs
View File

@@ -14,11 +14,12 @@ use tempfile::tempdir;
// See spawn.rs for details
fn start_mock_issuer() -> (SocketAddr, thread::JoinHandle<()>) {
fn start_mock_issuer(chatgpt_account_id: &str) -> (SocketAddr, thread::JoinHandle<()>) {
// Bind to a random available port
let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap();
let addr = listener.local_addr().unwrap();
let server = tiny_http::Server::from_listener(listener, None).unwrap();
let chatgpt_account_id = chatgpt_account_id.to_string();
let handle = thread::spawn(move || {
while let Ok(mut req) = server.recv() {
@@ -41,7 +42,7 @@ fn start_mock_issuer() -> (SocketAddr, thread::JoinHandle<()>) {
"email": "user@example.com",
"https://api.openai.com/auth": {
"chatgpt_plan_type": "pro",
"chatgpt_account_id": "acc-123"
"chatgpt_account_id": chatgpt_account_id,
}
});
let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b);
@@ -80,7 +81,8 @@ fn start_mock_issuer() -> (SocketAddr, thread::JoinHandle<()>) {
async fn end_to_end_login_flow_persists_auth_json() -> Result<()> {
skip_if_no_network!(Ok(()));
let (issuer_addr, issuer_handle) = start_mock_issuer();
let chatgpt_account_id = "12345678-0000-0000-0000-000000000000";
let (issuer_addr, issuer_handle) = start_mock_issuer(chatgpt_account_id);
let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port());
let tmp = tempdir()?;
@@ -113,8 +115,15 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> {
port: 0,
open_browser: false,
force_state: Some(state),
forced_chatgpt_workspace_id: Some(chatgpt_account_id.to_string()),
};
let server = run_login_server(opts)?;
assert!(
server
.auth_url
.contains(format!("allowed_workspace_id={chatgpt_account_id}").as_str()),
"auth URL should include forced workspace parameter"
);
let login_port = server.actual_port;
// Simulate browser callback, and follow redirect to /success
@@ -138,7 +147,7 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> {
assert_eq!(json["OPENAI_API_KEY"], "access-123");
assert_eq!(json["tokens"]["access_token"], "access-123");
assert_eq!(json["tokens"]["refresh_token"], "refresh-123");
assert_eq!(json["tokens"]["account_id"], "acc-123");
assert_eq!(json["tokens"]["account_id"], chatgpt_account_id);
// Stop mock issuer
drop(issuer_handle);
@@ -149,7 +158,7 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> {
async fn creates_missing_codex_home_dir() -> Result<()> {
skip_if_no_network!(Ok(()));
let (issuer_addr, _issuer_handle) = start_mock_issuer();
let (issuer_addr, _issuer_handle) = start_mock_issuer("org-123");
let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port());
let tmp = tempdir()?;
@@ -166,6 +175,7 @@ async fn creates_missing_codex_home_dir() -> Result<()> {
port: 0,
open_browser: false,
force_state: Some(state),
forced_chatgpt_workspace_id: None,
};
let server = run_login_server(opts)?;
let login_port = server.actual_port;
@@ -185,11 +195,67 @@ async fn creates_missing_codex_home_dir() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> {
skip_if_no_network!(Ok(()));
let (issuer_addr, _issuer_handle) = start_mock_issuer("org-actual");
let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port());
let tmp = tempdir()?;
let codex_home = tmp.path().to_path_buf();
let state = "state-mismatch".to_string();
let opts = ServerOptions {
codex_home: codex_home.clone(),
client_id: codex_login::CLIENT_ID.to_string(),
issuer,
port: 0,
open_browser: false,
force_state: Some(state.clone()),
forced_chatgpt_workspace_id: Some("org-required".to_string()),
};
let server = run_login_server(opts)?;
assert!(
server
.auth_url
.contains("allowed_workspace_id=org-required"),
"auth URL should include forced workspace parameter"
);
let login_port = server.actual_port;
let client = reqwest::Client::new();
let url = format!("http://127.0.0.1:{login_port}/auth/callback?code=abc&state={state}");
let resp = client.get(&url).send().await?;
assert!(resp.status().is_success());
let body = resp.text().await?;
assert!(
body.contains("Login is restricted to workspace id org-required"),
"error body should mention workspace restriction"
);
let result = server.block_until_done().await;
assert!(
result.is_err(),
"login should fail due to workspace mismatch"
);
let err = result.unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::PermissionDenied);
let auth_path = codex_home.join("auth.json");
assert!(
!auth_path.exists(),
"auth.json should not be written when the workspace mismatches"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> {
skip_if_no_network!(Ok(()));
let (issuer_addr, _issuer_handle) = start_mock_issuer();
let (issuer_addr, _issuer_handle) = start_mock_issuer("org-123");
let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port());
let first_tmp = tempdir()?;
@@ -202,6 +268,7 @@ async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> {
port: 0,
open_browser: false,
force_state: Some("cancel_state".to_string()),
forced_chatgpt_workspace_id: None,
};
let first_server = run_login_server(first_opts)?;
@@ -220,6 +287,7 @@ async fn cancels_previous_login_server_when_port_is_in_use() -> Result<()> {
port: login_port,
open_browser: false,
force_state: Some("cancel_state_2".to_string()),
forced_chatgpt_workspace_id: None,
};
let second_server = run_login_server(second_opts)?;