From 9db2cbe18872f8edb0312707be9cf10df344c756 Mon Sep 17 00:00:00 2001
From: chrarnoldus <12196001+chrarnoldus@users.noreply.github.com>
Date: Mon, 31 Aug 2026 09:08:29 +0000
Subject: [PATCH] feat(byok): support Bedrock API keys through Vercel

Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>
---
 .../organizations/byok/BYOKKeysManager.tsx    |  25 ++-
 .../openrouter/inference-provider-id.test.ts  |  35 ++++
 .../openrouter/inference-provider-id.ts       |  12 +-
 .../ai-gateway/providers/openrouter/types.ts  |   4 +-
 .../ai-gateway/providers/vercel/index.test.ts |  54 ++++++
 .../lib/ai-gateway/providers/vercel/index.ts  |   8 +-
 .../ai-gateway/rewriteModelResponse.test.ts   |   7 +-
 apps/web/src/routers/byok-router.test.ts      | 171 +++++++++++++++++-
 8 files changed, 298 insertions(+), 18 deletions(-)

diff --git a/apps/web/src/components/organizations/byok/BYOKKeysManager.tsx b/apps/web/src/components/organizations/byok/BYOKKeysManager.tsx
index ed96468d5a..a2b417c6a5 100644
--- a/apps/web/src/components/organizations/byok/BYOKKeysManager.tsx
+++ b/apps/web/src/components/organizations/byok/BYOKKeysManager.tsx
@@ -43,7 +43,7 @@ import {
   DirectUserByokInferenceProviderIdSchema,
   UserByokProviderIdSchema,
   VercelUserByokInferenceProviderIdSchema,
-  AwsCredentialsSchema,
+  BedrockCredentialsSchema,
   VertexCredentialsSchema,
   type VercelUserByokInferenceProviderId,
 } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id';
@@ -280,7 +280,7 @@ export function BYOKKeysManager({ organizationId }: BYOKKeysManagerProps) {
     if (!value) return null;
     const schema =
       providerId === VercelUserByokInferenceProviderIdSchema.enum.bedrock
-        ? AwsCredentialsSchema
+        ? BedrockCredentialsSchema
         : providerId === VercelUserByokInferenceProviderIdSchema.enum.vertex
           ? VertexCredentialsSchema
           : null;
@@ -293,11 +293,10 @@ export function BYOKKeysManager({ organizationId }: BYOKKeysManagerProps) {
     }
     const result = schema.safeParse(parsed);
     if (!result.success) {
-      const providerName =
-        providerId === VercelUserByokInferenceProviderIdSchema.enum.bedrock
-          ? 'AWS'
-          : 'Google Vertex';
-      return `Invalid ${providerName} credentials:\n${z.prettifyError(result.error)}`;
+      if (providerId === VercelUserByokInferenceProviderIdSchema.enum.bedrock) {
+        return 'Enter JSON with apiKey and region, or accessKeyId, secretAccessKey, and region. Use only one authentication method.';
+      }
+      return `Invalid Google Vertex credentials:\n${z.prettifyError(result.error)}`;
     }
     return null;
   };
@@ -581,7 +580,7 @@ export function BYOKKeysManager({ organizationId }: BYOKKeysManagerProps) {
               <div className="space-y-2">
                 <Label htmlFor="apiKey">
                   {selectedProvider === VercelUserByokInferenceProviderIdSchema.enum.bedrock
-                    ? 'AWS Credentials'
+                    ? 'AWS Bedrock Credentials'
                     : selectedProvider === VercelUserByokInferenceProviderIdSchema.enum.vertex
                       ? 'Google Vertex Credentials'
                       : 'API Key'}
@@ -650,7 +649,15 @@ export function BYOKKeysManager({ organizationId }: BYOKKeysManagerProps) {
                   <Alert>
                     <Info className="size-4" />
                     <AlertDescription>
-                      <p>Enter your AWS credentials as JSON:</p>
+                      <p>Enter a Bedrock API key and AWS region as JSON:</p>
+                      <code className="mt-1 block text-xs break-all">
+                        {'{"apiKey": "...", "region": "us-east-1"}'}
+                      </code>
+                      <p className="mt-1">
+                        Generate an API key in the AWS Bedrock console. Use a region where your key
+                        and model are available, and replace the key before it expires.
+                      </p>
+                      <p className="mt-1">Or enter IAM credentials as JSON:</p>
                       <code className="mt-1 block text-xs break-all">
                         {'{"accessKeyId": "...", "secretAccessKey": "...", "region": "us-east-1"}'}
                       </code>
diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.test.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.test.ts
index 8968f16606..f26832acf5 100644
--- a/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.test.ts
+++ b/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.test.ts
@@ -1,4 +1,5 @@
 import {
+  BedrockCredentialsSchema,
   DirectUserByokInferenceProviderIdSchema,
   getVercelUserByokProviderIdForEndpoint,
   normalizeVercelInferenceProviderIdForRouting,
@@ -9,6 +10,40 @@ import {
   VercelUserByokInferenceProviderIdSchema,
 } from './inference-provider-id';
 
+describe('BedrockCredentialsSchema', () => {
+  test.each([
+    { accessKeyId: 'AKIAEXAMPLE', secretAccessKey: 'secret', region: 'us-east-1' },
+    { apiKey: 'bedrock-api-key', region: 'eu-west-1' },
+  ])('accepts Bedrock credentials: %j', credentials => {
+    expect(BedrockCredentialsSchema.parse(credentials)).toEqual(credentials);
+  });
+
+  test.each([
+    null,
+    'bedrock-api-key',
+    {},
+    { region: 'us-east-1' },
+    { apiKey: 'bedrock-api-key' },
+    { apiKey: '', region: 'us-east-1' },
+    { apiKey: '   ', region: 'us-east-1' },
+    { apiKey: 123, region: 'us-east-1' },
+    { apiKey: 'bedrock-api-key', region: '' },
+    { apiKey: 'bedrock-api-key', region: '   ' },
+    { apiKey: 'bedrock-api-key', region: 123 },
+    { apiKey: 'bedrock-api-key', region: 'us-east-1', accessKeyId: 'AKIAEXAMPLE' },
+    { apiKey: 'bedrock-api-key', region: 'us-east-1', secretAccessKey: 'secret' },
+    {
+      apiKey: 'bedrock-api-key',
+      region: 'us-east-1',
+      accessKeyId: 'AKIAEXAMPLE',
+      secretAccessKey: 'secret',
+    },
+    { accessKeyId: 'AKIAEXAMPLE', region: 'us-east-1' },
+  ])('rejects incomplete, invalid, or mixed credentials: %j', credentials => {
+    expect(BedrockCredentialsSchema.safeParse(credentials).success).toBe(false);
+  });
+});
+
 describe('inference provider ids', () => {
   test('direct BYOK provider ids do not overlap with OpenRouter provider ids', () => {
     const overlappingProviderIds = DirectUserByokInferenceProviderIdSchema.options.filter(
diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.ts
index 85bcc5a1f2..1daf812b72 100644
--- a/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.ts
+++ b/apps/web/src/lib/ai-gateway/providers/openrouter/inference-provider-id.ts
@@ -293,7 +293,17 @@ export const AwsCredentialsSchema = z.object({
   region: z.string(),
 });
 
-export type AwsCredentials = z.infer<typeof AwsCredentialsSchema>;
+export const BedrockCredentialsSchema = z.union([
+  AwsCredentialsSchema.extend({ apiKey: z.never().optional() }),
+  z.object({
+    apiKey: z.string().trim().min(1),
+    region: z.string().trim().min(1),
+    accessKeyId: z.never().optional(),
+    secretAccessKey: z.never().optional(),
+  }),
+]);
+
+export type BedrockCredentials = z.infer<typeof BedrockCredentialsSchema>;
 
 export const VertexCredentialsSchema = z.object({
   project: z.string().min(1),
diff --git a/apps/web/src/lib/ai-gateway/providers/openrouter/types.ts b/apps/web/src/lib/ai-gateway/providers/openrouter/types.ts
index b4851986cd..4d5b07f3a4 100644
--- a/apps/web/src/lib/ai-gateway/providers/openrouter/types.ts
+++ b/apps/web/src/lib/ai-gateway/providers/openrouter/types.ts
@@ -3,7 +3,7 @@ import type { GatewayProviderOptions } from '@ai-sdk/gateway';
 import type { AnthropicProviderOptions } from '@ai-sdk/anthropic';
 import type { ReasoningDetailUnion } from '@/lib/ai-gateway/custom-llm/reasoning-details';
 import type {
-  AwsCredentials,
+  BedrockCredentials,
   VertexCredentials,
 } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id';
 import type Anthropic from '@anthropic-ai/sdk';
@@ -26,7 +26,7 @@ export function isOpenRouterProviderConfig(value: unknown): value is OpenRouterP
 
 export type VercelInferenceProviderConfig =
   | { apiKey: string; baseURL?: string }
-  | AwsCredentials
+  | BedrockCredentials
   | VertexCredentials;
 
 export type VercelProviderConfig = {
diff --git a/apps/web/src/lib/ai-gateway/providers/vercel/index.test.ts b/apps/web/src/lib/ai-gateway/providers/vercel/index.test.ts
index 51fda74501..36a1f7d0e5 100644
--- a/apps/web/src/lib/ai-gateway/providers/vercel/index.test.ts
+++ b/apps/web/src/lib/ai-gateway/providers/vercel/index.test.ts
@@ -372,6 +372,60 @@ describe('applyVercelSettings BYOK pinning', () => {
     });
   });
 
+  it.each<GatewayRequest>([
+    {
+      kind: 'chat_completions',
+      body: {
+        model: 'anthropic/claude-sonnet-4.5',
+        messages: [{ role: 'user', content: 'hello' }],
+      },
+    },
+    {
+      kind: 'messages',
+      body: {
+        model: 'anthropic/claude-sonnet-4.5',
+        messages: [{ role: 'user', content: 'hello' }],
+        max_tokens: 100,
+      },
+    },
+    {
+      kind: 'responses',
+      body: { model: 'anthropic/claude-sonnet-4.5', input: 'hello' },
+    },
+  ])('forwards Bedrock API key credentials for $kind', async request => {
+    const credentials = { apiKey: 'bedrock-api-key', region: 'eu-west-1' };
+    await applyVercelSettings('anthropic/claude-sonnet-4.5', request, [
+      { decryptedAPIKey: JSON.stringify(credentials), providerId: 'bedrock' },
+    ]);
+
+    expect(request.body.providerOptions?.gateway?.byok).toEqual({ bedrock: [credentials] });
+    expect(request.body.providerOptions?.gateway?.only).toEqual(['bedrock']);
+  });
+
+  it('retains a Bedrock API key when the caller ignores its only BYOK provider', async () => {
+    const request = byokRequest(['amazon-bedrock']);
+    const credentials = { apiKey: 'bedrock-api-key', region: 'us-east-1' };
+    await applyVercelSettings('anthropic/claude-sonnet-4.5', request, [
+      { decryptedAPIKey: JSON.stringify(credentials), providerId: 'bedrock' },
+    ]);
+
+    expect(request.body.providerOptions?.gateway?.byok).toEqual({ bedrock: [credentials] });
+    expect(request.body.providerOptions?.gateway?.only).toEqual(['bedrock']);
+  });
+
+  it.each([
+    'bedrock-secret',
+    '{"apiKey":"bedrock-secret"',
+    '{"apiKey":"bedrock-secret"}',
+    '{"apiKey":"bedrock-secret","region":"us-east-1","accessKeyId":"AKIAEXAMPLE","secretAccessKey":"secret"}',
+  ])('rejects malformed Bedrock credentials without exposing secrets: %s', async credentials => {
+    await expect(
+      applyVercelSettings('anthropic/claude-sonnet-4.5', byokRequest([]), [
+        { decryptedAPIKey: credentials, providerId: 'bedrock' },
+      ])
+    ).rejects.toEqual(new Error('Failed to parse AWS credentials'));
+  });
+
   it('uses one Vertex credential key for Anthropic models served by Vertex', async () => {
     const request = byokRequest([]);
 
diff --git a/apps/web/src/lib/ai-gateway/providers/vercel/index.ts b/apps/web/src/lib/ai-gateway/providers/vercel/index.ts
index fbadabf222..b140ad627b 100644
--- a/apps/web/src/lib/ai-gateway/providers/vercel/index.ts
+++ b/apps/web/src/lib/ai-gateway/providers/vercel/index.ts
@@ -2,7 +2,7 @@ import type { BYOKResult } from '@/lib/ai-gateway/providers/types';
 import type { VercelUserByokInferenceProviderId } from '@/lib/ai-gateway/providers/openrouter/inference-provider-id';
 import {
   DirectUserByokInferenceProviderIdSchema,
-  AwsCredentialsSchema,
+  BedrockCredentialsSchema,
   normalizeVercelInferenceProviderIdForRouting,
   openRouterToVercelInferenceProviderId,
   VertexCredentialsSchema,
@@ -172,9 +172,9 @@ export function convertProviderOptions(
   };
 }
 
-function parseAwsCredentials(input: string) {
+function parseBedrockCredentials(input: string) {
   try {
-    return AwsCredentialsSchema.parse(JSON.parse(input));
+    return BedrockCredentialsSchema.parse(JSON.parse(input));
   } catch {
     throw new Error('Failed to parse AWS credentials');
   }
@@ -228,7 +228,7 @@ export function getVercelInferenceProviderConfigForUserByok(
   }
 
   if (key === VercelUserByokInferenceProviderIdSchema.enum.bedrock) {
-    list.push(parseAwsCredentials(provider.decryptedAPIKey));
+    list.push(parseBedrockCredentials(provider.decryptedAPIKey));
   } else if (key === VercelUserByokInferenceProviderIdSchema.enum.vertex) {
     list.push(parseVertexCredentials(provider.decryptedAPIKey));
   } else {
diff --git a/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts b/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts
index e6e22d49b0..564a1e2a27 100644
--- a/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts
+++ b/apps/web/src/lib/ai-gateway/rewriteModelResponse.test.ts
@@ -1100,6 +1100,7 @@ describe('sanitizeApiRequestLogRequest', () => {
                   secretAccessKey: 'bedrock-secret',
                   region: 'us-east-1',
                 },
+                { apiKey: 'bedrock-api-key', region: 'us-west-2' },
               ],
             },
           },
@@ -1116,7 +1117,7 @@ describe('sanitizeApiRequestLogRequest', () => {
           order: ['friendli', 'novita'],
           byok: {
             friendli: [{ apiKey: '[redacted]' }],
-            bedrock: [{ apiKey: '[redacted]' }],
+            bedrock: [{ apiKey: '[redacted]' }, { apiKey: '[redacted]' }],
           },
         },
         anthropic: { effort: 'high' },
@@ -1129,6 +1130,10 @@ describe('sanitizeApiRequestLogRequest', () => {
       accessKeyId: 'AKIAEXAMPLE',
       secretAccessKey: 'bedrock-secret',
     });
+    expect(request.body.providerOptions.gateway.byok.bedrock[1]).toEqual({
+      apiKey: 'bedrock-api-key',
+      region: 'us-west-2',
+    });
   });
 });
 
diff --git a/apps/web/src/routers/byok-router.test.ts b/apps/web/src/routers/byok-router.test.ts
index 43a50f62f0..47809313aa 100644
--- a/apps/web/src/routers/byok-router.test.ts
+++ b/apps/web/src/routers/byok-router.test.ts
@@ -12,7 +12,7 @@ import {
 } from '@kilocode/db/schema';
 import { eq, and } from 'drizzle-orm';
 import type { User, Organization } from '@kilocode/db/schema';
-import { encryptApiKey } from '@/lib/ai-gateway/byok/encryption';
+import { decryptApiKey, encryptApiKey } from '@/lib/ai-gateway/byok/encryption';
 import { BYOK_ENCRYPTION_KEY } from '@/lib/config.server';
 
 const VERTEX_CREDENTIALS = JSON.stringify({
@@ -24,6 +24,17 @@ const VERTEX_CREDENTIALS = JSON.stringify({
   },
 });
 
+const BEDROCK_IAM_CREDENTIALS = {
+  accessKeyId: 'AKIAEXAMPLE',
+  secretAccessKey: 'bedrock-test-secret',
+  region: 'us-east-1',
+};
+
+const BEDROCK_API_KEY_CREDENTIALS = {
+  apiKey: 'bedrock-test-api-key',
+  region: 'us-west-2',
+};
+
 describe('BYOK Router', () => {
   let ownerUser: User;
   let memberUser: User;
@@ -526,6 +537,164 @@ describe('BYOK Router', () => {
     });
   });
 
+  describe('Bedrock credentials', () => {
+    test.each(['user', 'organization'] as const)(
+      'encrypts a %s API key credential and returns only metadata from create and list',
+      async owner => {
+        const caller = await createCallerForUser(ownerUser.id);
+        const scope = owner === 'organization' ? { organizationId: organizationA.id } : {};
+        const credential = JSON.stringify(BEDROCK_API_KEY_CREDENTIALS);
+
+        const created = await caller.byok.create({
+          ...scope,
+          provider_id: 'bedrock',
+          api_key: credential,
+        });
+
+        expect(created).toEqual({
+          id: expect.any(String),
+          provider_id: 'bedrock',
+          provider_name: 'bedrock',
+          management_source: 'user',
+          is_enabled: true,
+          created_by: ownerUser.id,
+          created_at: expect.any(String),
+          updated_at: expect.any(String),
+        });
+        await expect(caller.byok.list(scope)).resolves.toEqual([created]);
+
+        const [stored] = await db
+          .select()
+          .from(byok_api_keys)
+          .where(eq(byok_api_keys.id, created.id));
+        expect(stored.organization_id).toBe(owner === 'organization' ? organizationA.id : null);
+        expect(stored.kilo_user_id).toBe(owner === 'user' ? ownerUser.id : null);
+        expect(stored.encrypted_api_key).toEqual({
+          iv: expect.any(String),
+          data: expect.any(String),
+          authTag: expect.any(String),
+        });
+        expect(JSON.stringify(stored.encrypted_api_key)).not.toContain(
+          BEDROCK_API_KEY_CREDENTIALS.apiKey
+        );
+        expect(decryptApiKey(stored.encrypted_api_key, BYOK_ENCRYPTION_KEY)).toBe(credential);
+      }
+    );
+
+    test.each([
+      {
+        owner: 'user',
+        from: 'IAM',
+        to: 'API key',
+        initial: BEDROCK_IAM_CREDENTIALS,
+        replacement: BEDROCK_API_KEY_CREDENTIALS,
+      },
+      {
+        owner: 'organization',
+        from: 'API key',
+        to: 'IAM',
+        initial: BEDROCK_API_KEY_CREDENTIALS,
+        replacement: BEDROCK_IAM_CREDENTIALS,
+      },
+    ])(
+      'replaces a $owner credential from $from to $to',
+      async ({ owner, initial, replacement }) => {
+        const caller = await createCallerForUser(ownerUser.id);
+        const scope = owner === 'organization' ? { organizationId: organizationA.id } : {};
+        const created = await caller.byok.create({
+          ...scope,
+          provider_id: 'bedrock',
+          api_key: JSON.stringify(initial),
+        });
+        const [before] = await db
+          .select()
+          .from(byok_api_keys)
+          .where(eq(byok_api_keys.id, created.id));
+        expect(decryptApiKey(before.encrypted_api_key, BYOK_ENCRYPTION_KEY)).toBe(
+          JSON.stringify(initial)
+        );
+
+        const updated = await caller.byok.update({
+          ...scope,
+          id: created.id,
+          api_key: JSON.stringify(replacement),
+        });
+
+        expect(updated).toEqual({ ...created, updated_at: expect.any(String) });
+        await expect(caller.byok.list(scope)).resolves.toEqual([updated]);
+        const [after] = await db
+          .select()
+          .from(byok_api_keys)
+          .where(eq(byok_api_keys.id, created.id));
+        expect(after.organization_id).toBe(before.organization_id);
+        expect(after.kilo_user_id).toBe(before.kilo_user_id);
+        expect(after.encrypted_api_key).not.toEqual(before.encrypted_api_key);
+        expect(decryptApiKey(after.encrypted_api_key, BYOK_ENCRYPTION_KEY)).toBe(
+          JSON.stringify(replacement)
+        );
+      }
+    );
+
+    test.each([
+      { name: 'missing API key', credentials: { region: 'us-east-1' } },
+      { name: 'empty API key', credentials: { apiKey: '', region: 'us-east-1' } },
+      { name: 'missing region', credentials: { apiKey: 'bedrock-test-api-key' } },
+      { name: 'empty region', credentials: { apiKey: 'bedrock-test-api-key', region: '' } },
+      {
+        name: 'mixed IAM and API key',
+        credentials: { ...BEDROCK_IAM_CREDENTIALS, apiKey: 'bedrock-test-api-key' },
+      },
+    ])('rejects $name credentials before create or update writes', async ({ credentials }) => {
+      const caller = await createCallerForUser(ownerUser.id);
+      const scope = { organizationId: organizationA.id };
+      const created = await caller.byok.create({
+        ...scope,
+        provider_id: 'bedrock',
+        api_key: JSON.stringify(BEDROCK_IAM_CREDENTIALS),
+      });
+      const keysBefore = await db
+        .select()
+        .from(byok_api_keys)
+        .where(eq(byok_api_keys.organization_id, organizationA.id));
+      const logsBefore = await db
+        .select()
+        .from(organization_audit_logs)
+        .where(eq(organization_audit_logs.organization_id, organizationA.id))
+        .orderBy(organization_audit_logs.id);
+      const api_key = JSON.stringify(credentials);
+
+      await expect(
+        caller.byok.create({ ...scope, provider_id: 'bedrock', api_key })
+      ).rejects.toMatchObject({
+        code: 'BAD_REQUEST',
+        message: 'Invalid credentials for provider: bedrock',
+      });
+      const keysAfterCreate = await db
+        .select()
+        .from(byok_api_keys)
+        .where(eq(byok_api_keys.organization_id, organizationA.id));
+      expect(keysAfterCreate).toEqual(keysBefore);
+
+      await expect(caller.byok.update({ ...scope, id: created.id, api_key })).rejects.toMatchObject(
+        {
+          code: 'BAD_REQUEST',
+          message: 'Invalid credentials for provider: bedrock',
+        }
+      );
+      const keysAfterUpdate = await db
+        .select()
+        .from(byok_api_keys)
+        .where(eq(byok_api_keys.organization_id, organizationA.id));
+      expect(keysAfterUpdate).toEqual(keysBefore);
+      const logsAfter = await db
+        .select()
+        .from(organization_audit_logs)
+        .where(eq(organization_audit_logs.organization_id, organizationA.id))
+        .orderBy(organization_audit_logs.id);
+      expect(logsAfter).toEqual(logsBefore);
+    });
+  });
+
   describe('deprecated codestral provider', () => {
     test('declines to test a legacy codestral key with a deprecation message', async () => {
       const caller = await createCallerForUser(ownerUser.id);