fix: issue with aws credentials not being passed in correctly

[seawatts]

Issue

Lendflow caught a bug where they were getting the error The security token included in the request is invalid when calling the baml client functions within the playground as well as their local environment while trying to use an assumed role. They were setting the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY but did not have the ability to also set AWS_SESSION_TOKEN which is required in the aws credentials chain.

Root cause and fix

When assuming an IAM role, it’s required by the AWS credentials chain to supply the AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, as well as AWS_SESSION_TOKEN. Previously, we did not pass the AWS_SESSION_TOKEN into the credentials chain. I updated the code to do so.

Additionally, I fixed another issue with AWS_PROFILE and SSO. There is now separate logic for when AWS_PROFILE is passed in that will use the credentials chain to grab the SSO token from disk (not available in WASM).

---

Important

Fix AWS credentials handling by including AWS_SESSION_TOKEN and adding AWS_PROFILE logic for SSO, with updates to integration tests.

• Behavior:
  • Include AWS_SESSION_TOKEN in credentials chain in aws_bedrock.rs and aws_client.rs.
  • Add logic for AWS_PROFILE handling for SSO in aws_client.rs.
• Code Changes:
  • Update UnresolvedAwsBedrock and ResolvedAwsBedrock structs to include session_token and profile.
  • Modify AwsClient::client_anyhow() to set credentials provider with session_token.
• Tests:
  • Add integration tests for invalid AWS credentials and profiles in integ-tests/typescript/tests/integ-tests.test.ts.
• Misc:
  • Remove patchelf dependency from pyproject.toml.

^{This description was created by for 9a3efc3. 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	Dec 31, 2024 5:40pm

[ellipsis-dev]

👍 Looks good to me! Reviewed everything up to bf6fc64 in 13 seconds

More details

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

1. integ-tests/python/pyproject.toml:23

• Draft comment:
  The patchelf dependency was removed, but this change is not mentioned in the PR description. Ensure that this removal is intentional and won't affect the build or tests.
• Reason this comment was not posted:
  Comment did not seem useful.

Workflow ID: wflow_3s2dnsisE8SqGCOR

---

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

[ellipsis-dev]

👍 Looks good to me! Incremental review on b7a18c2 in 47 seconds

More details

• Looked at 413 lines of code in 6 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:100

• Draft comment:
  Consider using unwrap_or instead of unwrap_or_else for string literals to improve readability. This applies to other similar instances in this file.
• Reason this comment was not posted:
  Confidence changes required: 20%
  The code in aws_bedrock.rs and aws_client.rs has multiple instances where unwrap_or_else is used with a closure that returns a string. This can be simplified using unwrap_or. This is a minor optimization but improves readability.

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

• Draft comment:
  Remove the redundant comma after await in the loader.credentials_provider call.
• Reason this comment was not posted:
  Confidence changes required: 10%
  In aws_client.rs, the client_anyhow function has a redundant comma in the await call. This is a minor issue but should be corrected for consistency.

Workflow ID: wflow_lRFF8wsMhgDmak34

---

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

[ellipsis-dev]

👍 Looks good to me! Incremental review on 9a3efc3 in 53 seconds

More details

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

1. integ-tests/typescript/tests/integ-tests.test.ts:45

• Draft comment:
  The test for handling invalid AWS profile is currently skipped. Ensure there's a valid reason for this, or consider enabling it to verify the functionality.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The test case for handling invalid AWS profile is currently skipped. This might be intentional, but it's important to ensure that all test cases are active unless there's a specific reason to skip them. The test should be reviewed to determine if it can be enabled or if there's a valid reason for it to remain skipped.

Workflow ID: wflow_PnhJNYHLiMXRAAxL

---

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

[ellipsis-dev]

Skipped PR review on 6370553 because no changed files had a supported extension. If you think this was in error, please contact us and we'll fix it right away.