Reset ChatGPT Codex auth during OAuth setup

crates/goose/src/providers/chatgpt_codex.rs

@@ -806,6 +806,10 @@ impl ChatGptCodexAuthProvider {
         }
     }
 
+    fn clear_cached_tokens(&self) {
+        self.cache.clear();
+    }
+
     async fn get_valid_token(&self) -> Result<TokenData> {
         if let Some(mut token_data) = self.cache.load() {
             if token_data.expires_at > Utc::now() + chrono::Duration::seconds(60) {
@@ -865,6 +869,11 @@ pub struct ChatGptCodexProvider {
 }
 
 impl ChatGptCodexProvider {
+    pub async fn cleanup() -> Result<()> {
+        TokenCache::new().clear();
+        Ok(())
+    }
+
     pub async fn from_env(model: ModelConfig) -> Result<Self> {
         let auth_provider = Arc::new(ChatGptCodexAuthProvider::new(
             ChatGptCodexAuthState::instance(),
@@ -997,10 +1006,23 @@ impl Provider for ChatGptCodexProvider {
     }
 
     async fn configure_oauth(&self) -> Result<(), ProviderError> {
-        self.auth_provider
-            .get_valid_token()
+        let previous_token = self.auth_provider.cache.load();
+        self.auth_provider.clear_cached_tokens();
+
+        let result = perform_oauth_flow(self.auth_provider.state.as_ref())
             .await
-            .map_err(|e| ProviderError::Authentication(format!("OAuth flow failed: {}", e)))?;
+            .and_then(|token_data| self.auth_provider.cache.save(&token_data));
+
+        if let Err(e) = result {
+            if let Some(previous_token) = previous_token.as_ref() {
+                let _ = self.auth_provider.cache.save(previous_token);
+            }
+            return Err(ProviderError::Authentication(format!(
+                "OAuth flow failed: {}",
+                e
+            )));
+        }
+
         Ok(())
     }
 

crates/goose/src/providers/init.rs

@@ -94,6 +94,10 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
         "databricks",
         Arc::new(|| Box::pin(DatabricksProvider::cleanup())),
     );
+    registry.set_cleanup(
+        "chatgpt_codex",
+        Arc::new(|| Box::pin(ChatGptCodexProvider::cleanup())),
+    );
 
     if let Err(e) = load_custom_providers_into_registry(&mut registry) {
         tracing::warn!("Failed to load custom providers: {}", e);