From ca9069e0fa5daebbaa8274822717278934809d65 Mon Sep 17 00:00:00 2001 From: Brian Mcclaugherty Date: Mon, 25 Aug 2025 10:32:59 -0400 Subject: [PATCH] Add AWS Bedrock sessionToken (#478) * add aws bedrock sessionToken * update changelog for bedrock sessionToken --------- Co-authored-by: Brendan Kellam --- CHANGELOG.md | 1 + .../language-model-providers.mdx | 3 + docs/snippets/schemas/v3/index.schema.mdx | 62 +++++++++++++++++++ .../schemas/v3/languageModel.schema.mdx | 62 +++++++++++++++++++ packages/schemas/src/v3/index.schema.ts | 62 +++++++++++++++++++ packages/schemas/src/v3/index.type.ts | 16 +++++ .../schemas/src/v3/languageModel.schema.ts | 62 +++++++++++++++++++ packages/schemas/src/v3/languageModel.type.ts | 16 +++++ packages/web/src/env.mjs | 1 + packages/web/src/features/chat/actions.ts | 3 + schemas/v3/languageModel.json | 4 ++ 11 files changed, 292 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c369861..d3d3d717 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - Updated NextJS to version 15. [#477](https://github.com/sourcebot-dev/sourcebot/pull/477) +- Add `sessionToken` as optional Bedrock configuration parameter. [#478](https://github.com/sourcebot-dev/sourcebot/pull/478) ## [4.6.4] - 2025-08-11 diff --git a/docs/docs/configuration/language-model-providers.mdx b/docs/docs/configuration/language-model-providers.mdx index d10951c7..44865e6d 100644 --- a/docs/docs/configuration/language-model-providers.mdx +++ b/docs/docs/configuration/language-model-providers.mdx @@ -67,6 +67,9 @@ For a detailed description of all the providers, please refer to the [schema](ht "accessKeySecret": { "env": "AWS_SECRET_ACCESS_KEY" }, + "sessionToken": { + "env": "AWS_SESSION_TOKEN" + }, "region": "YOUR_REGION_HERE", // defaults to the AWS_REGION env var if not set "baseUrl": "OPTIONAL_BASE_URL" } diff --git a/docs/snippets/schemas/v3/index.schema.mdx b/docs/snippets/schemas/v3/index.schema.mdx index 51600da2..291cc3b2 100644 --- a/docs/snippets/schemas/v3/index.schema.mdx +++ b/docs/snippets/schemas/v3/index.schema.mdx @@ -1284,6 +1284,37 @@ } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", @@ -2628,6 +2659,37 @@ } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", diff --git a/docs/snippets/schemas/v3/languageModel.schema.mdx b/docs/snippets/schemas/v3/languageModel.schema.mdx index 22785e7c..c24c17ac 100644 --- a/docs/snippets/schemas/v3/languageModel.schema.mdx +++ b/docs/snippets/schemas/v3/languageModel.schema.mdx @@ -81,6 +81,37 @@ } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", @@ -1425,6 +1456,37 @@ } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", diff --git a/packages/schemas/src/v3/index.schema.ts b/packages/schemas/src/v3/index.schema.ts index 804093fb..efa924f6 100644 --- a/packages/schemas/src/v3/index.schema.ts +++ b/packages/schemas/src/v3/index.schema.ts @@ -1283,6 +1283,37 @@ const schema = { } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", @@ -2627,6 +2658,37 @@ const schema = { } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", diff --git a/packages/schemas/src/v3/index.type.ts b/packages/schemas/src/v3/index.type.ts index 85dc22b2..ef225f46 100644 --- a/packages/schemas/src/v3/index.type.ts +++ b/packages/schemas/src/v3/index.type.ts @@ -488,6 +488,22 @@ export interface AmazonBedrockLanguageModel { */ env: string; }; + /** + * Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable. + */ + sessionToken?: + | { + /** + * The name of the secret that contains the token. + */ + secret: string; + } + | { + /** + * The name of the environment variable that contains the token. Only supported in declarative connection configs. + */ + env: string; + }; /** * The AWS region. Defaults to the `AWS_REGION` environment variable. */ diff --git a/packages/schemas/src/v3/languageModel.schema.ts b/packages/schemas/src/v3/languageModel.schema.ts index 8142fcae..2e8ec4bd 100644 --- a/packages/schemas/src/v3/languageModel.schema.ts +++ b/packages/schemas/src/v3/languageModel.schema.ts @@ -80,6 +80,37 @@ const schema = { } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", @@ -1424,6 +1455,37 @@ const schema = { } ] }, + "sessionToken": { + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable.", + "anyOf": [ + { + "type": "object", + "properties": { + "secret": { + "type": "string", + "description": "The name of the secret that contains the token." + } + }, + "required": [ + "secret" + ], + "additionalProperties": false + }, + { + "type": "object", + "properties": { + "env": { + "type": "string", + "description": "The name of the environment variable that contains the token. Only supported in declarative connection configs." + } + }, + "required": [ + "env" + ], + "additionalProperties": false + } + ] + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.", diff --git a/packages/schemas/src/v3/languageModel.type.ts b/packages/schemas/src/v3/languageModel.type.ts index d420d11b..7de00803 100644 --- a/packages/schemas/src/v3/languageModel.type.ts +++ b/packages/schemas/src/v3/languageModel.type.ts @@ -59,6 +59,22 @@ export interface AmazonBedrockLanguageModel { */ env: string; }; + /** + * Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable. + */ + sessionToken?: + | { + /** + * The name of the secret that contains the token. + */ + secret: string; + } + | { + /** + * The name of the environment variable that contains the token. Only supported in declarative connection configs. + */ + env: string; + }; /** * The AWS region. Defaults to the `AWS_REGION` environment variable. */ diff --git a/packages/web/src/env.mjs b/packages/web/src/env.mjs index 808395a7..b96beb7c 100644 --- a/packages/web/src/env.mjs +++ b/packages/web/src/env.mjs @@ -121,6 +121,7 @@ export const env = createEnv({ AWS_ACCESS_KEY_ID: z.string().optional(), AWS_SECRET_ACCESS_KEY: z.string().optional(), + AWS_SESSION_TOKEN: z.string().optional(), AWS_REGION: z.string().optional(), SOURCEBOT_CHAT_MODEL_TEMPERATURE: numberSchema.default(0.3), diff --git a/packages/web/src/features/chat/actions.ts b/packages/web/src/features/chat/actions.ts index 9ca1e826..b2449282 100644 --- a/packages/web/src/features/chat/actions.ts +++ b/packages/web/src/features/chat/actions.ts @@ -389,6 +389,9 @@ export const _getAISDKLanguageModelAndOptions = async (config: LanguageModel, or secretAccessKey: config.accessKeySecret ? await getTokenFromConfig(config.accessKeySecret, orgId, prisma) : env.AWS_SECRET_ACCESS_KEY, + sessionToken: config.sessionToken + ? await getTokenFromConfig(config.sessionToken, orgId, prisma) + : env.AWS_SESSION_TOKEN, headers: config.headers ? await extractLanguageModelHeaders(config.headers, orgId, prisma) : undefined, diff --git a/schemas/v3/languageModel.json b/schemas/v3/languageModel.json index 660cd44f..d362fac5 100644 --- a/schemas/v3/languageModel.json +++ b/schemas/v3/languageModel.json @@ -25,6 +25,10 @@ "$ref": "./shared.json#/definitions/Token", "description": "Optional secret access key to use with the model. Defaults to the `AWS_SECRET_ACCESS_KEY` environment variable." }, + "sessionToken": { + "$ref": "./shared.json#/definitions/Token", + "description": "Optional session token to use with the model. Defaults to the `AWS_SESSION_TOKEN` environment variable." + }, "region": { "type": "string", "description": "The AWS region. Defaults to the `AWS_REGION` environment variable.",