mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
af8a9d2d2b
Stacked on #16508. This removes the temporary `codex-core` / `codex-login` re-export shims from the ownership split and rewrites callsites to import directly from `codex-model-provider-info`, `codex-models-manager`, `codex-api`, `codex-protocol`, `codex-feedback`, and `codex-response-debug-context`. No behavior change intended; this is the mechanical import cleanup layer split out from the ownership move. --------- Co-authored-by: Codex <noreply@openai.com>
827 lines
25 KiB
Rust
827 lines
25 KiB
Rust
use super::*;
|
|
use crate::auth::storage::FileAuthStorage;
|
|
use crate::auth::storage::get_auth_file;
|
|
use crate::token_data::IdTokenInfo;
|
|
use codex_app_server_protocol::AuthMode;
|
|
use codex_protocol::account::PlanType as AccountPlanType;
|
|
use codex_protocol::auth::KnownPlan as InternalKnownPlan;
|
|
use codex_protocol::auth::PlanType as InternalPlanType;
|
|
|
|
use base64::Engine;
|
|
use codex_protocol::config_types::ForcedLoginMethod;
|
|
use codex_protocol::config_types::ModelProviderAuthInfo;
|
|
use pretty_assertions::assert_eq;
|
|
use serde::Serialize;
|
|
use serde_json::json;
|
|
use std::sync::Arc;
|
|
use tempfile::TempDir;
|
|
use tempfile::tempdir;
|
|
|
|
#[tokio::test]
|
|
async fn refresh_without_id_token() {
|
|
let codex_home = tempdir().unwrap();
|
|
let fake_jwt = write_auth_file(
|
|
AuthFileParams {
|
|
openai_api_key: None,
|
|
chatgpt_plan_type: Some("pro".to_string()),
|
|
chatgpt_account_id: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let storage = create_auth_storage(
|
|
codex_home.path().to_path_buf(),
|
|
AuthCredentialsStoreMode::File,
|
|
);
|
|
let updated = super::persist_tokens(
|
|
&storage,
|
|
/*id_token*/ None,
|
|
Some("new-access-token".to_string()),
|
|
Some("new-refresh-token".to_string()),
|
|
)
|
|
.expect("update_tokens should succeed");
|
|
|
|
let tokens = updated.tokens.expect("tokens should exist");
|
|
assert_eq!(tokens.id_token.raw_jwt, fake_jwt);
|
|
assert_eq!(tokens.access_token, "new-access-token");
|
|
assert_eq!(tokens.refresh_token, "new-refresh-token");
|
|
}
|
|
|
|
#[test]
|
|
fn login_with_api_key_overwrites_existing_auth_json() {
|
|
let dir = tempdir().unwrap();
|
|
let auth_path = dir.path().join("auth.json");
|
|
let stale_auth = json!({
|
|
"OPENAI_API_KEY": "sk-old",
|
|
"tokens": {
|
|
"id_token": "stale.header.payload",
|
|
"access_token": "stale-access",
|
|
"refresh_token": "stale-refresh",
|
|
"account_id": "stale-acc"
|
|
}
|
|
});
|
|
std::fs::write(
|
|
&auth_path,
|
|
serde_json::to_string_pretty(&stale_auth).unwrap(),
|
|
)
|
|
.unwrap();
|
|
|
|
super::login_with_api_key(dir.path(), "sk-new", AuthCredentialsStoreMode::File)
|
|
.expect("login_with_api_key should succeed");
|
|
|
|
let storage = FileAuthStorage::new(dir.path().to_path_buf());
|
|
let auth = storage
|
|
.try_read_auth_json(&auth_path)
|
|
.expect("auth.json should parse");
|
|
assert_eq!(auth.openai_api_key.as_deref(), Some("sk-new"));
|
|
assert!(auth.tokens.is_none(), "tokens should be cleared");
|
|
}
|
|
|
|
#[test]
|
|
fn missing_auth_json_returns_none() {
|
|
let dir = tempdir().unwrap();
|
|
let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File)
|
|
.expect("call should succeed");
|
|
assert_eq!(auth, None);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial(codex_api_key)]
|
|
async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
|
|
let codex_home = tempdir().unwrap();
|
|
let fake_jwt = write_auth_file(
|
|
AuthFileParams {
|
|
openai_api_key: None,
|
|
chatgpt_plan_type: Some("pro".to_string()),
|
|
chatgpt_account_id: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.unwrap()
|
|
.unwrap();
|
|
assert_eq!(None, auth.api_key());
|
|
assert_eq!(AuthMode::Chatgpt, auth.auth_mode());
|
|
assert_eq!(auth.get_chatgpt_user_id().as_deref(), Some("user-12345"));
|
|
|
|
let auth_dot_json = auth
|
|
.get_current_auth_json()
|
|
.expect("AuthDotJson should exist");
|
|
let last_refresh = auth_dot_json
|
|
.last_refresh
|
|
.expect("last_refresh should be recorded");
|
|
|
|
assert_eq!(
|
|
AuthDotJson {
|
|
auth_mode: None,
|
|
openai_api_key: None,
|
|
tokens: Some(TokenData {
|
|
id_token: IdTokenInfo {
|
|
email: Some("user@example.com".to_string()),
|
|
chatgpt_plan_type: Some(InternalPlanType::Known(InternalKnownPlan::Pro)),
|
|
chatgpt_user_id: Some("user-12345".to_string()),
|
|
chatgpt_account_id: None,
|
|
raw_jwt: fake_jwt,
|
|
},
|
|
access_token: "test-access-token".to_string(),
|
|
refresh_token: "test-refresh-token".to_string(),
|
|
account_id: None,
|
|
}),
|
|
last_refresh: Some(last_refresh),
|
|
},
|
|
auth_dot_json
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial(codex_api_key)]
|
|
async fn loads_api_key_from_auth_json() {
|
|
let dir = tempdir().unwrap();
|
|
let auth_file = dir.path().join("auth.json");
|
|
std::fs::write(
|
|
auth_file,
|
|
r#"{"OPENAI_API_KEY":"sk-test-key","tokens":null,"last_refresh":null}"#,
|
|
)
|
|
.unwrap();
|
|
|
|
let auth = super::load_auth(
|
|
dir.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.unwrap()
|
|
.unwrap();
|
|
assert_eq!(auth.auth_mode(), AuthMode::ApiKey);
|
|
assert_eq!(auth.api_key(), Some("sk-test-key"));
|
|
|
|
assert!(auth.get_token_data().is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn logout_removes_auth_file() -> Result<(), std::io::Error> {
|
|
let dir = tempdir()?;
|
|
let auth_dot_json = AuthDotJson {
|
|
auth_mode: Some(ApiAuthMode::ApiKey),
|
|
openai_api_key: Some("sk-test-key".to_string()),
|
|
tokens: None,
|
|
last_refresh: None,
|
|
};
|
|
super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?;
|
|
let auth_file = get_auth_file(dir.path());
|
|
assert!(auth_file.exists());
|
|
assert!(logout(dir.path(), AuthCredentialsStoreMode::File)?);
|
|
assert!(!auth_file.exists());
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn unauthorized_recovery_reports_mode_and_step_names() {
|
|
let dir = tempdir().unwrap();
|
|
let manager = AuthManager::shared(
|
|
dir.path().to_path_buf(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
);
|
|
let managed = UnauthorizedRecovery {
|
|
manager: Arc::clone(&manager),
|
|
step: UnauthorizedRecoveryStep::Reload,
|
|
expected_account_id: None,
|
|
mode: UnauthorizedRecoveryMode::Managed,
|
|
};
|
|
assert_eq!(managed.mode_name(), "managed");
|
|
assert_eq!(managed.step_name(), "reload");
|
|
|
|
let external = UnauthorizedRecovery {
|
|
manager,
|
|
step: UnauthorizedRecoveryStep::ExternalRefresh,
|
|
expected_account_id: None,
|
|
mode: UnauthorizedRecoveryMode::External,
|
|
};
|
|
assert_eq!(external.mode_name(), "external");
|
|
assert_eq!(external.step_name(), "external_refresh");
|
|
}
|
|
|
|
#[test]
|
|
fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
|
|
let codex_home = tempdir().unwrap();
|
|
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 auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("load auth")
|
|
.expect("auth available");
|
|
let mut updated_auth_dot_json = auth
|
|
.get_current_auth_json()
|
|
.expect("AuthDotJson should exist");
|
|
let updated_tokens = updated_auth_dot_json
|
|
.tokens
|
|
.as_mut()
|
|
.expect("tokens should exist");
|
|
updated_tokens.access_token = "new-access-token".to_string();
|
|
updated_tokens.refresh_token = "new-refresh-token".to_string();
|
|
let updated_auth = CodexAuth::from_auth_dot_json(
|
|
codex_home.path(),
|
|
updated_auth_dot_json,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("updated auth should parse");
|
|
|
|
let manager = AuthManager::from_auth_for_testing(auth.clone());
|
|
let error = RefreshTokenFailedError::new(
|
|
RefreshTokenFailedReason::Exhausted,
|
|
"refresh token already used",
|
|
);
|
|
manager.record_permanent_refresh_failure_if_unchanged(&auth, &error);
|
|
|
|
assert_eq!(manager.refresh_failure_for_auth(&auth), Some(error));
|
|
assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None);
|
|
}
|
|
|
|
#[test]
|
|
fn external_auth_tokens_without_chatgpt_metadata_cannot_seed_chatgpt_auth() {
|
|
let err = AuthDotJson::from_external_tokens(&ExternalAuthTokens::access_token_only(
|
|
"test-access-token",
|
|
))
|
|
.expect_err("bearer-only external auth should not seed ChatGPT auth");
|
|
|
|
assert_eq!(
|
|
err.to_string(),
|
|
"external auth tokens are missing ChatGPT metadata"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn external_bearer_only_auth_manager_uses_cached_provider_token() {
|
|
let script = ProviderAuthScript::new(&["provider-token", "next-token"]).unwrap();
|
|
let manager = AuthManager::external_bearer_only(script.auth_config());
|
|
|
|
let first = manager
|
|
.auth()
|
|
.await
|
|
.and_then(|auth| auth.api_key().map(str::to_string));
|
|
let second = manager
|
|
.auth()
|
|
.await
|
|
.and_then(|auth| auth.api_key().map(str::to_string));
|
|
|
|
assert_eq!(first.as_deref(), Some("provider-token"));
|
|
assert_eq!(second.as_deref(), Some("provider-token"));
|
|
assert_eq!(manager.auth_mode(), Some(AuthMode::ApiKey));
|
|
assert_eq!(manager.get_api_auth_mode(), Some(ApiAuthMode::ApiKey));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn external_bearer_only_auth_manager_disables_auto_refresh_when_interval_is_zero() {
|
|
let script = ProviderAuthScript::new(&["provider-token", "next-token"]).unwrap();
|
|
let mut auth_config = script.auth_config();
|
|
auth_config.refresh_interval_ms = 0;
|
|
let manager = AuthManager::external_bearer_only(auth_config);
|
|
|
|
let first = manager
|
|
.auth()
|
|
.await
|
|
.and_then(|auth| auth.api_key().map(str::to_string));
|
|
let second = manager
|
|
.auth()
|
|
.await
|
|
.and_then(|auth| auth.api_key().map(str::to_string));
|
|
|
|
assert_eq!(first.as_deref(), Some("provider-token"));
|
|
assert_eq!(second.as_deref(), Some("provider-token"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn external_bearer_only_auth_manager_returns_none_when_command_fails() {
|
|
let script = ProviderAuthScript::new_failing().unwrap();
|
|
let manager = AuthManager::external_bearer_only(script.auth_config());
|
|
|
|
assert_eq!(manager.auth().await, None);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn unauthorized_recovery_uses_external_refresh_for_bearer_manager() {
|
|
let script = ProviderAuthScript::new(&["provider-token", "refreshed-provider-token"]).unwrap();
|
|
let mut auth_config = script.auth_config();
|
|
auth_config.refresh_interval_ms = 0;
|
|
let manager = AuthManager::external_bearer_only(auth_config);
|
|
let initial_token = manager
|
|
.auth()
|
|
.await
|
|
.and_then(|auth| auth.api_key().map(str::to_string));
|
|
let mut recovery = manager.unauthorized_recovery();
|
|
|
|
assert!(recovery.has_next());
|
|
assert_eq!(recovery.mode_name(), "external");
|
|
assert_eq!(recovery.step_name(), "external_refresh");
|
|
|
|
let result = recovery
|
|
.next()
|
|
.await
|
|
.expect("external refresh should succeed");
|
|
|
|
assert_eq!(result.auth_state_changed(), Some(true));
|
|
let refreshed_token = manager
|
|
.auth()
|
|
.await
|
|
.and_then(|auth| auth.api_key().map(str::to_string));
|
|
assert_eq!(initial_token.as_deref(), Some("provider-token"));
|
|
assert_eq!(refreshed_token.as_deref(), Some("refreshed-provider-token"));
|
|
}
|
|
|
|
struct ProviderAuthScript {
|
|
tempdir: TempDir,
|
|
command: String,
|
|
args: Vec<String>,
|
|
}
|
|
|
|
impl ProviderAuthScript {
|
|
fn new(tokens: &[&str]) -> std::io::Result<Self> {
|
|
let tempdir = tempfile::tempdir()?;
|
|
let token_file = tempdir.path().join("tokens.txt");
|
|
// `cmd.exe`'s `set /p` treats LF-only input as one line, so use CRLF on Windows.
|
|
let token_line_ending = if cfg!(windows) { "\r\n" } else { "\n" };
|
|
let mut token_file_contents = String::new();
|
|
for token in tokens {
|
|
token_file_contents.push_str(token);
|
|
token_file_contents.push_str(token_line_ending);
|
|
}
|
|
std::fs::write(&token_file, token_file_contents)?;
|
|
|
|
#[cfg(unix)]
|
|
let (command, args) = {
|
|
let script_path = tempdir.path().join("print-token.sh");
|
|
std::fs::write(
|
|
&script_path,
|
|
r#"#!/bin/sh
|
|
first_line=$(sed -n '1p' tokens.txt)
|
|
printf '%s\n' "$first_line"
|
|
tail -n +2 tokens.txt > tokens.next
|
|
mv tokens.next tokens.txt
|
|
"#,
|
|
)?;
|
|
let mut permissions = std::fs::metadata(&script_path)?.permissions();
|
|
{
|
|
use std::os::unix::fs::PermissionsExt;
|
|
permissions.set_mode(0o755);
|
|
}
|
|
std::fs::set_permissions(&script_path, permissions)?;
|
|
("./print-token.sh".to_string(), Vec::new())
|
|
};
|
|
|
|
#[cfg(windows)]
|
|
let (command, args) = {
|
|
let script_path = tempdir.path().join("print-token.cmd");
|
|
std::fs::write(
|
|
&script_path,
|
|
r#"@echo off
|
|
setlocal EnableExtensions DisableDelayedExpansion
|
|
set "first_line="
|
|
<tokens.txt set /p "first_line="
|
|
if not defined first_line exit /b 1
|
|
setlocal EnableDelayedExpansion
|
|
echo(!first_line!
|
|
endlocal
|
|
more +1 tokens.txt > tokens.next
|
|
move /y tokens.next tokens.txt >nul
|
|
"#,
|
|
)?;
|
|
(
|
|
"cmd.exe".to_string(),
|
|
vec![
|
|
"/d".to_string(),
|
|
"/s".to_string(),
|
|
"/c".to_string(),
|
|
".\\print-token.cmd".to_string(),
|
|
],
|
|
)
|
|
};
|
|
|
|
Ok(Self {
|
|
tempdir,
|
|
command,
|
|
args,
|
|
})
|
|
}
|
|
|
|
fn new_failing() -> std::io::Result<Self> {
|
|
let tempdir = tempfile::tempdir()?;
|
|
|
|
#[cfg(unix)]
|
|
let (command, args) = {
|
|
let script_path = tempdir.path().join("fail.sh");
|
|
std::fs::write(
|
|
&script_path,
|
|
r#"#!/bin/sh
|
|
exit 1
|
|
"#,
|
|
)?;
|
|
let mut permissions = std::fs::metadata(&script_path)?.permissions();
|
|
{
|
|
use std::os::unix::fs::PermissionsExt;
|
|
permissions.set_mode(0o755);
|
|
}
|
|
std::fs::set_permissions(&script_path, permissions)?;
|
|
("./fail.sh".to_string(), Vec::new())
|
|
};
|
|
|
|
#[cfg(windows)]
|
|
let (command, args) = (
|
|
"cmd.exe".to_string(),
|
|
vec![
|
|
"/d".to_string(),
|
|
"/s".to_string(),
|
|
"/c".to_string(),
|
|
"exit /b 1".to_string(),
|
|
],
|
|
);
|
|
|
|
Ok(Self {
|
|
tempdir,
|
|
command,
|
|
args,
|
|
})
|
|
}
|
|
|
|
fn auth_config(&self) -> ModelProviderAuthInfo {
|
|
serde_json::from_value(json!({
|
|
"command": self.command,
|
|
"args": self.args,
|
|
// Process startup can be slow on loaded Windows CI workers, so leave enough slack to
|
|
// avoid turning these auth-cache assertions into a process-launch timing test.
|
|
"timeout_ms": 10_000,
|
|
"refresh_interval_ms": 60000,
|
|
"cwd": self.tempdir.path(),
|
|
}))
|
|
.expect("provider auth config should deserialize")
|
|
}
|
|
}
|
|
|
|
struct AuthFileParams {
|
|
openai_api_key: Option<String>,
|
|
chatgpt_plan_type: Option<String>,
|
|
chatgpt_account_id: Option<String>,
|
|
}
|
|
|
|
fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result<String> {
|
|
let fake_jwt = fake_jwt_for_auth_file_params(¶ms)?;
|
|
let auth_file = get_auth_file(codex_home);
|
|
let auth_json_data = json!({
|
|
"OPENAI_API_KEY": params.openai_api_key,
|
|
"tokens": {
|
|
"id_token": fake_jwt,
|
|
"access_token": "test-access-token",
|
|
"refresh_token": "test-refresh-token"
|
|
},
|
|
"last_refresh": Utc::now(),
|
|
});
|
|
let auth_json = serde_json::to_string_pretty(&auth_json_data)?;
|
|
std::fs::write(auth_file, auth_json)?;
|
|
Ok(fake_jwt)
|
|
}
|
|
|
|
fn fake_jwt_for_auth_file_params(params: &AuthFileParams) -> std::io::Result<String> {
|
|
#[derive(Serialize)]
|
|
struct Header {
|
|
alg: &'static str,
|
|
typ: &'static str,
|
|
}
|
|
|
|
let header = Header {
|
|
alg: "none",
|
|
typ: "JWT",
|
|
};
|
|
let mut auth_payload = serde_json::json!({
|
|
"chatgpt_user_id": "user-12345",
|
|
"user_id": "user-12345",
|
|
});
|
|
|
|
if let Some(chatgpt_plan_type) = params.chatgpt_plan_type.as_ref() {
|
|
auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type.clone());
|
|
}
|
|
|
|
if let Some(chatgpt_account_id) = params.chatgpt_account_id.as_ref() {
|
|
auth_payload["chatgpt_account_id"] = serde_json::Value::String(chatgpt_account_id.clone());
|
|
}
|
|
|
|
let payload = serde_json::json!({
|
|
"email": "user@example.com",
|
|
"email_verified": true,
|
|
"https://api.openai.com/auth": auth_payload,
|
|
});
|
|
let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b);
|
|
let header_b64 = b64(&serde_json::to_vec(&header)?);
|
|
let payload_b64 = b64(&serde_json::to_vec(&payload)?);
|
|
let signature_b64 = b64(b"sig");
|
|
Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
|
|
}
|
|
|
|
async fn build_config(
|
|
codex_home: &Path,
|
|
forced_login_method: Option<ForcedLoginMethod>,
|
|
forced_chatgpt_workspace_id: Option<String>,
|
|
) -> AuthConfig {
|
|
AuthConfig {
|
|
codex_home: codex_home.to_path_buf(),
|
|
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
|
|
forced_login_method,
|
|
forced_chatgpt_workspace_id,
|
|
}
|
|
}
|
|
|
|
/// Use sparingly.
|
|
/// TODO (gpeal): replace this with an injectable env var provider.
|
|
#[cfg(test)]
|
|
struct EnvVarGuard {
|
|
key: &'static str,
|
|
original: Option<std::ffi::OsString>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl EnvVarGuard {
|
|
fn set(key: &'static str, value: &str) -> Self {
|
|
let original = env::var_os(key);
|
|
unsafe {
|
|
env::set_var(key, value);
|
|
}
|
|
Self { key, original }
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
impl Drop for EnvVarGuard {
|
|
fn drop(&mut self) {
|
|
unsafe {
|
|
match &self.original {
|
|
Some(value) => env::set_var(self.key, value),
|
|
None => env::remove_var(self.key),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
|
|
let codex_home = tempdir().unwrap();
|
|
login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
|
|
.expect("seed api key");
|
|
|
|
let config = build_config(
|
|
codex_home.path(),
|
|
Some(ForcedLoginMethod::Chatgpt),
|
|
/*forced_chatgpt_workspace_id*/ None,
|
|
)
|
|
.await;
|
|
|
|
let err =
|
|
super::enforce_login_restrictions(&config).expect_err("expected method mismatch to error");
|
|
assert!(err.to_string().contains("ChatGPT login is required"));
|
|
assert!(
|
|
!codex_home.path().join("auth.json").exists(),
|
|
"auth.json should be removed on mismatch"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial(codex_api_key)]
|
|
async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
|
|
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_another_org".to_string()),
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let config = build_config(
|
|
codex_home.path(),
|
|
/*forced_login_method*/ None,
|
|
Some("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!(
|
|
!codex_home.path().join("auth.json").exists(),
|
|
"auth.json should be removed on mismatch"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial(codex_api_key)]
|
|
async fn enforce_login_restrictions_allows_matching_workspace() {
|
|
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("org_mine".to_string()),
|
|
)
|
|
.await;
|
|
|
|
super::enforce_login_restrictions(&config).expect("matching workspace should succeed");
|
|
assert!(
|
|
codex_home.path().join("auth.json").exists(),
|
|
"auth.json should remain when restrictions pass"
|
|
);
|
|
}
|
|
|
|
#[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();
|
|
login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
|
|
.expect("seed api key");
|
|
|
|
let config = build_config(
|
|
codex_home.path(),
|
|
/*forced_login_method*/ None,
|
|
Some("org_mine".to_string()),
|
|
)
|
|
.await;
|
|
|
|
super::enforce_login_restrictions(&config).expect("matching workspace should succeed");
|
|
assert!(
|
|
codex_home.path().join("auth.json").exists(),
|
|
"auth.json should remain when restrictions pass"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[serial(codex_api_key)]
|
|
async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
|
|
let _guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env");
|
|
let codex_home = tempdir().unwrap();
|
|
|
|
let config = build_config(
|
|
codex_home.path(),
|
|
Some(ForcedLoginMethod::Chatgpt),
|
|
/*forced_chatgpt_workspace_id*/ None,
|
|
)
|
|
.await;
|
|
|
|
let err = super::enforce_login_restrictions(&config)
|
|
.expect_err("environment API key should not satisfy forced ChatGPT login");
|
|
assert!(
|
|
err.to_string()
|
|
.contains("ChatGPT login is required, but an API key is currently being used.")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn plan_type_maps_known_plan() {
|
|
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: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("load auth")
|
|
.expect("auth available");
|
|
|
|
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Pro));
|
|
}
|
|
|
|
#[test]
|
|
fn plan_type_maps_self_serve_business_usage_based_plan() {
|
|
let codex_home = tempdir().unwrap();
|
|
let _jwt = write_auth_file(
|
|
AuthFileParams {
|
|
openai_api_key: None,
|
|
chatgpt_plan_type: Some("self_serve_business_usage_based".to_string()),
|
|
chatgpt_account_id: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("load auth")
|
|
.expect("auth available");
|
|
|
|
pretty_assertions::assert_eq!(
|
|
auth.account_plan_type(),
|
|
Some(AccountPlanType::SelfServeBusinessUsageBased)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn plan_type_maps_enterprise_cbp_usage_based_plan() {
|
|
let codex_home = tempdir().unwrap();
|
|
let _jwt = write_auth_file(
|
|
AuthFileParams {
|
|
openai_api_key: None,
|
|
chatgpt_plan_type: Some("enterprise_cbp_usage_based".to_string()),
|
|
chatgpt_account_id: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("load auth")
|
|
.expect("auth available");
|
|
|
|
pretty_assertions::assert_eq!(
|
|
auth.account_plan_type(),
|
|
Some(AccountPlanType::EnterpriseCbpUsageBased)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn plan_type_maps_unknown_to_unknown() {
|
|
let codex_home = tempdir().unwrap();
|
|
let _jwt = write_auth_file(
|
|
AuthFileParams {
|
|
openai_api_key: None,
|
|
chatgpt_plan_type: Some("mystery-tier".to_string()),
|
|
chatgpt_account_id: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("load auth")
|
|
.expect("auth available");
|
|
|
|
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown));
|
|
}
|
|
|
|
#[test]
|
|
fn missing_plan_type_maps_to_unknown() {
|
|
let codex_home = tempdir().unwrap();
|
|
let _jwt = write_auth_file(
|
|
AuthFileParams {
|
|
openai_api_key: None,
|
|
chatgpt_plan_type: None,
|
|
chatgpt_account_id: None,
|
|
},
|
|
codex_home.path(),
|
|
)
|
|
.expect("failed to write auth file");
|
|
|
|
let auth = super::load_auth(
|
|
codex_home.path(),
|
|
/*enable_codex_api_key_env*/ false,
|
|
AuthCredentialsStoreMode::File,
|
|
)
|
|
.expect("load auth")
|
|
.expect("auth available");
|
|
|
|
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown));
|
|
}
|