diff --git a/site/src/pages/AISettingsPage/ProvidersPage/components/ProviderForm.tsx b/site/src/pages/AISettingsPage/ProvidersPage/components/ProviderForm.tsx index 99a1d24a92..48b560add7 100644 --- a/site/src/pages/AISettingsPage/ProvidersPage/components/ProviderForm.tsx +++ b/site/src/pages/AISettingsPage/ProvidersPage/components/ProviderForm.tsx @@ -25,6 +25,7 @@ export type ProviderFormValues = { smallFastModel: string; accessKey: string; accessKeySecret: string; + roleArn: string; apiKey: string; enabled: boolean; }; @@ -66,6 +67,7 @@ const defaultInitialValues: ProviderFormValues = { smallFastModel: "", accessKey: "", accessKeySecret: "", + roleArn: "", apiKey: "", enabled: true, }; @@ -526,6 +528,16 @@ export const ProviderForm: FC = ({ View docs

+ +

+ Optional. When a role ARN is set, the gateway assumes that role + (using the base identity) before calling Bedrock. +

)} diff --git a/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.test.ts b/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.test.ts index 1ae786008c..a721fd6e82 100644 --- a/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.test.ts +++ b/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.test.ts @@ -29,6 +29,7 @@ const baseOpenAIFormValues: ProviderFormValues = { smallFastModel: "", accessKey: "", accessKeySecret: "", + roleArn: "", apiKey: "sk-test", enabled: true, }; @@ -42,6 +43,7 @@ const baseBedrockFormValues: ProviderFormValues = { smallFastModel: "anthropic.claude-haiku-4-5", accessKey: "AKIA-test", accessKeySecret: "secret", + roleArn: "", apiKey: "", enabled: true, }; @@ -55,6 +57,7 @@ const baseCopilotFormValues: ProviderFormValues = { smallFastModel: "", accessKey: "", accessKeySecret: "", + roleArn: "", apiKey: "", enabled: true, }; @@ -407,6 +410,30 @@ describe("providerFormValuesToCreate", () => { }); expect(req.api_keys).toBeUndefined(); }); + + it("includes role_arn when a role ARN is provided", () => { + const req = providerFormValuesToCreate({ + ...baseBedrockFormValues, + roleArn: "arn:aws:iam::123456789012:role/BedrockRole", + }); + const s = req.settings as unknown as Record; + expect(s.role_arn).toBe("arn:aws:iam::123456789012:role/BedrockRole"); + }); + + it("omits role_arn when the form value is blank", () => { + const req = providerFormValuesToCreate(baseBedrockFormValues); + const s = req.settings as unknown as Record; + expect(s.role_arn).toBeUndefined(); + }); + + it("trims whitespace around the role ARN", () => { + const req = providerFormValuesToCreate({ + ...baseBedrockFormValues, + roleArn: " arn:aws:iam::123456789012:role/BedrockRole ", + }); + const s = req.settings as unknown as Record; + expect(s.role_arn).toBe("arn:aws:iam::123456789012:role/BedrockRole"); + }); }); describe("Copilot", () => { @@ -559,6 +586,32 @@ describe("providerFormValuesToUpdate", () => { expect(s.access_key).toBeUndefined(); expect(s.access_key_secret).toBeUndefined(); }); + + it("sends role_arn even when the access keys are kept", () => { + const req = providerFormValuesToUpdate( + { + ...baseBedrockFormValues, + accessKey: SAVED_CREDENTIAL_MASK, + accessKeySecret: SAVED_CREDENTIAL_MASK, + roleArn: "arn:aws:iam::123456789012:role/BedrockRole", + }, + MockAIProviderBedrock, + ); + const s = req.settings as unknown as Record; + expect(s.role_arn).toBe("arn:aws:iam::123456789012:role/BedrockRole"); + }); + + it("omits role_arn when the field is cleared", () => { + const req = providerFormValuesToUpdate( + { + ...baseBedrockFormValues, + roleArn: "", + }, + MockAIProviderBedrock, + ); + const s = req.settings as unknown as Record; + expect(s.role_arn).toBeUndefined(); + }); }); describe("Copilot", () => { @@ -635,6 +688,23 @@ describe("aiProviderToFormValues", () => { expect(values.accessKeySecret).toBe(""); }); + it("round-trips role_arn back into the form", () => { + const provider: AIProvider = { + ...MockAIProviderBedrock, + settings: settings({ + _type: "bedrock", + role_arn: "arn:aws:iam::123456789012:role/BedrockRole", + }), + }; + const values = aiProviderToFormValues(provider); + expect(values.roleArn).toBe("arn:aws:iam::123456789012:role/BedrockRole"); + }); + + it("seeds an empty role ARN when the provider has none", () => { + const values = aiProviderToFormValues(MockAIProviderBedrock); + expect(values.roleArn).toBe(""); + }); + it("seeds Copilot form values without a credential field", () => { const values = aiProviderToFormValues(MockAIProviderCopilot); expect(values.type).toBe("copilot"); diff --git a/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.ts b/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.ts index 67eec7e4d9..2eaaa65a4f 100644 --- a/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.ts +++ b/site/src/pages/AISettingsPage/ProvidersPage/components/providerFormApiMap.ts @@ -109,6 +109,7 @@ const buildBedrockSettings = ( smallFastModel: string, accessKey: string, accessKeySecret: string, + roleArn: string, ): BedrockSettingsWire => ({ _type: BEDROCK_SETTINGS_TYPE, _version: BEDROCK_SETTINGS_VERSION, @@ -117,6 +118,7 @@ const buildBedrockSettings = ( small_fast_model: smallFastModel, ...(accessKey ? { access_key: accessKey } : {}), ...(accessKeySecret ? { access_key_secret: accessKeySecret } : {}), + ...(roleArn ? { role_arn: roleArn } : {}), }); // Bedrock credentials live in `settings`; openai/anthropic keys go in @@ -141,6 +143,7 @@ export const providerFormValuesToCreate = ( values.smallFastModel.trim(), sanitizeCredential(values.accessKey), sanitizeCredential(values.accessKeySecret), + values.roleArn.trim(), ); return { type: "anthropic", @@ -215,6 +218,7 @@ export const providerFormValuesToUpdate = ( values.smallFastModel.trim(), credentialsChanged ? newAccessKey : "", credentialsChanged ? newAccessKeySecret : "", + values.roleArn.trim(), ); return { ...base, settings: settings as AIProviderSettings }; @@ -238,6 +242,7 @@ export const aiProviderToFormValues = ( smallFastModel: s.small_fast_model ?? "", accessKey: "", accessKeySecret: "", + roleArn: s.role_arn ?? "", enabled: provider.enabled, }; }