Add user based hosted key throttling

Author: Theodore Li

apps/sim/lib/core/hosted-key-throttler/index.ts

@@ -0,0 +1,16 @@
+export {
+  getHostedKeyThrottler,
+  HostedKeyThrottler,
+  resetHostedKeyThrottler,
+} from './throttler'
+export {
+  DEFAULT_BURST_MULTIPLIER,
+  THROTTLE_WINDOW_MS,
+  toTokenBucketConfig,
+  type AcquireKeyResult,
+  type CustomThrottle,
+  type PerRequestThrottle,
+  type ThrottleConfig,
+  type ThrottleDimension,
+  type ThrottleMode,
+} from './types'

apps/sim/lib/core/hosted-key-throttler/throttler.test.ts

@@ -0,0 +1,132 @@
+import { loggerMock } from '@sim/testing'
+import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
+import { HostedKeyThrottler } from './throttler'
+import type { PerRequestThrottle } from './types'
+import type { ConsumeResult, RateLimitStorageAdapter } from '@/lib/core/rate-limiter/storage'
+
+vi.mock('@sim/logger', () => loggerMock)
+
+interface MockAdapter {
+  consumeTokens: Mock
+  getTokenStatus: Mock
+  resetBucket: Mock
+}
+
+const createMockAdapter = (): MockAdapter => ({
+  consumeTokens: vi.fn(),
+  getTokenStatus: vi.fn(),
+  resetBucket: vi.fn(),
+})
+
+describe('HostedKeyThrottler', () => {
+  const testProvider = 'exa'
+  const envKeys = ['EXA_API_KEY_1', 'EXA_API_KEY_2', 'EXA_API_KEY_3']
+  let mockAdapter: MockAdapter
+  let throttler: HostedKeyThrottler
+  let originalEnv: NodeJS.ProcessEnv
+
+  const perRequestThrottle: PerRequestThrottle = {
+    mode: 'per_request',
+    userRequestsPerMinute: 10,
+  }
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockAdapter = createMockAdapter()
+    throttler = new HostedKeyThrottler(mockAdapter as RateLimitStorageAdapter)
+
+    originalEnv = { ...process.env }
+    process.env.EXA_API_KEY_1 = 'test-key-1'
+    process.env.EXA_API_KEY_2 = 'test-key-2'
+    process.env.EXA_API_KEY_3 = 'test-key-3'
+  })
+
+  afterEach(() => {
+    process.env = originalEnv
+  })
+
+  describe('acquireKey', () => {
+    it('should return error when no keys are configured', async () => {
+      delete process.env.EXA_API_KEY_1
+      delete process.env.EXA_API_KEY_2
+      delete process.env.EXA_API_KEY_3
+
+      const result = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle)
+
+      expect(result.success).toBe(false)
+      expect(result.error).toContain('No hosted keys configured')
+    })
+
+    it('should throttle user when they exceed their rate limit', async () => {
+      const throttledResult: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 30000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(throttledResult)
+
+      const result = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle, 'user-123')
+
+      expect(result.success).toBe(false)
+      expect(result.userThrottled).toBe(true)
+      expect(result.retryAfterMs).toBeDefined()
+      expect(result.error).toContain('Rate limit exceeded')
+    })
+
+    it('should allow user within their rate limit', async () => {
+      const allowedResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
+
+      const result = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle, 'user-123')
+
+      expect(result.success).toBe(true)
+      expect(result.userThrottled).toBeUndefined()
+      expect(result.key).toBe('test-key-1')
+    })
+
+    it('should distribute requests across keys round-robin style', async () => {
+      const allowedResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
+
+      const r1 = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle, 'user-1')
+      const r2 = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle, 'user-2')
+      const r3 = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle, 'user-3')
+      const r4 = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle, 'user-4')
+
+      expect(r1.keyIndex).toBe(0)
+      expect(r2.keyIndex).toBe(1)
+      expect(r3.keyIndex).toBe(2)
+      expect(r4.keyIndex).toBe(0) // Wraps back
+    })
+
+    it('should work without userId (no per-user throttling)', async () => {
+      const result = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle)
+
+      expect(result.success).toBe(true)
+      expect(result.key).toBe('test-key-1')
+      expect(mockAdapter.consumeTokens).not.toHaveBeenCalled()
+    })
+
+    it('should handle partial key availability', async () => {
+      delete process.env.EXA_API_KEY_2
+
+      const result = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle)
+
+      expect(result.success).toBe(true)
+      expect(result.key).toBe('test-key-1')
+      expect(result.envVarName).toBe('EXA_API_KEY_1')
+
+      const r2 = await throttler.acquireKey(testProvider, envKeys, perRequestThrottle)
+      expect(r2.keyIndex).toBe(2) // Skips missing key 1
+      expect(r2.envVarName).toBe('EXA_API_KEY_3')
+    })
+  })
+})

apps/sim/lib/core/hosted-key-throttler/throttler.ts

@@ -0,0 +1,202 @@
+import { createLogger } from '@sim/logger'
+import {
+  createStorageAdapter,
+  type RateLimitStorageAdapter,
+  type TokenBucketConfig,
+} from '@/lib/core/rate-limiter/storage'
+import {
+  DEFAULT_BURST_MULTIPLIER,
+  THROTTLE_WINDOW_MS,
+  toTokenBucketConfig,
+  type AcquireKeyResult,
+  type PerRequestThrottle,
+  type ThrottleConfig,
+} from './types'
+
+const logger = createLogger('HostedKeyThrottler')
+
+/** Dimension name for per-user rate limiting */
+const USER_REQUESTS_DIMENSION = 'user_requests'
+
+/**
+ * Information about an available hosted key
+ */
+interface AvailableKey {
+  key: string
+  keyIndex: number
+  envVarName: string
+}
+
+/**
+ * HostedKeyThrottler provides:
+ * 1. Per-user rate limiting (enforced - blocks users who exceed their limit)
+ * 2. Least-loaded key selection (distributes requests evenly across keys)
+ */
+export class HostedKeyThrottler {
+  private storage: RateLimitStorageAdapter
+  /** In-memory request counters per key: "provider:keyIndex" -> count */
+  private keyRequestCounts = new Map<string, number>()
+
+  constructor(storage?: RateLimitStorageAdapter) {
+    this.storage = storage ?? createStorageAdapter()
+  }
+
+  /**
+   * Build storage key for per-user rate limiting
+   */
+  private buildUserStorageKey(provider: string, userId: string): string {
+    return `hosted:${provider}:user:${userId}:${USER_REQUESTS_DIMENSION}`
+  }
+
+  /**
+   * Get available keys from environment variables
+   */
+  private getAvailableKeys(envKeys: string[]): AvailableKey[] {
+    const keys: AvailableKey[] = []
+    for (let i = 0; i < envKeys.length; i++) {
+      const envVarName = envKeys[i]
+      const key = process.env[envVarName]
+      if (key) {
+        keys.push({ key, keyIndex: i, envVarName })
+      }
+    }
+    return keys
+  }
+
+  /**
+   * Get user rate limit config from throttle config
+   */
+  private getUserRateLimitConfig(throttle: ThrottleConfig): TokenBucketConfig | null {
+    if (throttle.mode !== 'per_request' || !throttle.userRequestsPerMinute) {
+      return null
+    }
+    return toTokenBucketConfig(
+      throttle.userRequestsPerMinute,
+      throttle.burstMultiplier ?? DEFAULT_BURST_MULTIPLIER,
+      THROTTLE_WINDOW_MS
+    )
+  }
+
+  /**
+   * Check and consume user rate limit. Returns null if allowed, or retry info if throttled.
+   */
+  private async checkUserRateLimit(
+    provider: string,
+    userId: string,
+    throttle: ThrottleConfig
+  ): Promise<{ throttled: true; retryAfterMs: number } | null> {
+    const config = this.getUserRateLimitConfig(throttle)
+    if (!config) return null
+
+    const storageKey = this.buildUserStorageKey(provider, userId)
+
+    try {
+      const result = await this.storage.consumeTokens(storageKey, 1, config)
+      if (!result.allowed) {
+        const retryAfterMs = Math.max(0, result.resetAt.getTime() - Date.now())
+        logger.info(`User ${userId} throttled for ${provider}`, {
+          provider,
+          userId,
+          retryAfterMs,
+          tokensRemaining: result.tokensRemaining,
+        })
+        return { throttled: true, retryAfterMs }
+      }
+      return null
+    } catch (error) {
+      logger.error(`Error checking user rate limit for ${provider}`, { error, userId })
+      return null // Allow on error
+    }
+  }
+
+  /**
+   * Acquire the best available key.
+   *
+   * 1. Per-user throttling (enforced): Users exceeding their limit get blocked
+   * 2. Least-loaded key selection: Picks the key with fewest requests
+   */
+  async acquireKey(
+    provider: string,
+    envKeys: string[],
+    throttle: ThrottleConfig,
+    userId?: string
+  ): Promise<AcquireKeyResult> {
+    if (userId && throttle.mode === 'per_request' && throttle.userRequestsPerMinute) {
+      const userThrottleResult = await this.checkUserRateLimit(provider, userId, throttle)
+      if (userThrottleResult) {
+        return {
+          success: false,
+          userThrottled: true,
+          retryAfterMs: userThrottleResult.retryAfterMs,
+          error: `Rate limit exceeded. Please wait ${Math.ceil(userThrottleResult.retryAfterMs / 1000)} seconds.`,
+        }
+      }
+    }
+
+    const availableKeys = this.getAvailableKeys(envKeys)
+
+    if (availableKeys.length === 0) {
+      logger.warn(`No hosted keys configured for provider ${provider}`)
+      return {
+        success: false,
+        error: `No hosted keys configured for ${provider}`,
+      }
+    }
+
+    // Select the key with fewest requests
+    let leastLoaded = availableKeys[0]
+    let minCount = this.getKeyCount(provider, leastLoaded.keyIndex)
+
+    for (let i = 1; i < availableKeys.length; i++) {
+      const count = this.getKeyCount(provider, availableKeys[i].keyIndex)
+      if (count < minCount) {
+        minCount = count
+        leastLoaded = availableKeys[i]
+      }
+    }
+
+    this.incrementKeyCount(provider, leastLoaded.keyIndex)
+
+    logger.debug(`Selected hosted key for ${provider}`, {
+      provider,
+      keyIndex: leastLoaded.keyIndex,
+      envVarName: leastLoaded.envVarName,
+      requestCount: minCount + 1,
+    })
+
+    return {
+      success: true,
+      key: leastLoaded.key,
+      keyIndex: leastLoaded.keyIndex,
+      envVarName: leastLoaded.envVarName,
+    }
+  }
+
+  private getKeyCount(provider: string, keyIndex: number): number {
+    return this.keyRequestCounts.get(`${provider}:${keyIndex}`) ?? 0
+  }
+
+  private incrementKeyCount(provider: string, keyIndex: number): void {
+    const key = `${provider}:${keyIndex}`
+    this.keyRequestCounts.set(key, (this.keyRequestCounts.get(key) ?? 0) + 1)
+  }
+}
+
+let cachedThrottler: HostedKeyThrottler | null = null
+
+/**
+ * Get the singleton HostedKeyThrottler instance
+ */
+export function getHostedKeyThrottler(): HostedKeyThrottler {
+  if (!cachedThrottler) {
+    cachedThrottler = new HostedKeyThrottler()
+  }
+  return cachedThrottler
+}
+
+/**
+ * Reset the cached throttler (for testing)
+ */
+export function resetHostedKeyThrottler(): void {
+  cachedThrottler = null
+}