feat: provide saner semantics around aws_session_token

[sxlijin]

Change aws credential loading logic:

• if any of access_key_id, secret_access_key, or session_token are set, all 3 are loaded explicitly (either from the .baml client definition or the dynamic client properties)
• if, and only if, none of the 3 are set, all 3 are loaded from, respectively, AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN

This most closely matches the behavior of the AWS SDKs (Python, TS, and Rust). See slack thread which is copied below:

OK, so chris and i figured out what happened with bedrock/ethan:

in #1266, chris correctly added support for aws session token so that if a user set it in aws.baml as properties { session_token env.AWS_SESSION_TOKEN }, baml would respect that (prior to #1266 baml would not)

• however, 1266 also introduced an implicit default: if AWS_SESSION_TOKEN is set in the process' environment, but the user only set properties { access_key_id ... ; secret_access_key ... ; } then baml would construct the aws creds using access_key_id secret_access_key and session_token
• in ethan's case this is problematic, because he uses custom values for the access_key_id secret_access_key pair from his lambda secrets, but the aws lambda environment also sets AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
• as a result, after updating past fix: issue with aws credentials not being passed in correctly #1266, his (access_key_id, secret_access_key) did not agree with his AWS_SESSION_TOKEN and caused a runtime failure
• he also has no way to opt out of this behavior, because we do not currently provide a way to force session_token to null: it is always inferred from the environment by default

to solve this, we're going to use the following logic for aws creds:

• if any of access_key_id, secret_access_key, or session_token are set in baml client properties, we will never magically infer a value for any of the 3 from the environment
  • but if none of the 3 are set, we will read all 3 from the env
• this behavior feels most in line with how credential init in the aws sdk normally works
  • in ts, an AwsCredentialIdentityProvider is any function that returns { accessKeyId: string, secretAccessKey: string, sessionToken?: string, expiration?: Date} (docs)
  • in python, if you set any of the 3, a creds object is constructed using the explicitly provided values of all 3 (impl callsite, impl source code)
  • in rust, this is what constructing Credentials::new does when you override the credentials_loader

stepping back, this is a mix of the two approaches: (1) session_token: unset defaults to reading AWS_SESSION_TOKEN from the env, and user is allowed to explicitly set session_token: null or (2) session_token: unset never reads from the env, and the user must always set it

NB: this does not explain why multiple other customers are complaining about not being able to figure out how to use aws bedrock. so @Vaibhav Gupta we still need to see what those other complaints are

---

Important

Updates AWS credential loading logic to use explicit credentials if set, otherwise defaults to environment variables, aligning with AWS SDK behavior.

• Behavior:
  • Updates AWS credential loading logic in aws_bedrock.rs and aws_client.rs.
  • If any of access_key_id, secret_access_key, or session_token are set, all are loaded explicitly.
  • If none are set, all are loaded from environment variables AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN.
• Error Handling:
  • Adds error checks for environment variable placeholders in aws_client.rs.
• Misc:
  • Adjusts credential provider logic in aws_client.rs to use DefaultCredentialsChain when no credentials are provided.

^{This description was created by for 8199507. It will automatically update as commits are pushed.}

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
baml	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Jan 7, 2025 2:36am

[ellipsis-dev]

👍 Looks good to me! Reviewed everything up to 3fa4835 in 1 minute and 42 seconds

More details

• Looked at 2869 lines of code in 37 files
• Skipped 2 files when reviewing.
• Skipped posting 1 drafted comments based on config settings.

1. engine/baml-lib/llm-client/src/clients/aws_bedrock.rs:205

• Draft comment:
  Consider refactoring the repeated pattern of fetching environment variables into a helper function to reduce redundancy and improve maintainability.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The code in aws_bedrock.rs has a repeated pattern for fetching environment variables. This can be refactored to a helper function to reduce redundancy and improve maintainability.

Workflow ID: wflow_0WDxqhPg3pw13jFk

---

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

[sxlijin]

@seawatts i suspect this can be collapsed as

Suggested change

	(None, None, None) => {
	// If no credentials provided, get them all from env vars
	let access_key_id = match ctx.get_env_var("AWS_ACCESS_KEY_ID") {
	Ok(key) if !key.is_empty() => Some(key),
	_ => None,
	};
	let secret_access_key = match ctx.get_env_var("AWS_SECRET_ACCESS_KEY") {
	Ok(key) if !key.is_empty() => Some(key),
	_ => None,
	};
	let session_token = match ctx.get_env_var("AWS_SESSION_TOKEN") {
	Ok(token) if !token.is_empty() => Some(token),
	_ => None,
	};
	(access_key_id, secret_access_key, session_token)
	}
	(None, None, None) => {
	// If no credentials provided, get them all from env vars
	(ctx.get_env_var("AWS_ACCESS_KEY_ID").ok(), ctx.get_env_var("AWS_SECRET_ACCESS_KEY").ok(), ctx.get_env_var("AWS_SESSION_TOKEN").ok())
	}

but i don't want to put this through another round of testing right now so i'm leaving it as-is

(i'm not sure if the empty string difference is particularly significant here, i think it just triggers the fallback-to-env-var logic twice via two different paths

[ellipsis-dev]

👍 Looks good to me! Reviewed everything up to 8199507 in 17 seconds

More details

• Looked at 182 lines of code in 2 files
• Skipped 0 files when reviewing.
• Skipped posting 2 drafted comments based on config settings.

1. engine/baml-lib/llm-client/src/clients/aws_bedrock.rs:205

• Draft comment:
  The logic for loading AWS credentials has been updated to ensure that if any of the credentials (access_key_id, secret_access_key, session_token) are set, they are used explicitly, and if none are set, they are loaded from environment variables. This change aligns with the behavior of AWS SDKs.
• Reason this comment was not posted:
  Confidence changes required: 10%
  The code in aws_bedrock.rs and aws_client.rs has been updated to change the logic for loading AWS credentials. The new logic ensures that if any of the credentials (access_key_id, secret_access_key, session_token) are set, they are used explicitly, and if none are set, they are loaded from environment variables. This change aligns with the behavior of AWS SDKs.

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

• Draft comment:
  The logic for loading AWS credentials has been updated to ensure that if any of the credentials (access_key_id, secret_access_key, session_token) are set, they are used explicitly, and if none are set, they are loaded from environment variables. This change aligns with the behavior of AWS SDKs.
• Reason this comment was not posted:
  Confidence changes required: 10%
  The code in aws_client.rs has been updated to change the logic for loading AWS credentials. The new logic ensures that if any of the credentials (access_key_id, secret_access_key, session_token) are set, they are used explicitly, and if none are set, they are loaded from environment variables. This change aligns with the behavior of AWS SDKs.

Workflow ID: wflow_LPzSp0Hb5O6Ei43Z

---

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.