From dd2756b570697b0fcf1426ca5203f6f9e0938d0c Mon Sep 17 00:00:00 2001
From: Timothy Carambat <rambat1010@gmail.com>
Date: Tue, 29 Oct 2024 16:34:52 -0700
Subject: [PATCH] add `sessionToken` validation connection auth for AWSbedrock
 (#2554)

---
 .../AwsBedrockLLMOptions/index.jsx            | 61 +++++++++++++++++++
 server/models/systemSettings.js               |  3 +
 server/utils/AiProviders/bedrock/index.js     | 25 ++++++++
 .../utils/agents/aibitat/providers/bedrock.js | 16 +++++
 server/utils/helpers/updateENV.js             | 11 ++++
 5 files changed, 116 insertions(+)

diff --git a/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx b/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
index 00f44a35ff..569ec4395b 100644
--- a/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
+++ b/frontend/src/components/LLMSelection/AwsBedrockLLMOptions/index.jsx
@@ -1,7 +1,12 @@
 import { ArrowSquareOut, Info } from "@phosphor-icons/react";
 import { AWS_REGIONS } from "./regions";
+import { useState } from "react";
 
 export default function AwsBedrockLLMOptions({ settings }) {
+  const [useSessionToken, setUseSessionToken] = useState(
+    settings?.AwsBedrockLLMConnectionMethod === "sessionToken"
+  );
+
   return (
     <div className="w-full flex flex-col">
       {!settings?.credentialsOnly && (
@@ -24,6 +29,43 @@ export default function AwsBedrockLLMOptions({ settings }) {
         </div>
       )}
 
+      <div className="flex flex-col gap-y-2">
+        <input
+          type="hidden"
+          name="AwsBedrockLLMConnectionMethod"
+          value={useSessionToken ? "sessionToken" : "iam"}
+        />
+        <div className="flex flex-col w-full">
+          <label className="text-white text-sm font-semibold block mb-3">
+            Use session token
+          </label>
+          <p className="text-white/50 text-sm">
+            Select the method to authenticate with AWS Bedrock.
+          </p>
+        </div>
+        <div className="flex items-center justify-start gap-x-4 bg-zinc-900 p-2.5 rounded-lg w-fit">
+          <span
+            className={`text-sm ${!useSessionToken ? "text-white" : "text-white/50"}`}
+          >
+            IAM
+          </span>
+          <label className="relative inline-flex items-center cursor-pointer">
+            <input
+              type="checkbox"
+              className="sr-only peer"
+              checked={useSessionToken}
+              onChange={(e) => setUseSessionToken(e.target.checked)}
+            />
+            <div className="w-11 h-6 bg-zinc-700 peer-focus:outline-none rounded-full peer peer-checked:after:translate-x-full after:content-[''] after:absolute after:top-[2px] after:left-[2px] after:bg-white after:rounded-full after:h-5 after:w-5 after:transition-all peer-checked:bg-primary-button"></div>
+          </label>
+          <span
+            className={`text-sm ${useSessionToken ? "text-white" : "text-white/50"}`}
+          >
+            Session Token
+          </span>
+        </div>
+      </div>
+
       <div className="w-full flex items-center gap-[36px] my-1.5">
         <div className="flex flex-col w-60">
           <label className="text-white text-sm font-semibold block mb-3">
@@ -59,6 +101,25 @@ export default function AwsBedrockLLMOptions({ settings }) {
             spellCheck={false}
           />
         </div>
+        {useSessionToken && (
+          <div className="flex flex-col w-60">
+            <label className="text-white text-sm font-semibold block mb-3">
+              AWS Bedrock Session Token
+            </label>
+            <input
+              type="password"
+              name="AwsBedrockLLMSessionToken"
+              className="border-none bg-zinc-900 text-white placeholder:text-white/20 text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+              placeholder="AWS Bedrock Session Token"
+              defaultValue={
+                settings?.AwsBedrockLLMSessionToken ? "*".repeat(20) : ""
+              }
+              required={true}
+              autoComplete="off"
+              spellCheck={false}
+            />
+          </div>
+        )}
         <div className="flex flex-col w-60">
           <label className="text-white text-sm font-semibold block mb-3">
             AWS region
diff --git a/server/models/systemSettings.js b/server/models/systemSettings.js
index 55569be078..41dfc92931 100644
--- a/server/models/systemSettings.js
+++ b/server/models/systemSettings.js
@@ -505,8 +505,11 @@ const SystemSettings = {
       GenericOpenAiKey: !!process.env.GENERIC_OPEN_AI_API_KEY,
       GenericOpenAiMaxTokens: process.env.GENERIC_OPEN_AI_MAX_TOKENS,
 
+      AwsBedrockLLMConnectionMethod:
+        process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam",
       AwsBedrockLLMAccessKeyId: !!process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
       AwsBedrockLLMAccessKey: !!process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+      AwsBedrockLLMSessionToken: !!process.env.AWS_BEDROCK_LLM_SESSION_TOKEN,
       AwsBedrockLLMRegion: process.env.AWS_BEDROCK_LLM_REGION,
       AwsBedrockLLMModel: process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE,
       AwsBedrockLLMTokenLimit: process.env.AWS_BEDROCK_LLM_MODEL_TOKEN_LIMIT,
diff --git a/server/utils/AiProviders/bedrock/index.js b/server/utils/AiProviders/bedrock/index.js
index c271f7297e..8d9be2d392 100644
--- a/server/utils/AiProviders/bedrock/index.js
+++ b/server/utils/AiProviders/bedrock/index.js
@@ -31,6 +31,14 @@ class AWSBedrockLLM {
     if (!process.env.AWS_BEDROCK_LLM_REGION)
       throw new Error("No AWS Bedrock LLM region was set.");
 
+    if (
+      process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD === "sessionToken" &&
+      !process.env.AWS_BEDROCK_LLM_SESSION_TOKEN
+    )
+      throw new Error(
+        "No AWS Bedrock LLM session token was set while using session token as the authentication method."
+      );
+
     this.model =
       modelPreference || process.env.AWS_BEDROCK_LLM_MODEL_PREFERENCE;
     this.limits = {
@@ -41,6 +49,20 @@ class AWSBedrockLLM {
 
     this.embedder = embedder ?? new NativeEmbedder();
     this.defaultTemp = 0.7;
+    this.#log(
+      `Loaded with model: ${this.model}. Will communicate with AWS Bedrock using ${this.authMethod} authentication.`
+    );
+  }
+
+  /**
+   * Get the authentication method for the AWS Bedrock LLM.
+   * There are only two valid values for this setting - anything else will default to "iam".
+   * @returns {"iam"|"sessionToken"}
+   */
+  get authMethod() {
+    const method = process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam";
+    if (!["iam", "sessionToken"].includes(method)) return "iam";
+    return method;
   }
 
   #bedrockClient({ temperature = 0.7 }) {
@@ -51,6 +73,9 @@ class AWSBedrockLLM {
       credentials: {
         accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
         secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+        ...(this.authMethod === "sessionToken"
+          ? { sessionToken: process.env.AWS_BEDROCK_LLM_SESSION_TOKEN }
+          : {}),
       },
       temperature,
     });
diff --git a/server/utils/agents/aibitat/providers/bedrock.js b/server/utils/agents/aibitat/providers/bedrock.js
index 819dfc4053..6cdb92897a 100644
--- a/server/utils/agents/aibitat/providers/bedrock.js
+++ b/server/utils/agents/aibitat/providers/bedrock.js
@@ -22,6 +22,11 @@ class AWSBedrockProvider extends InheritMultiple([Provider, UnTooled]) {
       credentials: {
         accessKeyId: process.env.AWS_BEDROCK_LLM_ACCESS_KEY_ID,
         secretAccessKey: process.env.AWS_BEDROCK_LLM_ACCESS_KEY,
+        // If we're using a session token, we need to pass it in as a credential
+        // otherwise we must omit it so it does not conflict if using IAM auth
+        ...(this.authMethod === "sessionToken"
+          ? { sessionToken: process.env.AWS_BEDROCK_LLM_SESSION_TOKEN }
+          : {}),
       },
       model,
     });
@@ -31,6 +36,17 @@ class AWSBedrockProvider extends InheritMultiple([Provider, UnTooled]) {
     this.verbose = true;
   }
 
+  /**
+   * Get the authentication method for the AWS Bedrock LLM.
+   * There are only two valid values for this setting - anything else will default to "iam".
+   * @returns {"iam"|"sessionToken"}
+   */
+  get authMethod() {
+    const method = process.env.AWS_BEDROCK_LLM_CONNECTION_METHOD || "iam";
+    if (!["iam", "sessionToken"].includes(method)) return "iam";
+    return method;
+  }
+
   get client() {
     return this._client;
   }
diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js
index a884f1324b..6081159a52 100644
--- a/server/utils/helpers/updateENV.js
+++ b/server/utils/helpers/updateENV.js
@@ -213,6 +213,13 @@ const KEY_MAPPING = {
   },
 
   // AWS Bedrock LLM InferenceSettings
+  AwsBedrockLLMConnectionMethod: {
+    envKey: "AWS_BEDROCK_LLM_CONNECTION_METHOD",
+    checks: [
+      (input) =>
+        ["iam", "sessionToken"].includes(input) ? null : "Invalid value",
+    ],
+  },
   AwsBedrockLLMAccessKeyId: {
     envKey: "AWS_BEDROCK_LLM_ACCESS_KEY_ID",
     checks: [isNotEmpty],
@@ -221,6 +228,10 @@ const KEY_MAPPING = {
     envKey: "AWS_BEDROCK_LLM_ACCESS_KEY",
     checks: [isNotEmpty],
   },
+  AwsBedrockLLMSessionToken: {
+    envKey: "AWS_BEDROCK_LLM_SESSION_TOKEN",
+    checks: [],
+  },
   AwsBedrockLLMRegion: {
     envKey: "AWS_BEDROCK_LLM_REGION",
     checks: [isNotEmpty],