fix(gcp): use refreshed gcloud token after reauth (retry Vertex AI on 401/403) (#9849)

Signed-off-by: Michael Neale <michael.neale@gmail.com>
Signed-off-by: Douwe M Osinga <douwe@sidewalklabs.com>
Co-authored-by: Douwe M Osinga <douwe@sidewalklabs.com>

crates/goose/src/providers/gcpauth.rs

@@ -324,7 +324,7 @@ struct TokenResponse {
 #[derive(Debug)]
 pub struct GcpAuth {
     /// The loaded credentials (service account or authorized user)
-    credentials: AdcCredentials,
+    credentials: RwLock<AdcCredentials>,
     /// HTTP client for making token exchange requests
     client: reqwest::Client,
     /// Thread-safe cache for the current token
@@ -348,12 +348,19 @@ impl GcpAuth {
     /// * `Result<Self, AuthError>` - A new GcpAuth instance or an error if initialization fails
     pub async fn new() -> Result<Self, AuthError> {
         Ok(Self {
-            credentials: AdcCredentials::load().await?,
+            credentials: RwLock::new(AdcCredentials::load().await?),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(None)),
         })
     }
 
+    pub async fn refresh_credentials(&self) -> Result<(), AuthError> {
+        let reloaded = AdcCredentials::load().await?;
+        *self.credentials.write().await = reloaded;
+        *self.cached_token.write().await = None;
+        Ok(())
+    }
+
     /// Retrieves a valid authentication token.
     ///
     /// This method implements an efficient token management strategy:
@@ -386,7 +393,7 @@ impl GcpAuth {
         }
 
         // Get new token
-        let token_response = match &self.credentials {
+        let token_response = match &*self.credentials.read().await {
             AdcCredentials::ServiceAccount(creds) => self.get_service_account_token(creds).await?,
             AdcCredentials::AuthorizedUser(creds) => self.get_authorized_user_token(creds).await?,
             AdcCredentials::DefaultAccount(creds) => self.get_default_access_token(creds).await?,
@@ -687,7 +694,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     // Helper function to create a test GcpAuth instance with credentials
     async fn create_test_auth_with_creds(creds: AdcCredentials) -> GcpAuth {
         GcpAuth {
-            credentials: creds,
+            credentials: RwLock::new(creds),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(None)),
         }
@@ -696,7 +703,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     #[tokio::test]
     async fn test_token_caching() {
         let auth = GcpAuth {
-            credentials: AdcCredentials::ServiceAccount(mock_service_account()),
+            credentials: RwLock::new(AdcCredentials::ServiceAccount(mock_service_account())),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(Some(CachedToken {
                 token: AuthToken {
@@ -719,7 +726,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     #[tokio::test]
     async fn test_token_expiration() {
         let auth = GcpAuth {
-            credentials: AdcCredentials::ServiceAccount(mock_service_account()),
+            credentials: RwLock::new(AdcCredentials::ServiceAccount(mock_service_account())),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(Some(CachedToken {
                 token: AuthToken {
@@ -757,7 +764,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     #[tokio::test]
     async fn test_concurrent_token_access() {
         let auth = Arc::new(GcpAuth {
-            credentials: AdcCredentials::ServiceAccount(mock_service_account()),
+            credentials: RwLock::new(AdcCredentials::ServiceAccount(mock_service_account())),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(Some(CachedToken {
                 token: AuthToken {
@@ -788,7 +795,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     #[tokio::test]
     async fn test_token_refresh_race_condition() {
         let auth = Arc::new(GcpAuth {
-            credentials: AdcCredentials::ServiceAccount(mock_service_account()),
+            credentials: RwLock::new(AdcCredentials::ServiceAccount(mock_service_account())),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(Some(CachedToken {
                 token: AuthToken {
@@ -841,7 +848,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     #[tokio::test]
     async fn test_authorized_user_token() {
         let auth = GcpAuth {
-            credentials: AdcCredentials::AuthorizedUser(mock_authorized_user()),
+            credentials: RwLock::new(AdcCredentials::AuthorizedUser(mock_authorized_user())),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(None)),
         };
@@ -858,7 +865,7 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
     #[tokio::test]
     async fn test_service_account_jwt_creation() {
         let auth = GcpAuth {
-            credentials: AdcCredentials::ServiceAccount(mock_service_account()),
+            credentials: RwLock::new(AdcCredentials::ServiceAccount(mock_service_account())),
             client: reqwest::Client::new(),
             cached_token: Arc::new(RwLock::new(None)),
         };
@@ -1124,4 +1131,9 @@ iXVBc2YmAuU8hiOFUPxtyQfNzG5fQ0rhJSewdtyWxIadJSLj6fsK+AEsNQ==
         .await;
         assert!(matches!(result, Err(AuthError::Credentials(_))));
     }
+
+    // Note: there is intentionally no test that exercises refresh_credentials() end to end.
+    // Doing so would require pointing GOOGLE_APPLICATION_CREDENTIALS at a temp file via
+    // std::env::set_var, which is unsafe in the 2024 edition and races with any other thread
+    // reading the environment. We don't mutate process-global env state in tests for that.
 }

crates/goose/src/providers/gcpvertexai.rs

@@ -284,6 +284,7 @@ impl GcpVertexAIProvider {
         let mut overloaded_attempts = 0;
         let mut last_error = None;
         let max_retries = self.retry_config.max_retries;
+        let mut retried_auth = false;
 
         loop {
             if rate_limit_attempts > max_retries && overloaded_attempts > max_retries {
@@ -295,10 +296,21 @@ impl GcpVertexAIProvider {
                 );
             }
 
-            let auth_header = self
-                .get_auth_header()
-                .await
-                .map_err(|e| ProviderError::Authentication(e.to_string()))?;
+            let auth_header = match self.get_auth_header().await {
+                Ok(header) => header,
+                Err(e) => {
+                    if !retried_auth {
+                        retried_auth = true;
+                        if self.auth.refresh_credentials().await.is_ok() {
+                            tracing::info!(
+                                "gcloud token exchange failed ({e}); reloaded credentials and retrying"
+                            );
+                            continue;
+                        }
+                    }
+                    return Err(ProviderError::Authentication(e.to_string()));
+                }
+            };
 
             let mut request = self
                 .client
@@ -355,6 +367,17 @@ impl GcpVertexAIProvider {
             } else if status == StatusCode::OK {
                 return Ok(response);
             } else if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN {
+                if !retried_auth {
+                    retried_auth = true;
+                    if let Err(e) = self.auth.refresh_credentials().await {
+                        tracing::warn!("Failed to reload gcloud credentials after {status}: {e}");
+                    } else {
+                        tracing::info!(
+                            "Vertex AI returned {status}; reloaded gcloud credentials and retrying"
+                        );
+                        continue;
+                    }
+                }
                 return Err(ProviderError::Authentication(format!(
                     "Authentication failed with status: {status}"
                 )));