From 9dd3444b78bea323e058d1be80083eeb088cb399 Mon Sep 17 00:00:00 2001 From: Michael Zeng Date: Fri, 15 May 2026 17:19:34 -0700 Subject: [PATCH] exec-server: minimize remote auth selection fix --- codex-rs/login/src/auth/auth_tests.rs | 89 --------------------------- codex-rs/login/src/auth/manager.rs | 63 +++++++++---------- 2 files changed, 28 insertions(+), 124 deletions(-) diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 83c36b9109..63ab3e0c42 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -209,7 +209,6 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -266,7 +265,6 @@ async fn loads_api_key_from_auth_json() { let auth = super::load_auth( dir.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -344,7 +342,6 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -746,7 +743,6 @@ async fn load_auth_reads_access_token_from_env() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, Some(&chatgpt_base_url), ) @@ -778,7 +774,6 @@ async fn load_auth_keeps_codex_api_key_env_precedence() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ true, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -789,85 +784,6 @@ async fn load_auth_keeps_codex_api_key_env_precedence() { assert_eq!(auth.api_key(), Some("sk-env")); } -#[tokio::test] -#[serial(codex_auth_env)] -async fn load_auth_can_ignore_access_token_env_for_persisted_chatgpt() { - let codex_home = tempdir().unwrap(); - 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); - 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, - /*enable_codex_access_token_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - ) - .await - .expect("persisted auth should load") - .expect("persisted auth should be present"); - - assert!(auth.is_chatgpt_auth()); -} - -#[tokio::test] -#[serial(codex_auth_env)] -async fn load_auth_can_ignore_api_key_env_for_agent_identity_selection() { - let codex_home = tempdir().unwrap(); - 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"); - 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 _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"); - - let chatgpt_base_url = format!("{}/backend-api", server.uri()); - let _authapi_guard = - EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url); - let auth = super::load_auth( - codex_home.path(), - /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, - AuthCredentialsStoreMode::File, - Some(&chatgpt_base_url), - ) - .await - .expect("env auth should load") - .expect("env auth should be present"); - - let CodexAuth::AgentIdentity(agent_identity) = auth else { - panic!("env auth should load as agent identity"); - }; - assert_eq!(agent_identity.record(), &expected_record); - assert_eq!(agent_identity.process_task_id(), "task-123"); - server.verify().await; -} - #[tokio::test] #[serial(codex_auth_env)] async fn enforce_login_restrictions_logs_out_for_method_mismatch() { @@ -1266,7 +1182,6 @@ async fn plan_type_maps_known_plan() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -1295,7 +1210,6 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -1327,7 +1241,6 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -1359,7 +1272,6 @@ async fn plan_type_maps_unknown_to_unknown() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) @@ -1388,7 +1300,6 @@ async fn missing_plan_type_maps_to_unknown() { let auth = super::load_auth( codex_home.path(), /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, ) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 0a79ec3a61..7fe37f13b0 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -246,7 +246,6 @@ impl CodexAuth { load_auth( codex_home, /*enable_codex_api_key_env*/ false, - /*enable_codex_access_token_env*/ true, auth_credentials_store_mode, chatgpt_base_url, ) @@ -619,7 +618,6 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< let Some(auth) = load_auth( &config.codex_home, /*enable_codex_api_key_env*/ true, - /*enable_codex_access_token_env*/ true, config.auth_credentials_store_mode, config.chatgpt_base_url.as_deref(), ) @@ -731,6 +729,22 @@ fn logout_all_stores( } async fn load_auth( + codex_home: &Path, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, + chatgpt_base_url: Option<&str>, +) -> std::io::Result> { + load_auth_with_env( + codex_home, + enable_codex_api_key_env, + /*enable_codex_access_token_env*/ true, + auth_credentials_store_mode, + chatgpt_base_url, + ) + .await +} + +async fn load_auth_with_env( codex_home: &Path, enable_codex_api_key_env: bool, enable_codex_access_token_env: bool, @@ -1292,10 +1306,6 @@ impl Debug for AuthManager { .field("codex_home", &self.codex_home) .field("inner", &self.inner) .field("enable_codex_api_key_env", &self.enable_codex_api_key_env) - .field( - "enable_codex_access_token_env", - &self.enable_codex_access_token_env, - ) .field( "auth_credentials_store_mode", &self.auth_credentials_store_mode, @@ -1338,7 +1348,7 @@ impl AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option, ) -> Self { - let managed_auth = load_auth( + let managed_auth = load_auth_with_env( &codex_home, enable_codex_api_key_env, enable_codex_access_token_env, @@ -1547,7 +1557,7 @@ impl AuthManager { } async fn load_auth_from_storage(&self) -> Option { - load_auth( + load_auth_with_env( &self.codex_home, self.enable_codex_api_key_env, self.enable_codex_access_token_env, @@ -1635,25 +1645,6 @@ impl AuthManager { ) } - async fn shared_with_env( - codex_home: PathBuf, - enable_codex_api_key_env: bool, - enable_codex_access_token_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - chatgpt_base_url: Option, - ) -> Arc { - Arc::new( - Self::new_with_env( - codex_home, - enable_codex_api_key_env, - enable_codex_access_token_env, - auth_credentials_store_mode, - chatgpt_base_url, - ) - .await, - ) - } - /// Convenience constructor returning an `Arc` wrapper from resolved config. pub async fn shared_from_config( config: &impl AuthManagerConfig, @@ -1674,14 +1665,16 @@ impl AuthManager { enable_codex_api_key_env: bool, enable_codex_access_token_env: bool, ) -> Arc { - let auth_manager = Self::shared_with_env( - config.codex_home(), - enable_codex_api_key_env, - enable_codex_access_token_env, - config.cli_auth_credentials_store_mode(), - Some(config.chatgpt_base_url()), - ) - .await; + let auth_manager = Arc::new( + Self::new_with_env( + config.codex_home(), + enable_codex_api_key_env, + enable_codex_access_token_env, + config.cli_auth_credentials_store_mode(), + Some(config.chatgpt_base_url()), + ) + .await, + ); auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id()); auth_manager }