Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Abhijay007 merged 6 commits into from
Jun 18, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
24489e6
fix(gcp): retry Vertex AI request once with refreshed token on 401/403
michaelneale Jun 17, 2026
bdd7638
Merge remote-tracking branch 'origin/main' into fix/gcp-vertex-token-…
Jun 18, 2026
3a12606
fix(gcp): reload gcloud credentials on auth failure, not just retry
Jun 18, 2026
67f139a
fix(gcp): also reload credentials on token-exchange auth failure
Jun 18, 2026
d47207d
docs(gcp): drop verbose doc comment on refresh_credentials
Jun 18, 2026
9e52335
test(gcp): drop refresh_credentials test that mutated process env
Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions crates/goose/src/providers/gcpauth.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
Comment on lines +359 to +360

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Acquire auth locks in a consistent order

In the concurrent reauth path, this takes credentials.write() and then waits for cached_token.write(), while get_token() takes cached_token.write() first and then awaits credentials.read() for the token exchange. If a refresh is triggered while one request is refreshing a token and another is queued for the cache lock, the queued get_token() can acquire the cache lock before this refresh does, then block on the credentials writer that is itself waiting for the cache lock, leaving both requests stuck. Avoid holding both locks in opposite orders, or clone the credentials under the read lock and drop it before the token-exchange await.

Useful? React with 👍 / 👎.

Ok(())
}

/// Retrieves a valid authentication token.
///
/// This method implements an efficient token management strategy:
Expand Down Expand Up @@ -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?,
Expand Down Expand Up @@ -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)),
}
Expand All @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)),
};
Expand All @@ -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)),
};
Expand Down Expand Up @@ -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.
}
31 changes: 27 additions & 4 deletions crates/goose/src/providers/gcpvertexai.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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 {

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Refresh credentials before returning token-exchange auth errors

This retry only runs after Vertex returns 401/403, so it misses the reauth path where the cached access token has already expired and the in-memory authorized-user refresh token is stale. In that case get_auth_header() fails during the OAuth token exchange before any Vertex request is sent, returning at the auth-header step and never reaching this branch; since refresh_credentials() is only called here, the freshly rewritten ADC file still is not picked up and the user must restart. Please also reload/clear credentials once on auth-header token-exchange failures before surfacing the authentication error.

Useful? React with 👍 / 👎.

@DOsinga DOsinga Jun 18, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — addressed in 67f139a. get_auth_header() failures (the expired-refresh-token / invalid_rapt case in #9689, where the OAuth token exchange fails before any Vertex request is sent) now also reload credentials and retry once, sharing the same single-retry budget as the 401/403 path.

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}"
)));
Expand Down
Loading