diff --git a/crates/goose/src/providers/chatgpt_codex.rs b/crates/goose/src/providers/chatgpt_codex.rs index 6f4ab1ca..05c62553 100644 --- a/crates/goose/src/providers/chatgpt_codex.rs +++ b/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 { 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 { let auth_provider = Arc::new(ChatGptCodexAuthProvider::new( ChatGptCodexAuthState::instance(), @@ -997,10 +1006,25 @@ 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() { + if self.auth_provider.cache.load().is_none() { + let _ = self.auth_provider.cache.save(previous_token); + } + } + return Err(ProviderError::Authentication(format!( + "OAuth flow failed: {}", + e + ))); + } + Ok(()) } diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index c0c228bb..6e864251 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -100,6 +100,10 @@ async fn init_registry() -> RwLock { "kimi_code", Arc::new(|| Box::pin(KimiCodeProvider::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);