Skip to content

Commit 7d6b4cb

Browse files
authored
fix: make vertex work in vscode playground again (#1645)
For some reason, when we do the gcp oauth exchange, we were asking for access tokens that would be valid for 60h. GCP probably at some recent point started denying these requests (because access tokens should not live for 60h, that's absurd), so we trim the expiration time down to 1h, which appears to fix things. (We request a new oauth token for every single request in wasm, so the lifetime of the token could literally be measured in seconds + clock skew.) <img width="521" alt="image" src="https://github.com/user-attachments/assets/9bfde855-3c08-4928-8ef4-d77e3b5d1095" /> <img width="1092" alt="image" src="https://github.com/user-attachments/assets/629cb64c-f0c9-46f3-b9a4-562e9f503a27" /> Also improve the error message when this does fail: <img width="571" alt="image" src="https://github.com/user-attachments/assets/122a25c0-040c-4509-8619-ec28e2b3cdd3" />
1 parent 1c1fdb3 commit 7d6b4cb

1 file changed

Lines changed: 8 additions & 35 deletions

File tree

  • engine/baml-runtime/src/internal/llm_client/primitive/vertex

engine/baml-runtime/src/internal/llm_client/primitive/vertex/wasm_auth.rs

Lines changed: 8 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -19,20 +19,25 @@ impl VertexAuth {
1919
pub async fn new(auth_strategy: &ResolvedGcpAuthStrategy) -> Result<VertexAuth> {
2020
Ok(match auth_strategy {
2121
ResolvedGcpAuthStrategy::FilePath(path) => {
22-
anyhow::bail!("Failed to auth - cannot load credentials from files in WASM")
22+
anyhow::bail!(
23+
"Failed to auth - cannot load credentials from a file in WASM (path='{}...', path.len={})",
24+
path.chars().take(5).collect::<String>(),
25+
path.len()
26+
)
2327
}
2428
ResolvedGcpAuthStrategy::JsonString(json) => {
2529
log::debug!("Attempting to auth using JsonString strategy");
2630
Self(serde_json::from_str(&json).context("Failed to parse service account credentials as GCP service account creds (are you using JSON format creds?)")?)
2731
}
2832
ResolvedGcpAuthStrategy::JsonObject(json) => {
33+
// NB: this should never happen in WASM, there's no way to pass a JSON object in
2934
log::debug!("Attempting to auth using JsonObject strategy");
3035
Self(serde_json::from_value(
3136
serde_json::to_value(&json).context("Failed to parse service account credentials as GCP service account creds (issue during serialization)")?).context("Failed to parse service account credentials as GCP service account creds (are you using JSON format creds?)")?)
3237
}
3338
ResolvedGcpAuthStrategy::SystemDefault => {
3439
anyhow::bail!(
35-
"Failed to auth - cannot load GCP application default credentials in WASM"
40+
"Failed to auth - failed to load default credentials in WASM (please set env.GOOGLE_APPLICATION_CREDENTIALS, see https://docs.boundaryml.com/ref/llm-client-providers/google-vertex#using-a-vertex-ai-client-in-the-playground)"
3641
)
3742
}
3843
})
@@ -95,7 +100,7 @@ impl Claims {
95100
iss: service_account.client_email.clone(),
96101
scope: DEFAULT_SCOPE.to_string(),
97102
aud: service_account.token_uri.clone(),
98-
exp: (now + chrono::Duration::hours(60)).timestamp(),
103+
exp: (now + chrono::Duration::hours(1)).timestamp(),
99104
iat: now.timestamp(),
100105
}
101106
}
@@ -108,35 +113,3 @@ pub struct ServiceAccount {
108113
pub client_email: String,
109114
pub private_key: String,
110115
}
111-
112-
async fn get_access_token(service_account: &ServiceAccount) -> Result<String> {
113-
// Create the JWT
114-
let claims = Claims::from_service_account(service_account);
115-
116-
let jwt = encode_jwt(&serde_json::to_value(claims)?, &service_account.private_key)
117-
.await
118-
.map_err(|e| anyhow::anyhow!(format!("{e:?}")))?;
119-
120-
// Make the token request
121-
let client = reqwest::Client::new();
122-
let params = [
123-
("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer"),
124-
("assertion", &jwt),
125-
];
126-
let res: serde_json::Value = client
127-
.post(&service_account.token_uri)
128-
.form(&params)
129-
.send()
130-
.await?
131-
.json()
132-
.await?;
133-
134-
Ok(res
135-
.as_object()
136-
.context("Token exchange did not return a JSON object")?
137-
.get("access_token")
138-
.context("Access token not found in response")?
139-
.as_str()
140-
.context("Access token is not a string")?
141-
.to_string())
142-
}

0 commit comments

Comments
 (0)