From 4775449e191d52f3cfb926c13236a5bd7b36c0e2 Mon Sep 17 00:00:00 2001 From: rreichel3-oai Date: Thu, 16 Apr 2026 13:24:25 -0400 Subject: [PATCH] Support multiple forced ChatGPT workspaces --- .../codex_app_server_protocol.schemas.json | 5 +- .../codex_app_server_protocol.v2.schemas.json | 5 +- .../schema/json/v2/ConfigReadResponse.json | 5 +- .../schema/typescript/v2/Config.ts | 3 +- .../app-server-protocol/src/protocol/v1.rs | 2 +- .../src/protocol/v2/config.rs | 2 +- .../request_processors/account_processor.rs | 9 +- codex-rs/cli/src/login.rs | 2 +- codex-rs/config/src/config_toml.rs | 16 +++- codex-rs/core/config.schema.json | 22 ++++- codex-rs/core/src/config/mod.rs | 28 +++--- codex-rs/login/src/auth/auth_tests.rs | 96 ++++++++++++++++++- codex-rs/login/src/auth/manager.rs | 70 +++++++------- codex-rs/login/src/server.rs | 24 +++-- .../login/tests/suite/device_code_login.rs | 2 +- .../login/tests/suite/login_server_e2e.rs | 45 ++++++++- codex-rs/tui/src/local_chatgpt_auth.rs | 17 ++-- .../src/openai_codex/generated/v2_all.py | 2 +- 18 files changed, 270 insertions(+), 85 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 6da572813f..deb9707b4c 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -7126,8 +7126,11 @@ ] }, "forced_chatgpt_workspace_id": { + "items": { + "type": "string" + }, "type": [ - "string", + "array", "null" ] }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index e338f4151f..91d9b1f551 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -3515,8 +3515,11 @@ ] }, "forced_chatgpt_workspace_id": { + "items": { + "type": "string" + }, "type": [ - "string", + "array", "null" ] }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json index 2e0a05725c..821b323eef 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json @@ -235,8 +235,11 @@ ] }, "forced_chatgpt_workspace_id": { + "items": { + "type": "string" + }, "type": [ - "string", + "array", "null" ] }, diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts b/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts index cc7e340ea3..1b04d60d28 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts @@ -4,6 +4,7 @@ import type { ForcedLoginMethod } from "../ForcedLoginMethod"; import type { ReasoningEffort } from "../ReasoningEffort"; import type { ReasoningSummary } from "../ReasoningSummary"; +import type { ServiceTier } from "../ServiceTier"; import type { Verbosity } from "../Verbosity"; import type { WebSearchMode } from "../WebSearchMode"; import type { JsonValue } from "../serde_json/JsonValue"; @@ -19,4 +20,4 @@ export type Config = {model: string | null, review_model: string | null, model_c * [UNSTABLE] Optional default for where approval requests are routed for * review. */ -approvals_reviewer: ApprovalsReviewer | null, sandbox_mode: SandboxMode | null, sandbox_workspace_write: SandboxWorkspaceWrite | null, forced_chatgpt_workspace_id: string | null, forced_login_method: ForcedLoginMethod | null, web_search: WebSearchMode | null, tools: ToolsV2 | null, profile: string | null, profiles: { [key in string]?: ProfileV2 }, instructions: string | null, developer_instructions: string | null, compact_prompt: string | null, model_reasoning_effort: ReasoningEffort | null, model_reasoning_summary: ReasoningSummary | null, model_verbosity: Verbosity | null, service_tier: string | null, analytics: AnalyticsConfig | null} & ({ [key in string]?: number | string | boolean | Array | { [key in string]?: JsonValue } | null }); +approvals_reviewer: ApprovalsReviewer | null, sandbox_mode: SandboxMode | null, sandbox_workspace_write: SandboxWorkspaceWrite | null, forced_chatgpt_workspace_id: Array | null, forced_login_method: ForcedLoginMethod | null, web_search: WebSearchMode | null, tools: ToolsV2 | null, profile: string | null, profiles: { [key in string]?: ProfileV2 }, instructions: string | null, developer_instructions: string | null, compact_prompt: string | null, model_reasoning_effort: ReasoningEffort | null, model_reasoning_summary: ReasoningSummary | null, model_verbosity: Verbosity | null, service_tier: ServiceTier | null, analytics: AnalyticsConfig | null} & ({ [key in string]?: number | string | boolean | Array | { [key in string]?: JsonValue } | null }); diff --git a/codex-rs/app-server-protocol/src/protocol/v1.rs b/codex-rs/app-server-protocol/src/protocol/v1.rs index f8e9afdd12..e35e0011fa 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -201,7 +201,7 @@ pub struct UserSavedConfig { pub approval_policy: Option, pub sandbox_mode: Option, pub sandbox_settings: Option, - pub forced_chatgpt_workspace_id: Option, + pub forced_chatgpt_workspace_id: Option>, pub forced_login_method: Option, pub model: Option, pub model_reasoning_effort: Option, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/config.rs b/codex-rs/app-server-protocol/src/protocol/v2/config.rs index 28a2ad01f9..6e1db1bb7f 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/config.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/config.rs @@ -232,7 +232,7 @@ pub struct Config { pub approvals_reviewer: Option, pub sandbox_mode: Option, pub sandbox_workspace_write: Option, - pub forced_chatgpt_workspace_id: Option, + pub forced_chatgpt_workspace_id: Option>, pub forced_login_method: Option, pub web_search: Option, pub tools: Option, diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index f32f87ed15..71c7e4b3e2 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -570,11 +570,14 @@ impl AccountRequestProcessor { } } - if let Some(expected_workspace) = self.config.forced_chatgpt_workspace_id.as_deref() - && chatgpt_account_id != expected_workspace + if let Some(expected_workspaces) = self.config.forced_chatgpt_workspace_id.as_deref() + && !expected_workspaces + .iter() + .any(|expected_workspace| chatgpt_account_id == *expected_workspace) { return Err(invalid_request(format!( - "External auth must use workspace {expected_workspace}, but received {chatgpt_account_id:?}." + "External auth must use one of workspace(s) {:?}, but received {chatgpt_account_id:?}.", + expected_workspaces ))); } diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 16add7ac90..6fdc62fdb1 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -115,7 +115,7 @@ fn print_login_server_start(actual_port: u16, auth_url: &str) { pub async fn login_with_chatgpt( codex_home: PathBuf, - forced_chatgpt_workspace_id: Option, + forced_chatgpt_workspace_id: Option>, cli_auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result<()> { let opts = ServerOptions::new( diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index b627a83726..401d9fdf61 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -177,9 +177,17 @@ pub struct ConfigToml { /// Compact prompt used for history compaction. pub compact_prompt: Option, - /// When set, restricts ChatGPT login to a specific workspace identifier. + /// Optional commit attribution text for commit message co-author trailers. + /// This top-level setting only takes effect when `[features].codex_git_commit` + /// is enabled. + /// + /// When enabled and unset, Codex uses `Codex `. + /// Set to an empty string to disable automatic commit attribution. + pub commit_attribution: Option, + + /// When set, restricts ChatGPT login to one or more workspace identifiers. #[serde(default)] - pub forced_chatgpt_workspace_id: Option, + pub forced_chatgpt_workspace_id: Option, /// When set, restricts the login mechanism users may use. #[serde(default)] @@ -508,7 +516,9 @@ impl From for UserSavedConfig { approval_policy: config_toml.approval_policy, sandbox_mode: config_toml.sandbox_mode, sandbox_settings: config_toml.sandbox_workspace_write.map(From::from), - forced_chatgpt_workspace_id: config_toml.forced_chatgpt_workspace_id, + forced_chatgpt_workspace_id: config_toml + .forced_chatgpt_workspace_id + .map(ForcedChatgptWorkspaceIds::into_vec), forced_login_method: config_toml.forced_login_method, model: config_toml.model, model_reasoning_effort: config_toml.model_reasoning_effort, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 82fde87e78..dcb3533178 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -861,6 +861,20 @@ }, "type": "object" }, + "ForcedChatgptWorkspaceIds": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ], + "description": "Backward-compatible shape for workspace restrictions in config.toml." + }, "ForcedLoginMethod": { "enum": [ "chatgpt", @@ -4367,9 +4381,13 @@ "description": "Optional URI-based file opener. If set, citations to files in the model output will be hyperlinked using the specified URI scheme." }, "forced_chatgpt_workspace_id": { + "allOf": [ + { + "$ref": "#/definitions/ForcedChatgptWorkspaceIds" + } + ], "default": null, - "description": "When set, restricts ChatGPT login to a specific workspace identifier.", - "type": "string" + "description": "When set, restricts ChatGPT login to one or more workspace identifiers." }, "forced_login_method": { "allOf": [ diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 4aa3230332..e0d1a86576 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -746,15 +746,14 @@ pub struct Config { /// instructions inserted into developer messages when realtime becomes /// active. pub experimental_realtime_start_instructions: Option, - /// Experimental / do not use. When set, app-server fetches thread-scoped /// config from a remote service at this endpoint. pub experimental_thread_config_endpoint: Option, /// Experimental / do not use. Selects the thread persistence backend. pub experimental_thread_store: ThreadStoreConfig, - /// When set, restricts ChatGPT login to a specific workspace identifier. - pub forced_chatgpt_workspace_id: Option, + /// When set, restricts ChatGPT login to one or more workspace identifiers. + pub forced_chatgpt_workspace_id: Option>, /// When set, restricts the login mechanism users may use. pub forced_login_method: Option, @@ -881,7 +880,7 @@ impl AuthManagerConfig for Config { self.cli_auth_credentials_store_mode } - fn forced_chatgpt_workspace_id(&self) -> Option { + fn forced_chatgpt_workspace_id(&self) -> Option> { self.forced_chatgpt_workspace_id.clone() } @@ -2757,15 +2756,18 @@ impl Config { let include_apply_patch_tool_flag = features.enabled(Feature::ApplyPatchFreeform); let use_experimental_unified_exec_tool = features.enabled(Feature::UnifiedExec); - let forced_chatgpt_workspace_id = - cfg.forced_chatgpt_workspace_id.as_ref().and_then(|value| { - let trimmed = value.trim(); - if trimmed.is_empty() { - None - } else { - Some(trimmed.to_string()) - } - }); + let forced_chatgpt_workspace_id = cfg + .forced_chatgpt_workspace_id + .clone() + .map(codex_config::config_toml::ForcedChatgptWorkspaceIds::into_vec) + .map(|values| { + values + .into_iter() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .collect::>() + }) + .filter(|values| !values.is_empty()); let forced_login_method = cfg.forced_login_method; diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index fe57be06fa..d88a1eca99 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -592,6 +592,67 @@ exit 1 } } +#[tokio::test] +async fn auth_manager_notifies_when_auth_state_changes() { + let dir = tempdir().unwrap(); + let manager = AuthManager::shared( + dir.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + ); + let mut auth_state_rx = manager.subscribe_auth_state(); + + save_auth( + dir.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-test-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }, + AuthCredentialsStoreMode::File, + ) + .expect("save auth"); + + assert!( + manager.reload(), + "reload should report a changed auth state" + ); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("auth change notification should arrive") + .expect("auth state watch should remain open"); + + save_auth( + dir.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-updated-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }, + AuthCredentialsStoreMode::File, + ) + .expect("save updated auth"); + + assert!( + !manager.reload(), + "reload remains mode-stable even when the underlying credentials change" + ); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("auth reload notification should still arrive") + .expect("auth state watch should remain open"); + + manager.set_forced_chatgpt_workspace_id(Some(vec!["workspace-123".to_string()])); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("workspace change notification should arrive") + .expect("auth state watch should remain open"); +} + struct AuthFileParams { openai_api_key: Option, chatgpt_plan_type: Option, @@ -654,7 +715,7 @@ fn fake_jwt_for_auth_file_params(params: &AuthFileParams) -> std::io::Result, - forced_chatgpt_workspace_id: Option, + forced_chatgpt_workspace_id: Option>, ) -> AuthConfig { AuthConfig { codex_home: codex_home.to_path_buf(), @@ -823,14 +884,14 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { let config = build_config( codex_home.path(), /*forced_login_method*/ None, - Some("org_mine".to_string()), + Some(vec!["org_mine".to_string()]), ) .await; let err = super::enforce_login_restrictions(&config) .await .expect_err("expected workspace mismatch to error"); - assert!(err.to_string().contains("workspace org_mine")); + assert!(err.to_string().contains("workspace(s) org_mine")); assert!( !codex_home.path().join("auth.json").exists(), "auth.json should be removed on mismatch" @@ -855,7 +916,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() { let config = build_config( codex_home.path(), /*forced_login_method*/ None, - Some("org_mine".to_string()), + Some(vec!["org_mine".to_string()]), ) .await; @@ -870,6 +931,31 @@ async fn enforce_login_restrictions_allows_matching_workspace() { #[tokio::test] #[serial(codex_auth_env)] +async fn enforce_login_restrictions_allows_any_matching_workspace_in_list() { + let codex_home = tempdir().unwrap(); + let _jwt = write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("org_mine".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let config = build_config( + codex_home.path(), + /*forced_login_method*/ None, + Some(vec!["org_other".to_string(), "org_mine".to_string()]), + ) + .await; + + super::enforce_login_restrictions(&config) + .await + .expect("any matching workspace in the allowed list should succeed"); +} + +#[tokio::test] async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_forced_chatgpt_workspace_id_is_set() { let codex_home = tempdir().unwrap(); @@ -880,7 +966,7 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f let config = build_config( codex_home.path(), /*forced_login_method*/ None, - Some("org_mine".to_string()), + Some(vec!["org_mine".to_string()]), ) .await; diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 45bd55302b..def9b7965d 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -610,8 +610,8 @@ pub struct AuthConfig { pub codex_home: PathBuf, pub auth_credentials_store_mode: AuthCredentialsStoreMode, pub forced_login_method: Option, - pub forced_chatgpt_workspace_id: Option, pub chatgpt_base_url: Option, + pub forced_chatgpt_workspace_id: Option>, } pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { @@ -653,34 +653,38 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< } } - if let Some(expected_account_id) = config.forced_chatgpt_workspace_id.as_deref() { - // workspace is the external identifier for account id. - let chatgpt_account_id = match auth { - CodexAuth::ApiKey(_) => return Ok(()), - CodexAuth::AgentIdentity(_) => auth.get_account_id(), - CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { - let token_data = match auth.get_token_data() { - Ok(data) => data, - Err(err) => { - return logout_with_message( - &config.codex_home, - format!( - "Failed to load ChatGPT credentials while enforcing workspace restrictions: {err}. Logging out." - ), - config.auth_credentials_store_mode, - ); - } - }; - token_data.id_token.chatgpt_account_id + if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { + if !auth.is_chatgpt_auth() { + return Ok(()); + } + + let token_data = match auth.get_token_data() { + Ok(data) => data, + Err(err) => { + return logout_with_message( + &config.codex_home, + format!( + "Failed to load ChatGPT credentials while enforcing workspace restrictions: {err}. Logging out." + ), + config.auth_credentials_store_mode, + ); } }; - if chatgpt_account_id.as_deref() != Some(expected_account_id) { + + // workspace is the external identifier for account id. + let chatgpt_account_id = token_data.id_token.chatgpt_account_id.as_deref(); + if !chatgpt_account_id.is_some_and(|actual| { + expected_account_ids + .iter() + .any(|expected| expected == actual) + }) { + let expected_workspaces = expected_account_ids.join(", "); let message = match chatgpt_account_id { Some(actual) => format!( - "Login is restricted to workspace {expected_account_id}, but current credentials belong to {actual}. Logging out." + "Login is restricted to workspace(s) {expected_workspaces}, but current credentials belong to {actual}. Logging out." ), None => format!( - "Login is restricted to workspace {expected_account_id}, but current credentials lack a workspace identifier. Logging out." + "Login is restricted to workspace(s) {expected_workspaces}, but current credentials lack a workspace identifier. Logging out." ), }; return logout_with_message( @@ -1247,7 +1251,7 @@ pub struct AuthManager { inner: RwLock, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, - forced_chatgpt_workspace_id: RwLock>, + forced_chatgpt_workspace_id: RwLock>>, chatgpt_base_url: Option, refresh_lock: Semaphore, external_auth: RwLock>>, @@ -1266,8 +1270,8 @@ pub trait AuthManagerConfig { /// Returns the CLI auth credential storage mode for auth loading. fn cli_auth_credentials_store_mode(&self) -> AuthCredentialsStoreMode; - /// Returns the workspace ID that ChatGPT auth should be restricted to, if any. - fn forced_chatgpt_workspace_id(&self) -> Option; + /// Returns the workspace IDs that ChatGPT auth should be restricted to, if any. + fn forced_chatgpt_workspace_id(&self) -> Option>; /// Returns the ChatGPT backend base URL used for first-party backend authorization. fn chatgpt_base_url(&self) -> String; @@ -1548,7 +1552,7 @@ impl AuthManager { } } - pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { + pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option>) { if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() && *guard != workspace_id { @@ -1556,7 +1560,7 @@ impl AuthManager { } } - pub fn forced_chatgpt_workspace_id(&self) -> Option { + pub fn forced_chatgpt_workspace_id(&self) -> Option> { self.forced_chatgpt_workspace_id .read() .ok() @@ -1836,13 +1840,15 @@ impl AuthManager { "external auth refresh did not return ChatGPT metadata", ))); }; - if let Some(expected_workspace_id) = forced_chatgpt_workspace_id.as_deref() - && chatgpt_metadata.account_id != expected_workspace_id + if let Some(expected_workspace_ids) = forced_chatgpt_workspace_id.as_deref() + && !expected_workspace_ids + .iter() + .any(|expected| chatgpt_metadata.account_id == *expected) { return Err(RefreshTokenError::Transient(std::io::Error::other( format!( - "external auth refresh returned workspace {:?}, expected {expected_workspace_id:?}", - chatgpt_metadata.account_id, + "external auth refresh returned workspace {:?}, expected one of {:?}", + chatgpt_metadata.account_id, expected_workspace_ids, ), ))); } diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index ad0d6b5e5e..c137be9273 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -69,7 +69,7 @@ pub struct ServerOptions { pub port: u16, pub open_browser: bool, pub force_state: Option, - pub forced_chatgpt_workspace_id: Option, + pub forced_chatgpt_workspace_id: Option>, pub codex_streamlined_login: bool, pub cli_auth_credentials_store_mode: AuthCredentialsStoreMode, } @@ -79,7 +79,7 @@ impl ServerOptions { pub fn new( codex_home: PathBuf, client_id: String, - forced_chatgpt_workspace_id: Option, + forced_chatgpt_workspace_id: Option>, cli_auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> Self { Self { @@ -486,7 +486,7 @@ fn build_authorize_url( redirect_uri: &str, pkce: &PkceCodes, state: &str, - forced_chatgpt_workspace_id: Option<&str>, + forced_chatgpt_workspace_ids: Option<&[String]>, ) -> String { let mut query = vec![ ("response_type".to_string(), "code".to_string()), @@ -507,8 +507,13 @@ fn build_authorize_url( ("state".to_string(), state.to_string()), ("originator".to_string(), originator().value), ]; - if let Some(workspace_id) = forced_chatgpt_workspace_id { - query.push(("allowed_workspace_id".to_string(), workspace_id.to_string())); + if let Some(workspace_ids) = forced_chatgpt_workspace_ids { + query.extend( + workspace_ids + .iter() + .cloned() + .map(|workspace_id| ("allowed_workspace_id".to_string(), workspace_id)), + ); } let qs = query .into_iter() @@ -928,7 +933,7 @@ fn jwt_auth_claims(jwt: &str) -> serde_json::Map { /// Validates the ID token against an optional workspace restriction. pub(crate) fn ensure_workspace_allowed( - expected: Option<&str>, + expected: Option<&[String]>, id_token: &str, ) -> Result<(), String> { let Some(expected) = expected else { @@ -940,10 +945,13 @@ pub(crate) fn ensure_workspace_allowed( 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 { + if expected.iter().any(|workspace_id| workspace_id == actual) { Ok(()) } else { - Err(format!("Login is restricted to workspace id {expected}.")) + Err(format!( + "Login is restricted to workspace id(s) {}.", + expected.join(", ") + )) } } diff --git a/codex-rs/login/tests/suite/device_code_login.rs b/codex-rs/login/tests/suite/device_code_login.rs index bed94c7005..234acd3a43 100644 --- a/codex-rs/login/tests/suite/device_code_login.rs +++ b/codex-rs/login/tests/suite/device_code_login.rs @@ -183,7 +183,7 @@ async fn device_code_login_rejects_workspace_mismatch() -> anyhow::Result<()> { let issuer = mock_server.uri(); let mut opts = server_opts(&codex_home, issuer, AuthCredentialsStoreMode::File); - opts.forced_chatgpt_workspace_id = Some("org-required".to_string()); + opts.forced_chatgpt_workspace_id = Some(vec!["org-required".to_string()]); let err = run_device_code_login(opts) .await diff --git a/codex-rs/login/tests/suite/login_server_e2e.rs b/codex-rs/login/tests/suite/login_server_e2e.rs index ce58a51358..0bcc9a5241 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -121,7 +121,7 @@ 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()), + forced_chatgpt_workspace_id: Some(vec![chatgpt_account_id.to_string()]), codex_streamlined_login: false, }; let server = run_login_server(opts)?; @@ -204,6 +204,45 @@ async fn creates_missing_codex_home_dir() -> Result<()> { Ok(()) } +#[tokio::test] +async fn login_server_includes_all_forced_workspace_query_params() -> Result<()> { + skip_if_no_network!(Ok(())); + + let (issuer_addr, _issuer_handle) = start_mock_issuer("org-123"); + let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); + + let tmp = tempdir()?; + let codex_home = tmp.path().to_path_buf(); + let state = "state-multi".to_string(); + + let opts = ServerOptions { + codex_home, + cli_auth_credentials_store_mode: AuthCredentialsStoreMode::File, + client_id: codex_login::CLIENT_ID.to_string(), + issuer, + port: 0, + open_browser: false, + force_state: Some(state), + forced_chatgpt_workspace_id: Some(vec![ + "org-required-a".to_string(), + "org-required-b".to_string(), + ]), + }; + let server = run_login_server(opts)?; + assert!( + server + .auth_url + .contains("allowed_workspace_id=org-required-a") + ); + assert!( + server + .auth_url + .contains("allowed_workspace_id=org-required-b") + ); + + Ok(()) +} + #[tokio::test] async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { skip_if_no_network!(Ok(())); @@ -223,7 +262,7 @@ async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { port: 0, open_browser: false, force_state: Some(state.clone()), - forced_chatgpt_workspace_id: Some("org-required".to_string()), + forced_chatgpt_workspace_id: Some(vec!["org-required".to_string()]), codex_streamlined_login: false, }; let server = run_login_server(opts)?; @@ -241,7 +280,7 @@ async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { assert!(resp.status().is_success()); let body = resp.text().await?; assert!( - body.contains("Login is restricted to workspace id org-required"), + body.contains("Login is restricted to workspace id(s) org-required"), "error body should mention workspace restriction" ); diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index 1f84b289a7..c0b353434a 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -16,7 +16,7 @@ pub(crate) struct LocalChatgptAuth { pub(crate) fn load_local_chatgpt_auth( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, - forced_chatgpt_workspace_id: Option<&str>, + forced_chatgpt_workspace_id: Option<&[String]>, ) -> Result { let auth = load_auth_dot_json(codex_home, auth_credentials_store_mode) .map_err(|err| format!("failed to load local auth: {err}"))? @@ -33,11 +33,14 @@ pub(crate) fn load_local_chatgpt_auth( .account_id .or(tokens.id_token.chatgpt_account_id.clone()) .ok_or_else(|| "local ChatGPT auth is missing chatgpt account id".to_string())?; - if let Some(expected_workspace) = forced_chatgpt_workspace_id - && chatgpt_account_id != expected_workspace + if let Some(expected_workspaces) = forced_chatgpt_workspace_id + && !expected_workspaces + .iter() + .any(|expected_workspace| chatgpt_account_id == *expected_workspace) { return Err(format!( - "local ChatGPT auth must use workspace {expected_workspace}, but found {chatgpt_account_id:?}" + "local ChatGPT auth must use one of workspace(s) {:?}, but found {chatgpt_account_id:?}", + expected_workspaces )); } @@ -122,7 +125,7 @@ mod tests { let auth = load_local_chatgpt_auth( codex_home.path(), AuthCredentialsStoreMode::File, - Some("workspace-1"), + Some(&["workspace-1".to_string()]), ) .expect("chatgpt auth should load"); @@ -186,7 +189,7 @@ mod tests { let auth = load_local_chatgpt_auth( codex_home.path(), AuthCredentialsStoreMode::File, - Some("workspace-1"), + Some(&["workspace-1".to_string(), "workspace-2".to_string()]), ) .expect("managed auth should win"); @@ -202,7 +205,7 @@ mod tests { let auth = load_local_chatgpt_auth( codex_home.path(), AuthCredentialsStoreMode::File, - Some("workspace-1"), + Some(&["workspace-1".to_string()]), ) .expect("chatgpt auth should load"); diff --git a/sdk/python/src/openai_codex/generated/v2_all.py b/sdk/python/src/openai_codex/generated/v2_all.py index 363d87495c..882296e5a8 100644 --- a/sdk/python/src/openai_codex/generated/v2_all.py +++ b/sdk/python/src/openai_codex/generated/v2_all.py @@ -7449,7 +7449,7 @@ class Config(BaseModel): ] = None compact_prompt: str | None = None developer_instructions: str | None = None - forced_chatgpt_workspace_id: str | None = None + forced_chatgpt_workspace_id: list[str] | None = None forced_login_method: ForcedLoginMethod | None = None instructions: str | None = None model: str | None = None