Skip to content

Commit 103acdf

Browse files
Refactor external auth to use a single trait (#16356)
## Summary - Replace the separate external auth enum and refresher trait with a single `ExternalAuth` trait in login auth flow - Move bearer token auth behind `BearerTokenRefresher` and update `AuthManager` and app-server wiring to use the generic external auth API
1 parent 0fe873a commit 103acdf

6 files changed

Lines changed: 155 additions & 119 deletions

File tree

AGENTS.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ In the codex-rs folder where the rust code lives:
1717
- Do not add these comments for string or char literals unless the comment adds real clarity; those literals are intentionally exempt from the lint.
1818
- If you add one of these comments, the parameter name must exactly match the callee signature.
1919
- When possible, make `match` statements exhaustive and avoid wildcard arms.
20+
- Newly added traits should include doc comments that explain their role and how implementations are expected to use them.
2021
- When writing tests, prefer comparing the equality of entire objects over fields one by one.
2122
- When making a change that adds or changes an API, ensure that the documentation in the `docs/` folder is up to date if applicable.
2223
- If you change `ConfigToml` or nested config types, run `just write-config-schema` to update `codex-rs/core/config.schema.json`.

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,10 @@ use codex_core::models_manager::collaboration_mode_presets::CollaborationModesCo
7171
use codex_exec_server::EnvironmentManager;
7272
use codex_features::Feature;
7373
use codex_feedback::CodexFeedback;
74+
use codex_login::AuthMode as LoginAuthMode;
75+
use codex_login::auth::ExternalAuth;
7476
use codex_login::auth::ExternalAuthRefreshContext;
7577
use codex_login::auth::ExternalAuthRefreshReason;
76-
use codex_login::auth::ExternalAuthRefresher;
7778
use codex_login::auth::ExternalAuthTokens;
7879
use codex_protocol::ThreadId;
7980
use codex_protocol::protocol::SessionSource;
@@ -103,7 +104,11 @@ impl ExternalAuthRefreshBridge {
103104
}
104105

105106
#[async_trait]
106-
impl ExternalAuthRefresher for ExternalAuthRefreshBridge {
107+
impl ExternalAuth for ExternalAuthRefreshBridge {
108+
fn auth_mode(&self) -> LoginAuthMode {
109+
LoginAuthMode::Chatgpt
110+
}
111+
107112
async fn refresh(
108113
&self,
109114
context: ExternalAuthRefreshContext,
@@ -211,7 +216,7 @@ impl MessageProcessor {
211216
enable_codex_api_key_env,
212217
rpc_transport,
213218
} = args;
214-
let auth_manager = AuthManager::shared_with_external_chatgpt_auth_refresher(
219+
let auth_manager = AuthManager::shared_with_external_auth(
215220
config.codex_home.clone(),
216221
enable_codex_api_key_env,
217222
config.cli_auth_credentials_store_mode,
@@ -290,7 +295,7 @@ impl MessageProcessor {
290295
}
291296

292297
pub(crate) fn clear_runtime_references(&self) {
293-
self.auth_manager.clear_external_chatgpt_auth_refresher();
298+
self.auth_manager.clear_external_auth();
294299
}
295300

296301
pub(crate) async fn process_request(

codex-rs/login/src/auth/auth_tests.rs

Lines changed: 24 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,8 @@ async fn external_bearer_only_auth_manager_uses_cached_provider_token() {
283283

284284
assert_eq!(first.as_deref(), Some("provider-token"));
285285
assert_eq!(second.as_deref(), Some("provider-token"));
286+
assert_eq!(manager.auth_mode(), Some(crate::AuthMode::ApiKey));
287+
assert_eq!(manager.get_api_auth_mode(), Some(ApiAuthMode::ApiKey));
286288
}
287289

288290
#[tokio::test]
@@ -448,13 +450,29 @@ struct AuthFileParams {
448450
}
449451

450452
fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result<String> {
453+
let fake_jwt = fake_jwt_for_auth_file_params(&params)?;
451454
let auth_file = get_auth_file(codex_home);
452-
// Create a minimal valid JWT for the id_token field.
455+
let auth_json_data = json!({
456+
"OPENAI_API_KEY": params.openai_api_key,
457+
"tokens": {
458+
"id_token": fake_jwt,
459+
"access_token": "test-access-token",
460+
"refresh_token": "test-refresh-token"
461+
},
462+
"last_refresh": Utc::now(),
463+
});
464+
let auth_json = serde_json::to_string_pretty(&auth_json_data)?;
465+
std::fs::write(auth_file, auth_json)?;
466+
Ok(fake_jwt)
467+
}
468+
469+
fn fake_jwt_for_auth_file_params(params: &AuthFileParams) -> std::io::Result<String> {
453470
#[derive(Serialize)]
454471
struct Header {
455472
alg: &'static str,
456473
typ: &'static str,
457474
}
475+
458476
let header = Header {
459477
alg: "none",
460478
typ: "JWT",
@@ -464,13 +482,12 @@ fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result
464482
"user_id": "user-12345",
465483
});
466484

467-
if let Some(chatgpt_plan_type) = params.chatgpt_plan_type {
468-
auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type);
485+
if let Some(chatgpt_plan_type) = params.chatgpt_plan_type.as_ref() {
486+
auth_payload["chatgpt_plan_type"] = serde_json::Value::String(chatgpt_plan_type.clone());
469487
}
470488

471-
if let Some(chatgpt_account_id) = params.chatgpt_account_id {
472-
let org_value = serde_json::Value::String(chatgpt_account_id);
473-
auth_payload["chatgpt_account_id"] = org_value;
489+
if let Some(chatgpt_account_id) = params.chatgpt_account_id.as_ref() {
490+
auth_payload["chatgpt_account_id"] = serde_json::Value::String(chatgpt_account_id.clone());
474491
}
475492

476493
let payload = serde_json::json!({
@@ -482,20 +499,7 @@ fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result
482499
let header_b64 = b64(&serde_json::to_vec(&header)?);
483500
let payload_b64 = b64(&serde_json::to_vec(&payload)?);
484501
let signature_b64 = b64(b"sig");
485-
let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}");
486-
487-
let auth_json_data = json!({
488-
"OPENAI_API_KEY": params.openai_api_key,
489-
"tokens": {
490-
"id_token": fake_jwt,
491-
"access_token": "test-access-token",
492-
"refresh_token": "test-refresh-token"
493-
},
494-
"last_refresh": Utc::now(),
495-
});
496-
let auth_json = serde_json::to_string_pretty(&auth_json_data)?;
497-
std::fs::write(auth_file, auth_json)?;
498-
Ok(fake_jwt)
502+
Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
499503
}
500504

501505
async fn build_config(

codex-rs/login/src/auth/external_bearer.rs

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
use super::manager::ExternalAuth;
2+
use super::manager::ExternalAuthRefreshContext;
3+
use super::manager::ExternalAuthTokens;
4+
use async_trait::async_trait;
15
use codex_protocol::config_types::ModelProviderAuthInfo;
26
use std::fmt;
37
use std::io;
@@ -10,47 +14,61 @@ use tokio::process::Command;
1014
use tokio::sync::Mutex;
1115

1216
#[derive(Clone)]
13-
pub(crate) struct ExternalBearerAuth {
17+
pub(crate) struct BearerTokenRefresher {
1418
state: Arc<ExternalBearerAuthState>,
1519
}
1620

17-
impl ExternalBearerAuth {
21+
impl BearerTokenRefresher {
1822
pub(crate) fn new(config: ModelProviderAuthInfo) -> Self {
1923
Self {
2024
state: Arc::new(ExternalBearerAuthState::new(config)),
2125
}
2226
}
27+
}
2328

24-
pub(crate) async fn resolve_access_token(&self) -> io::Result<String> {
25-
let mut cached = self.state.cached_token.lock().await;
26-
if let Some(cached_token) = cached.as_ref()
27-
&& cached_token.fetched_at.elapsed() < self.state.config.refresh_interval()
28-
{
29-
return Ok(cached_token.access_token.clone());
30-
}
29+
#[async_trait]
30+
impl ExternalAuth for BearerTokenRefresher {
31+
fn auth_mode(&self) -> crate::AuthMode {
32+
crate::AuthMode::ApiKey
33+
}
3134

32-
let access_token = run_provider_auth_command(&self.state.config).await?;
33-
*cached = Some(CachedExternalBearerToken {
34-
access_token: access_token.clone(),
35-
fetched_at: Instant::now(),
36-
});
37-
Ok(access_token)
35+
async fn resolve(&self) -> io::Result<Option<ExternalAuthTokens>> {
36+
let access_token = {
37+
let mut cached = self.state.cached_token.lock().await;
38+
if let Some(cached_token) = cached.as_ref()
39+
&& cached_token.fetched_at.elapsed() < self.state.config.refresh_interval()
40+
{
41+
cached_token.access_token.clone()
42+
} else {
43+
let access_token = run_provider_auth_command(&self.state.config).await?;
44+
*cached = Some(CachedExternalBearerToken {
45+
access_token: access_token.clone(),
46+
fetched_at: Instant::now(),
47+
});
48+
access_token
49+
}
50+
};
51+
Ok(Some(ExternalAuthTokens::access_token_only(access_token)))
3852
}
3953

40-
pub(crate) async fn refresh_after_unauthorized(&self) -> io::Result<()> {
54+
async fn refresh(
55+
&self,
56+
_context: ExternalAuthRefreshContext,
57+
) -> io::Result<ExternalAuthTokens> {
4158
let access_token = run_provider_auth_command(&self.state.config).await?;
4259
let mut cached = self.state.cached_token.lock().await;
4360
*cached = Some(CachedExternalBearerToken {
44-
access_token,
61+
access_token: access_token.clone(),
4562
fetched_at: Instant::now(),
4663
});
47-
Ok(())
64+
Ok(ExternalAuthTokens::access_token_only(access_token))
4865
}
4966
}
5067

51-
impl fmt::Debug for ExternalBearerAuth {
68+
impl fmt::Debug for BearerTokenRefresher {
5269
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53-
f.debug_struct("ExternalBearerAuth").finish_non_exhaustive()
70+
f.debug_struct("BearerTokenRefresher")
71+
.finish_non_exhaustive()
5472
}
5573
}
5674

0 commit comments

Comments
 (0)