From 987275c64a95a50ac9fa0946ffc8de00c452281f Mon Sep 17 00:00:00 2001
From: Daniel Tabuenca <dtabuenc@gmail.com>
Date: Sat, 8 Feb 2025 11:40:20 -0600
Subject: [PATCH] feat: bedrock temporary credentials support (#1207)

---
 README.md                        | 4 +++-
 lua/avante/providers/bedrock.lua | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 18ab5c1..a9979de 100644
--- a/README.md
+++ b/README.md
@@ -471,8 +471,10 @@ Given its early stage, `avante.nvim` currently supports the following basic func
 > For Amazon Bedrock:
 >
 > ```sh
-> export BEDROCK_KEYS=aws_access_key_id,aws_secret_access_key,aws_region
+> export BEDROCK_KEYS=aws_access_key_id,aws_secret_access_key,aws_region[,aws_session_token]
+>
 > ```
+> Note: The aws_session_token is optional and only needed when using temporary AWS credentials
 
 1. Open a code file in Neovim.
 2. Use the `:AvanteAsk` command to query the AI about the code.
diff --git a/lua/avante/providers/bedrock.lua b/lua/avante/providers/bedrock.lua
index 2b0cd7e..49720f3 100644
--- a/lua/avante/providers/bedrock.lua
+++ b/lua/avante/providers/bedrock.lua
@@ -64,6 +64,7 @@ M.parse_curl_args = function(provider, prompt_opts)
   local aws_access_key_id = parts[1]
   local aws_secret_access_key = parts[2]
   local aws_region = parts[3]
+  local aws_session_token = parts[4]
 
   local endpoint = string.format(
     "https://bedrock-runtime.%s.amazonaws.com/model/%s/invoke-with-response-stream",
@@ -75,6 +76,8 @@ M.parse_curl_args = function(provider, prompt_opts)
     ["Content-Type"] = "application/json",
   }
 
+  if aws_session_token and aws_session_token ~= "" then headers["x-amz-security-token"] = aws_session_token end
+
   local body_payload = M.build_bedrock_payload(prompt_opts, body_opts)
 
   local rawArgs = {