From 02a720525015f3913a7ffa5d46834d5f50868f9e Mon Sep 17 00:00:00 2001 From: rreichel3-oai Date: Thu, 14 May 2026 17:11:36 -0400 Subject: [PATCH] [codex] Support multiple forced ChatGPT workspaces (#18161) ## Summary This change lets `forced_chatgpt_workspace_id` accept multiple workspace IDs instead of a single value. It keeps the existing config key name, adds backward-compatible parsing for a single string in `config.toml`, and normalizes the setting into an allowed workspace list across login enforcement, app-server config surfaces, and local ChatGPT auth helpers. ## Why Workspace-restricted deployments may need to allow more than one ChatGPT workspace without dropping the guardrail entirely. ## Server-side impact Codex's local server and app-server protocol needed changes because they previously assumed a single workspace ID. The local login flow now matches the auth backend interface by sending the allowed workspace list as a single comma-separated `allowed_workspace_id` query parameter. ## Validation This was tested with: - A single workspace config - With multi-workspace configs - With multiple workspaces in the config - The user only being a part of a subset of them All were successful. Automated coverage: - `cargo test -p codex-login` - `cargo test -p codex-app-server-protocol` - `cargo test -p codex-tui local_chatgpt_auth` - `cargo test --locked -p codex-app-server login_account_chatgpt_includes_forced_workspace_allowlist_query_param` --- .../codex_app_server_protocol.schemas.json | 24 +++- .../codex_app_server_protocol.v2.schemas.json | 24 +++- .../schema/json/v2/ConfigReadResponse.json | 24 +++- .../schema/typescript/v2/Config.ts | 3 +- .../v2/ForcedChatgptWorkspaceIds.ts | 8 ++ .../schema/typescript/v2/index.ts | 1 + .../app-server-protocol/src/protocol/v1.rs | 3 +- .../src/protocol/v2/config.rs | 20 ++- .../request_processors/account_processor.rs | 6 +- codex-rs/app-server/tests/suite/v2/account.rs | 111 ++++++++++++---- .../app-server/tests/suite/v2/config_rpc.rs | 81 ++++++++++++ codex-rs/cli/src/login.rs | 2 +- codex-rs/config/src/config_toml.rs | 112 ++++++++++++++++- codex-rs/core/config.schema.json | 22 +++- codex-rs/core/src/config/config_tests.rs | 49 ++++++++ codex-rs/core/src/config/mod.rs | 28 +++-- codex-rs/login/src/auth/auth_tests.rs | 119 ++++++++++++++++-- codex-rs/login/src/auth/manager.rs | 39 +++--- codex-rs/login/src/server.rs | 19 +-- .../login/tests/suite/device_code_login.rs | 13 +- .../login/tests/suite/login_server_e2e.rs | 68 ++++++++-- codex-rs/tui/src/local_chatgpt_auth.rs | 14 +-- 22 files changed, 673 insertions(+), 117 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/ForcedChatgptWorkspaceIds.ts 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 066e1f4870..e996ca5417 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,9 +7126,13 @@ ] }, "forced_chatgpt_workspace_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "$ref": "#/definitions/v2/ForcedChatgptWorkspaceIds" + }, + { + "type": "null" + } ] }, "forced_login_method": { @@ -8697,6 +8701,20 @@ ], "type": "object" }, + "ForcedChatgptWorkspaceIds": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ], + "description": "Backward-compatible API shape for ChatGPT workspace login restrictions." + }, "ForcedLoginMethod": { "enum": [ "chatgpt", 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 95ce7e4aef..7f73f9fe07 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,9 +3515,13 @@ ] }, "forced_chatgpt_workspace_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "$ref": "#/definitions/ForcedChatgptWorkspaceIds" + }, + { + "type": "null" + } ] }, "forced_login_method": { @@ -5086,6 +5090,20 @@ ], "type": "object" }, + "ForcedChatgptWorkspaceIds": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ], + "description": "Backward-compatible API shape for ChatGPT workspace login restrictions." + }, "ForcedLoginMethod": { "enum": [ "chatgpt", 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 cb9176e560..c21d146a4a 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ConfigReadResponse.json @@ -235,9 +235,13 @@ ] }, "forced_chatgpt_workspace_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "$ref": "#/definitions/ForcedChatgptWorkspaceIds" + }, + { + "type": "null" + } ] }, "forced_login_method": { @@ -581,6 +585,20 @@ } ] }, + "ForcedChatgptWorkspaceIds": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ], + "description": "Backward-compatible API shape for ChatGPT workspace login restrictions." + }, "ForcedLoginMethod": { "enum": [ "chatgpt", 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..50d652cea5 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/Config.ts @@ -10,6 +10,7 @@ import type { JsonValue } from "../serde_json/JsonValue"; import type { AnalyticsConfig } from "./AnalyticsConfig"; import type { ApprovalsReviewer } from "./ApprovalsReviewer"; import type { AskForApproval } from "./AskForApproval"; +import type { ForcedChatgptWorkspaceIds } from "./ForcedChatgptWorkspaceIds"; import type { ProfileV2 } from "./ProfileV2"; import type { SandboxMode } from "./SandboxMode"; import type { SandboxWorkspaceWrite } from "./SandboxWorkspaceWrite"; @@ -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: ForcedChatgptWorkspaceIds | 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 }); diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ForcedChatgptWorkspaceIds.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ForcedChatgptWorkspaceIds.ts new file mode 100644 index 0000000000..d0582c8f4f --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ForcedChatgptWorkspaceIds.ts @@ -0,0 +1,8 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +/** + * Backward-compatible API shape for ChatGPT workspace login restrictions. + */ +export type ForcedChatgptWorkspaceIds = string | Array; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index ab6eaefb5a..984154ba04 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -113,6 +113,7 @@ export type { FileSystemPath } from "./FileSystemPath"; export type { FileSystemSandboxEntry } from "./FileSystemSandboxEntry"; export type { FileSystemSpecialPath } from "./FileSystemSpecialPath"; export type { FileUpdateChange } from "./FileUpdateChange"; +export type { ForcedChatgptWorkspaceIds } from "./ForcedChatgptWorkspaceIds"; export type { FsChangedNotification } from "./FsChangedNotification"; export type { FsCopyParams } from "./FsCopyParams"; export type { FsCopyResponse } from "./FsCopyResponse"; diff --git a/codex-rs/app-server-protocol/src/protocol/v1.rs b/codex-rs/app-server-protocol/src/protocol/v1.rs index f8e9afdd12..3c45c20b8f 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -22,6 +22,7 @@ use serde::Serialize; use ts_rs::TS; use crate::protocol::common::AuthMode; +use crate::protocol::v2::ForcedChatgptWorkspaceIds; #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema, TS)] #[serde(rename_all = "camelCase")] @@ -201,7 +202,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 108ce9252f..ca929b02a9 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/config.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/config.rs @@ -225,6 +225,24 @@ pub struct AppsConfig { pub apps: HashMap, } +/// Backward-compatible API shape for ChatGPT workspace login restrictions. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(untagged)] +#[ts(export_to = "v2/")] +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, + } + } +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, ExperimentalApi)] #[serde(rename_all = "snake_case")] #[ts(export_to = "v2/")] @@ -242,7 +260,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..17d78fb20b 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,11 @@ 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.contains(&chatgpt_account_id) { 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) {expected_workspaces:?}, but received {chatgpt_account_id:?}.", ))); } diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 50c365d633..a15b46abd9 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -44,20 +44,30 @@ use std::path::Path; use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; +use url::Url; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; -const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); const LOGIN_ISSUER_ENV_VAR: &str = "CODEX_APP_SERVER_LOGIN_ISSUER"; +const WORKSPACE_ID_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174000"; +const WORKSPACE_ID_SECOND_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174001"; +const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; +const WORKSPACE_ID_EMBEDDED: &str = "123e4567-e89b-42d3-a456-426614174010"; +const WORKSPACE_ID_INITIAL: &str = "123e4567-e89b-42d3-a456-426614174011"; +const WORKSPACE_ID_REFRESHED: &str = "123e4567-e89b-42d3-a456-426614174012"; +const WORKSPACE_ID_DEVICE: &str = "123e4567-e89b-42d3-a456-426614174013"; +const WORKSPACE_ID_STALE: &str = "123e4567-e89b-42d3-a456-426614174014"; // Helper to create a minimal config.toml for the app server #[derive(Default)] struct CreateConfigTomlParams { forced_method: Option, forced_workspace_id: Option, + forced_workspace_ids: Option>, requires_openai_auth: Option, base_url: Option, model_provider_id: Option, @@ -76,6 +86,13 @@ fn create_config_toml(codex_home: &Path, params: CreateConfigTomlParams) -> std: }; let forced_workspace_line = if let Some(ws) = params.forced_workspace_id { format!("forced_chatgpt_workspace_id = \"{ws}\"\n") + } else if let Some(workspaces) = params.forced_workspace_ids { + let workspaces = workspaces + .into_iter() + .map(|workspace_id| format!("\"{workspace_id}\"")) + .collect::>() + .join(", "); + format!("forced_chatgpt_workspace_id = [{workspaces}]\n") } else { String::new() }; @@ -248,7 +265,7 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { &ChatGptIdTokenClaims::new() .email("embedded@example.com") .plan_type("pro") - .chatgpt_account_id("org-embedded"), + .chatgpt_account_id(WORKSPACE_ID_EMBEDDED), )?; let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; @@ -257,7 +274,7 @@ async fn set_auth_token_updates_account_and_notifies() -> Result<()> { let set_id = mcp .send_chatgpt_auth_tokens_login_request( access_token, - "org-embedded".to_string(), + WORKSPACE_ID_EMBEDDED.to_string(), Some("pro".to_string()), ) .await?; @@ -322,7 +339,7 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { &ChatGptIdTokenClaims::new() .email("embedded@example.com") .plan_type("pro") - .chatgpt_account_id("org-embedded"), + .chatgpt_account_id(WORKSPACE_ID_EMBEDDED), )?; let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; @@ -331,7 +348,7 @@ async fn account_read_refresh_token_is_noop_in_external_mode() -> Result<()> { let set_id = mcp .send_chatgpt_auth_tokens_login_request( access_token, - "org-embedded".to_string(), + WORKSPACE_ID_EMBEDDED.to_string(), Some("pro".to_string()), ) .await?; @@ -441,13 +458,13 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") - .chatgpt_account_id("org-initial"), + .chatgpt_account_id(WORKSPACE_ID_INITIAL), )?; let refreshed_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("refreshed@example.com") .plan_type("pro") - .chatgpt_account_id("org-refreshed"), + .chatgpt_account_id(WORKSPACE_ID_REFRESHED), )?; let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; @@ -456,7 +473,7 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { let set_id = mcp .send_chatgpt_auth_tokens_login_request( initial_access_token.clone(), - "org-initial".to_string(), + WORKSPACE_ID_INITIAL.to_string(), Some("pro".to_string()), ) .await?; @@ -499,7 +516,7 @@ async fn external_auth_refreshes_on_unauthorized() -> Result<()> { respond_to_refresh_request( &mut mcp, &refreshed_access_token, - "org-refreshed", + WORKSPACE_ID_REFRESHED, Some("pro"), ) .await?; @@ -553,7 +570,7 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") - .chatgpt_account_id("org-initial"), + .chatgpt_account_id(WORKSPACE_ID_INITIAL), )?; let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; @@ -562,7 +579,7 @@ async fn external_auth_refresh_error_fails_turn() -> Result<()> { let set_id = mcp .send_chatgpt_auth_tokens_login_request( initial_access_token, - "org-initial".to_string(), + WORKSPACE_ID_INITIAL.to_string(), Some("pro".to_string()), ) .await?; @@ -651,7 +668,7 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { create_config_toml( codex_home.path(), CreateConfigTomlParams { - forced_workspace_id: Some("org-expected".to_string()), + forced_workspace_id: Some(WORKSPACE_ID_ALLOWED.to_string()), requires_openai_auth: Some(true), base_url: Some(format!("{}/v1", mock_server.uri())), ..Default::default() @@ -669,13 +686,13 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") - .chatgpt_account_id("org-expected"), + .chatgpt_account_id(WORKSPACE_ID_ALLOWED), )?; let refreshed_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("refreshed@example.com") .plan_type("pro") - .chatgpt_account_id("org-other"), + .chatgpt_account_id(WORKSPACE_ID_DISALLOWED), )?; let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; @@ -684,7 +701,7 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { let set_id = mcp .send_chatgpt_auth_tokens_login_request( initial_access_token, - "org-expected".to_string(), + WORKSPACE_ID_ALLOWED.to_string(), Some("pro".to_string()), ) .await?; @@ -738,7 +755,7 @@ async fn external_auth_refresh_mismatched_workspace_fails_turn() -> Result<()> { request_id, serde_json::to_value(ChatgptAuthTokensRefreshResponse { access_token: refreshed_access_token, - chatgpt_account_id: "org-other".to_string(), + chatgpt_account_id: WORKSPACE_ID_DISALLOWED.to_string(), chatgpt_plan_type: Some("pro".to_string()), })?, ) @@ -790,7 +807,7 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { &ChatGptIdTokenClaims::new() .email("initial@example.com") .plan_type("pro") - .chatgpt_account_id("org-initial"), + .chatgpt_account_id(WORKSPACE_ID_INITIAL), )?; let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; @@ -799,7 +816,7 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { let set_id = mcp .send_chatgpt_auth_tokens_login_request( initial_access_token, - "org-initial".to_string(), + WORKSPACE_ID_INITIAL.to_string(), Some("pro".to_string()), ) .await?; @@ -853,7 +870,7 @@ async fn external_auth_refresh_invalid_access_token_fails_turn() -> Result<()> { request_id, serde_json::to_value(ChatgptAuthTokensRefreshResponse { access_token: "not-a-jwt".to_string(), - chatgpt_account_id: "org-initial".to_string(), + chatgpt_account_id: WORKSPACE_ID_INITIAL.to_string(), chatgpt_plan_type: Some("pro".to_string()), })?, ) @@ -1062,7 +1079,7 @@ async fn login_account_chatgpt_device_code_succeeds_and_notifies() -> Result<()> &ChatGptIdTokenClaims::new() .email("device@example.com") .plan_type("pro") - .chatgpt_account_id("org-device"), + .chatgpt_account_id(WORKSPACE_ID_DEVICE), )?; mock_device_code_oauth_token(&mock_server, &id_token).await; @@ -1378,14 +1395,14 @@ async fn set_auth_token_cancels_active_chatgpt_login() -> Result<()> { &ChatGptIdTokenClaims::new() .email("embedded@example.com") .plan_type("pro") - .chatgpt_account_id("org-embedded"), + .chatgpt_account_id(WORKSPACE_ID_EMBEDDED), )?; // Set an external auth token instead of completing the ChatGPT login flow. // This should cancel the active login attempt. let set_id = mcp .send_chatgpt_auth_tokens_login_request( access_token, - "org-embedded".to_string(), + WORKSPACE_ID_EMBEDDED.to_string(), Some("pro".to_string()), ) .await?; @@ -1428,7 +1445,7 @@ async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result create_config_toml( codex_home.path(), CreateConfigTomlParams { - forced_workspace_id: Some("ws-forced".to_string()), + forced_workspace_id: Some(WORKSPACE_ID_ALLOWED.to_string()), ..Default::default() }, )?; @@ -1448,12 +1465,56 @@ async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result bail!("unexpected login response: {login:?}"); }; assert!( - auth_url.contains("allowed_workspace_id=ws-forced"), + auth_url.contains(&format!("allowed_workspace_id={WORKSPACE_ID_ALLOWED}")), "auth URL should include forced workspace" ); Ok(()) } +#[tokio::test] +// Serialize tests that launch the login server since it binds to a fixed port. +#[serial(login_port)] +async fn login_account_chatgpt_includes_forced_workspace_allowlist_query_param() -> Result<()> { + let codex_home = TempDir::new()?; + create_config_toml( + codex_home.path(), + CreateConfigTomlParams { + forced_workspace_ids: Some(vec![ + WORKSPACE_ID_ALLOWED.to_string(), + WORKSPACE_ID_SECOND_ALLOWED.to_string(), + ]), + ..Default::default() + }, + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp.send_login_account_chatgpt_request().await?; + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + + let login: LoginAccountResponse = to_response(resp)?; + let LoginAccountResponse::Chatgpt { auth_url, .. } = login else { + bail!("unexpected login response: {login:?}"); + }; + let auth_url = Url::parse(&auth_url)?; + let allowed_workspace_ids = auth_url + .query_pairs() + .filter_map(|(key, value)| (key == "allowed_workspace_id").then(|| value.into_owned())) + .collect::>(); + assert_eq!( + allowed_workspace_ids, + vec![format!( + "{WORKSPACE_ID_ALLOWED},{WORKSPACE_ID_SECOND_ALLOWED}" + )] + ); + Ok(()) +} + #[tokio::test] async fn get_account_no_auth() -> Result<()> { let codex_home = TempDir::new()?; @@ -1662,7 +1723,7 @@ async fn get_account_omits_chatgpt_after_permanent_refresh_failure() -> Result<( codex_home.path(), ChatGptAuthFixture::new("stale-access-token") .refresh_token("stale-refresh-token") - .account_id("acct_123") + .account_id(WORKSPACE_ID_STALE) .email("user@example.com") .plan_type("pro") .last_refresh(Some(Utc::now() - ChronoDuration::days(9))), diff --git a/codex-rs/app-server/tests/suite/v2/config_rpc.rs b/codex-rs/app-server/tests/suite/v2/config_rpc.rs index 2cbf476f1a..f95b9fc6dc 100644 --- a/codex-rs/app-server/tests/suite/v2/config_rpc.rs +++ b/codex-rs/app-server/tests/suite/v2/config_rpc.rs @@ -14,6 +14,7 @@ use codex_app_server_protocol::ConfigReadParams; use codex_app_server_protocol::ConfigReadResponse; use codex_app_server_protocol::ConfigValueWriteParams; use codex_app_server_protocol::ConfigWriteResponse; +use codex_app_server_protocol::ForcedChatgptWorkspaceIds; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::MergeStrategy; @@ -164,6 +165,86 @@ allowed_domains = ["example.com"] Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn config_read_accepts_legacy_forced_chatgpt_workspace_id() -> Result<()> { + const WORKSPACE_ID: &str = "123e4567-e89b-42d3-a456-426614174000"; + + let codex_home = TempDir::new()?; + write_config( + &codex_home, + &format!( + r#" +forced_chatgpt_workspace_id = "{WORKSPACE_ID}" +"# + ), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_config_read_request(ConfigReadParams { + include_layers: false, + cwd: None, + }) + .await?; + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ConfigReadResponse { config, .. } = to_response(resp)?; + + assert_eq!( + config.forced_chatgpt_workspace_id, + Some(ForcedChatgptWorkspaceIds::Single(WORKSPACE_ID.to_string())) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn config_read_accepts_forced_chatgpt_workspace_id_list() -> Result<()> { + const WORKSPACE_ID_A: &str = "123e4567-e89b-42d3-a456-426614174000"; + const WORKSPACE_ID_B: &str = "123e4567-e89b-42d3-a456-426614174001"; + + let codex_home = TempDir::new()?; + write_config( + &codex_home, + &format!( + r#" +forced_chatgpt_workspace_id = ["{WORKSPACE_ID_A}", "{WORKSPACE_ID_B}"] +"# + ), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_config_read_request(ConfigReadParams { + include_layers: false, + cwd: None, + }) + .await?; + let resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ConfigReadResponse { config, .. } = to_response(resp)?; + + assert_eq!( + config.forced_chatgpt_workspace_id, + Some(ForcedChatgptWorkspaceIds::Multiple(vec![ + WORKSPACE_ID_A.to_string(), + WORKSPACE_ID_B.to_string(), + ])) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn config_read_includes_nested_web_search_tool_config() -> Result<()> { let codex_home = TempDir::new()?; 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 591b259676..055df69529 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -27,6 +27,7 @@ use crate::types::ToolSuggestConfig; use crate::types::Tui; use crate::types::UriBasedFileOpener; use crate::types::WindowsToml; +use codex_app_server_protocol::ForcedChatgptWorkspaceIds as ApiForcedChatgptWorkspaceIds; use codex_app_server_protocol::Tools; use codex_app_server_protocol::UserSavedConfig; use codex_features::FeaturesToml; @@ -56,6 +57,7 @@ use schemars::JsonSchema; use serde::Deserialize; use serde::Deserializer; use serde::Serialize; +use serde::de::Error as SerdeError; const RESERVED_MODEL_PROVIDER_IDS: [&str; 4] = [ AMAZON_BEDROCK_PROVIDER_ID, @@ -86,6 +88,55 @@ const fn default_hide_agent_reasoning() -> Option { Some(false) } +/// Backward-compatible shape for ChatGPT workspace login restrictions in config.toml. +#[derive(Serialize, 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, + } + } + + pub fn into_api(self) -> ApiForcedChatgptWorkspaceIds { + match self { + Self::Single(value) => ApiForcedChatgptWorkspaceIds::Single(value), + Self::Multiple(values) => ApiForcedChatgptWorkspaceIds::Multiple(values), + } + } +} + +impl<'de> Deserialize<'de> for ForcedChatgptWorkspaceIds { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + #[derive(Deserialize)] + #[serde(untagged)] + enum Repr { + Single(String), + Multiple(Vec), + } + + match Repr::deserialize(deserializer)? { + Repr::Single(value) if value.contains(',') => Err(D::Error::custom( + "forced_chatgpt_workspace_id must be a single workspace ID string or a TOML list \ +of strings; comma-separated strings are not supported. Use \ +`forced_chatgpt_workspace_id = [\"123e4567-e89b-42d3-a456-426614174000\", \ +\"123e4567-e89b-42d3-a456-426614174001\"]` instead.", + )), + Repr::Single(value) => Ok(Self::Single(value)), + Repr::Multiple(values) => Ok(Self::Multiple(values)), + } + } +} + /// Base config deserialized from ~/.codex/config.toml. #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, JsonSchema)] #[schemars(deny_unknown_fields)] @@ -177,9 +228,9 @@ pub struct ConfigToml { /// Compact prompt used for history compaction. pub compact_prompt: 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)] @@ -507,7 +558,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_api), forced_login_method: config_toml.forced_login_method, model: config_toml.model, model_reasoning_effort: config_toml.model_reasoning_effort, @@ -950,3 +1003,56 @@ pub fn validate_oss_provider(provider: &str) -> std::io::Result<()> { )), } } + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + const WORKSPACE_ID_A: &str = "123e4567-e89b-42d3-a456-426614174000"; + const WORKSPACE_ID_B: &str = "123e4567-e89b-42d3-a456-426614174001"; + + #[test] + fn forced_chatgpt_workspace_id_accepts_single_string() { + let config: ConfigToml = toml::from_str(&format!( + r#"forced_chatgpt_workspace_id = "{WORKSPACE_ID_A}""# + )) + .expect("single workspace id should deserialize"); + + assert_eq!( + config + .forced_chatgpt_workspace_id + .expect("workspace id should be set") + .into_vec(), + vec![WORKSPACE_ID_A.to_string()] + ); + } + + #[test] + fn forced_chatgpt_workspace_id_accepts_string_list() { + let config: ConfigToml = toml::from_str(&format!( + r#"forced_chatgpt_workspace_id = ["{WORKSPACE_ID_A}", "{WORKSPACE_ID_B}"]"# + )) + .expect("workspace id list should deserialize"); + + assert_eq!( + config + .forced_chatgpt_workspace_id + .expect("workspace ids should be set") + .into_vec(), + vec![WORKSPACE_ID_A.to_string(), WORKSPACE_ID_B.to_string()] + ); + } + + #[test] + fn forced_chatgpt_workspace_id_rejects_comma_separated_string() { + let err = toml::from_str::(&format!( + r#"forced_chatgpt_workspace_id = "{WORKSPACE_ID_A},{WORKSPACE_ID_B}""# + )) + .expect_err("comma-separated string should be rejected"); + + let message = err.to_string(); + assert!(message.contains("TOML list of strings")); + assert!(message.contains("comma-separated strings are not supported")); + } +} diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 9fc95349d2..84ce98ffc5 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -852,6 +852,20 @@ }, "type": "object" }, + "ForcedChatgptWorkspaceIds": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ], + "description": "Backward-compatible shape for ChatGPT workspace login restrictions in config.toml." + }, "ForcedLoginMethod": { "enum": [ "chatgpt", @@ -4383,9 +4397,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/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index ff28e50436..c2ef25e6cb 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -2725,6 +2725,55 @@ async fn runtime_config_resolves_terminal_resize_reflow_defaults_and_overrides() ); } +#[tokio::test] +async fn forced_chatgpt_workspace_id_empty_values_disable_runtime_restriction() +-> std::io::Result<()> { + let cases: Vec<(&str, &str, Option>)> = vec![ + ("unset", "", None), + ("empty string", r#"forced_chatgpt_workspace_id = """#, None), + ( + "whitespace string", + r#"forced_chatgpt_workspace_id = " ""#, + None, + ), + ("empty list", r#"forced_chatgpt_workspace_id = []"#, None), + ( + "blank list entries", + r#"forced_chatgpt_workspace_id = ["", " "]"#, + None, + ), + ( + "mixed list entries", + r#"forced_chatgpt_workspace_id = ["", " 123e4567-e89b-42d3-a456-426614174000 ", "123e4567-e89b-42d3-a456-426614174001"]"#, + Some(vec![ + "123e4567-e89b-42d3-a456-426614174000", + "123e4567-e89b-42d3-a456-426614174001", + ]), + ), + ]; + + for (name, toml, expected) in cases { + let cfg_toml: ConfigToml = toml::from_str(toml) + .unwrap_or_else(|err| panic!("{name} should parse forced_chatgpt_workspace_id: {err}")); + let config = Config::load_from_base_config_with_overrides( + cfg_toml, + ConfigOverrides::default(), + tempdir().expect("tempdir").abs(), + ) + .await?; + + let expected = expected.map(|values| { + values + .into_iter() + .map(ToString::to_string) + .collect::>() + }); + assert_eq!(config.forced_chatgpt_workspace_id, expected, "{name}"); + } + + Ok(()) +} + #[tokio::test] async fn legacy_remote_thread_store_endpoint_is_rejected() { let cfg: ConfigToml = diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index cc5203251c..16d8a57b30 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -752,15 +752,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, @@ -891,7 +890,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() } @@ -2849,15 +2848,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..63ab3e0c42 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -22,6 +22,10 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; +const WORKSPACE_ID_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174000"; +const WORKSPACE_ID_SECOND_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174001"; +const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; + #[tokio::test] async fn refresh_without_id_token() { let codex_home = tempdir().unwrap(); @@ -87,7 +91,7 @@ fn login_with_api_key_overwrites_existing_auth_json() { async fn login_with_access_token_writes_only_token() { let dir = tempdir().unwrap(); let auth_path = dir.path().join("auth.json"); - let record = agent_identity_record("account-123"); + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); let agent_identity = signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); let server = MockServer::start().await; @@ -145,7 +149,7 @@ async fn login_with_access_token_rejects_invalid_jwt() { #[tokio::test] async fn login_with_access_token_rejects_unsigned_jwt() { let dir = tempdir().unwrap(); - let record = agent_identity_record("account-123"); + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) @@ -329,7 +333,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { AuthFileParams { openai_api_key: None, chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: Some("org_mine".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), }, codex_home.path(), ) @@ -654,7 +658,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(), @@ -712,7 +716,7 @@ fn remove_access_token_env_var() -> EnvVarGuard { #[serial(codex_auth_env)] async fn load_auth_reads_access_token_from_env() { let codex_home = tempdir().unwrap(); - let expected_record = agent_identity_record("account-123"); + let expected_record = agent_identity_record(WORKSPACE_ID_ALLOWED); let agent_identity = signed_agent_identity_jwt(&expected_record, json!(expected_record.plan_type)) .expect("signed agent identity"); @@ -762,7 +766,7 @@ async fn load_auth_reads_access_token_from_env() { #[serial(codex_auth_env)] async fn load_auth_keeps_codex_api_key_env_precedence() { let codex_home = tempdir().unwrap(); - let record = agent_identity_record("account-123"); + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity"); let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity); let _api_key_guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env"); @@ -814,7 +818,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { AuthFileParams { openai_api_key: None, chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: Some("org_another_org".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_DISALLOWED.to_string()), }, codex_home.path(), ) @@ -823,14 +827,17 @@ 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![WORKSPACE_ID_ALLOWED.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(&format!("workspace(s) {WORKSPACE_ID_ALLOWED}")) + ); assert!( !codex_home.path().join("auth.json").exists(), "auth.json should be removed on mismatch" @@ -846,7 +853,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() { AuthFileParams { openai_api_key: None, chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: Some("org_mine".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), }, codex_home.path(), ) @@ -855,7 +862,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![WORKSPACE_ID_ALLOWED.to_string()]), ) .await; @@ -870,6 +877,94 @@ 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(WORKSPACE_ID_ALLOWED.to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let config = build_config( + codex_home.path(), + /*forced_login_method*/ None, + Some(vec![ + WORKSPACE_ID_SECOND_ALLOWED.to_string(), + WORKSPACE_ID_ALLOWED.to_string(), + ]), + ) + .await; + + super::enforce_login_restrictions(&config) + .await + .expect("any matching workspace in the allowed list should succeed"); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismatch() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + let record = agent_identity_record(WORKSPACE_ID_DISALLOWED); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/agent-identities/jwks")) + .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/agent-runtime-id/task/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "task_id": "task-123", + }))) + .expect(1) + .mount(&server) + .await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + let _authapi_guard = + EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url); + save_auth( + codex_home.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::AgentIdentity), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: Some(agent_identity), + }, + AuthCredentialsStoreMode::File, + ) + .expect("seed agent identity auth"); + + let config = AuthConfig { + codex_home: codex_home.path().to_path_buf(), + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + forced_login_method: None, + forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), + chatgpt_base_url: Some(chatgpt_base_url), + }; + + let err = super::enforce_login_restrictions(&config) + .await + .expect_err("expected workspace mismatch to error"); + assert!(err.to_string().contains(&format!( + "current credentials belong to {WORKSPACE_ID_DISALLOWED}" + ))); + assert!( + !codex_home.path().join("auth.json").exists(), + "auth.json should be removed on mismatch" + ); + server.verify().await; +} + +#[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 +975,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![WORKSPACE_ID_ALLOWED.to_string()]), ) .await; diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 45bd55302b..a2e4e8e0d8 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,9 +653,8 @@ 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 { + if let Some(expected_account_ids) = config.forced_chatgpt_workspace_id.as_deref() { + let chatgpt_account_id = match &auth { CodexAuth::ApiKey(_) => return Ok(()), CodexAuth::AgentIdentity(_) => auth.get_account_id(), CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { @@ -674,13 +673,21 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< token_data.id_token.chatgpt_account_id } }; - if chatgpt_account_id.as_deref() != Some(expected_account_id) { + + // workspace is the external identifier for account id. + let chatgpt_account_id = 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 +1254,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 +1273,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 +1555,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 +1563,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 +1843,13 @@ 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.contains(&chatgpt_metadata.account_id) { 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..b72bc946f2 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,8 @@ 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.push(("allowed_workspace_id".to_string(), workspace_ids.join(","))); } let qs = query .into_iter() @@ -928,7 +928,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 +940,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..e575e6b8c7 100644 --- a/codex-rs/login/tests/suite/device_code_login.rs +++ b/codex-rs/login/tests/suite/device_code_login.rs @@ -23,6 +23,9 @@ use core_test_support::skip_if_no_network; // ---------- Small helpers ---------- +const WORKSPACE_ID_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174000"; +const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; + fn make_jwt(payload: serde_json::Value) -> String { let header = json!({ "alg": "none", "typ": "JWT" }); let header_b64 = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).unwrap()); @@ -131,7 +134,7 @@ async fn device_code_login_integration_succeeds() -> anyhow::Result<()> { let jwt = make_jwt(json!({ "https://api.openai.com/auth": { - "chatgpt_account_id": "acct_321" + "chatgpt_account_id": WORKSPACE_ID_ALLOWED } })); @@ -152,7 +155,7 @@ async fn device_code_login_integration_succeeds() -> anyhow::Result<()> { assert_eq!(tokens.access_token, "access-token-123"); assert_eq!(tokens.refresh_token, "refresh-token-123"); assert_eq!(tokens.id_token.raw_jwt, jwt); - assert_eq!(tokens.account_id.as_deref(), Some("acct_321")); + assert_eq!(tokens.account_id.as_deref(), Some(WORKSPACE_ID_ALLOWED)); Ok(()) } @@ -174,8 +177,8 @@ async fn device_code_login_rejects_workspace_mismatch() -> anyhow::Result<()> { let jwt = make_jwt(json!({ "https://api.openai.com/auth": { - "chatgpt_account_id": "acct_321", - "organization_id": "org-actual" + "chatgpt_account_id": WORKSPACE_ID_DISALLOWED, + "organization_id": WORKSPACE_ID_DISALLOWED } })); @@ -183,7 +186,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![WORKSPACE_ID_ALLOWED.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..bc7bb401e8 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -12,10 +12,15 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_login::ServerOptions; use codex_login::run_login_server; use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; use tempfile::tempdir; +use url::Url; const DEFAULT_LOGIN_PORT: u16 = 1455; const FALLBACK_LOGIN_PORT: u16 = 1457; +const WORKSPACE_ID_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174000"; +const WORKSPACE_ID_SECOND_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174001"; +const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; // See spawn.rs for details @@ -121,7 +126,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)?; @@ -165,7 +170,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("org-123"); + let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -204,11 +209,52 @@ async fn creates_missing_codex_home_dir() -> Result<()> { Ok(()) } +#[tokio::test] +async fn login_server_includes_forced_workspaces_as_one_query_param() -> Result<()> { + skip_if_no_network!(Ok(())); + + let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); + 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![ + WORKSPACE_ID_ALLOWED.to_string(), + WORKSPACE_ID_SECOND_ALLOWED.to_string(), + ]), + codex_streamlined_login: false, + }; + let server = run_login_server(opts)?; + let auth_url = Url::parse(&server.auth_url)?; + let allowed_workspace_ids = auth_url + .query_pairs() + .filter_map(|(key, value)| (key == "allowed_workspace_id").then(|| value.into_owned())) + .collect::>(); + assert_eq!( + allowed_workspace_ids, + vec![format!( + "{WORKSPACE_ID_ALLOWED},{WORKSPACE_ID_SECOND_ALLOWED}" + )] + ); + + 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_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_DISALLOWED); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -223,14 +269,14 @@ 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![WORKSPACE_ID_ALLOWED.to_string()]), codex_streamlined_login: false, }; let server = run_login_server(opts)?; assert!( server .auth_url - .contains("allowed_workspace_id=org-required"), + .contains(&format!("allowed_workspace_id={WORKSPACE_ID_ALLOWED}")), "auth URL should include forced workspace parameter" ); let login_port = server.actual_port; @@ -241,7 +287,9 @@ 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(&format!( + "Login is restricted to workspace id(s) {WORKSPACE_ID_ALLOWED}" + )), "error body should mention workspace restriction" ); @@ -266,7 +314,7 @@ async fn forced_chatgpt_workspace_id_mismatch_blocks_login() -> Result<()> { async fn oauth_access_denied_missing_entitlement_blocks_login_with_clear_error() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer("org-123"); + let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -334,7 +382,7 @@ async fn oauth_access_denied_missing_entitlement_blocks_login_with_clear_error() async fn oauth_access_denied_unknown_reason_uses_generic_error_page() -> Result<()> { skip_if_no_network!(Ok(())); - let (issuer_addr, _issuer_handle) = start_mock_issuer("org-123"); + let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -442,7 +490,7 @@ async fn falls_back_to_registered_fallback_port_when_default_port_is_in_use() -> }) }; - let (issuer_addr, _issuer_handle) = start_mock_issuer("org-123"); + let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let tmp = tempdir()?; @@ -480,7 +528,7 @@ async fn falls_back_to_registered_fallback_port_when_default_port_is_in_use() -> 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("org-123"); + let (issuer_addr, _issuer_handle) = start_mock_issuer(WORKSPACE_ID_ALLOWED); let issuer = format!("http://{}:{}", issuer_addr.ip(), issuer_addr.port()); let first_tmp = tempdir()?; diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index 1f84b289a7..270d64f5f7 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,11 @@ 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.contains(&chatgpt_account_id) { 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) {expected_workspaces:?}, but found {chatgpt_account_id:?}", )); } @@ -122,7 +122,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 +186,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 +202,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");