From 270ed782e5ca5dc6f76cd82c62631f4da08a2e37 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 | 2 +- .../app-server-protocol/src/protocol/v1.rs | 2 +- .../app-server-protocol/src/protocol/v2.rs | 2 +- .../app-server/src/codex_message_processor.rs | 9 ++-- codex-rs/cli/src/login.rs | 2 +- codex-rs/config/src/config_toml.rs | 25 +++++++++-- codex-rs/core/config.schema.json | 22 ++++++++- codex-rs/core/src/agent_identity.rs | 21 +++++---- codex-rs/core/src/config/mod.rs | 28 ++++++------ codex-rs/login/src/auth/auth_tests.rs | 37 ++++++++++++--- codex-rs/login/src/auth/manager.rs | 35 +++++++++------ 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/codex_app_server/generated/v2_all.py | 2 +- 19 files changed, 214 insertions(+), 76 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 2cb5fb84df..be3a2445f4 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 @@ -6574,8 +6574,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 4ff918e642..c911b12dfb 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 @@ -3185,8 +3185,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 58f8186266..0047d78680 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json @@ -234,8 +234,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 508fe84e92..1b04d60d28 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts @@ -20,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: ServiceTier | 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 d642e7fab9..8bb3490801 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -198,7 +198,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.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 8ac6651ae9..cb2e4488c6 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -776,7 +776,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/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 7a2ada0765..b83cff4055 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1677,13 +1677,16 @@ impl CodexMessageProcessor { } } - 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) { let error = JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: 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 ), data: None, }; diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index fd0dfee3af..8480fcd477 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -112,7 +112,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 0c8e14c16d..b1036a6088 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -64,6 +64,23 @@ const RESERVED_MODEL_PROVIDER_IDS: [&str; 4] = [ LMSTUDIO_OSS_PROVIDER_ID, ]; +/// Backward-compatible shape for workspace restrictions in config.toml. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema)] +#[serde(untagged)] +pub enum ForcedChatgptWorkspaceIds { + Single(String), + Multiple(Vec), +} + +impl ForcedChatgptWorkspaceIds { + pub fn into_vec(self) -> Vec { + match self { + Self::Single(value) => vec![value], + Self::Multiple(values) => values, + } + } +} + /// Base config deserialized from ~/.codex/config.toml. #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] #[schemars(deny_unknown_fields)] @@ -151,9 +168,9 @@ pub struct ConfigToml { /// Set to an empty string to disable automatic commit attribution. pub commit_attribution: Option, - /// When set, restricts ChatGPT login to a specific workspace identifier. + /// 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)] @@ -413,7 +430,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 b0faea2b56..b253ed4d74 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -731,6 +731,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", @@ -2564,9 +2578,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/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index 3a4c6787f2..260dfe965d 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -454,19 +454,24 @@ impl AgentIdentityBinding { } } - fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { + fn from_auth(auth: &CodexAuth, forced_workspace_id: Option>) -> Option { if !auth.is_chatgpt_auth() { return None; } let token_data = auth.get_token_data().ok()?; - let resolved_account_id = - forced_workspace_id - .filter(|value| !value.is_empty()) - .or(token_data - .account_id - .clone() - .filter(|value| !value.is_empty()))?; + let token_account_id = token_data + .account_id + .clone() + .filter(|value| !value.is_empty()); + let forced_account_id = forced_workspace_id.and_then(|values| { + let mut values = values.into_iter().filter(|value| !value.is_empty()); + match token_account_id.as_ref() { + Some(account_id) => values.find(|value| value == account_id), + None => values.next(), + } + }); + let resolved_account_id = forced_account_id.or(token_account_id)?; Some(Self { binding_id: format!("chatgpt-account-{resolved_account_id}"), diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index a6240fdcb3..54f3802090 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -530,12 +530,11 @@ 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 uses a remote thread /// store at this endpoint instead of the local filesystem/SQLite store. pub experimental_thread_store_endpoint: Option, - /// 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, @@ -634,7 +633,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() } @@ -1946,15 +1945,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 326adea6e4..d4d695d325 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -605,7 +605,7 @@ async fn auth_manager_notifies_when_auth_state_changes() { .expect("auth reload notification should still arrive") .expect("auth state watch should remain open"); - manager.set_forced_chatgpt_workspace_id(Some("workspace-123".to_string())); + 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") @@ -674,7 +674,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(), @@ -754,13 +754,13 @@ 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) .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" @@ -784,7 +784,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; @@ -795,6 +795,31 @@ async fn enforce_login_restrictions_allows_matching_workspace() { ); } +#[tokio::test] +#[serial(codex_api_key)] +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) + .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() { @@ -805,7 +830,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 29733f13da..682fe5e022 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -582,7 +582,7 @@ 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 forced_chatgpt_workspace_id: Option>, } pub fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { @@ -620,7 +620,7 @@ pub fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { } } - if let Some(expected_account_id) = config.forced_chatgpt_workspace_id.as_deref() { + if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { if !auth.is_chatgpt_auth() { return Ok(()); } @@ -640,13 +640,18 @@ pub fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { // 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 != Some(expected_account_id) { + 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( @@ -1197,7 +1202,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: RwLock>, background_agent_task_auth_mode: RwLock, refresh_lock: Semaphore, @@ -1218,8 +1223,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) -> Option { @@ -1512,7 +1517,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 { @@ -1521,7 +1526,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() @@ -1913,13 +1918,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 0c7b810184..0947340e59 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -64,7 +64,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 cli_auth_credentials_store_mode: AuthCredentialsStoreMode, } @@ -73,7 +73,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 { @@ -471,7 +471,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()), @@ -492,8 +492,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() @@ -870,7 +875,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 { @@ -882,10 +887,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 9522f5b0b0..cc203601cc 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -117,7 +117,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()]), }; let server = run_login_server(opts)?; assert!( @@ -198,6 +198,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(())); @@ -217,7 +256,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()]), }; let server = run_login_server(opts)?; assert!( @@ -234,7 +273,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/codex_app_server/generated/v2_all.py b/sdk/python/src/codex_app_server/generated/v2_all.py index 21d4968c33..fa97a61a64 100644 --- a/sdk/python/src/codex_app_server/generated/v2_all.py +++ b/sdk/python/src/codex_app_server/generated/v2_all.py @@ -6281,7 +6281,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