Make billing actor id required for throttling

Author: Theodore Li

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/hooks/use-editor-subblock-layout.ts

@@ -109,7 +109,7 @@ export function useEditorSubblockLayout(
       // Check required feature if specified - declarative feature gating
       if (!isSubBlockFeatureEnabled(block)) return false
 
-      // Hide tool API key fields when hosted key is available
+      // Hide tool API key fields when hosted
       if (isSubBlockHiddenByHostedKey(block)) return false
 
       // Special handling for trigger-config type (legacy trigger configuration UI)

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

@@ -1,12 +1,12 @@
 import { loggerMock } from '@sim/testing'
 import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
-import { HostedKeyRateLimiter } from './hosted-key-rate-limiter'
-import type { CustomRateLimit, PerRequestRateLimit } from './types'
 import type {
   ConsumeResult,
   RateLimitStorageAdapter,
   TokenStatus,
 } from '@/lib/core/rate-limiter/storage'
+import { HostedKeyRateLimiter } from './hosted-key-rate-limiter'
+import type { CustomRateLimit, PerRequestRateLimit } from './types'
 
 vi.mock('@sim/logger', () => loggerMock)
 
@@ -52,12 +52,24 @@ describe('HostedKeyRateLimiter', () => {
 
   describe('acquireKey', () => {
     it('should return error when no keys are configured', async () => {
-      delete process.env.EXA_API_KEY_COUNT
-      delete process.env.EXA_API_KEY_1
-      delete process.env.EXA_API_KEY_2
-      delete process.env.EXA_API_KEY_3
+      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)
+      const result = await rateLimiter.acquireKey(
+        testProvider,
+        envKeyPrefix,
+        perRequestRateLimit,
+        'workspace-1'
+      )
 
       expect(result.success).toBe(false)
       expect(result.error).toContain('No hosted keys configured')
@@ -143,24 +155,33 @@ describe('HostedKeyRateLimiter', () => {
       expect(r4.keyIndex).toBe(0) // Wraps back
     })
 
-    it('should work without billingActorId (no rate limiting)', async () => {
-      const result = await rateLimiter.acquireKey(testProvider, envKeyPrefix, perRequestRateLimit)
-
-      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 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)
+      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)
+      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')
     })
@@ -256,14 +277,6 @@ describe('HostedKeyRateLimiter', () => {
       expect(result.error).toContain('tokens')
     })
 
-    it('should skip dimension pre-check without billingActorId', async () => {
-      const result = await rateLimiter.acquireKey(testProvider, envKeyPrefix, customRateLimit)
-
-      expect(result.success).toBe(true)
-      expect(mockAdapter.consumeTokens).not.toHaveBeenCalled()
-      expect(mockAdapter.getTokenStatus).not.toHaveBeenCalled()
-    })
-
     it('should pre-check all dimensions and block on first depleted one', async () => {
       const multiDimensionConfig: CustomRateLimit = {
         mode: 'custom',

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

@@ -5,13 +5,13 @@ import {
   type TokenBucketConfig,
 } from '@/lib/core/rate-limiter/storage'
 import {
-  DEFAULT_BURST_MULTIPLIER,
-  DEFAULT_WINDOW_MS,
-  toTokenBucketConfig,
   type AcquireKeyResult,
   type CustomRateLimit,
+  DEFAULT_BURST_MULTIPLIER,
+  DEFAULT_WINDOW_MS,
   type HostedKeyRateLimitConfig,
   type ReportUsageResult,
+  toTokenBucketConfig,
 } from './types'
 
 const logger = createLogger('HostedKeyRateLimiter')
@@ -21,7 +21,7 @@ const logger = createLogger('HostedKeyRateLimiter')
  * E.g. with `EXA_API_KEY_COUNT=5`, returns `['EXA_API_KEY_1', ..., 'EXA_API_KEY_5']`.
  */
 function resolveEnvKeys(prefix: string): string[] {
-  const count = parseInt(process.env[`${prefix}_COUNT`] || '0', 10)
+  const count = Number.parseInt(process.env[`${prefix}_COUNT`] || '0', 10)
   const names: string[] = []
   for (let i = 1; i <= count; i++) {
     names.push(`${prefix}_${i}`)
@@ -192,9 +192,9 @@ export class HostedKeyRateLimiter {
     provider: string,
     envKeyPrefix: string,
     config: HostedKeyRateLimitConfig,
-    billingActorId?: string
+    billingActorId: string
   ): Promise<AcquireKeyResult> {
-    if (billingActorId && config.requestsPerMinute) {
+    if (config.requestsPerMinute) {
       const rateLimitResult = await this.checkActorRateLimit(provider, billingActorId, config)
       if (rateLimitResult) {
         return {
@@ -206,7 +206,7 @@ export class HostedKeyRateLimiter {
       }
     }
 
-    if (billingActorId && config.mode === 'custom' && config.dimensions.length > 0) {
+    if (config.mode === 'custom' && config.dimensions.length > 0) {
       const dimensionResult = await this.preCheckDimensions(provider, billingActorId, config)
       if (dimensionResult) {
         return {

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

@@ -4,14 +4,14 @@ export {
   resetHostedKeyRateLimiter,
 } from './hosted-key-rate-limiter'
 export {
-  DEFAULT_BURST_MULTIPLIER,
-  DEFAULT_WINDOW_MS,
-  toTokenBucketConfig,
   type AcquireKeyResult,
   type CustomRateLimit,
+  DEFAULT_BURST_MULTIPLIER,
+  DEFAULT_WINDOW_MS,
   type HostedKeyRateLimitConfig,
   type HostedKeyRateLimitMode,
   type PerRequestRateLimit,
   type RateLimitDimension,
   type ReportUsageResult,
+  toTokenBucketConfig,
 } from './types'

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

@@ -1,22 +1,20 @@
-export type { RateLimitResult, RateLimitStatus } from './rate-limiter'
-export { RateLimiter } from './rate-limiter'
-export type { RateLimitStorageAdapter, TokenBucketConfig } from './storage'
-export type { RateLimitConfig, SubscriptionPlan, TriggerType } from './types'
-export { RATE_LIMITS, RateLimitError } from './types'
-export {
-  getHostedKeyRateLimiter,
-  HostedKeyRateLimiter,
-  resetHostedKeyRateLimiter,
-} from './hosted-key'
 export {
-  DEFAULT_BURST_MULTIPLIER,
-  DEFAULT_WINDOW_MS,
-  toTokenBucketConfig,
   type AcquireKeyResult,
   type CustomRateLimit,
+  DEFAULT_BURST_MULTIPLIER,
+  DEFAULT_WINDOW_MS,
+  getHostedKeyRateLimiter,
   type HostedKeyRateLimitConfig,
+  HostedKeyRateLimiter,
   type HostedKeyRateLimitMode,
   type PerRequestRateLimit,
   type RateLimitDimension,
   type ReportUsageResult,
+  resetHostedKeyRateLimiter,
+  toTokenBucketConfig,
 } from './hosted-key'
+export type { RateLimitResult, RateLimitStatus } from './rate-limiter'
+export { RateLimiter } from './rate-limiter'
+export type { RateLimitStorageAdapter, TokenBucketConfig } from './storage'
+export type { RateLimitConfig, SubscriptionPlan, TriggerType } from './types'
+export { RATE_LIMITS, RateLimitError } from './types'

apps/sim/tools/index.ts

@@ -80,7 +80,17 @@ async function injectHostedKeyIfNeeded(
   const provider = byokProviderId || tool.id
   const billingActorId = executionContext?.workspaceId
 
-  const acquireResult = await rateLimiter.acquireKey(provider, envKeyPrefix, rateLimit, billingActorId)
+  if (!billingActorId) {
+    logger.error(`[${requestId}] No workspace ID available for hosted key rate limiting`)
+    return { isUsingHostedKey: false }
+  }
+
+  const acquireResult = await rateLimiter.acquireKey(
+    provider,
+    envKeyPrefix,
+    rateLimit,
+    billingActorId
+  )
 
   if (!acquireResult.success && acquireResult.billingActorRateLimited) {
     logger.warn(`[${requestId}] Billing actor ${billingActorId} rate limited for ${tool.id}`, {
@@ -304,10 +314,10 @@ async function reportCustomDimensionUsage(
 
     for (const dim of result.dimensions) {
       if (!dim.allowed) {
-        logger.warn(
-          `[${requestId}] Dimension ${dim.name} overdrawn after ${tool.id} execution`,
-          { consumed: dim.consumed, tokensRemaining: dim.tokensRemaining }
-        )
+        logger.warn(`[${requestId}] Dimension ${dim.name} overdrawn after ${tool.id} execution`, {
+          consumed: dim.consumed,
+          tokensRemaining: dim.tokensRemaining,
+        })
       }
     }
   } catch (error) {
@@ -736,7 +746,13 @@ export async function executeTool(
 
       // Post-execution: report custom dimension usage and calculate cost
       if (hostedKeyInfo.isUsingHostedKey && finalResult.success) {
-        await reportCustomDimensionUsage(tool, contextParams, finalResult.output, executionContext, requestId)
+        await reportCustomDimensionUsage(
+          tool,
+          contextParams,
+          finalResult.output,
+          executionContext,
+          requestId
+        )
 
         const { cost: hostedKeyCost, metadata } = await processHostedKeyCost(
           tool,
@@ -804,7 +820,13 @@ export async function executeTool(
 
     // Post-execution: report custom dimension usage and calculate cost
     if (hostedKeyInfo.isUsingHostedKey && finalResult.success) {
-      await reportCustomDimensionUsage(tool, contextParams, finalResult.output, executionContext, requestId)
+      await reportCustomDimensionUsage(
+        tool,
+        contextParams,
+        finalResult.output,
+        executionContext,
+        requestId
+      )
 
       const { cost: hostedKeyCost, metadata } = await processHostedKeyCost(
         tool,