feat (provider/amazon-bedrock): add support for session tokens (#2755)

Co-authored-by: jdavis <45986634+shoopapa@users.noreply.github.com>
Co-authored-by: Joe Davis <joedavis29@gmail.com>

.changeset/shy-beans-retire.md

@@ -0,0 +1,6 @@
+---
+'@ai-sdk/amazon-bedrock': patch
+'@ai-sdk/provider-utils': patch
+---
+
+feat (provider/amazon-bedrock): add support for session tokens

content/providers/01-ai-sdk-providers/08-amazon-bedrock.mdx

@@ -88,6 +88,17 @@ const bedrock = createAmazonBedrock({
   region: 'us-east-1',
   accessKeyId: 'xxxxxxxxx',
   secretAccessKey: 'xxxxxxxxx',
+  sessionToken: 'xxxxxxxxx',
+});
+
+// or with bedrockOptions
+const bedrock = createAmazonBedrock({
+  bedrockOptions: {
+    region: 'us-east-1',
+    credentials: {
+      // ...
+    },
+  },
 });
 ```
 
@@ -108,9 +119,14 @@ You can use the following optional settings to customize the Amazon Bedrock prov
   The AWS secret access key that you want to use for the API calls.
   It uses the `AWS_SECRET_ACCESS_KEY` environment variable by default.
 
+- **sessionToken** _string_
+
+  Optional. The AWS session token that you want to use for the API calls.
+  It uses the `AWS_SESSION_TOKEN` environment variable by default.
+
 - **bedrockOptions** _object_
 
-  Optional. The configuration options used by the [Amazon Bedrock Library](https://docs.aws.amazon.com/AWSJavaScraiptSDK/v3/latest/Package/-aws-sdk-client-bedrock-runtime/)
+  Optional. The configuration options used by the [Amazon Bedrock Library](https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/Package/-aws-sdk-client-bedrock-runtime/)
   (`BedrockRuntimeClientConfig`), including:
 
   - **region** _string_

packages/amazon-bedrock/src/bedrock-chat-language-model.test.ts

@@ -24,6 +24,7 @@ const provider = createAmazonBedrock({
   region: 'us-east-1',
   accessKeyId: 'test-access-key',
   secretAccessKey: 'test-secret-key',
+  sessionToken: 'test-token-key',
 });
 
 const model = provider('anthropic.claude-3-haiku-20240307-v1:0');

packages/amazon-bedrock/src/bedrock-provider.ts

@@ -1,4 +1,8 @@
-import { generateId, loadSetting } from '@ai-sdk/provider-utils';
+import {
+  generateId,
+  loadOptionalSetting,
+  loadSetting,
+} from '@ai-sdk/provider-utils';
 import {
   BedrockRuntimeClient,
   BedrockRuntimeClientConfig,
@@ -13,6 +17,7 @@ export interface AmazonBedrockProviderSettings {
   region?: string;
   accessKeyId?: string;
   secretAccessKey?: string;
+  sessionToken?: string;
 
   /**
    * Complete Bedrock configuration for setting advanced authentication and
@@ -65,6 +70,10 @@ export function createAmazonBedrock(
             environmentVariableName: 'AWS_SECRET_ACCESS_KEY',
             description: 'AWS secret access key',
           }),
+          sessionToken: loadOptionalSetting({
+            settingValue: options.sessionToken,
+            environmentVariableName: 'AWS_SESSION_TOKEN',
+          }),
         },
       },
     );

packages/provider-utils/src/index.ts

@@ -6,7 +6,8 @@ export * from './generate-id';
 export * from './get-error-message';
 export * from './is-abort-error';
 export * from './load-api-key';
-export * from './load-setting';
+export { loadSetting } from './load-setting';
+export { loadOptionalSetting } from './load-optional-setting';
 export * from './parse-json';
 export * from './post-to-api';
 export * from './response-handler';

packages/provider-utils/src/load-optional-setting.ts

@@ -0,0 +1,30 @@
+/**
+ * Loads an optional `string` setting from the environment or a parameter.
+ *
+ * @param settingValue - The setting value.
+ * @param environmentVariableName - The environment variable name.
+ * @returns The setting value.
+ */
+export function loadOptionalSetting({
+  settingValue,
+  environmentVariableName,
+}: {
+  settingValue: string | undefined;
+  environmentVariableName: string;
+}): string | undefined {
+  if (typeof settingValue === 'string') {
+    return settingValue;
+  }
+
+  if (settingValue != null || typeof process === 'undefined') {
+    return undefined;
+  }
+
+  settingValue = process.env[environmentVariableName];
+
+  if (settingValue == null || typeof settingValue !== 'string') {
+    return undefined;
+  }
+
+  return settingValue;
+}

packages/provider-utils/src/load-setting.ts

@@ -1,5 +1,14 @@
 import { LoadSettingError } from '@ai-sdk/provider';
 
+/**
+ * Loads a `string` setting from the environment or a parameter.
+ *
+ * @param settingValue - The setting value.
+ * @param environmentVariableName - The environment variable name.
+ * @param settingName - The setting name.
+ * @param description - The description of the setting.
+ * @returns The setting value.
+ */
 export function loadSetting({
   settingValue,
   environmentVariableName,
@@ -23,21 +32,29 @@ export function loadSetting({
 
   if (typeof process === 'undefined') {
     throw new LoadSettingError({
-      message: `${description} setting is missing. Pass it using the '${settingName}' parameter. Environment variables is not supported in this environment.`,
+      message:
+        `${description} setting is missing. ` +
+        `Pass it using the '${settingName}' parameter. ` +
+        `Environment variables is not supported in this environment.`,
     });
   }
 
   settingValue = process.env[environmentVariableName];
 
   if (settingValue == null) {
     throw new LoadSettingError({
-      message: `${description} setting is missing. Pass it using the '${settingName}' parameter or the ${environmentVariableName} environment variable.`,
+      message:
+        `${description} setting is missing. ` +
+        `Pass it using the '${settingName}' parameter ` +
+        `or the ${environmentVariableName} environment variable.`,
     });
   }
 
   if (typeof settingValue !== 'string') {
     throw new LoadSettingError({
-      message: `${description} setting must be a string. The value of the ${environmentVariableName} environment variable is not a string.`,
+      message:
+        `${description} setting must be a string. ` +
+        `The value of the ${environmentVariableName} environment variable is not a string.`,
     });
   }