Skip to content

Commit 4a42c4e

Browse files
authored
[Auth] Choose which auth storage to use based on config (#5792)
This PR is a follow-up to #5591. It allows users to choose which auth storage mode they want by using the new `cli_auth_credentials_store_mode` config.
1 parent 66a4b89 commit 4a42c4e

File tree

30 files changed

+361
-80
lines changed

30 files changed

+361
-80
lines changed

codex-rs/app-server/src/codex_message_processor.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,11 @@ impl CodexMessageProcessor {
323323
}
324324
}
325325

326-
match login_with_api_key(&self.config.codex_home, &params.api_key) {
326+
match login_with_api_key(
327+
&self.config.codex_home,
328+
&params.api_key,
329+
self.config.cli_auth_credentials_store_mode,
330+
) {
327331
Ok(()) => {
328332
self.auth_manager.reload();
329333
self.outgoing
@@ -367,6 +371,7 @@ impl CodexMessageProcessor {
367371
config.codex_home.clone(),
368372
CLIENT_ID.to_string(),
369373
config.forced_chatgpt_workspace_id.clone(),
374+
config.cli_auth_credentials_store_mode,
370375
)
371376
};
372377

codex-rs/app-server/src/message_processor.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,11 @@ impl MessageProcessor {
3737
feedback: CodexFeedback,
3838
) -> Self {
3939
let outgoing = Arc::new(outgoing);
40-
let auth_manager = AuthManager::shared(config.codex_home.clone(), false);
40+
let auth_manager = AuthManager::shared(
41+
config.codex_home.clone(),
42+
false,
43+
config.cli_auth_credentials_store_mode,
44+
);
4145
let conversation_manager = Arc::new(ConversationManager::new(
4246
auth_manager.clone(),
4347
SessionSource::VSCode,

codex-rs/app-server/tests/common/auth_fixtures.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use base64::Engine;
66
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
77
use chrono::DateTime;
88
use chrono::Utc;
9+
use codex_core::auth::AuthCredentialsStoreMode;
910
use codex_core::auth::AuthDotJson;
1011
use codex_core::auth::save_auth;
1112
use codex_core::token_data::TokenData;
@@ -108,7 +109,11 @@ pub fn encode_id_token(claims: &ChatGptIdTokenClaims) -> Result<String> {
108109
Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
109110
}
110111

111-
pub fn write_chatgpt_auth(codex_home: &Path, fixture: ChatGptAuthFixture) -> Result<()> {
112+
pub fn write_chatgpt_auth(
113+
codex_home: &Path,
114+
fixture: ChatGptAuthFixture,
115+
cli_auth_credentials_store_mode: AuthCredentialsStoreMode,
116+
) -> Result<()> {
112117
let id_token_raw = encode_id_token(&fixture.claims)?;
113118
let id_token = parse_id_token(&id_token_raw).context("parse id token")?;
114119
let tokens = TokenData {
@@ -126,5 +131,5 @@ pub fn write_chatgpt_auth(codex_home: &Path, fixture: ChatGptAuthFixture) -> Res
126131
last_refresh,
127132
};
128133

129-
save_auth(codex_home, &auth).context("write auth.json")
134+
save_auth(codex_home, &auth, cli_auth_credentials_store_mode).context("write auth.json")
130135
}

codex-rs/app-server/tests/suite/login.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use codex_app_server_protocol::JSONRPCResponse;
1212
use codex_app_server_protocol::LoginChatGptResponse;
1313
use codex_app_server_protocol::LogoutChatGptResponse;
1414
use codex_app_server_protocol::RequestId;
15+
use codex_core::auth::AuthCredentialsStoreMode;
1516
use codex_login::login_with_api_key;
1617
use serial_test::serial;
1718
use tempfile::TempDir;
@@ -45,7 +46,12 @@ stream_max_retries = 0
4546
async fn logout_chatgpt_removes_auth() {
4647
let codex_home = TempDir::new().unwrap_or_else(|e| panic!("create tempdir: {e}"));
4748
create_config_toml(codex_home.path()).expect("write config.toml");
48-
login_with_api_key(codex_home.path(), "sk-test-key").expect("seed api key");
49+
login_with_api_key(
50+
codex_home.path(),
51+
"sk-test-key",
52+
AuthCredentialsStoreMode::File,
53+
)
54+
.expect("seed api key");
4955
assert!(codex_home.path().join("auth.json").exists());
5056

5157
let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)])

codex-rs/app-server/tests/suite/rate_limits.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use codex_app_server_protocol::JSONRPCError;
99
use codex_app_server_protocol::JSONRPCResponse;
1010
use codex_app_server_protocol::LoginApiKeyParams;
1111
use codex_app_server_protocol::RequestId;
12+
use codex_core::auth::AuthCredentialsStoreMode;
1213
use codex_protocol::protocol::RateLimitSnapshot;
1314
use codex_protocol::protocol::RateLimitWindow;
1415
use pretty_assertions::assert_eq;
@@ -106,6 +107,7 @@ async fn get_account_rate_limits_returns_snapshot() -> Result<()> {
106107
ChatGptAuthFixture::new("chatgpt-token")
107108
.account_id("account-123")
108109
.plan_type("pro"),
110+
AuthCredentialsStoreMode::File,
109111
)
110112
.context("write chatgpt auth")?;
111113

codex-rs/app-server/tests/suite/user_info.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use app_test_support::write_chatgpt_auth;
77
use codex_app_server_protocol::JSONRPCResponse;
88
use codex_app_server_protocol::RequestId;
99
use codex_app_server_protocol::UserInfoResponse;
10+
use codex_core::auth::AuthCredentialsStoreMode;
1011
use pretty_assertions::assert_eq;
1112
use tempfile::TempDir;
1213
use tokio::time::timeout;
@@ -22,6 +23,7 @@ async fn user_info_returns_email_from_auth_json() {
2223
ChatGptAuthFixture::new("access")
2324
.refresh_token("refresh")
2425
.email("[email protected]"),
26+
AuthCredentialsStoreMode::File,
2527
)
2628
.expect("write chatgpt auth");
2729

codex-rs/chatgpt/src/apply_command.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ pub async fn run_apply_command(
3232
)
3333
.await?;
3434

35-
init_chatgpt_token_from_auth(&config.codex_home).await?;
35+
init_chatgpt_token_from_auth(&config.codex_home, config.cli_auth_credentials_store_mode)
36+
.await?;
3637

3738
let task_response = get_task(&config, apply_cli.task_id).await?;
3839
apply_diff_from_task(task_response, cwd).await

codex-rs/chatgpt/src/chatgpt_client.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@ pub(crate) async fn chatgpt_get_request<T: DeserializeOwned>(
1313
path: String,
1414
) -> anyhow::Result<T> {
1515
let chatgpt_base_url = &config.chatgpt_base_url;
16-
init_chatgpt_token_from_auth(&config.codex_home).await?;
16+
init_chatgpt_token_from_auth(&config.codex_home, config.cli_auth_credentials_store_mode)
17+
.await?;
1718

1819
// Make direct HTTP request to ChatGPT backend API with the token
1920
let client = create_client();

codex-rs/chatgpt/src/chatgpt_token.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use std::path::Path;
33
use std::sync::LazyLock;
44
use std::sync::RwLock;
55

6+
use codex_core::auth::AuthCredentialsStoreMode;
67
use codex_core::token_data::TokenData;
78

89
static CHATGPT_TOKEN: LazyLock<RwLock<Option<TokenData>>> = LazyLock::new(|| RwLock::new(None));
@@ -18,8 +19,11 @@ pub fn set_chatgpt_token_data(value: TokenData) {
1819
}
1920

2021
/// Initialize the ChatGPT token from auth.json file
21-
pub async fn init_chatgpt_token_from_auth(codex_home: &Path) -> std::io::Result<()> {
22-
let auth = CodexAuth::from_auth_storage(codex_home)?;
22+
pub async fn init_chatgpt_token_from_auth(
23+
codex_home: &Path,
24+
auth_credentials_store_mode: AuthCredentialsStoreMode,
25+
) -> std::io::Result<()> {
26+
let auth = CodexAuth::from_auth_storage(codex_home, auth_credentials_store_mode)?;
2327
if let Some(auth) = auth {
2428
let token_data = auth.get_token_data().await?;
2529
set_chatgpt_token_data(token_data);

codex-rs/cli/src/login.rs

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use codex_app_server_protocol::AuthMode;
22
use codex_common::CliConfigOverrides;
33
use codex_core::CodexAuth;
4+
use codex_core::auth::AuthCredentialsStoreMode;
45
use codex_core::auth::CLIENT_ID;
56
use codex_core::auth::login_with_api_key;
67
use codex_core::auth::logout;
@@ -17,11 +18,13 @@ use std::path::PathBuf;
1718
pub async fn login_with_chatgpt(
1819
codex_home: PathBuf,
1920
forced_chatgpt_workspace_id: Option<String>,
21+
cli_auth_credentials_store_mode: AuthCredentialsStoreMode,
2022
) -> std::io::Result<()> {
2123
let opts = ServerOptions::new(
2224
codex_home,
2325
CLIENT_ID.to_string(),
2426
forced_chatgpt_workspace_id,
27+
cli_auth_credentials_store_mode,
2528
);
2629
let server = run_login_server(opts)?;
2730

@@ -43,7 +46,13 @@ pub async fn run_login_with_chatgpt(cli_config_overrides: CliConfigOverrides) ->
4346

4447
let forced_chatgpt_workspace_id = config.forced_chatgpt_workspace_id.clone();
4548

46-
match login_with_chatgpt(config.codex_home, forced_chatgpt_workspace_id).await {
49+
match login_with_chatgpt(
50+
config.codex_home,
51+
forced_chatgpt_workspace_id,
52+
config.cli_auth_credentials_store_mode,
53+
)
54+
.await
55+
{
4756
Ok(_) => {
4857
eprintln!("Successfully logged in");
4958
std::process::exit(0);
@@ -66,7 +75,11 @@ pub async fn run_login_with_api_key(
6675
std::process::exit(1);
6776
}
6877

69-
match login_with_api_key(&config.codex_home, &api_key) {
78+
match login_with_api_key(
79+
&config.codex_home,
80+
&api_key,
81+
config.cli_auth_credentials_store_mode,
82+
) {
7083
Ok(_) => {
7184
eprintln!("Successfully logged in");
7285
std::process::exit(0);
@@ -121,6 +134,7 @@ pub async fn run_login_with_device_code(
121134
config.codex_home,
122135
client_id.unwrap_or(CLIENT_ID.to_string()),
123136
forced_chatgpt_workspace_id,
137+
config.cli_auth_credentials_store_mode,
124138
);
125139
if let Some(iss) = issuer_base_url {
126140
opts.issuer = iss;
@@ -140,7 +154,7 @@ pub async fn run_login_with_device_code(
140154
pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! {
141155
let config = load_config_or_exit(cli_config_overrides).await;
142156

143-
match CodexAuth::from_auth_storage(&config.codex_home) {
157+
match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode) {
144158
Ok(Some(auth)) => match auth.mode {
145159
AuthMode::ApiKey => match auth.get_token().await {
146160
Ok(api_key) => {
@@ -171,7 +185,7 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! {
171185
pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! {
172186
let config = load_config_or_exit(cli_config_overrides).await;
173187

174-
match logout(&config.codex_home) {
188+
match logout(&config.codex_home, config.cli_auth_credentials_store_mode) {
175189
Ok(true) => {
176190
eprintln!("Successfully logged out");
177191
std::process::exit(0);

0 commit comments

Comments
 (0)