feat: make gcp auth work seamlessly from vscode (#1860)

Implement loading GCP oauth token via a JS callback using the gcp node
sdk, saving users from having to create a service account and put the
secret key in an env var to use the playground.

Also convert from using a DI-esque pattern where things are wired
through RuntimeContext to instead using a singleton pattern which uses
channels to mimic RPCs across the wasm/js boundary to drive the
callbacks (necessary because of the `Send` requirements that tokio
imposes on `js_sys::Value` which is not `Send`).

DONE:

- [x] - test the errors - what happens when the user is not gcloud
auth'd? is the error message good?
- [x] - restore the vertex wasm_auth path when env vars are explicitly
set
- [x] - fix this error (explicit env var should never resolve to file in
wasm) "Failed to auth - cannot load credentials from a file in WASM
(path='$INTE...', path.len=51)",
- [x] - fix this error (explicit env var should have a useful error)
"Failed to auth - cannot load credentials from a file in WASM
(path='$INTE...', path.len=51)",
- [x] - try to drop the custom result enums
- [x] - rename aws* to js-callback-bridge*
- [x] - update docs (both this and the vertex docs saying that claude is
not supported)

Fixes #1857 
<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Implement seamless GCP OAuth token loading from VSCode using JS
callbacks, refactor AWS credential handling, and update documentation.
> 
>   - **Behavior**:
> - Implement GCP OAuth token loading via JS callback using GCP Node
SDK, removing need for service account keys in env vars.
> - Refactor AWS credential handling to use JS callback bridge pattern.
> - Update `VertexAuth` in `vertex/std_auth.rs` and
`vertex/wasm_auth.rs` to support new auth strategies.
>   - **JS Callback Bridge**:
> - Introduce `JsCallbackProvider` in `js_callback_provider.rs` for AWS
and GCP credential requests.
> - Implement `init_js_callback_bridge()` in `js_callback_bridge.rs` to
set up JS callbacks for credential loading.
>   - **Documentation**:
> - Update `vertex.mdx` to reflect new GCP auth method and deprecate
`credentials_content`.
>   - **Misc**:
> - Rename `AuthStrategy::JsonString` and `AuthStrategy::JsonFile` to
`StringContainingJson` and `MaybeFilePath` respectively.
>     - Add `derive-new` dependency in `Cargo.toml` and `Cargo.lock`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for daee78a. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: sxlijin

engine/Cargo.lock

@@ -50,6 +50,7 @@ bytes = "1.6.0"
 cfg-if = "1.0.0"
 clap = { version = "4.4.6", features = ["cargo", "derive"] }
 dashmap = "5.5.3"
+derive-new = "0.7"
 derive_builder = "0.20.0"
 derive_more = { version = "0.99.18", features = ["constructor"] }
 either = "1.8.1"

engine/Cargo.toml

@@ -9,8 +9,6 @@ use anyhow::Result;
 use baml_types::{GetEnvVar, StringOr, UnresolvedValue};
 use either::Either;
 use indexmap::IndexMap;
-use secrecy::SecretString;
-use serde::Deserialize;
 
 use super::helpers::{Error, PropertyHandler, UnresolvedUrl};
 
@@ -26,19 +24,12 @@ enum UnresolvedGcpAuthStrategy<Meta> {
     SystemDefault,
 }
 
-#[derive(Debug, Deserialize)]
-pub struct ServiceAccount {
-    pub token_uri: String,
-    pub project_id: String,
-    pub client_email: String,
-    pub private_key: SecretString,
-}
-
 pub enum ResolvedGcpAuthStrategy {
-    /// This is the mechanism that GCP SDKs usually support for GOOGLE_APPLICATION_CREDENTIALS
-    FilePath(String),
+    /// GCP SDKs usually support passing in GOOGLE_APPLICATION_CREDENTIALS as a file path
+    /// In WASM, however, we treat both StringContainingJson and MaybeFilePath as a string
+    MaybeFilePath(String),
     /// Because the WASM playground needs a way to pass in credentials.
-    JsonString(String),
+    StringContainingJson(String),
     /// JsonObject was implemented for a user: https://github.com/BoundaryML/baml/issues/1001
     JsonObject(IndexMap<String, String>),
     /// The normal GCP application default credentials flow, after checking
@@ -94,8 +85,8 @@ impl<Meta> UnresolvedGcpAuthStrategy<Meta> {
             UnresolvedGcpAuthStrategy::CredentialsString(s) => {
                 let s = s.resolve(ctx)?;
                 match serde_json::from_str::<serde_json::Value>(&s) {
-                    Ok(_) => ResolvedGcpAuthStrategy::JsonString(s),
-                    Err(_) => ResolvedGcpAuthStrategy::FilePath(s),
+                    Ok(_) => ResolvedGcpAuthStrategy::StringContainingJson(s),
+                    Err(_) => ResolvedGcpAuthStrategy::MaybeFilePath(s),
                 }
             }
             UnresolvedGcpAuthStrategy::CredentialsJsonObject(m) => {
@@ -107,7 +98,7 @@ impl<Meta> UnresolvedGcpAuthStrategy<Meta> {
             }
             UnresolvedGcpAuthStrategy::CredentialsContentString(s) => {
                 let s = s.resolve(ctx)?;
-                ResolvedGcpAuthStrategy::JsonString(s)
+                ResolvedGcpAuthStrategy::StringContainingJson(s)
             }
             UnresolvedGcpAuthStrategy::SystemDefault => {
                 log::debug!("Neither options.credentials nor options.credentials_content are set, falling back to env vars");
@@ -125,16 +116,16 @@ impl<Meta> UnresolvedGcpAuthStrategy<Meta> {
                             log::warn!("Resolving GOOGLE_APPLICATION_CREDENTIALS from env, but it is an empty string");
                         }
                         match serde_json::from_str::<serde_json::Value>(&credentials) {
-                            Ok(_) => ResolvedGcpAuthStrategy::JsonString(credentials),
-                            Err(_) => ResolvedGcpAuthStrategy::FilePath(credentials),
+                            Ok(_) => ResolvedGcpAuthStrategy::StringContainingJson(credentials),
+                            Err(_) => ResolvedGcpAuthStrategy::MaybeFilePath(credentials),
                         }
                     }
                     (None, Some(credentials_content)) => {
                         log::debug!("Using GOOGLE_APPLICATION_CREDENTIALS_CONTENT from env");
                         if credentials_content.is_empty() {
                             log::warn!("Resolving GOOGLE_APPLICATION_CREDENTIALS_CONTENT from env, but it is an empty string");
                         }
-                        ResolvedGcpAuthStrategy::JsonString(credentials_content)
+                        ResolvedGcpAuthStrategy::StringContainingJson(credentials_content)
                     }
                     (None, None) => {
                         log::debug!("Using UseSystemDefault strategy");

engine/baml-lib/llm-client/src/clients/vertex.rs

@@ -26,6 +26,7 @@ cfg-if.workspace = true
 clap.workspace = true
 colored = "2.1.0"
 dashmap.workspace = true
+derive-new.workspace = true
 derive_more.workspace = true
 dunce = "1.0.4"
 either.workspace = true

engine/baml-runtime/Cargo.toml

@@ -311,8 +311,7 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
             Err(e) => return e.into_response(),
         };
 
-        let ctx_mgr =
-            RuntimeContextManager::new_from_env_vars(std::env::vars().collect(), None, None);
+        let ctx_mgr = RuntimeContextManager::new_from_env_vars(std::env::vars().collect(), None);
         let client_registry = b_options.and_then(|options| options.client_registry);
 
         let locked = self.b.read().await;
@@ -400,7 +399,7 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
 
         tokio::spawn(async move {
             let ctx_mgr =
-                RuntimeContextManager::new_from_env_vars(std::env::vars().collect(), None, None);
+                RuntimeContextManager::new_from_env_vars(std::env::vars().collect(), None);
 
             let result_stream = self.b.read().await.stream_function(
                 b_fn,

engine/baml-runtime/src/cli/serve/mod.rs

@@ -51,7 +51,7 @@ use crate::internal::llm_client::{
     ModelFeatures, ResolveMediaUrls,
 };
 use crate::tracingv2::storage::storage::BAML_TRACER;
-use crate::{json_body, AwsCredProvider, JsonBodyInput, RenderCurlSettings, RuntimeContext};
+use crate::{json_body, JsonBodyInput, RenderCurlSettings, RuntimeContext};
 // See https://github.com/awslabs/aws-sdk-rust/issues/169
 use super::custom_http_client;
 #[cfg(target_arch = "wasm32")]
@@ -307,7 +307,6 @@ impl AwsClient {
         &self,
         span_id: Option<Uuid>,
         http_request_id: &HttpRequestId,
-        aws_cred_provider: AwsCredProvider,
     ) -> Result<bedrock::Client> {
         #[cfg(target_arch = "wasm32")]
         let loader = super::wasm::load_aws_config();
@@ -324,7 +323,6 @@ impl AwsClient {
                 #[cfg(target_arch = "wasm32")]
                 {
                     loader.credentials_provider(WasmAwsCreds {
-                        aws_cred_provider: aws_cred_provider.clone(),
                         profile: self.properties.profile.clone(),
                     })
                 }
@@ -564,11 +562,7 @@ impl WithStreamChat for AwsClient {
         let prompt = internal_baml_jinja::RenderedPrompt::Chat(chat_messages.to_vec());
 
         let aws_client = match self
-            .client_anyhow(
-                ctx.span_id.clone(),
-                &http_request_id,
-                ctx.aws_cred_provider.clone(),
-            )
+            .client_anyhow(ctx.span_id.clone(), &http_request_id)
             .await
         {
             Ok(c) => c,
@@ -864,11 +858,7 @@ impl WithChat for AwsClient {
         let prompt = internal_baml_jinja::RenderedPrompt::Chat(chat_messages.to_vec());
 
         let aws_client = match self
-            .client_anyhow(
-                ctx.span_id.clone(),
-                &http_request_id,
-                ctx.aws_cred_provider.clone(),
-            )
+            .client_anyhow(ctx.span_id.clone(), &http_request_id)
             .await
         {
             Ok(c) => c,

engine/baml-runtime/src/internal/llm_client/primitive/aws/aws_client.rs

@@ -38,7 +38,7 @@ use std::sync::Arc;
 use std::time::SystemTime;
 use time::OffsetDateTime;
 
-use crate::{AwsCredProvider, AwsCredProviderImpl, AwsCredResult};
+use crate::{js_callback_provider::get_js_callback_provider, AwsCredResult, JsCallbackProvider};
 
 pub fn load_aws_config() -> ConfigLoader {
     log::debug!("Loading AWS config for wasm specifically");
@@ -160,8 +160,6 @@ impl HttpClient for BrowserHttp2 {
 }
 
 pub(super) struct WasmAwsCreds {
-    // pub default_chain: aws_config::default_provider::credentials::DefaultCredentialsChain,
-    pub aws_cred_provider: AwsCredProvider,
     pub profile: Option<String>,
 }
 
@@ -176,59 +174,25 @@ impl std::fmt::Debug for WasmAwsCreds {
 
 impl WasmAwsCreds {
     async fn provide_credentials_impl(&self) -> aws_credential_types::provider::Result {
-        match self.aws_cred_provider.clone() {
-            Some(AwsCredProviderImpl {
-                req_tx,
-                mut resp_rx,
-            }) => {
-                if let Err(e) = req_tx.send(self.profile.clone()).await {
-                    log::error!(
-                        "Failed to send AWS cred request across WASM bridge: {:?}",
-                        e
-                    );
-                    return Err(CredentialsError::unhandled(e));
-                };
-                let creds = match resp_rx.recv().await {
-                    Ok(Ok(creds)) => creds,
-                    Ok(Err(e)) => {
-                        log::error!("Error in AWS cred provider: {:?}", e);
-                        return Err(CredentialsError::unhandled(e));
-                    }
-                    Err(e) => {
-                        log::error!(
-                            "Failed to recv AWS cred response across WASM bridge: {:?}",
-                            e
-                        );
-                        return Err(CredentialsError::unhandled(e));
-                    }
-                };
-
-                match creds {
-                    AwsCredResult::Ok {
-                        access_key_id,
-                        secret_access_key,
-                        session_token,
-                        expiration,
-                        ..
-                    } => Ok(Credentials::new(
-                        access_key_id,
-                        secret_access_key,
-                        session_token,
-                        match expiration {
-                            Some(expiration) => match expiration.parse::<DateTime<Utc>>() {
-                                Ok(dt) => Some(dt.into()),
-                                Err(_) => None,
-                            },
-                            None => None,
-                        },
-                        "baml-playground-wasm-bridge",
-                    )),
-                    AwsCredResult::Err { name, message } => Err(CredentialsError::unhandled(
-                        format!("{}: {}", name, message),
-                    )),
-                }
+        let cred_provider = get_js_callback_provider().map_err(CredentialsError::unhandled)?;
+        match cred_provider.aws_req(self.profile.clone()).await {
+            Err(e) => {
+                log::error!("Error calling AWS cred provider: {:?}", e);
+                return Err(CredentialsError::unhandled(e));
             }
-            None => Err(CredentialsError::not_loaded_no_source()),
+            Ok(aws_creds) => Ok(Credentials::new(
+                aws_creds.access_key_id,
+                aws_creds.secret_access_key,
+                aws_creds.session_token,
+                match aws_creds.expiration {
+                    Some(expiration) => match expiration.parse::<DateTime<Utc>>() {
+                        Ok(dt) => Some(dt.into()),
+                        Err(_) => None,
+                    },
+                    None => None,
+                },
+                "baml-playground-wasm-bridge",
+            )),
         }
     }
 }

engine/baml-runtime/src/internal/llm_client/primitive/aws/wasm.rs

@@ -14,14 +14,30 @@ pub enum VertexAuth {
 impl VertexAuth {
     pub async fn new(auth_strategy: &ResolvedGcpAuthStrategy) -> Result<VertexAuth> {
         match auth_strategy {
-            ResolvedGcpAuthStrategy::FilePath(path) => {
+            ResolvedGcpAuthStrategy::MaybeFilePath(path) => {
                 log::debug!("Attempting to auth using JsonFile strategy");
-                let authz_user = gcp_auth::CustomServiceAccount::from_file(&path)?;
+                let authz_user =
+                    gcp_auth::CustomServiceAccount::from_file(&path).context(format!(
+                        "Failed to parse credentials as JSON file: {}",
+                        serde_json::to_string(&path)
+                            .expect("Serialization of string should always succeed")
+                    ))?;
                 Ok(VertexAuth::CustomServiceAccount(authz_user))
             }
-            ResolvedGcpAuthStrategy::JsonString(s) => {
+            ResolvedGcpAuthStrategy::StringContainingJson(s) => {
                 log::debug!("Attempting to auth using JsonString strategy");
-                let authz_user = gcp_auth::CustomServiceAccount::from_json(&s)?;
+                let authz_user = gcp_auth::CustomServiceAccount::from_json(&s).context(format!(
+                    "Failed to parse credentials as JSON string: {}",
+                    {
+                        let s = serde_json::to_string(&s)
+                            .expect("Serialization of string should always succeed");
+                        if s.len() > 8 {
+                            format!("{}...{}", &s[..4], &s[s.len() - 4..])
+                        } else {
+                            s
+                        }
+                    }
+                ))?;
                 Ok(VertexAuth::CustomServiceAccount(authz_user))
             }
             ResolvedGcpAuthStrategy::JsonObject(o) => {

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

@@ -30,9 +30,7 @@ use chrono::{Duration, Utc};
 use futures::StreamExt;
 #[cfg(not(target_arch = "wasm32"))]
 use gcp_auth::TokenProvider;
-use internal_llm_client::vertex::{
-    BaseUrlOrLocation, ResolvedGcpAuthStrategy, ResolvedVertex, ServiceAccount,
-};
+use internal_llm_client::vertex::{BaseUrlOrLocation, ResolvedGcpAuthStrategy, ResolvedVertex};
 use internal_llm_client::{
     AllowedRoleMetadata, ClientProvider, ResolvedClientProperty, UnresolvedClientProperty,
 };