Skip to content

Commit

Permalink
aws is crashing - "time not implemented on this platform"
Browse files Browse the repository at this point in the history
  • Loading branch information
@sxlijin
sxlijin committed Feb 24, 2025
1 parent c2b014e commit 7a4ecce
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 71 deletions.
158 changes: 90 additions & 68 deletions engine/baml-runtime/src/internal/llm_client/primitive/aws/aws_client.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::sync::Arc;

use aws_config::Region;
use aws_config::{identity::IdentityCache, retry::RetryConfig, BehaviorVersion, ConfigLoader};
Expand Down Expand Up @@ -140,82 +141,103 @@ impl AwsClient {
// cURL previews.
async fn client_anyhow(&self, ctx: &RuntimeContext) -> Result<bedrock::Client> {
#[cfg(target_arch = "wasm32")]
let loader = super::wasm::load_aws_config();
let mut loader = {
let loader = super::wasm::load_aws_config();
let mut builder =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder();
if let Some(profile) = self.properties.profile.as_ref() {
builder = builder.profile_name(profile);
}
log::debug!("Building wasm aws credentials chain - UNCONDITIONALLY BECAUSE BULLSHIT $ENV_VAR UNCONDITIONAL SUBSTITUTION");
loader.credentials_provider(WasmAwsCreds {
default_chain: builder.build().await,
// aws_cred_provider: ctx.aws_cred_provider.clone(),
})
};
#[cfg(not(target_arch = "wasm32"))]
let loader = aws_config::defaults(BehaviorVersion::latest());

// Set credentials provider
let mut loader = match (
self.properties.access_key_id.as_ref(),
self.properties.secret_access_key.as_ref(),
self.properties.session_token.as_ref(),
) {
(None, None, None) => {
let mut builder =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder();
if let Some(profile) = self.properties.profile.as_ref() {
builder = builder.profile_name(profile);
}
log::debug!("Building wasm aws credentials chain - none of access key id / secret access key / session token provided");
// is it because of the 'static lifetime requirement?
#[cfg(target_arch = "wasm32")]
{
loader.credentials_provider(WasmAwsCreds {
default_chain: builder.build().await,
// aws_cred_provider: ctx.aws_cred_provider.clone(),
})
}
let loader = {
aws_config::defaults(BehaviorVersion::latest());

// Set credentials provider
let mut loader = match (
self.properties.access_key_id.as_ref(),
self.properties.secret_access_key.as_ref(),
self.properties.session_token.as_ref(),
) {
(None, None, None) => {
let mut builder =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder(
);
if let Some(profile) = self.properties.profile.as_ref() {
builder = builder.profile_name(profile);
}
log::debug!("Building wasm aws credentials chain - none of access key id / secret access key / session token provided");
// is it because of the 'static lifetime requirement?
#[cfg(target_arch = "wasm32")]
{
loader.credentials_provider(WasmAwsCreds {
default_chain: builder.build().await,
// aws_cred_provider: ctx.aws_cred_provider.clone(),
})
}

#[cfg(not(target_arch = "wasm32"))]
{
loader.credentials_provider(builder.build().await)
}
}
_ => {
log::debug!("Building wasm aws credentials chain - at least one of access key id / secret access key / session token provided");
if let Some(aws_access_key_id) = self.properties.access_key_id.as_ref() {
if aws_access_key_id.starts_with("$") {
return Err(anyhow::anyhow!(
"AWS access key id expected, please set: env.{}",
&aws_access_key_id[1..]
));
#[cfg(not(target_arch = "wasm32"))]
{
loader.credentials_provider(builder.build().await)
}
}
if let Some(aws_secret_access_key) = self.properties.secret_access_key.as_ref() {
// Exposing the secret key here is relatively safe. First, we expose it only
// to check if it starts with $. If so, the remainer should be an env
// var name, which is also safe to expose.
if aws_secret_access_key
.api_key
.expose_secret()
.starts_with("$")
_ => {
log::debug!(
"Building wasm aws credentials chain - at least one was provided {:?} {} {:?}",
self.properties.access_key_id,
self.properties.secret_access_key.is_some(),
self.properties.session_token,
);
if let Some(aws_access_key_id) = self.properties.access_key_id.as_ref() {
if aws_access_key_id.starts_with("$") {
return Err(anyhow::anyhow!(
"AWS access key id expected, please set: env.{}",
&aws_access_key_id[1..]
));
}
}
if let Some(aws_secret_access_key) = self.properties.secret_access_key.as_ref()
{
return Err(anyhow::anyhow!(
"AWS secret access key expected, please set: env.{}",
&aws_secret_access_key.api_key.expose_secret()[1..]
));
// Exposing the secret key here is relatively safe. First, we expose it only
// to check if it starts with $. If so, the remainer should be an env
// var name, which is also safe to expose.
if aws_secret_access_key
.api_key
.expose_secret()
.starts_with("$")
{
return Err(anyhow::anyhow!(
"AWS secret access key expected, please set: env.{}",
&aws_secret_access_key.api_key.expose_secret()[1..]
));
}
}
}
if let Some(aws_session_token) = self.properties.session_token.as_ref() {
if aws_session_token.starts_with("$") {
return Err(anyhow::anyhow!(
"AWS session token expected, please set: env.{}",
&aws_session_token[1..]
));
if let Some(aws_session_token) = self.properties.session_token.as_ref() {
if aws_session_token.starts_with("$") {
return Err(anyhow::anyhow!(
"AWS session token expected, please set: env.{}",
&aws_session_token[1..]
));
}
}
loader.credentials_provider(Credentials::new(
self.properties.access_key_id.clone().unwrap_or("".into()),
self.properties
.secret_access_key
.as_ref()
.map_or("", |key| key.api_key.expose_secret())
.to_string(),
self.properties.session_token.clone(),
None,
"baml-runtime",
))
}
loader.credentials_provider(Credentials::new(
self.properties.access_key_id.clone().unwrap_or("".into()),
self.properties
.secret_access_key
.as_ref()
.map_or("", |key| key.api_key.expose_secret())
.to_string(),
self.properties.session_token.clone(),
None,
"baml-runtime",
))
}
};
};

// Set region if specified
Expand Down
25 changes: 23 additions & 2 deletions engine/baml-runtime/src/internal/llm_client/primitive/aws/wasm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// SPDX-License-Identifier: Apache-2.0

use aws_config::ConfigLoader;
use aws_credential_types::{provider::future::ProvideCredentials, Credentials};
use aws_smithy_async::{
rt::sleep::{AsyncSleep, Sleep},
time::TimeSource,
Expand All @@ -22,6 +23,7 @@ use aws_smithy_runtime_api::{
use aws_smithy_types::body::SdkBody;

use aws_config::{BehaviorVersion, SdkConfig};
use chrono::{DateTime, Utc};
use core::pin::Pin;
use core::task::{Context, Poll};
use futures::Stream;
Expand Down Expand Up @@ -168,12 +170,20 @@ impl HttpClient for BrowserHttp2 {
}
}

#[derive(Debug)]
pub(super) struct WasmAwsCreds {
pub default_chain: aws_config::default_provider::credentials::DefaultCredentialsChain,
// pub aws_cred_provider: Arc<AwsCredProvider>,
}

impl std::fmt::Debug for WasmAwsCreds {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WasmAwsCreds")
.field("default_chain", &self.default_chain)
.field("aws_cred_provider", &"<no-repr-available>")
.finish()
}
}

impl aws_credential_types::provider::ProvideCredentials for WasmAwsCreds {
fn provide_credentials<'a>(
&'a self,
Expand All @@ -182,6 +192,17 @@ impl aws_credential_types::provider::ProvideCredentials for WasmAwsCreds {
Self: 'a,
{
log::debug!("Providing AWS credentials for wasm");
self.default_chain.provide_credentials()
// self.default_chain.provide_credentials()
// let datetime_str = "2025-02-24T19:38:05.000Z";
// let datetime: DateTime<Utc> = datetime_str.parse().expect("Invalid datetime format");
// let expires_after =
// web_time::UNIX_EPOCH + web_time::Duration::from_secs(datetime.timestamp() as u64);
ProvideCredentials::ready(Ok(Credentials::new(
"fake-access-key-id".to_string(),
"fake-secret-access-key".to_string(),
None,
None,
"hardcoded-boundaryml-dev",
)))
}
}
2 changes: 1 addition & 1 deletion engine/baml-runtime/src/types/runtime_context.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ cfg_if::cfg_if!(
} else {
use futures::future::BoxFuture;
pub type BamlSrcReader = Option<Box<fn(&str) -> BoxFuture<'static, Result<Vec<u8>>>>>;
pub type AwsCredProvider = Option<Box<fn(Option<&str>) -> BoxFuture<'static, Result<HashMap<String, String>>>>>;
pub type AwsCredProvider = Option<Box<fn(Option<&str>) -> BoxFuture<'static, Result<HashMap<String, String>>>>>>;
}
);

Expand Down
3 changes: 3 additions & 0 deletions typescript/playground-common/src/shared/baml-project-panel/atoms.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,9 @@ export const runtimeAtom = unwrap(
const envVars = get(envVarsAtom)
const vscodeEnv = (await vscode.loadEnv())?.envVars ?? {}

const awsCreds = await vscode.loadAwsCreds()
console.log('awsCreds', awsCreds)

if (wasm === undefined || project === undefined) {
return { rt: undefined, diags: undefined }
}
Expand Down

0 comments on commit 7a4ecce

Please sign in to comment.