chore: import upstream snapshot with attribution

apps/sim/lib/core/rate-limiter/hosted-key/hosted-key-rate-limiter.test.ts

@@ -0,0 +1,864 @@
+import { sleep } from '@sim/utils/helpers'
+import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
+import type {
+  ConsumeResult,
+  RateLimitStorageAdapter,
+  TokenStatus,
+} from '@/lib/core/rate-limiter/storage'
+import { HostedKeyRateLimiter } from './hosted-key-rate-limiter'
+import { HEARTBEAT_REFRESH_INTERVAL_MS, type HostedKeyQueue } from './queue'
+import type { CustomRateLimit, PerRequestRateLimit } from './types'
+
+/** Force the queue wait to give up on the first iteration by reporting a retry time
+ *  larger than the 5-minute MAX_QUEUE_WAIT_MS cap. */
+const RETRY_PAST_CAP_MS = 6 * 60 * 1000
+
+interface MockAdapter {
+  consumeTokens: Mock
+  getTokenStatus: Mock
+  resetBucket: Mock
+}
+
+const createMockAdapter = (): MockAdapter => ({
+  consumeTokens: vi.fn(),
+  getTokenStatus: vi.fn(),
+  resetBucket: vi.fn(),
+})
+
+interface MockQueue {
+  enqueue: Mock
+  checkHead: Mock
+  refreshHeartbeat: Mock
+  dequeue: Mock
+}
+
+/** Stub queue that defaults to "you're at the head, no waiting" — i.e. acts as if the
+ *  queue is empty or Redis is unavailable. Tests override per-call to simulate ordering. */
+const createMockQueue = (): MockQueue => {
+  const queue: MockQueue = {
+    enqueue: vi.fn().mockResolvedValue({ position: 0, enabled: true }),
+    checkHead: vi.fn().mockResolvedValue('head'),
+    refreshHeartbeat: vi.fn().mockResolvedValue(undefined),
+    dequeue: vi.fn().mockResolvedValue(undefined),
+  }
+  return queue
+}
+
+describe('HostedKeyRateLimiter', () => {
+  const testProvider = 'exa'
+  const envKeyPrefix = 'EXA_API_KEY'
+  let mockAdapter: MockAdapter
+  let mockQueue: MockQueue
+  let rateLimiter: HostedKeyRateLimiter
+  let originalEnv: NodeJS.ProcessEnv
+
+  const perRequestRateLimit: PerRequestRateLimit = {
+    mode: 'per_request',
+    requestsPerMinute: 10,
+  }
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockAdapter = createMockAdapter()
+    mockQueue = createMockQueue()
+    rateLimiter = new HostedKeyRateLimiter(
+      mockAdapter as RateLimitStorageAdapter,
+      mockQueue as unknown as HostedKeyQueue
+    )
+
+    originalEnv = { ...process.env }
+    process.env.EXA_API_KEY_COUNT = '3'
+    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 () => {
+      const allowedResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
+
+      process.env.EXA_API_KEY_COUNT = undefined
+      process.env.EXA_API_KEY_1 = undefined
+      process.env.EXA_API_KEY_2 = undefined
+      process.env.EXA_API_KEY_3 = undefined
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(false)
+      expect(result.error).toContain('No hosted keys configured')
+    })
+
+    it('should rate limit billing actor when wait exceeds the queue cap', async () => {
+      // resetAt past the 5-minute cap forces the wait loop to bail immediately.
+      const rateLimitedResult: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(rateLimitedResult)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-123'
+      )
+
+      expect(result.success).toBe(false)
+      expect(result.billingActorRateLimited).toBe(true)
+      expect(result.retryAfterMs).toBeDefined()
+      expect(result.error).toContain('Rate limit exceeded')
+    })
+
+    it('should wait for capacity then succeed when bucket refills within the cap', async () => {
+      // First call: bucket empty, refills in 100ms (well under cap).
+      // Second call: bucket has capacity, consumed.
+      const blocked: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 100),
+      }
+      const allowed: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValueOnce(blocked).mockResolvedValueOnce(allowed)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-wait'
+      )
+
+      expect(result.success).toBe(true)
+      expect(result.key).toBe('test-key-1')
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(2)
+    })
+
+    it('should allow billing actor 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 rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-123'
+      )
+
+      expect(result.success).toBe(true)
+      expect(result.billingActorRateLimited).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 rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1'
+      )
+      const r2 = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-2'
+      )
+      const r3 = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-3'
+      )
+      const r4 = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-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 handle partial key availability', async () => {
+      const allowedResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
+
+      process.env.EXA_API_KEY_2 = undefined
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(true)
+      expect(result.key).toBe('test-key-1')
+      expect(result.envVarName).toBe('EXA_API_KEY_1')
+
+      const r2 = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-2'
+      )
+      expect(r2.keyIndex).toBe(2) // Skips missing key 1
+      expect(r2.envVarName).toBe('EXA_API_KEY_3')
+    })
+  })
+
+  describe('FIFO queue ordering', () => {
+    const allowed: ConsumeResult = {
+      allowed: true,
+      tokensRemaining: 9,
+      resetAt: new Date(Date.now() + 60000),
+    }
+
+    it('enqueues every call onto the per-workspace+provider queue', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+
+      await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit, 'workspace-1')
+
+      expect(mockQueue.enqueue).toHaveBeenCalledWith(
+        testProvider,
+        'workspace-1',
+        expect.any(String)
+      )
+    })
+
+    it('always dequeues at the end of a successful acquisition', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+
+      await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit, 'workspace-1')
+
+      expect(mockQueue.dequeue).toHaveBeenCalledWith(
+        testProvider,
+        'workspace-1',
+        expect.any(String)
+      )
+    })
+
+    it('always dequeues even when the call fails (no keys configured)', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+      process.env.EXA_API_KEY_COUNT = '0'
+
+      await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit, 'workspace-1')
+
+      expect(mockQueue.dequeue).toHaveBeenCalled()
+    })
+
+    it('waits at the head of the queue before consuming from the bucket', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+      // First two checkHead calls say we're waiting; third says we're up.
+      mockQueue.checkHead
+        .mockResolvedValueOnce('waiting')
+        .mockResolvedValueOnce('waiting')
+        .mockResolvedValueOnce('head')
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(true)
+      expect(mockQueue.checkHead).toHaveBeenCalledTimes(3)
+      // Bucket is only consumed once we reach the head.
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(1)
+    })
+
+    it('refreshes the heartbeat while waiting at the head of the queue', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+
+      // We need the wait loop to iterate long enough for HEARTBEAT_REFRESH_INTERVAL_MS
+      // to elapse. Use fake timers so we don't actually sleep.
+      vi.useFakeTimers()
+      try {
+        // Queue says we're waiting forever — except after some time we're at head.
+        mockQueue.checkHead.mockImplementation(async () => {
+          // Advance past the heartbeat interval each time we poll, then say we're up.
+          vi.advanceTimersByTime(15_000)
+          return mockQueue.checkHead.mock.calls.length >= 2 ? 'head' : 'waiting'
+        })
+
+        const promise = rateLimiter.acquireKey(
+          testProvider,
+          envKeyPrefix,
+          perRequestRateLimit,
+          'workspace-1'
+        )
+        // Drain pending timers so the sleep() resolves.
+        await vi.runAllTimersAsync()
+        await promise
+
+        expect(mockQueue.refreshHeartbeat).toHaveBeenCalled()
+      } finally {
+        vi.useRealTimers()
+      }
+    })
+
+    it('returns 429 when the queue wait exceeds the cap', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+      mockQueue.checkHead.mockResolvedValue('waiting')
+
+      vi.useFakeTimers()
+      try {
+        const promise = rateLimiter.acquireKey(
+          testProvider,
+          envKeyPrefix,
+          perRequestRateLimit,
+          'workspace-1'
+        )
+        // Burn past the 5-minute cap.
+        await vi.advanceTimersByTimeAsync(6 * 60 * 1000)
+        const result = await promise
+
+        expect(result.success).toBe(false)
+        expect(result.billingActorRateLimited).toBe(true)
+      } finally {
+        vi.useRealTimers()
+      }
+    })
+
+    it('treats "missing" status as proceed (queue evicted, fall through to bucket race)', async () => {
+      mockAdapter.consumeTokens.mockResolvedValue(allowed)
+      mockQueue.checkHead.mockResolvedValueOnce('missing')
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(true)
+    })
+  })
+
+  describe('execution-budget-bounded waits', () => {
+    it('bails immediately when the execution signal is already aborted', async () => {
+      const blocked: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 100),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(blocked)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1',
+        AbortSignal.abort()
+      )
+
+      expect(result.success).toBe(false)
+      expect(result.billingActorRateLimited).toBe(true)
+      // Aborted budget => give up on the first bucket check rather than looping.
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(1)
+    })
+
+    it('stops waiting promptly when the signal aborts mid-sleep', async () => {
+      // Bucket reports a long refill, so the wait sleeps up to the heartbeat cap (10s).
+      // Aborting mid-sleep must wake the wait within a tick, not after the full interval.
+      const blocked: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 10_000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(blocked)
+
+      const controller = new AbortController()
+      const start = Date.now()
+      const promise = rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1',
+        controller.signal
+      )
+      // Let the first bucket check run and the sleep begin, then abort.
+      await sleep(20)
+      controller.abort()
+      const result = await promise
+
+      expect(result.success).toBe(false)
+      expect(result.billingActorRateLimited).toBe(true)
+      // Resolved well before the 10s capped sleep would otherwise have elapsed.
+      expect(Date.now() - start).toBeLessThan(2000)
+    })
+
+    it('keeps waiting past the no-signal fallback cap while the signal is live', async () => {
+      // A live (non-aborted) signal means the run still has budget, so the wait must not
+      // 429 at the 5-minute MAX_QUEUE_WAIT_MS fallback. The bucket frees up after ~7 min.
+      const blocked: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 10_000),
+      }
+      const allowedResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60_000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(blocked)
+
+      vi.useFakeTimers()
+      try {
+        const promise = rateLimiter.acquireKey(
+          testProvider,
+          envKeyPrefix,
+          perRequestRateLimit,
+          'workspace-1',
+          new AbortController().signal
+        )
+        // Burn well past the 5-minute fallback cap — without a signal this would have 429'd.
+        await vi.advanceTimersByTimeAsync(7 * 60 * 1000)
+        mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
+        await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_INTERVAL_MS)
+        const result = await promise
+
+        expect(result.success).toBe(true)
+        expect(result.key).toBe('test-key-1')
+      } finally {
+        vi.useRealTimers()
+      }
+    })
+
+    it('refreshes the heartbeat during a long low-RPM bucket wait', async () => {
+      // Provider with a long refill (retryAfterMs >> heartbeat TTL). The sleep must be
+      // capped so the heartbeat is renewed and the head is not reaped mid-wait.
+      const blocked: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 60_000),
+      }
+      const allowedResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 60_000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(blocked)
+
+      vi.useFakeTimers()
+      try {
+        const promise = rateLimiter.acquireKey(
+          testProvider,
+          envKeyPrefix,
+          perRequestRateLimit,
+          'workspace-1',
+          new AbortController().signal
+        )
+        await vi.advanceTimersByTimeAsync(3 * HEARTBEAT_REFRESH_INTERVAL_MS)
+        mockAdapter.consumeTokens.mockResolvedValue(allowedResult)
+        await vi.advanceTimersByTimeAsync(HEARTBEAT_REFRESH_INTERVAL_MS)
+        await promise
+
+        expect(mockQueue.refreshHeartbeat).toHaveBeenCalled()
+      } finally {
+        vi.useRealTimers()
+      }
+    })
+  })
+
+  describe('acquireKey with custom rate limit', () => {
+    const customRateLimit: CustomRateLimit = {
+      mode: 'custom',
+      requestsPerMinute: 5,
+      dimensions: [
+        {
+          name: 'tokens',
+          limitPerMinute: 1000,
+          extractUsage: (_params, response) => (response.tokenCount as number) ?? 0,
+        },
+      ],
+    }
+
+    it('should enforce requestsPerMinute for custom mode when wait exceeds the cap', async () => {
+      const rateLimitedResult: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(rateLimitedResult)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        customRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(false)
+      expect(result.billingActorRateLimited).toBe(true)
+      expect(result.error).toContain('Rate limit exceeded')
+    })
+
+    it('should allow request when actor request limit and dimensions have budget', async () => {
+      const allowedConsume: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 4,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedConsume)
+
+      const budgetAvailable: TokenStatus = {
+        tokensAvailable: 500,
+        maxTokens: 2000,
+        lastRefillAt: new Date(),
+        nextRefillAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.getTokenStatus.mockResolvedValue(budgetAvailable)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        customRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(true)
+      expect(result.key).toBe('test-key-1')
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(1)
+      expect(mockAdapter.getTokenStatus).toHaveBeenCalledTimes(1)
+    })
+
+    it('should block request when a dimension wait exceeds the cap', async () => {
+      const allowedConsume: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 4,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedConsume)
+
+      const depleted: TokenStatus = {
+        tokensAvailable: 0,
+        maxTokens: 2000,
+        lastRefillAt: new Date(),
+        nextRefillAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
+      }
+      mockAdapter.getTokenStatus.mockResolvedValue(depleted)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        customRateLimit,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(false)
+      expect(result.billingActorRateLimited).toBe(true)
+      expect(result.error).toContain('tokens')
+    })
+
+    it('should wait for dimension capacity then succeed when budget refills', async () => {
+      const allowedConsume: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 4,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedConsume)
+
+      const depleted: TokenStatus = {
+        tokensAvailable: 0,
+        maxTokens: 2000,
+        lastRefillAt: new Date(),
+        nextRefillAt: new Date(Date.now() + 100),
+      }
+      const refilled: TokenStatus = {
+        tokensAvailable: 500,
+        maxTokens: 2000,
+        lastRefillAt: new Date(),
+        nextRefillAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.getTokenStatus.mockResolvedValueOnce(depleted).mockResolvedValueOnce(refilled)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        customRateLimit,
+        'workspace-dim-wait'
+      )
+
+      expect(result.success).toBe(true)
+      expect(mockAdapter.getTokenStatus).toHaveBeenCalledTimes(2)
+    })
+
+    it('should pre-check all dimensions and block on first depleted one', async () => {
+      const multiDimensionConfig: CustomRateLimit = {
+        mode: 'custom',
+        requestsPerMinute: 10,
+        dimensions: [
+          {
+            name: 'tokens',
+            limitPerMinute: 1000,
+            extractUsage: (_p, r) => (r.tokenCount as number) ?? 0,
+          },
+          {
+            name: 'search_units',
+            limitPerMinute: 50,
+            extractUsage: (_p, r) => (r.searchUnits as number) ?? 0,
+          },
+        ],
+      }
+
+      const allowedConsume: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 9,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(allowedConsume)
+
+      const tokensBudget: TokenStatus = {
+        tokensAvailable: 500,
+        maxTokens: 2000,
+        lastRefillAt: new Date(),
+        nextRefillAt: new Date(Date.now() + 60000),
+      }
+      const searchUnitsDepleted: TokenStatus = {
+        tokensAvailable: 0,
+        maxTokens: 100,
+        lastRefillAt: new Date(),
+        nextRefillAt: new Date(Date.now() + RETRY_PAST_CAP_MS),
+      }
+      mockAdapter.getTokenStatus
+        .mockResolvedValueOnce(tokensBudget)
+        .mockResolvedValueOnce(searchUnitsDepleted)
+
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        multiDimensionConfig,
+        'workspace-1'
+      )
+
+      expect(result.success).toBe(false)
+      expect(result.billingActorRateLimited).toBe(true)
+      expect(result.error).toContain('search_units')
+    })
+  })
+
+  describe('reportUsage', () => {
+    const customConfig: CustomRateLimit = {
+      mode: 'custom',
+      requestsPerMinute: 5,
+      dimensions: [
+        {
+          name: 'tokens',
+          limitPerMinute: 1000,
+          extractUsage: (_params, response) => (response.tokenCount as number) ?? 0,
+        },
+      ],
+    }
+
+    it('should consume actual tokens from dimension bucket after execution', async () => {
+      const consumeResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 850,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(consumeResult)
+
+      const result = await rateLimiter.reportUsage(
+        testProvider,
+        'workspace-1',
+        customConfig,
+        {},
+        { tokenCount: 150 }
+      )
+
+      expect(result.dimensions).toHaveLength(1)
+      expect(result.dimensions[0].name).toBe('tokens')
+      expect(result.dimensions[0].consumed).toBe(150)
+      expect(result.dimensions[0].allowed).toBe(true)
+      expect(result.dimensions[0].tokensRemaining).toBe(850)
+
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledWith(
+        'hosted:exa:actor:workspace-1:tokens',
+        150,
+        expect.objectContaining({ maxTokens: 2000, refillRate: 1000 })
+      )
+    })
+
+    it('should handle overdrawn bucket gracefully (optimistic concurrency)', async () => {
+      const overdrawnResult: ConsumeResult = {
+        allowed: false,
+        tokensRemaining: 0,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(overdrawnResult)
+
+      const result = await rateLimiter.reportUsage(
+        testProvider,
+        'workspace-1',
+        customConfig,
+        {},
+        { tokenCount: 500 }
+      )
+
+      expect(result.dimensions[0].allowed).toBe(false)
+      expect(result.dimensions[0].consumed).toBe(500)
+    })
+
+    it('should skip consumption when extractUsage returns 0', async () => {
+      const result = await rateLimiter.reportUsage(
+        testProvider,
+        'workspace-1',
+        customConfig,
+        {},
+        { tokenCount: 0 }
+      )
+
+      expect(result.dimensions).toHaveLength(1)
+      expect(result.dimensions[0].consumed).toBe(0)
+      expect(mockAdapter.consumeTokens).not.toHaveBeenCalled()
+    })
+
+    it('should handle multiple dimensions independently', async () => {
+      const multiConfig: CustomRateLimit = {
+        mode: 'custom',
+        requestsPerMinute: 10,
+        dimensions: [
+          {
+            name: 'tokens',
+            limitPerMinute: 1000,
+            extractUsage: (_p, r) => (r.tokenCount as number) ?? 0,
+          },
+          {
+            name: 'search_units',
+            limitPerMinute: 50,
+            extractUsage: (_p, r) => (r.searchUnits as number) ?? 0,
+          },
+        ],
+      }
+
+      const tokensConsumed: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 800,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      const searchConsumed: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 47,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens
+        .mockResolvedValueOnce(tokensConsumed)
+        .mockResolvedValueOnce(searchConsumed)
+
+      const result = await rateLimiter.reportUsage(
+        testProvider,
+        'workspace-1',
+        multiConfig,
+        {},
+        { tokenCount: 200, searchUnits: 3 }
+      )
+
+      expect(result.dimensions).toHaveLength(2)
+      expect(result.dimensions[0]).toEqual({
+        name: 'tokens',
+        consumed: 200,
+        allowed: true,
+        tokensRemaining: 800,
+      })
+      expect(result.dimensions[1]).toEqual({
+        name: 'search_units',
+        consumed: 3,
+        allowed: true,
+        tokensRemaining: 47,
+      })
+
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(2)
+    })
+
+    it('should continue with remaining dimensions if extractUsage throws', async () => {
+      const throwingConfig: CustomRateLimit = {
+        mode: 'custom',
+        requestsPerMinute: 10,
+        dimensions: [
+          {
+            name: 'broken',
+            limitPerMinute: 100,
+            extractUsage: () => {
+              throw new Error('extraction failed')
+            },
+          },
+          {
+            name: 'tokens',
+            limitPerMinute: 1000,
+            extractUsage: (_p, r) => (r.tokenCount as number) ?? 0,
+          },
+        ],
+      }
+
+      const consumeResult: ConsumeResult = {
+        allowed: true,
+        tokensRemaining: 900,
+        resetAt: new Date(Date.now() + 60000),
+      }
+      mockAdapter.consumeTokens.mockResolvedValue(consumeResult)
+
+      const result = await rateLimiter.reportUsage(
+        testProvider,
+        'workspace-1',
+        throwingConfig,
+        {},
+        { tokenCount: 100 }
+      )
+
+      expect(result.dimensions).toHaveLength(1)
+      expect(result.dimensions[0].name).toBe('tokens')
+      expect(mockAdapter.consumeTokens).toHaveBeenCalledTimes(1)
+    })
+
+    it('should handle storage errors gracefully', async () => {
+      mockAdapter.consumeTokens.mockRejectedValue(new Error('db connection lost'))
+
+      const result = await rateLimiter.reportUsage(
+        testProvider,
+        'workspace-1',
+        customConfig,
+        {},
+        { tokenCount: 100 }
+      )
+
+      expect(result.dimensions).toHaveLength(0)
+    })
+  })
+})

apps/sim/lib/core/rate-limiter/hosted-key/hosted-key-rate-limiter.ts

@@ -0,0 +1,695 @@
+import { createLogger } from '@sim/logger'
+import { sleep } from '@sim/utils/helpers'
+import { generateShortId } from '@sim/utils/id'
+import { getMaxExecutionTimeout } from '@/lib/core/execution-limits'
+import {
+  createStorageAdapter,
+  type RateLimitStorageAdapter,
+  type TokenBucketConfig,
+} from '@/lib/core/rate-limiter/storage'
+import { PlatformEvents } from '@/lib/core/telemetry'
+import { getHostedKeyQueue, HEARTBEAT_REFRESH_INTERVAL_MS, type HostedKeyQueue } from './queue'
+import {
+  type AcquireKeyResult,
+  type CustomRateLimit,
+  DEFAULT_BURST_MULTIPLIER,
+  DEFAULT_WINDOW_MS,
+  type HostedKeyRateLimitConfig,
+  type ReportUsageResult,
+  toTokenBucketConfig,
+} from './types'
+
+const logger = createLogger('HostedKeyRateLimiter')
+
+/**
+ * Fallback ceiling on how long a hosted-key acquisition waits for the per-workspace
+ * bucket to refill when no execution `AbortSignal` is available to bound the wait (e.g.
+ * a caller without a wired execution deadline, or Redis no-op mode). When a signal IS
+ * provided, the wait is instead bounded by the surrounding execution budget — the signal
+ * fires when the run hits its plan timeout or is cancelled — and this constant no longer
+ * applies (see {@link ABSOLUTE_MAX_QUEUE_WAIT_MS} for the backstop in that case).
+ */
+const MAX_QUEUE_WAIT_MS = 5 * 60 * 1000
+
+/**
+ * Hard safety ceiling applied even when an execution `AbortSignal` is present, in case the
+ * signal never fires. Matches the longest possible execution budget (enterprise async) so
+ * it never truncates a legitimately long-running background run before its own deadline.
+ */
+const ABSOLUTE_MAX_QUEUE_WAIT_MS = getMaxExecutionTimeout()
+
+/**
+ * Floor on per-iteration sleep when the bucket reports `retryAfterMs <= 0`,
+ * which can happen due to clock skew or sub-millisecond resets. Prevents a
+ * tight retry loop hammering the storage adapter.
+ */
+const MIN_QUEUE_RETRY_DELAY_MS = 50
+
+/**
+ * Poll interval while waiting to reach the head of the FIFO queue. 200ms balances
+ * acquisition latency (worst-case wait for advancement is one poll period) against
+ * Redis load — at this cadence, N waiters generate N×5 EVAL/sec, which is fine for
+ * the typical low-tens contention. Revisit if telemetry shows hot Redis under load.
+ */
+const QUEUE_HEAD_POLL_MS = 200
+
+/**
+ * Sleep for `ms`, resolving early if `signal` aborts. Cleans up its own timer and listener
+ * so neither leaks. Callers don't need to distinguish an early (aborted) return from a normal
+ * one — the surrounding wait loop re-checks its budget immediately after and bails when the
+ * signal has fired. Falls back to a plain sleep when no signal is provided.
+ */
+function interruptibleSleep(ms: number, signal?: AbortSignal): Promise<void> {
+  if (!signal) return sleep(ms)
+  if (signal.aborted) return Promise.resolve()
+  return new Promise<void>((resolve) => {
+    const onAbort = () => {
+      clearTimeout(timer)
+      signal.removeEventListener('abort', onAbort)
+      resolve()
+    }
+    const timer = setTimeout(() => {
+      signal.removeEventListener('abort', onAbort)
+      resolve()
+    }, ms)
+    signal.addEventListener('abort', onAbort, { once: true })
+    // Catch an abort that fired between the guard above and addEventListener.
+    if (signal.aborted) onAbort()
+  })
+}
+
+/**
+ * Resolves env var names for a numbered key prefix using a `{PREFIX}_COUNT` env var.
+ * E.g. with `EXA_API_KEY_COUNT=5`, returns `['EXA_API_KEY_1', ..., 'EXA_API_KEY_5']`.
+ */
+function resolveEnvKeys(prefix: string): string[] {
+  const count = Number.parseInt(process.env[`${prefix}_COUNT`] || '0', 10)
+  const names: string[] = []
+  for (let i = 1; i <= count; i++) {
+    names.push(`${prefix}_${i}`)
+  }
+  return names
+}
+
+/** Dimension name for per-billing-actor request rate limiting */
+const ACTOR_REQUESTS_DIMENSION = 'actor_requests'
+
+/**
+ * Information about an available hosted key
+ */
+interface AvailableKey {
+  key: string
+  keyIndex: number
+  envVarName: string
+}
+
+/**
+ * Mutable heartbeat bookkeeping shared across every wait phase of a single `acquireKey`
+ * call. Carrying one `lastHeartbeatAt` across the queue-head wait and the bucket waits
+ * ensures the heartbeat-refresh cadence reflects the *actual* last write to Redis rather
+ * than resetting per phase, which could otherwise let the ticket heartbeat lapse and get
+ * reaped mid-wait (breaking FIFO ordering).
+ */
+interface WaitState {
+  /** Epoch ms of the last heartbeat write to Redis. */
+  lastHeartbeatAt: number
+}
+
+/**
+ * HostedKeyRateLimiter provides:
+ * 1. Per-billing-actor rate limiting (enforced - blocks actors who exceed their limit)
+ * 2. Round-robin key selection (distributes requests evenly across keys)
+ * 3. Post-execution dimension usage tracking for custom rate limits
+ *
+ * The billing actor is typically a workspace ID, meaning rate limits are shared
+ * across all users within the same workspace.
+ */
+export class HostedKeyRateLimiter {
+  private storage: RateLimitStorageAdapter
+  private queue: HostedKeyQueue
+  /** Round-robin counter per provider for even key distribution */
+  private roundRobinCounters = new Map<string, number>()
+
+  constructor(storage?: RateLimitStorageAdapter, queue?: HostedKeyQueue) {
+    this.storage = storage ?? createStorageAdapter()
+    this.queue = queue ?? getHostedKeyQueue()
+  }
+
+  private buildActorStorageKey(provider: string, billingActorId: string): string {
+    return `hosted:${provider}:actor:${billingActorId}:${ACTOR_REQUESTS_DIMENSION}`
+  }
+
+  private buildDimensionStorageKey(
+    provider: string,
+    billingActorId: string,
+    dimensionName: string
+  ): string {
+    return `hosted:${provider}:actor:${billingActorId}:${dimensionName}`
+  }
+
+  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
+  }
+
+  /**
+   * Build a token bucket config for the per-billing-actor request rate limit.
+   * Works for both `per_request` and `custom` modes since both define `requestsPerMinute`.
+   */
+  private getActorRateLimitConfig(config: HostedKeyRateLimitConfig): TokenBucketConfig | null {
+    if (!config.requestsPerMinute) return null
+    return toTokenBucketConfig(
+      config.requestsPerMinute,
+      config.burstMultiplier ?? DEFAULT_BURST_MULTIPLIER,
+      DEFAULT_WINDOW_MS
+    )
+  }
+
+  /**
+   * Check and consume billing actor request rate limit. Returns null if allowed, or retry info if blocked.
+   */
+  private async checkActorRateLimit(
+    provider: string,
+    billingActorId: string,
+    config: HostedKeyRateLimitConfig
+  ): Promise<{ rateLimited: true; retryAfterMs: number } | null> {
+    const bucketConfig = this.getActorRateLimitConfig(config)
+    if (!bucketConfig) return null
+
+    const storageKey = this.buildActorStorageKey(provider, billingActorId)
+
+    try {
+      const result = await this.storage.consumeTokens(storageKey, 1, bucketConfig)
+      if (!result.allowed) {
+        const retryAfterMs = Math.max(0, result.resetAt.getTime() - Date.now())
+        logger.info(`Billing actor ${billingActorId} rate limited for ${provider}`, {
+          provider,
+          billingActorId,
+          retryAfterMs,
+          tokensRemaining: result.tokensRemaining,
+        })
+        return { rateLimited: true, retryAfterMs }
+      }
+      return null
+    } catch (error) {
+      logger.error(`Error checking billing actor rate limit for ${provider}`, {
+        error,
+        billingActorId,
+      })
+      return null
+    }
+  }
+
+  /**
+   * Pre-check that the billing actor has available budget in all custom dimensions.
+   * Does NOT consume tokens -- just verifies the actor isn't already depleted.
+   * Returns retry info for the most restrictive exhausted dimension, or null if all pass.
+   */
+  private async preCheckDimensions(
+    provider: string,
+    billingActorId: string,
+    config: CustomRateLimit
+  ): Promise<{ rateLimited: true; retryAfterMs: number; dimension: string } | null> {
+    for (const dimension of config.dimensions) {
+      const storageKey = this.buildDimensionStorageKey(provider, billingActorId, dimension.name)
+      const bucketConfig = toTokenBucketConfig(
+        dimension.limitPerMinute,
+        dimension.burstMultiplier ?? DEFAULT_BURST_MULTIPLIER,
+        DEFAULT_WINDOW_MS
+      )
+
+      try {
+        const status = await this.storage.getTokenStatus(storageKey, bucketConfig)
+        if (status.tokensAvailable < 1) {
+          const retryAfterMs = Math.max(0, status.nextRefillAt.getTime() - Date.now())
+          logger.info(
+            `Billing actor ${billingActorId} exhausted dimension ${dimension.name} for ${provider}`,
+            {
+              provider,
+              billingActorId,
+              dimension: dimension.name,
+              tokensAvailable: status.tokensAvailable,
+              retryAfterMs,
+            }
+          )
+          return { rateLimited: true, retryAfterMs, dimension: dimension.name }
+        }
+      } catch (error) {
+        logger.error(`Error pre-checking dimension ${dimension.name} for ${provider}`, {
+          error,
+          billingActorId,
+        })
+      }
+    }
+    return null
+  }
+
+  /**
+   * Acquire an available key via round-robin selection.
+   *
+   * For both modes:
+   *   1. Per-billing-actor request rate limiting (enforced): the call enqueues itself
+   *      onto a per-workspace+provider FIFO queue. Only the head of the queue attempts
+   *      to consume from the token bucket, guaranteeing strict ordering across callers
+   *      within a workspace. Different workspaces have independent queues and don't
+   *      block each other.
+   *   2. Round-robin key selection: cycles through available keys for even distribution
+   *
+   * For `custom` mode additionally:
+   *   3. Pre-checks dimension budgets: head waits on dimension refill the same way it
+   *      waits on actor request capacity.
+   *
+   * The wait is bounded by the surrounding execution budget: when `signal` is provided it
+   * fires at the run's plan timeout (or on cancellation), so a queued call uses its full
+   * available budget rather than a flat cap. When no signal is available the wait falls
+   * back to `MAX_QUEUE_WAIT_MS`. On exhaustion the call returns today's 429 result. The
+   * ticket is removed from the queue on exit regardless of success or failure.
+   *
+   * @param envKeyPrefix - Env var prefix (e.g. 'EXA_API_KEY'). Keys resolved via `{prefix}_COUNT`.
+   * @param billingActorId - The billing actor (typically workspace ID) to rate limit against
+   * @param signal - Optional execution `AbortSignal`; bounds the queue wait to the run's budget.
+   */
+  async acquireKey(
+    provider: string,
+    envKeyPrefix: string,
+    config: HostedKeyRateLimitConfig,
+    billingActorId: string,
+    signal?: AbortSignal
+  ): Promise<AcquireKeyResult> {
+    const ticketId = generateShortId()
+    const startedAt = Date.now()
+    const waitState: WaitState = { lastHeartbeatAt: startedAt }
+    const enqueueResult = await this.queue.enqueue(provider, billingActorId, ticketId)
+
+    try {
+      // Wait for our turn at the head of the queue (no-op when Redis unavailable).
+      const headStatus = await this.waitForQueueHead(
+        provider,
+        billingActorId,
+        ticketId,
+        startedAt,
+        waitState,
+        signal
+      )
+      if (headStatus.timedOut) {
+        PlatformEvents.hostedKeyQueueWaitExceeded({
+          provider,
+          workspaceId: billingActorId,
+          waitedMs: Date.now() - startedAt,
+          reason: 'queue_position',
+        })
+        return {
+          success: false,
+          billingActorRateLimited: true,
+          retryAfterMs: MAX_QUEUE_WAIT_MS,
+          error: `Rate limit exceeded — request waited too long in the queue. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
+        }
+      }
+
+      let dimensionWaited = false
+
+      if (config.requestsPerMinute) {
+        const rateLimitResult = await this.waitForActorCapacity(
+          provider,
+          billingActorId,
+          ticketId,
+          config,
+          startedAt,
+          waitState,
+          signal
+        )
+        if (rateLimitResult.rateLimited) {
+          return {
+            success: false,
+            billingActorRateLimited: true,
+            retryAfterMs: rateLimitResult.retryAfterMs,
+            error: `Rate limit exceeded. Please wait ${Math.ceil(rateLimitResult.retryAfterMs / 1000)} seconds. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
+          }
+        }
+      }
+
+      if (config.mode === 'custom' && config.dimensions.length > 0) {
+        const dimensionResult = await this.waitForDimensionCapacity(
+          provider,
+          billingActorId,
+          ticketId,
+          config,
+          startedAt,
+          waitState,
+          signal
+        )
+        if (dimensionResult.rateLimited) {
+          return {
+            success: false,
+            billingActorRateLimited: true,
+            retryAfterMs: dimensionResult.retryAfterMs,
+            error: `Rate limit exceeded for ${dimensionResult.dimension}. Please wait ${Math.ceil(dimensionResult.retryAfterMs / 1000)} seconds. If you're getting throttled frequently, consider adding your own API key under Settings > BYOK to avoid shared rate limits.`,
+          }
+        }
+        dimensionWaited = dimensionResult.waited
+      }
+
+      const totalWaitedMs = Date.now() - startedAt
+      if (enqueueResult.enabled && (enqueueResult.position > 0 || totalWaitedMs > 100)) {
+        // Attribute the wait to its dominant cause: queue depth takes precedence (it's
+        // reported alongside queuePosition), otherwise the bucket phase that actually slept.
+        const reason: 'queue_position' | 'actor_requests' | 'dimension' =
+          enqueueResult.position > 0
+            ? 'queue_position'
+            : dimensionWaited
+              ? 'dimension'
+              : 'actor_requests'
+        PlatformEvents.hostedKeyQueueWaited({
+          provider,
+          workspaceId: billingActorId,
+          waitedMs: totalWaitedMs,
+          attempts: 1,
+          reason,
+          queuePosition: enqueueResult.position,
+        })
+      }
+
+      const envKeys = resolveEnvKeys(envKeyPrefix)
+      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}`,
+        }
+      }
+
+      const counter = this.roundRobinCounters.get(provider) ?? 0
+      const selected = availableKeys[counter % availableKeys.length]
+      this.roundRobinCounters.set(provider, counter + 1)
+
+      logger.debug(`Selected hosted key for ${provider}`, {
+        provider,
+        keyIndex: selected.keyIndex,
+        envVarName: selected.envVarName,
+      })
+
+      return {
+        success: true,
+        key: selected.key,
+        keyIndex: selected.keyIndex,
+        envVarName: selected.envVarName,
+      }
+    } finally {
+      // Always remove our ticket so the next caller can advance, regardless of whether
+      // we succeeded, hit the cap, or threw. Best-effort; safe to call multiple times.
+      await this.queue.dequeue(provider, billingActorId, ticketId)
+    }
+  }
+
+  /**
+   * Remaining time budget for waiting, in milliseconds. When an execution `AbortSignal` is
+   * present it governs the wait: the budget is exhausted the moment the signal aborts (the
+   * run hit its plan timeout or was cancelled), with {@link ABSOLUTE_MAX_QUEUE_WAIT_MS} as a
+   * backstop should the signal never fire. Without a signal we fall back to the flat
+   * {@link MAX_QUEUE_WAIT_MS} ceiling.
+   */
+  private remainingWaitBudgetMs(startedAt: number, signal?: AbortSignal): number {
+    if (signal?.aborted) return 0
+    const ceiling = signal ? ABSOLUTE_MAX_QUEUE_WAIT_MS : MAX_QUEUE_WAIT_MS
+    return ceiling - (Date.now() - startedAt)
+  }
+
+  /** Refresh the ticket heartbeat if the refresh interval has elapsed since the last write. */
+  private async maybeRefreshHeartbeat(
+    provider: string,
+    billingActorId: string,
+    ticketId: string,
+    waitState: WaitState
+  ): Promise<void> {
+    if (Date.now() - waitState.lastHeartbeatAt >= HEARTBEAT_REFRESH_INTERVAL_MS) {
+      await this.queue.refreshHeartbeat(provider, billingActorId, ticketId)
+      waitState.lastHeartbeatAt = Date.now()
+    }
+  }
+
+  /**
+   * Sleep before the next bucket re-check, refreshing the heartbeat first if due. The sleep
+   * is capped at {@link HEARTBEAT_REFRESH_INTERVAL_MS} so that no single wait can outlive the
+   * heartbeat TTL — even when the bucket reports a long `retryAfterMs` (e.g. low-RPM
+   * providers). Without this cap a multi-second sleep could let the heartbeat lapse, the
+   * head get reaped as dead, and a second caller advance and race us for the bucket. The
+   * sleep also resolves early if `signal` aborts, so a cancelled/timed-out run stops waiting
+   * promptly rather than overshooting by up to the cap.
+   */
+  private async heartbeatAwareSleep(
+    provider: string,
+    billingActorId: string,
+    ticketId: string,
+    desiredMs: number,
+    waitState: WaitState,
+    signal?: AbortSignal
+  ): Promise<void> {
+    await this.maybeRefreshHeartbeat(provider, billingActorId, ticketId, waitState)
+    const sleepMs = Math.min(
+      Math.max(MIN_QUEUE_RETRY_DELAY_MS, desiredMs),
+      HEARTBEAT_REFRESH_INTERVAL_MS
+    )
+    await interruptibleSleep(sleepMs, signal)
+  }
+
+  /**
+   * Block until our ticket reaches the head of the queue. Refreshes the heartbeat on a
+   * regular cadence so we don't get reaped as dead. Returns `timedOut: true` once the wait
+   * budget is exhausted before reaching the head (see {@link remainingWaitBudgetMs}).
+   *
+   * No-op when Redis is unavailable (queue.enqueue returns enabled=false and checkHead
+   * always returns 'head').
+   */
+  private async waitForQueueHead(
+    provider: string,
+    billingActorId: string,
+    ticketId: string,
+    startedAt: number,
+    waitState: WaitState,
+    signal?: AbortSignal
+  ): Promise<{ timedOut: boolean }> {
+    while (true) {
+      const status = await this.queue.checkHead(provider, billingActorId, ticketId)
+      if (status === 'head') return { timedOut: false }
+
+      // 'missing' shouldn't normally happen — the queue list TTL (10min) outlives a typical
+      // wait — but if it does (e.g. Redis flushed mid-wait), treat as "you're up" so the
+      // caller proceeds to the bucket race rather than hanging forever.
+      if (status === 'missing') return { timedOut: false }
+
+      if (this.remainingWaitBudgetMs(startedAt, signal) <= 0) {
+        return { timedOut: true }
+      }
+
+      await this.maybeRefreshHeartbeat(provider, billingActorId, ticketId, waitState)
+      await interruptibleSleep(QUEUE_HEAD_POLL_MS, signal)
+    }
+  }
+
+  /**
+   * Wait for actor request-rate capacity. Called once we're at the head of the FIFO
+   * queue, so other callers can't race us for the next token — they're blocked behind us
+   * at queue level. Re-checks the bucket until the wait budget is exhausted (accounting for
+   * time already spent waiting in the queue). `waited` reports whether the loop ever slept.
+   */
+  private async waitForActorCapacity(
+    provider: string,
+    billingActorId: string,
+    ticketId: string,
+    config: HostedKeyRateLimitConfig,
+    startedAt: number,
+    waitState: WaitState,
+    signal?: AbortSignal
+  ): Promise<
+    { rateLimited: false; waited: boolean } | { rateLimited: true; retryAfterMs: number }
+  > {
+    let waited = false
+
+    while (true) {
+      const result = await this.checkActorRateLimit(provider, billingActorId, config)
+      if (!result) return { rateLimited: false, waited }
+
+      const remaining = this.remainingWaitBudgetMs(startedAt, signal)
+      if (remaining <= 0 || result.retryAfterMs > remaining) {
+        PlatformEvents.hostedKeyQueueWaitExceeded({
+          provider,
+          workspaceId: billingActorId,
+          waitedMs: Date.now() - startedAt,
+          reason: 'actor_requests',
+        })
+        return { rateLimited: true, retryAfterMs: result.retryAfterMs }
+      }
+
+      waited = true
+      await this.heartbeatAwareSleep(
+        provider,
+        billingActorId,
+        ticketId,
+        result.retryAfterMs,
+        waitState,
+        signal
+      )
+    }
+  }
+
+  /**
+   * Wait for custom-mode dimension capacity. `preCheckDimensions` is read-only — it does
+   * not consume — so re-running it after a sleep is safe and does not double-charge.
+   * Post-execution `reportUsage` performs the actual consumption. `waited` reports whether
+   * the loop ever slept.
+   */
+  private async waitForDimensionCapacity(
+    provider: string,
+    billingActorId: string,
+    ticketId: string,
+    config: CustomRateLimit,
+    startedAt: number,
+    waitState: WaitState,
+    signal?: AbortSignal
+  ): Promise<
+    | { rateLimited: false; waited: boolean }
+    | { rateLimited: true; retryAfterMs: number; dimension: string }
+  > {
+    let waited = false
+
+    while (true) {
+      const result = await this.preCheckDimensions(provider, billingActorId, config)
+      if (!result) return { rateLimited: false, waited }
+
+      const remaining = this.remainingWaitBudgetMs(startedAt, signal)
+      if (remaining <= 0 || result.retryAfterMs > remaining) {
+        PlatformEvents.hostedKeyQueueWaitExceeded({
+          provider,
+          workspaceId: billingActorId,
+          waitedMs: Date.now() - startedAt,
+          reason: 'dimension',
+          dimension: result.dimension,
+        })
+        return {
+          rateLimited: true,
+          retryAfterMs: result.retryAfterMs,
+          dimension: result.dimension,
+        }
+      }
+
+      waited = true
+      await this.heartbeatAwareSleep(
+        provider,
+        billingActorId,
+        ticketId,
+        result.retryAfterMs,
+        waitState,
+        signal
+      )
+    }
+  }
+
+  /**
+   * Report actual usage after successful tool execution (custom mode only).
+   * Calls `extractUsage` on each dimension and consumes the actual token count.
+   * This is the "post-execution" phase of the optimistic two-phase approach.
+   */
+  async reportUsage(
+    provider: string,
+    billingActorId: string,
+    config: CustomRateLimit,
+    params: Record<string, unknown>,
+    response: Record<string, unknown>
+  ): Promise<ReportUsageResult> {
+    const results: ReportUsageResult['dimensions'] = []
+
+    for (const dimension of config.dimensions) {
+      let usage: number
+      try {
+        usage = dimension.extractUsage(params, response)
+      } catch (error) {
+        logger.error(`Failed to extract usage for dimension ${dimension.name}`, {
+          provider,
+          billingActorId,
+          error,
+        })
+        continue
+      }
+
+      if (usage <= 0) {
+        results.push({
+          name: dimension.name,
+          consumed: 0,
+          allowed: true,
+          tokensRemaining: 0,
+        })
+        continue
+      }
+
+      const storageKey = this.buildDimensionStorageKey(provider, billingActorId, dimension.name)
+      const bucketConfig = toTokenBucketConfig(
+        dimension.limitPerMinute,
+        dimension.burstMultiplier ?? DEFAULT_BURST_MULTIPLIER,
+        DEFAULT_WINDOW_MS
+      )
+
+      try {
+        const consumeResult = await this.storage.consumeTokens(storageKey, usage, bucketConfig)
+
+        results.push({
+          name: dimension.name,
+          consumed: usage,
+          allowed: consumeResult.allowed,
+          tokensRemaining: consumeResult.tokensRemaining,
+        })
+
+        if (!consumeResult.allowed) {
+          logger.warn(
+            `Dimension ${dimension.name} overdrawn for ${provider} (optimistic concurrency)`,
+            { provider, billingActorId, usage, tokensRemaining: consumeResult.tokensRemaining }
+          )
+        }
+
+        logger.debug(`Consumed ${usage} from dimension ${dimension.name} for ${provider}`, {
+          provider,
+          billingActorId,
+          usage,
+          allowed: consumeResult.allowed,
+          tokensRemaining: consumeResult.tokensRemaining,
+        })
+      } catch (error) {
+        logger.error(`Failed to consume tokens for dimension ${dimension.name}`, {
+          provider,
+          billingActorId,
+          usage,
+          error,
+        })
+      }
+    }
+
+    return { dimensions: results }
+  }
+}
+
+let cachedInstance: HostedKeyRateLimiter | null = null
+
+/**
+ * Get the singleton HostedKeyRateLimiter instance
+ */
+export function getHostedKeyRateLimiter(): HostedKeyRateLimiter {
+  if (!cachedInstance) {
+    cachedInstance = new HostedKeyRateLimiter()
+  }
+  return cachedInstance
+}
+
+/**
+ * Reset the cached rate limiter (for testing)
+ */
+export function resetHostedKeyRateLimiter(): void {
+  cachedInstance = null
+}

apps/sim/lib/core/rate-limiter/hosted-key/index.ts

@@ -0,0 +1,11 @@
+export {
+  getHostedKeyRateLimiter,
+  HostedKeyRateLimiter,
+  resetHostedKeyRateLimiter,
+} from './hosted-key-rate-limiter'
+export {
+  DEFAULT_BURST_MULTIPLIER,
+  DEFAULT_WINDOW_MS,
+  type HostedKeyRateLimitConfig,
+  toTokenBucketConfig,
+} from './types'

apps/sim/lib/core/rate-limiter/hosted-key/queue.test.ts

@@ -0,0 +1,234 @@
+import { redisConfigMock, redisConfigMockFns } from '@sim/testing'
+import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
+import { HostedKeyQueue } from './queue'
+
+vi.mock('@/lib/core/config/redis', () => redisConfigMock)
+
+interface MockPipeline {
+  rpush: Mock
+  expire: Mock
+  set: Mock
+  lrem: Mock
+  del: Mock
+  exec: Mock
+}
+
+interface MockRedis {
+  multi: Mock
+  set: Mock
+  eval: Mock
+  pipeline: MockPipeline
+}
+
+function createFakeRedis(): MockRedis {
+  const pipeline: MockPipeline = {
+    rpush: vi.fn(),
+    expire: vi.fn(),
+    set: vi.fn(),
+    lrem: vi.fn(),
+    del: vi.fn(),
+    exec: vi.fn(),
+  }
+  // Pipeline methods return the pipeline for chaining.
+  pipeline.rpush.mockReturnValue(pipeline)
+  pipeline.expire.mockReturnValue(pipeline)
+  pipeline.set.mockReturnValue(pipeline)
+  pipeline.lrem.mockReturnValue(pipeline)
+  pipeline.del.mockReturnValue(pipeline)
+
+  return {
+    multi: vi.fn(() => pipeline),
+    set: vi.fn(),
+    eval: vi.fn(),
+    pipeline,
+  }
+}
+
+const provider = 'exa'
+const workspaceId = 'workspace-1'
+const ticketId = 'ticket-1'
+
+describe('HostedKeyQueue', () => {
+  let queue: HostedKeyQueue
+  let mockRedis: MockRedis
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockRedis = createFakeRedis()
+    redisConfigMockFns.mockGetRedisClient.mockReturnValue(mockRedis)
+    queue = new HostedKeyQueue()
+  })
+
+  describe('enqueue', () => {
+    it('returns position 0 when first in line', async () => {
+      // RPUSH returns new list length; first push -> 1.
+      mockRedis.pipeline.exec.mockResolvedValueOnce([
+        [null, 1],
+        [null, 1],
+        [null, 'OK'],
+      ])
+
+      const result = await queue.enqueue(provider, workspaceId, ticketId)
+
+      expect(result).toEqual({ position: 0, enabled: true })
+      expect(mockRedis.pipeline.rpush).toHaveBeenCalledWith(
+        'hosted-queue:exa:workspace-1',
+        ticketId
+      )
+      expect(mockRedis.pipeline.set).toHaveBeenCalledWith(
+        'hosted-queue-tkt:exa:workspace-1:ticket-1',
+        '1',
+        'EX',
+        expect.any(Number)
+      )
+    })
+
+    it('returns higher position when others are ahead', async () => {
+      // Length 5 after push -> position 4.
+      mockRedis.pipeline.exec.mockResolvedValueOnce([
+        [null, 5],
+        [null, 1],
+        [null, 'OK'],
+      ])
+
+      const result = await queue.enqueue(provider, workspaceId, ticketId)
+
+      expect(result.position).toBe(4)
+    })
+
+    it('falls back to enabled=false when Redis is unavailable', async () => {
+      redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
+
+      const result = await queue.enqueue(provider, workspaceId, ticketId)
+
+      expect(result).toEqual({ position: 0, enabled: false })
+    })
+
+    it('falls back to enabled=false on Redis error', async () => {
+      mockRedis.pipeline.exec.mockRejectedValueOnce(new Error('connection lost'))
+
+      const result = await queue.enqueue(provider, workspaceId, ticketId)
+
+      expect(result.enabled).toBe(false)
+    })
+  })
+
+  describe('checkHead', () => {
+    it('returns "head" when our ticket is at the head', async () => {
+      mockRedis.eval.mockResolvedValueOnce('head')
+
+      const status = await queue.checkHead(provider, workspaceId, ticketId)
+
+      expect(status).toBe('head')
+    })
+
+    it('returns "waiting" when someone else is the head', async () => {
+      mockRedis.eval.mockResolvedValueOnce('waiting')
+
+      const status = await queue.checkHead(provider, workspaceId, ticketId)
+
+      expect(status).toBe('waiting')
+    })
+
+    it('returns "missing" when our ticket is not in the queue', async () => {
+      mockRedis.eval.mockResolvedValueOnce('missing')
+
+      const status = await queue.checkHead(provider, workspaceId, ticketId)
+
+      expect(status).toBe('missing')
+    })
+
+    it('passes queue list key, heartbeat prefix, and ticketId to the Lua script', async () => {
+      mockRedis.eval.mockResolvedValueOnce('head')
+
+      await queue.checkHead(provider, workspaceId, ticketId)
+
+      expect(mockRedis.eval).toHaveBeenCalledWith(
+        expect.stringContaining('lindex'),
+        1,
+        'hosted-queue:exa:workspace-1',
+        'hosted-queue-tkt:exa:workspace-1:',
+        ticketId
+      )
+    })
+
+    it('fails open to "head" on Redis error so callers do not hang', async () => {
+      mockRedis.eval.mockRejectedValueOnce(new Error('boom'))
+
+      const status = await queue.checkHead(provider, workspaceId, ticketId)
+
+      expect(status).toBe('head')
+    })
+
+    it('returns "head" no-op when Redis is unavailable', async () => {
+      redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
+
+      const status = await queue.checkHead(provider, workspaceId, ticketId)
+
+      expect(status).toBe('head')
+    })
+  })
+
+  describe('refreshHeartbeat', () => {
+    it('writes the heartbeat key with TTL and re-extends the queue list TTL', async () => {
+      mockRedis.pipeline.exec.mockResolvedValueOnce([
+        [null, 'OK'],
+        [null, 1],
+      ])
+
+      await queue.refreshHeartbeat(provider, workspaceId, ticketId)
+
+      expect(mockRedis.pipeline.set).toHaveBeenCalledWith(
+        'hosted-queue-tkt:exa:workspace-1:ticket-1',
+        '1',
+        'EX',
+        expect.any(Number)
+      )
+      expect(mockRedis.pipeline.expire).toHaveBeenCalledWith(
+        'hosted-queue:exa:workspace-1',
+        expect.any(Number)
+      )
+      expect(mockRedis.pipeline.exec).toHaveBeenCalledTimes(1)
+    })
+
+    it('is a no-op when Redis is unavailable', async () => {
+      redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
+
+      await expect(queue.refreshHeartbeat(provider, workspaceId, ticketId)).resolves.toBeUndefined()
+      expect(mockRedis.multi).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('dequeue', () => {
+    it('removes the ticket from the list and deletes the heartbeat', async () => {
+      mockRedis.pipeline.exec.mockResolvedValueOnce([
+        [null, 1],
+        [null, 1],
+      ])
+
+      await queue.dequeue(provider, workspaceId, ticketId)
+
+      expect(mockRedis.pipeline.lrem).toHaveBeenCalledWith(
+        'hosted-queue:exa:workspace-1',
+        1,
+        ticketId
+      )
+      expect(mockRedis.pipeline.del).toHaveBeenCalledWith(
+        'hosted-queue-tkt:exa:workspace-1:ticket-1'
+      )
+    })
+
+    it('is a no-op when Redis is unavailable', async () => {
+      redisConfigMockFns.mockGetRedisClient.mockReturnValueOnce(null)
+
+      await expect(queue.dequeue(provider, workspaceId, ticketId)).resolves.toBeUndefined()
+      expect(mockRedis.multi).not.toHaveBeenCalled()
+    })
+
+    it('swallows errors so callers do not throw on cleanup', async () => {
+      mockRedis.pipeline.exec.mockRejectedValueOnce(new Error('connection lost'))
+
+      await expect(queue.dequeue(provider, workspaceId, ticketId)).resolves.toBeUndefined()
+    })
+  })
+})

apps/sim/lib/core/rate-limiter/hosted-key/queue.ts

@@ -0,0 +1,209 @@
+import { createLogger } from '@sim/logger'
+import { toError } from '@sim/utils/errors'
+import { getRedisClient } from '@/lib/core/config/redis'
+
+const logger = createLogger('HostedKeyQueue')
+
+/**
+ * Per-ticket heartbeat TTL. Refreshed by the head while it's actively waiting
+ * on the bucket. If the holder crashes, the heartbeat key expires, and the next
+ * caller sees the head as dead and removes it (lazy cleanup).
+ */
+const TICKET_HEARTBEAT_TTL_SECONDS = 30
+
+/** How often the head should refresh its heartbeat while waiting. */
+export const HEARTBEAT_REFRESH_INTERVAL_MS = 10_000
+
+/**
+ * TTL on the queue list itself. Set on enqueue and re-extended by the head's heartbeat,
+ * so a long-waiting head can't let the list expire out from under the waiters behind it.
+ * Prevents abandoned queues from sticking around forever in Redis.
+ */
+const QUEUE_LIST_TTL_SECONDS = 600
+
+const queueListKey = (provider: string, billingActorId: string): string =>
+  `hosted-queue:${provider}:${billingActorId}`
+
+const heartbeatKey = (provider: string, billingActorId: string, ticketId: string): string =>
+  `hosted-queue-tkt:${provider}:${billingActorId}:${ticketId}`
+
+/**
+ * Atomically reap any dead head, then return our ticket's status. Combines what
+ * would otherwise be 3 round-trips (reap, LINDEX, LPOS) into one EVAL — meaningful
+ * because callers poll this every ~200ms while waiting in the queue.
+ *
+ * `KEYS[1]` = queue list key. `ARGV[1]` = heartbeat key prefix. `ARGV[2]` = our ticketId.
+ *
+ * Reaping is bounded: at most one dead head is removed per call. If multiple dead
+ * tickets pile up at the head, subsequent polls will clean them one by one. This
+ * keeps the script O(1) rather than O(N) and is sufficient because queue depth
+ * is bounded by concurrent callers per workspace (typically tens).
+ *
+ * Returns one of: "head", "waiting", "missing".
+ */
+const CHECK_HEAD_SCRIPT = `
+local head = redis.call("lindex", KEYS[1], 0)
+if head and redis.call("exists", ARGV[1] .. head) == 0 then
+  redis.call("lrem", KEYS[1], 1, head)
+  head = redis.call("lindex", KEYS[1], 0)
+end
+if not head then
+  return "missing"
+end
+if head == ARGV[2] then
+  return "head"
+end
+if redis.call("lpos", KEYS[1], ARGV[2]) == false then
+  return "missing"
+end
+return "waiting"
+`
+
+export interface EnqueueResult {
+  /** Position at the moment of enqueue (0 = head, you go next). */
+  position: number
+  /** Whether Redis was available — false means we're in no-op mode. */
+  enabled: boolean
+}
+
+/**
+ * Per-workspace+provider FIFO queue for hosted-key acquisitions.
+ *
+ * Callers `enqueue` to claim a position, then `waitForHead` until they're at
+ * the head, then attempt to consume from the token bucket. On success or cap
+ * exceeded, they `dequeue` to make room for the next caller.
+ *
+ * No-op when Redis is unavailable: every method returns "you're the head /
+ * empty / etc." so the rate limiter falls back to plain bucket racing.
+ */
+export class HostedKeyQueue {
+  /**
+   * Push a ticket onto the tail of the queue and write a heartbeat. Returns the
+   * position at enqueue time (0 = head, ready to proceed).
+   */
+  async enqueue(
+    provider: string,
+    billingActorId: string,
+    ticketId: string
+  ): Promise<EnqueueResult> {
+    const redis = getRedisClient()
+    if (!redis) {
+      return { position: 0, enabled: false }
+    }
+
+    const listKey = queueListKey(provider, billingActorId)
+    const hbKey = heartbeatKey(provider, billingActorId, ticketId)
+
+    try {
+      const pipeline = redis.multi()
+      pipeline.rpush(listKey, ticketId)
+      pipeline.expire(listKey, QUEUE_LIST_TTL_SECONDS)
+      pipeline.set(hbKey, '1', 'EX', TICKET_HEARTBEAT_TTL_SECONDS)
+      const results = await pipeline.exec()
+      // results[0] is the rpush response: [err, length]
+      const length = results?.[0] && typeof results[0][1] === 'number' ? results[0][1] : 1
+      // Position is length - 1 (just-pushed at the tail).
+      return { position: length - 1, enabled: true }
+    } catch (error) {
+      logger.warn(`Queue enqueue failed for ${listKey}`, { error: toError(error).message })
+      return { position: 0, enabled: false }
+    }
+  }
+
+  /**
+   * Check whether `ticketId` is currently at the head of the queue. If the head
+   * is a different ticket but its heartbeat has expired (caller crashed), reap
+   * it and re-check on the next poll.
+   *
+   * Returns:
+   *  - "head": you're at the head, proceed to consume from the bucket
+   *  - "waiting": someone else is the head and they're alive
+   *  - "missing": your ticket isn't in the queue at all (e.g. queue list TTL
+   *    expired); caller should re-enqueue or treat as enabled=false
+   */
+  async checkHead(
+    provider: string,
+    billingActorId: string,
+    ticketId: string
+  ): Promise<'head' | 'waiting' | 'missing'> {
+    const redis = getRedisClient()
+    if (!redis) {
+      return 'head'
+    }
+
+    const listKey = queueListKey(provider, billingActorId)
+    const hbPrefix = `hosted-queue-tkt:${provider}:${billingActorId}:`
+
+    try {
+      const result = (await redis.eval(CHECK_HEAD_SCRIPT, 1, listKey, hbPrefix, ticketId)) as
+        | 'head'
+        | 'waiting'
+        | 'missing'
+      return result
+    } catch (error) {
+      logger.warn(`Queue checkHead failed for ${listKey}`, { error: toError(error).message })
+      // Fail-open: treat as head so the caller proceeds rather than hanging.
+      return 'head'
+    }
+  }
+
+  /**
+   * Refresh the ticket's heartbeat so the head isn't reaped as dead while waiting on the
+   * bucket. Also re-extends the queue list TTL so a wait outliving {@link QUEUE_LIST_TTL_SECONDS}
+   * doesn't let the list expire and collapse FIFO ordering.
+   */
+  async refreshHeartbeat(
+    provider: string,
+    billingActorId: string,
+    ticketId: string
+  ): Promise<void> {
+    const redis = getRedisClient()
+    if (!redis) return
+
+    const listKey = queueListKey(provider, billingActorId)
+    const hbKey = heartbeatKey(provider, billingActorId, ticketId)
+    try {
+      const pipeline = redis.multi()
+      pipeline.set(hbKey, '1', 'EX', TICKET_HEARTBEAT_TTL_SECONDS)
+      pipeline.expire(listKey, QUEUE_LIST_TTL_SECONDS)
+      await pipeline.exec()
+    } catch (error) {
+      logger.warn(`Queue heartbeat refresh failed for ${hbKey}`, {
+        error: toError(error).message,
+      })
+    }
+  }
+
+  /**
+   * Remove a ticket from the queue and its heartbeat key. Best-effort; safe to
+   * call multiple times. LREM count=1 removes at most one matching entry.
+   */
+  async dequeue(provider: string, billingActorId: string, ticketId: string): Promise<void> {
+    const redis = getRedisClient()
+    if (!redis) return
+
+    const listKey = queueListKey(provider, billingActorId)
+    const hbKey = heartbeatKey(provider, billingActorId, ticketId)
+    try {
+      const pipeline = redis.multi()
+      pipeline.lrem(listKey, 1, ticketId)
+      pipeline.del(hbKey)
+      await pipeline.exec()
+    } catch (error) {
+      logger.warn(`Queue dequeue failed for ${listKey}`, { error: toError(error).message })
+    }
+  }
+}
+
+let cachedQueue: HostedKeyQueue | null = null
+
+export function getHostedKeyQueue(): HostedKeyQueue {
+  if (!cachedQueue) {
+    cachedQueue = new HostedKeyQueue()
+  }
+  return cachedQueue
+}
+
+export function resetHostedKeyQueue(): void {
+  cachedQueue = null
+}

apps/sim/lib/core/rate-limiter/hosted-key/types.ts

@@ -0,0 +1,108 @@
+import type { TokenBucketConfig } from '@/lib/core/rate-limiter/storage'
+
+export type HostedKeyRateLimitMode = 'per_request' | 'custom'
+
+/**
+ * Simple per-request rate limit configuration.
+ * Enforces per-billing-actor rate limiting and distributes requests across keys.
+ */
+export interface PerRequestRateLimit {
+  mode: 'per_request'
+  /** Maximum requests per minute per billing actor (enforced - blocks if exceeded) */
+  requestsPerMinute: number
+  /** Burst multiplier for token bucket max capacity. Default: 2 */
+  burstMultiplier?: number
+}
+
+/**
+ * Custom rate limit with multiple dimensions (e.g., tokens, search units).
+ * Allows tracking different usage metrics independently.
+ */
+export interface CustomRateLimit {
+  mode: 'custom'
+  /** Maximum requests per minute per billing actor (enforced - blocks if exceeded) */
+  requestsPerMinute: number
+  /** Multiple dimensions to track */
+  dimensions: RateLimitDimension[]
+  /** Burst multiplier for token bucket max capacity. Default: 2 */
+  burstMultiplier?: number
+}
+
+/**
+ * A single dimension for custom rate limiting.
+ * Each dimension has its own token bucket.
+ */
+interface RateLimitDimension {
+  /** Dimension name (e.g., 'tokens', 'search_units') - used in storage key */
+  name: string
+  /** Limit per minute for this dimension */
+  limitPerMinute: number
+  /** Burst multiplier for token bucket max capacity. Default: 2 */
+  burstMultiplier?: number
+  /**
+   * Extract usage amount from request params and response.
+   * Called after successful execution to consume the actual usage.
+   */
+  extractUsage: (params: Record<string, unknown>, response: Record<string, unknown>) => number
+}
+
+/** Union of all hosted key rate limit configuration types */
+export type HostedKeyRateLimitConfig = PerRequestRateLimit | CustomRateLimit
+
+/**
+ * Result from acquiring a key from the hosted key rate limiter
+ */
+export interface AcquireKeyResult {
+  /** Whether a key was successfully acquired */
+  success: boolean
+  /** The API key value (if success=true) */
+  key?: string
+  /** Index of the key in the envKeys array */
+  keyIndex?: number
+  /** Environment variable name of the selected key */
+  envVarName?: string
+  /** Error message if no key available */
+  error?: string
+  /** Whether the billing actor was rate limited (exceeded their limit) */
+  billingActorRateLimited?: boolean
+  /** Milliseconds until the billing actor's rate limit resets (if billingActorRateLimited=true) */
+  retryAfterMs?: number
+}
+
+/**
+ * Result from reporting post-execution usage for custom dimensions
+ */
+export interface ReportUsageResult {
+  /** Per-dimension consumption results */
+  dimensions: {
+    name: string
+    consumed: number
+    allowed: boolean
+    tokensRemaining: number
+  }[]
+}
+
+/**
+ * Convert rate limit config to token bucket config for a dimension
+ */
+export function toTokenBucketConfig(
+  limitPerMinute: number,
+  burstMultiplier = 2,
+  windowMs = 60000
+): TokenBucketConfig {
+  return {
+    maxTokens: limitPerMinute * burstMultiplier,
+    refillRate: limitPerMinute,
+    refillIntervalMs: windowMs,
+  }
+}
+
+/**
+ * Default rate limit window in milliseconds (1 minute)
+ */
+export const DEFAULT_WINDOW_MS = 60000
+
+/**
+ * Default burst multiplier
+ */
+export const DEFAULT_BURST_MULTIPLIER = 2