diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 50c365d633..c7115b1953 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -44,6 +44,7 @@ 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; @@ -58,6 +59,7 @@ const LOGIN_ISSUER_ENV_VAR: &str = "CODEX_APP_SERVER_LOGIN_ISSUER"; 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 +78,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() }; @@ -1454,6 +1463,45 @@ async fn login_account_chatgpt_includes_forced_workspace_query_param() -> Result 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!["ws-forced-a".to_string(), "ws-forced-b".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!["ws-forced-a,ws-forced-b".to_string()] + ); + Ok(()) +} + #[tokio::test] async fn get_account_no_auth() -> Result<()> { let codex_home = TempDir::new()?;