feat: distribute bedrock requests per regional limits (#37049)

apps/studio/lib/ai/bedrock.test.ts

@@ -1,73 +0,0 @@
-import { describe, it, expect } from 'vitest'
-import { selectBedrockRegion, bedrockRegionMap } from './bedrock'
-
-describe('selectBedrockRegion', () => {
-  it('should return a valid region for a given routing key', async () => {
-    const region = await selectBedrockRegion('test-key')
-    const validRegions = Object.keys(bedrockRegionMap)
-
-    expect(validRegions).toContain(region)
-  })
-
-  it('should return the same region for the same routing key', async () => {
-    const routingKey = 'consistent-key'
-    const region1 = await selectBedrockRegion(routingKey)
-    const region2 = await selectBedrockRegion(routingKey)
-
-    expect(region1).toBe(region2)
-  })
-
-  it('should distribute different keys across regions', async () => {
-    const keys = Array.from({ length: 100 }, (_, i) => `key-${i}`)
-    const regions = await Promise.all(keys.map((key) => selectBedrockRegion(key)))
-    const uniqueRegions = new Set(regions)
-    const validRegions = Object.keys(bedrockRegionMap)
-
-    // Should use all regions for 100 different keys
-    expect(uniqueRegions.size).toEqual(validRegions.length)
-  })
-
-  it('should distribute keys evenly across regions', async () => {
-    const numKeys = 3000
-    const keys = Array.from({ length: numKeys }, (_, i) => `key-${i}`)
-    const regions = await Promise.all(keys.map((key) => selectBedrockRegion(key)))
-    const validRegions = Object.keys(bedrockRegionMap)
-
-    // Count occurrences of each region
-    const regionCounts = regions.reduce<Record<string, number>>((acc, region) => {
-      acc[region] = (acc[region] ?? 0) + 1
-      return acc
-    }, {})
-
-    const expectedCountPerRegion = numKeys / validRegions.length
-    const tolerance = expectedCountPerRegion * 0.2 // Allow 20% deviation
-
-    // Each region should have roughly equal distribution
-    for (const count of Object.values(regionCounts)) {
-      expect(count).toBeGreaterThan(expectedCountPerRegion - tolerance)
-      expect(count).toBeLessThan(expectedCountPerRegion + tolerance)
-    }
-  })
-
-  it('should handle empty string', async () => {
-    const region = await selectBedrockRegion('')
-    const validRegions = Object.keys(bedrockRegionMap)
-
-    expect(validRegions).toContain(region)
-  })
-
-  it('should handle special characters in routing key', async () => {
-    const region = await selectBedrockRegion('key-with-special-chars!@#$%')
-    const validRegions = Object.keys(bedrockRegionMap)
-
-    expect(validRegions).toContain(region)
-  })
-
-  it('should return consistent results for unicode characters', async () => {
-    const routingKey = '🔑-unicode-key-测试'
-    const region1 = await selectBedrockRegion(routingKey)
-    const region2 = await selectBedrockRegion(routingKey)
-
-    expect(region1).toBe(region2)
-  })
-})

apps/studio/lib/ai/bedrock.ts

@@ -2,6 +2,7 @@ import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock'
 import { createCredentialChain, fromNodeProviderChain } from '@aws-sdk/credential-providers'
 import { CredentialsProviderError } from '@smithy/property-provider'
 import { awsCredentialsProvider } from '@vercel/functions/oidc'
+import { selectWeightedKey } from './util'
 
 const credentialProvider = createCredentialChain(
   // Vercel OIDC provider will be used for staging/production
@@ -34,40 +35,6 @@ async function vercelOidcProvider() {
   }
 }
 
-export const bedrockRegionMap = {
-  us1: 'us-east-1',
-  us3: 'us-west-2',
-} as const
-
-export type BedrockRegion = keyof typeof bedrockRegionMap
-
-export const bedrockForRegion = (region: BedrockRegion) =>
-  createAmazonBedrock({
-    credentialProvider,
-    region: bedrockRegionMap[region],
-  })
-
-/**
- * Selects a region based on a routing key using a consistent hashing algorithm.
- *
- * Ensures that the same key always maps to the same region
- * while distributing keys evenly across available regions.
- */
-export async function selectBedrockRegion(routingKey: string) {
-  const regions = Object.keys(bedrockRegionMap) as BedrockRegion[]
-  const encoder = new TextEncoder()
-  const data = encoder.encode(routingKey)
-  const hashBuffer = await crypto.subtle.digest('SHA-256', data)
-
-  // Use first 4 bytes (32 bit integer)
-  const hashInt = new DataView(hashBuffer).getUint32(0)
-
-  // Use modulo to map to available regions
-  const regionIndex = hashInt % regions.length
-
-  return regions[regionIndex]
-}
-
 export async function checkAwsCredentials() {
   try {
     const credentials = await credentialProvider()
@@ -76,3 +43,74 @@ export async function checkAwsCredentials() {
     return false
   }
 }
+
+export const bedrockRegionMap = {
+  use1: 'us-east-1',
+  use2: 'us-east-2',
+  usw2: 'us-west-2',
+  euc1: 'eu-central-1',
+} as const
+
+export type BedrockRegion = keyof typeof bedrockRegionMap
+
+export const regionPrefixMap: Record<BedrockRegion, string> = {
+  use1: 'us',
+  use2: 'us',
+  usw2: 'us',
+  euc1: 'eu',
+}
+
+export type BedrockModel =
+  | 'anthropic.claude-3-7-sonnet-20250219-v1:0'
+  | 'anthropic.claude-3-5-haiku-20241022-v1:0'
+
+export type RegionWeights = Record<BedrockRegion, number>
+
+/**
+ * Weights for distributing requests across Bedrock regions.
+ * Weights are proportional to our rate limits per model per region.
+ */
+const modelRegionWeights: Record<BedrockModel, RegionWeights> = {
+  ['anthropic.claude-3-7-sonnet-20250219-v1:0']: {
+    use1: 40,
+    use2: 10,
+    usw2: 10,
+    euc1: 10,
+  },
+  ['anthropic.claude-3-5-haiku-20241022-v1:0']: {
+    use1: 40,
+    use2: 0,
+    usw2: 40,
+    euc1: 0,
+  },
+}
+
+/**
+ * Creates a Bedrock client that routes requests to different regions
+ * based on a routing key.
+ *
+ * Used to load balance requests across multiple regions depending on
+ * their capacities.
+ */
+export function createRoutedBedrock(routingKey?: string) {
+  return async (modelId: BedrockModel) => {
+    const regionWeights = modelRegionWeights[modelId]
+
+    // Select the Bedrock region based on the routing key and the model
+    const bedrockRegion = routingKey
+      ? await selectWeightedKey(routingKey, regionWeights)
+      : // There's a few places where getModel is called without a routing key
+        // Will cause disproportionate load on use1 region
+        'use1'
+
+    const bedrock = createAmazonBedrock({
+      credentialProvider,
+      region: bedrockRegionMap[bedrockRegion],
+    })
+
+    // Cross-region models require the region prefix
+    const modelName = `${regionPrefixMap[bedrockRegion]}.${modelId}`
+
+    return bedrock(modelName)
+  }
+}

apps/studio/lib/ai/model.test.ts

@@ -7,10 +7,10 @@ vi.mock('@ai-sdk/openai', () => ({
   openai: vi.fn(() => 'openai-model'),
 }))
 
-vi.mock('./bedrock', () => ({
-  bedrockForRegion: vi.fn(() => () => 'bedrock-model'),
+vi.mock('./bedrock', async () => ({
+  ...(await vi.importActual('./bedrock')),
+  createRoutedBedrock: vi.fn(() => () => 'bedrock-model'),
   checkAwsCredentials: vi.fn(),
-  selectBedrockRegion: vi.fn(() => 'us'),
 }))
 
 describe('getModel', () => {
@@ -29,10 +29,7 @@ describe('getModel', () => {
 
     const { model, error } = await getModel()
 
-    console.log('Model:', model)
-
     expect(model).toEqual('bedrock-model')
-    expect(bedrockModule.bedrockForRegion).toHaveBeenCalledWith('us1')
     expect(error).toBeUndefined()
   })
 
@@ -40,7 +37,7 @@ describe('getModel', () => {
     vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(false)
     process.env.OPENAI_API_KEY = 'test-key'
 
-    const { model } = await getModel('test-key')
+    const { model } = await getModel()
 
     expect(model).toEqual('openai-model')
     expect(openai).toHaveBeenCalledWith('gpt-4.1-2025-04-14')

apps/studio/lib/ai/model.ts

@@ -1,23 +1,12 @@
 import { openai } from '@ai-sdk/openai'
 import { LanguageModel } from 'ai'
-import {
-  bedrockForRegion,
-  BedrockRegion,
-  checkAwsCredentials,
-  selectBedrockRegion,
-} from './bedrock'
-
-export const regionMap = {
-  us1: 'us',
-  us2: 'us',
-  us3: 'us',
-  eu: 'eu',
-}
+import { checkAwsCredentials, createRoutedBedrock } from './bedrock'
 
 // Default behaviour here is to be throttled (e.g if this env var is not available, IS_THROTTLED should be true, unless specified 'false')
 const IS_THROTTLED = process.env.IS_THROTTLED !== 'false'
-const PRO_MODEL = process.env.AI_PRO_MODEL ?? 'anthropic.claude-3-7-sonnet-20250219-v1:0'
-const NORMAL_MODEL = process.env.AI_NORMAL_MODEL ?? 'anthropic.claude-3-5-haiku-20241022-v1:0'
+
+const BEDROCK_PRO_MODEL = 'anthropic.claude-3-7-sonnet-20250219-v1:0'
+const BEDROCK_NORMAL_MODEL = 'anthropic.claude-3-5-haiku-20241022-v1:0'
 const OPENAI_MODEL = 'gpt-4.1-2025-04-14'
 
 export type ModelSuccess = {
@@ -46,14 +35,11 @@ export async function getModel(routingKey?: string, isLimited?: boolean): Promis
   const hasOpenAIKey = !!process.env.OPENAI_API_KEY
 
   if (hasAwsCredentials) {
-    // Select the Bedrock region based on the routing key
-    const bedrockRegion: BedrockRegion = routingKey ? await selectBedrockRegion(routingKey) : 'us1'
-    const bedrock = bedrockForRegion(bedrockRegion)
-    const model = IS_THROTTLED || isLimited ? NORMAL_MODEL : PRO_MODEL
-    const modelName = `${regionMap[bedrockRegion]}.${model}`
+    const bedrockModel = IS_THROTTLED || isLimited ? BEDROCK_NORMAL_MODEL : BEDROCK_PRO_MODEL
+    const bedrock = createRoutedBedrock(routingKey)
 
     return {
-      model: bedrock(modelName),
+      model: await bedrock(bedrockModel),
     }
   }
 

apps/studio/lib/ai/util.test.ts

@@ -0,0 +1,83 @@
+import { describe, it, expect } from 'vitest'
+import { selectWeightedKey } from './util'
+
+describe('selectWeightedKey', () => {
+  it('should return a valid key from the weights object', async () => {
+    const weights = { a: 10, b: 20, c: 30 }
+    const result = await selectWeightedKey('test-input', weights)
+
+    expect(Object.keys(weights)).toContain(result)
+  })
+
+  it('should return consistent results for the same input', async () => {
+    const weights = { region1: 40, region2: 10, region3: 20 }
+    const input = 'consistent-key'
+
+    const result1 = await selectWeightedKey(input, weights)
+    const result2 = await selectWeightedKey(input, weights)
+    const result3 = await selectWeightedKey(input, weights)
+
+    expect(result1).toBe(result2)
+    expect(result2).toBe(result3)
+  })
+
+  it('should distribute keys according to weights', async () => {
+    const weights = { a: 80, b: 10, c: 10 }
+    const numSamples = 10000
+    const samples = Array.from({ length: numSamples }, (_, i) => `sample-${i}`)
+
+    const results = await Promise.all(samples.map((sample) => selectWeightedKey(sample, weights)))
+
+    const counts = results.reduce<Record<string, number>>((acc, key) => {
+      acc[key] = (acc[key] ?? 0) + 1
+      return acc
+    }, {})
+
+    expect(counts.a / numSamples).toBeCloseTo(0.8, 1)
+    expect(counts.b / numSamples).toBeCloseTo(0.1, 1)
+    expect(counts.c / numSamples).toBeCloseTo(0.1, 1)
+  })
+
+  it('should handle equal weights', async () => {
+    const weights = { x: 25, y: 25, z: 25, w: 25 }
+    const numSamples = 8000
+    const samples = Array.from({ length: numSamples }, (_, i) => `equal-${i}`)
+
+    const results = await Promise.all(samples.map((sample) => selectWeightedKey(sample, weights)))
+
+    const counts = results.reduce<Record<string, number>>((acc, key) => {
+      acc[key] = (acc[key] ?? 0) + 1
+      return acc
+    }, {})
+
+    // Each key should get roughly 25% of the samples
+    Object.values(counts).forEach((count) => {
+      expect(count / numSamples).toBeCloseTo(0.25, 1)
+    })
+  })
+
+  it('should handle single key', async () => {
+    const weights = { only: 100 }
+    const result = await selectWeightedKey('any-input', weights)
+
+    expect(result).toBe('only')
+  })
+
+  it('should handle empty string input', async () => {
+    const weights = { a: 10, b: 20 }
+    const result = await selectWeightedKey('', weights)
+
+    expect(Object.keys(weights)).toContain(result)
+  })
+
+  it('should handle unicode characters in input', async () => {
+    const weights = { option1: 50, option2: 50 }
+    const unicodeInput = '🔑-unicode-key-测试'
+
+    const result1 = await selectWeightedKey(unicodeInput, weights)
+    const result2 = await selectWeightedKey(unicodeInput, weights)
+
+    expect(result1).toBe(result2)
+    expect(Object.keys(weights)).toContain(result1)
+  })
+})

apps/studio/lib/ai/util.ts

@@ -0,0 +1,42 @@
+/**
+ * Selects a key from weighted choices using consistent hashing
+ * on an input string.
+ *
+ * The same input always returns the same key, with distribution
+ * proportional to the provided weights.
+ *
+ * @example
+ * const region = await selectWeightedKey('my-unique-id', {
+ *   use1: 40,
+ *   use2: 10,
+ *   usw2: 10,
+ *   euc1: 10,
+ * })
+ * // Returns one of the keys based on the input and weights
+ */
+export async function selectWeightedKey<T extends string>(
+  input: string,
+  weights: Record<T, number>
+): Promise<T> {
+  const keys = Object.keys(weights) as T[]
+  const encoder = new TextEncoder()
+  const data = encoder.encode(input)
+  const hashBuffer = await crypto.subtle.digest('SHA-256', data)
+
+  // Use first 4 bytes (32 bit integer)
+  const hashInt = new DataView(hashBuffer).getUint32(0)
+
+  const totalWeight = keys.reduce((sum, key) => sum + weights[key], 0)
+
+  let cumulativeWeight = 0
+  const targetWeight = hashInt % totalWeight
+
+  for (const key of keys) {
+    cumulativeWeight += weights[key]
+    if (cumulativeWeight > targetWeight) {
+      return key
+    }
+  }
+
+  return keys[0]
+}