Immutable CodexAuth (#8857)

Historically we started with a CodexAuth that knew how to refresh it's
own tokens and then added AuthManager that did a different kind of
refresh (re-reading from disk).

I don't think it makes sense for both `CodexAuth` and `AuthManager` to
be mutable and contain behaviors.

Move all refresh logic into `AuthManager` and keep `CodexAuth` as a data
object.

codex-rs/core/tests/suite/auth_refresh.rs

@@ -3,7 +3,7 @@ use anyhow::Result;
 use base64::Engine;
 use chrono::Duration;
 use chrono::Utc;
-use codex_core::CodexAuth;
+use codex_core::AuthManager;
 use codex_core::auth::AuthCredentialsStoreMode;
 use codex_core::auth::AuthDotJson;
 use codex_core::auth::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR;
@@ -45,32 +45,149 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> {
         .await;
 
     let ctx = RefreshTokenTestContext::new(&server)?;
-    let auth = ctx.auth.clone();
+    let initial_last_refresh = Utc::now() - Duration::days(1);
+    let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN);
+    let initial_auth = AuthDotJson {
+        openai_api_key: None,
+        tokens: Some(initial_tokens.clone()),
+        last_refresh: Some(initial_last_refresh),
+    };
+    ctx.write_auth(&initial_auth)?;
 
-    let access = auth
+    let access = ctx
+        .auth_manager
         .refresh_token()
         .await
         .context("refresh should succeed")?;
-    assert_eq!(access, "new-access-token");
+    assert_eq!(access, Some("new-access-token".to_string()));
 
+    let refreshed_tokens = TokenData {
+        access_token: "new-access-token".to_string(),
+        refresh_token: "new-refresh-token".to_string(),
+        ..initial_tokens.clone()
+    };
     let stored = ctx.load_auth()?;
     let tokens = stored.tokens.as_ref().context("tokens should exist")?;
-    assert_eq!(tokens.access_token, "new-access-token");
-    assert_eq!(tokens.refresh_token, "new-refresh-token");
+    assert_eq!(tokens, &refreshed_tokens);
     let refreshed_at = stored
         .last_refresh
         .as_ref()
         .context("last_refresh should be recorded")?;
     assert!(
-        *refreshed_at >= ctx.initial_last_refresh,
+        *refreshed_at >= initial_last_refresh,
         "last_refresh should advance"
     );
 
-    let cached = auth
-        .get_token_data()
+    let cached_auth = ctx
+        .auth_manager
+        .auth()
         .await
+        .context("auth should be cached")?;
+    let cached = cached_auth
+        .get_token_data()
         .context("token data should be cached")?;
-    assert_eq!(cached.access_token, "new-access-token");
+    assert_eq!(cached, refreshed_tokens);
+
+    server.verify().await;
+    Ok(())
+}
+
+#[serial_test::serial(auth_refresh)]
+#[tokio::test]
+async fn returns_fresh_tokens_as_is() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "access_token": "new-access-token",
+            "refresh_token": "new-refresh-token"
+        })))
+        .mount(&server)
+        .await;
+
+    let ctx = RefreshTokenTestContext::new(&server)?;
+    let initial_last_refresh = Utc::now() - Duration::days(1);
+    let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN);
+    let initial_auth = AuthDotJson {
+        openai_api_key: None,
+        tokens: Some(initial_tokens.clone()),
+        last_refresh: Some(initial_last_refresh),
+    };
+    ctx.write_auth(&initial_auth)?;
+
+    let cached_auth = ctx
+        .auth_manager
+        .auth()
+        .await
+        .context("auth should be cached")?;
+    let cached = cached_auth
+        .get_token_data()
+        .context("token data should remain cached")?;
+    assert_eq!(cached, initial_tokens);
+
+    let stored = ctx.load_auth()?;
+    assert_eq!(stored, initial_auth);
+
+    let requests = server.received_requests().await.unwrap_or_default();
+    assert!(requests.is_empty(), "expected no refresh token requests");
+
+    Ok(())
+}
+
+#[serial_test::serial(auth_refresh)]
+#[tokio::test]
+async fn refreshes_token_when_last_refresh_is_stale() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "access_token": "new-access-token",
+            "refresh_token": "new-refresh-token"
+        })))
+        .expect(1)
+        .mount(&server)
+        .await;
+
+    let ctx = RefreshTokenTestContext::new(&server)?;
+    let stale_refresh = Utc::now() - Duration::days(9);
+    let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN);
+    let initial_auth = AuthDotJson {
+        openai_api_key: None,
+        tokens: Some(initial_tokens.clone()),
+        last_refresh: Some(stale_refresh),
+    };
+    ctx.write_auth(&initial_auth)?;
+
+    let cached_auth = ctx
+        .auth_manager
+        .auth()
+        .await
+        .context("auth should be cached")?;
+    let refreshed_tokens = TokenData {
+        access_token: "new-access-token".to_string(),
+        refresh_token: "new-refresh-token".to_string(),
+        ..initial_tokens.clone()
+    };
+    let cached = cached_auth
+        .get_token_data()
+        .context("token data should refresh")?;
+    assert_eq!(cached, refreshed_tokens);
+
+    let stored = ctx.load_auth()?;
+    let tokens = stored.tokens.as_ref().context("tokens should exist")?;
+    assert_eq!(tokens, &refreshed_tokens);
+    let refreshed_at = stored
+        .last_refresh
+        .as_ref()
+        .context("last_refresh should be recorded")?;
+    assert!(
+        *refreshed_at >= stale_refresh,
+        "last_refresh should advance"
+    );
 
     server.verify().await;
     Ok(())
@@ -94,9 +211,17 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re
         .await;
 
     let ctx = RefreshTokenTestContext::new(&server)?;
-    let auth = ctx.auth.clone();
+    let initial_last_refresh = Utc::now() - Duration::days(1);
+    let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN);
+    let initial_auth = AuthDotJson {
+        openai_api_key: None,
+        tokens: Some(initial_tokens.clone()),
+        last_refresh: Some(initial_last_refresh),
+    };
+    ctx.write_auth(&initial_auth)?;
 
-    let err = auth
+    let err = ctx
+        .auth_manager
         .refresh_token()
         .await
         .err()
@@ -104,16 +229,16 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re
     assert_eq!(err.failed_reason(), Some(RefreshTokenFailedReason::Expired));
 
     let stored = ctx.load_auth()?;
-    let tokens = stored.tokens.as_ref().context("tokens should remain")?;
-    assert_eq!(tokens.access_token, INITIAL_ACCESS_TOKEN);
-    assert_eq!(tokens.refresh_token, INITIAL_REFRESH_TOKEN);
-    assert_eq!(
-        *stored
-            .last_refresh
-            .as_ref()
-            .context("last_refresh should remain unchanged")?,
-        ctx.initial_last_refresh,
-    );
+    assert_eq!(stored, initial_auth);
+    let cached_auth = ctx
+        .auth_manager
+        .auth()
+        .await
+        .context("auth should remain cached")?;
+    let cached = cached_auth
+        .get_token_data()
+        .context("token data should remain cached")?;
+    assert_eq!(cached, initial_tokens);
 
     server.verify().await;
     Ok(())
@@ -135,9 +260,17 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()>
         .await;
 
     let ctx = RefreshTokenTestContext::new(&server)?;
-    let auth = ctx.auth.clone();
+    let initial_last_refresh = Utc::now() - Duration::days(1);
+    let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN);
+    let initial_auth = AuthDotJson {
+        openai_api_key: None,
+        tokens: Some(initial_tokens.clone()),
+        last_refresh: Some(initial_last_refresh),
+    };
+    ctx.write_auth(&initial_auth)?;
 
-    let err = auth
+    let err = ctx
+        .auth_manager
         .refresh_token()
         .await
         .err()
@@ -146,16 +279,16 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()>
     assert_eq!(err.failed_reason(), None);
 
     let stored = ctx.load_auth()?;
-    let tokens = stored.tokens.as_ref().context("tokens should remain")?;
-    assert_eq!(tokens.access_token, INITIAL_ACCESS_TOKEN);
-    assert_eq!(tokens.refresh_token, INITIAL_REFRESH_TOKEN);
-    assert_eq!(
-        *stored
-            .last_refresh
-            .as_ref()
-            .context("last_refresh should remain unchanged")?,
-        ctx.initial_last_refresh,
-    );
+    assert_eq!(stored, initial_auth);
+    let cached_auth = ctx
+        .auth_manager
+        .auth()
+        .await
+        .context("auth should remain cached")?;
+    let cached = cached_auth
+        .get_token_data()
+        .context("token data should remain cached")?;
+    assert_eq!(cached, initial_tokens);
 
     server.verify().await;
     Ok(())
@@ -163,44 +296,26 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()>
 
 struct RefreshTokenTestContext {
     codex_home: TempDir,
-    auth: CodexAuth,
-    initial_last_refresh: chrono::DateTime<Utc>,
+    auth_manager: AuthManager,
     _env_guard: EnvGuard,
 }
 
 impl RefreshTokenTestContext {
     fn new(server: &MockServer) -> Result<Self> {
         let codex_home = TempDir::new()?;
-        let initial_last_refresh = Utc::now() - Duration::days(1);
-        let mut id_token = IdTokenInfo::default();
-        id_token.raw_jwt = minimal_jwt();
-        let tokens = TokenData {
-            id_token,
-            access_token: INITIAL_ACCESS_TOKEN.to_string(),
-            refresh_token: INITIAL_REFRESH_TOKEN.to_string(),
-            account_id: Some("account-id".to_string()),
-        };
-        let auth_dot_json = AuthDotJson {
-            openai_api_key: None,
-            tokens: Some(tokens),
-            last_refresh: Some(initial_last_refresh),
-        };
-        save_auth(
-            codex_home.path(),
-            &auth_dot_json,
-            AuthCredentialsStoreMode::File,
-        )?;
 
         let endpoint = format!("{}/oauth/token", server.uri());
         let env_guard = EnvGuard::set(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, endpoint);
 
-        let auth = CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File)?
-            .context("auth should load from storage")?;
+        let auth_manager = AuthManager::new(
+            codex_home.path().to_path_buf(),
+            false,
+            AuthCredentialsStoreMode::File,
+        );
 
         Ok(Self {
             codex_home,
-            auth,
-            initial_last_refresh,
+            auth_manager,
             _env_guard: env_guard,
         })
     }
@@ -210,6 +325,16 @@ impl RefreshTokenTestContext {
             .context("load auth.json")?
             .context("auth.json should exist")
     }
+
+    fn write_auth(&self, auth_dot_json: &AuthDotJson) -> Result<()> {
+        save_auth(
+            self.codex_home.path(),
+            auth_dot_json,
+            AuthCredentialsStoreMode::File,
+        )?;
+        self.auth_manager.reload();
+        Ok(())
+    }
 }
 
 struct EnvGuard {
@@ -270,3 +395,14 @@ fn minimal_jwt() -> String {
     let signature_b64 = b64(b"sig");
     format!("{header_b64}.{payload_b64}.{signature_b64}")
 }
+
+fn build_tokens(access_token: &str, refresh_token: &str) -> TokenData {
+    let mut id_token = IdTokenInfo::default();
+    id_token.raw_jwt = minimal_jwt();
+    TokenData {
+        id_token,
+        access_token: access_token.to_string(),
+        refresh_token: refresh_token.to_string(),
+        account_id: Some("account-id".to_string()),
+    }
+}