fix(tools): check usage limits before running hosted-key tools via the API (#8602)

* fix(tools): check usage limits before running hosted-key tools via the API

* fix(tools): gate hosted-key calls whose key is an unresolved variable reference

* fix(tools): exempt variable references that resolve to the caller's own key

* fix(tools): resolve the key reference with the registry's own resolver

* fix(tools): decide hosted-key admission on registry-resolved params

Author: waleedlatif1

apps/docs/openapi-v2-resources.json

@@ -4601,7 +4601,7 @@
       "post": {
         "operationId": "executeTool",
         "summary": "Run Tool",
-        "description": "Run a built-in tool using published parameter IDs. Sim resolves `credentialId`, hosted keys, and whole-value `{{VAR_NAME}}` references for `user-only` parameters; other values pass through verbatim. Third-party refusal returns `200` with `status: \"failed\"`; the error envelope covers API failures. Hidden or missing tools return `404`; disallowed integrations return `403` with `error.details.code: INTEGRATION_NOT_ALLOWED`. Hosted-key use is billed to the workspace. Workspace API keys return `403`; use a personal API key or scoped OAuth token.\n\nOAuth scope: `api:write`.",
+        "description": "Run a built-in tool using published parameter IDs. Sim resolves `credentialId`, hosted keys, and whole-value `{{VAR_NAME}}` references for `user-only` parameters; other values pass through verbatim. Third-party refusal returns `200` with `status: \"failed\"`; the error envelope covers API failures. Hidden or missing tools return `404`; disallowed integrations return `403` with `error.details.code: INTEGRATION_NOT_ALLOWED`. Hosted-key use is billed to the workspace; a call that would use a hosted key from a workspace over its usage or billing limits returns `402`. Workspace API keys return `403`; use a personal API key or scoped OAuth token.\n\nOAuth scope: `api:write`.",
         "x-sim-operation": "tools.execute",
         "x-oauth-scope": "api:write",
         "tags": ["Catalog"],
@@ -4658,6 +4658,9 @@
           "401": {
             "$ref": "#/components/responses/Unauthorized"
           },
+          "402": {
+            "$ref": "#/components/responses/UsageLimitExceeded"
+          },
           "403": {
             "$ref": "#/components/responses/Forbidden"
           },
@@ -9192,6 +9195,22 @@
           }
         }
       },
+      "UsageLimitExceeded": {
+        "description": "The workspace has exceeded its usage or billing limits.",
+        "content": {
+          "application/json": {
+            "schema": {
+              "$ref": "#/components/schemas/V2Error"
+            },
+            "example": {
+              "error": {
+                "code": "USAGE_LIMIT_EXCEEDED",
+                "message": "Usage limit exceeded. Please upgrade your plan to continue."
+              }
+            }
+          }
+        }
+      },
       "Forbidden": {
         "description": "The caller lacks the rights this operation requires. When the cause is one a caller can act on, `error.details.code` names it. A resource in a workspace the caller cannot reach at all answers `404` instead, so absence and denial are indistinguishable.",
         "content": {

apps/sim/app/api/v2/tools/[toolId]/execute/route.test.ts

@@ -14,8 +14,10 @@ vi.mock('@/lib/api/server/routes/v2-api-key-auth', () => v2ApiKeyAuthModuleMock)
 vi.mock('@/lib/core/rate-limiter', () => v2RateLimiterModuleMock)
 vi.mock('@/lib/tool-execution/application/execute-tool', () => ({
   executeToolForCaller: { operation: { id: 'tools.execute' }, execute: mocks.execute },
+  ToolUsageLimitExceededError: class ToolUsageLimitExceededError extends Error {},
 }))
 
+import { ToolUsageLimitExceededError } from '@/lib/tool-execution/application/execute-tool'
 import { POST } from '@/app/api/v2/tools/[toolId]/execute/route'
 
 const WORKSPACE_ID = '11111111-2222-4333-8444-555555555555'
@@ -76,6 +78,15 @@ describe('POST /api/v2/tools/{toolId}/execute', () => {
     expect(body.data.error.message).toBe('Firecrawl returned 402')
   })
 
+  it('answers 402 when the workspace is over its usage limit', async () => {
+    mocks.execute.mockRejectedValue(new ToolUsageLimitExceededError('Usage limit exceeded'))
+
+    const response = await post({ workspaceId: WORKSPACE_ID })
+
+    expect(response.status).toBe(402)
+    expect((await response.json()).error.code).toBe('USAGE_LIMIT_EXCEEDED')
+  })
+
   it('rejects a timeout beyond the ceiling', async () => {
     const response = await post({ workspaceId: WORKSPACE_ID, timeoutSeconds: 100_000 })
 

apps/sim/app/api/v2/tools/[toolId]/execute/route.ts

@@ -1,12 +1,31 @@
 import { v2ExecuteToolContract } from '@/lib/api/contracts/v2/catalog'
-import { defineV2JsonRoute, v2ApiKeyAuth, v2RateLimits } from '@/lib/api/server/routes'
-import { executeToolForCaller } from '@/lib/tool-execution/application/execute-tool'
+import {
+  defineV2JsonRoute,
+  type V2ErrorPolicy,
+  v2ApiKeyAuth,
+  v2RateLimits,
+} from '@/lib/api/server/routes'
+import {
+  executeToolForCaller,
+  ToolUsageLimitExceededError,
+} from '@/lib/tool-execution/application/execute-tool'
 import { toolExecutionOperations } from '@/lib/tool-execution/application/operations'
 import { catalogErrorPolicy } from '@/app/api/v2/lib/catalog'
+import { v2Error } from '@/app/api/v2/lib/response'
 
 export const dynamic = 'force-dynamic'
 export const revalidate = 0
 
+/** {@link catalogErrorPolicy} plus the `402` a hosted-key call over its usage limit raises. */
+const executeToolErrorPolicy = {
+  render(error) {
+    if (error instanceof ToolUsageLimitExceededError) {
+      return v2Error('USAGE_LIMIT_EXCEEDED', error.message)
+    }
+    return catalogErrorPolicy.render(error)
+  },
+} satisfies V2ErrorPolicy
+
 /**
  * POST /api/v2/tools/{toolId}/execute — Run one built-in tool.
  *
@@ -20,7 +39,7 @@ export const POST = defineV2JsonRoute({
   operation: toolExecutionOperations.execute,
   auth: v2ApiKeyAuth,
   rateLimit: v2RateLimits.publicApi,
-  errorPolicy: catalogErrorPolicy,
+  errorPolicy: executeToolErrorPolicy,
   mapInput: ({ params, body }) => ({
     workspaceId: body.workspaceId,
     toolId: params.toolId,

apps/sim/lib/api/contracts/v2/openapi/resources.ts

@@ -2100,8 +2100,8 @@ const declaredRoutes = [
       applicationOperation: toolExecutionOperations.execute,
       operationId: 'executeTool',
       summary: 'Run Tool',
-      description: `Run a built-in tool using published parameter IDs. Sim resolves \`credentialId\`, hosted keys, and whole-value \`{{VAR_NAME}}\` references for \`user-only\` parameters; other values pass through verbatim. Third-party refusal returns \`200\` with \`status: "failed"\`; the error envelope covers API failures. Hidden or missing tools return \`404\`; disallowed integrations return \`403\` with \`error.details.code: INTEGRATION_NOT_ALLOWED\`. Hosted-key use is billed to the workspace. ${WORKSPACE_API_KEY_DENIED}`,
-      errors: RESOURCE_ERRORS,
+      description: `Run a built-in tool using published parameter IDs. Sim resolves \`credentialId\`, hosted keys, and whole-value \`{{VAR_NAME}}\` references for \`user-only\` parameters; other values pass through verbatim. Third-party refusal returns \`200\` with \`status: "failed"\`; the error envelope covers API failures. Hidden or missing tools return \`404\`; disallowed integrations return \`403\` with \`error.details.code: INTEGRATION_NOT_ALLOWED\`. Hosted-key use is billed to the workspace; a call that would use a hosted key from a workspace over its usage or billing limits returns \`402\`. ${WORKSPACE_API_KEY_DENIED}`,
+      errors: [...RESOURCE_ERRORS, 'UsageLimitExceeded'],
       success: { description: 'The outcome of the tool call.' },
     }),
     {

apps/sim/lib/api/mcp/generated/v2-operations.ts

@@ -1162,7 +1162,7 @@ export const V2_MCP_OPERATIONS = {
     contract: v2ExecuteToolContract,
     summary: 'Run Tool',
     description:
-      'Run a built-in tool using published parameter IDs. Sim resolves `credentialId`, hosted keys, and whole-value `{{VAR_NAME}}` references for `user-only` parameters; other values pass through verbatim. Third-party refusal returns `200` with `status: "failed"`; the error envelope covers API failures. Hidden or missing tools return `404`; disallowed integrations return `403` with `error.details.code: INTEGRATION_NOT_ALLOWED`. Hosted-key use is billed to the workspace. Workspace API keys return `403`; use a personal API key or scoped OAuth token.\n\nOAuth scope: `api:write`.',
+      'Run a built-in tool using published parameter IDs. Sim resolves `credentialId`, hosted keys, and whole-value `{{VAR_NAME}}` references for `user-only` parameters; other values pass through verbatim. Third-party refusal returns `200` with `status: "failed"`; the error envelope covers API failures. Hidden or missing tools return `404`; disallowed integrations return `403` with `error.details.code: INTEGRATION_NOT_ALLOWED`. Hosted-key use is billed to the workspace; a call that would use a hosted key from a workspace over its usage or billing limits returns `402`. Workspace API keys return `403`; use a personal API key or scoped OAuth token.\n\nOAuth scope: `api:write`.',
     workspaceKeyUnsupported: true,
     handler: () => import('@/app/api/v2/tools/[toolId]/execute/route').then((route) => route.POST),
   },

apps/sim/lib/tool-execution/application/execute-tool.test.ts

@@ -9,6 +9,10 @@ import {
   billingAttributionMock,
   billingAttributionMockFns,
 } from '@sim/testing/mocks/billing-attribution.mock'
+import {
+  billingUsageGateCacheMock,
+  billingUsageGateCacheMockFns,
+} from '@sim/testing/mocks/billing-usage-gate-cache.mock'
 import {
   billingUsageLogMock,
   billingUsageLogMockFns,
@@ -22,6 +26,7 @@ import {
   customBlockOperationsMockFns,
 } from '@sim/testing/mocks/custom-block-operations.mock'
 import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing/mocks/env-flags.mock'
+import { environmentUtilsMockFns } from '@sim/testing/mocks/environment-utils.mock'
 import {
   integrationsAvailabilityMock,
   integrationsAvailabilityMockFns,
@@ -87,11 +92,16 @@ vi.mock('@/lib/internal/file/operations', () => ({
 
 vi.mock('@/lib/billing/core/billing-attribution', () => billingAttributionMock)
 
+vi.mock('@/lib/billing/core/usage-gate-cache', () => billingUsageGateCacheMock)
+
 vi.mock('@/lib/billing/core/usage-log', () => billingUsageLogMock)
 
 import { executeFileTool } from '@/lib/internal/file/execute-tool'
 import type { InternalToolOperationContext } from '@/lib/internal/tool-operations/types'
-import { executeToolForCaller } from '@/lib/tool-execution/application/execute-tool'
+import {
+  executeToolForCaller,
+  ToolUsageLimitExceededError,
+} from '@/lib/tool-execution/application/execute-tool'
 import { getAllBlocks, getBlock, getBlockMeta } from '@/blocks/registry'
 import type { BlockConfig } from '@/blocks/types'
 import { fileReadTool } from '@/tools/file/get'
@@ -116,6 +126,18 @@ const TOOL_METADATA: Record<string, Record<string, unknown>> = {
     },
     hosting: { apiKeyParam: 'apiKey' },
   },
+  image_generate: {
+    id: 'image_generate',
+    name: 'Image Generate',
+    params: {
+      provider: { type: 'string', required: true, visibility: 'user-only' },
+      apiKey: { type: 'string', required: true, visibility: 'user-only' },
+    },
+    hosting: {
+      apiKeyParam: 'apiKey',
+      enabled: (params: { provider?: unknown }) => params.provider === 'falai',
+    },
+  },
   snowflake_execute_sql: {
     id: 'snowflake_execute_sql',
     name: 'Snowflake Execute SQL',
@@ -156,6 +178,7 @@ const mocks = {
   getAllBlocks: vi.mocked(getAllBlocks),
   executeRegistryTool: toolsMockFns.mockExecuteTool,
   resolveBillingAttribution: billingAttributionMockFns.mockResolveBillingAttribution,
+  checkUsageLimits: billingUsageGateCacheMockFns.mockCheckExecutionUsageLimits,
 }
 
 vi.mocked(getBlock).mockReturnValue(undefined as never)
@@ -193,6 +216,7 @@ function block(overrides: Partial<BlockConfig> & { type: string }): BlockConfig
 const fileBlock = block({ type: 'file_v5', tools: { access: ['file_read'] } })
 const slackBlock = block({ type: 'slack', tools: { access: ['slack_message'] } })
 const firecrawlBlock = block({ type: 'firecrawl', tools: { access: ['firecrawl_scrape'] } })
+const imageBlock = block({ type: 'image_generator', tools: { access: ['image_generate'] } })
 const previewBlock = block({
   type: 'preview_thing',
   preview: true,
@@ -234,6 +258,7 @@ describe('executeToolForCaller', () => {
       fileBlock,
       slackBlock,
       firecrawlBlock,
+      imageBlock,
       previewBlock,
       confluenceBlock,
       zendeskBlock,
@@ -242,6 +267,8 @@ describe('executeToolForCaller', () => {
     ])
     mocks.executeRegistryTool.mockResolvedValue({ success: true, output: { markdown: '# Hi' } })
     mocks.resolveBillingAttribution.mockResolvedValue({ workspaceId: WORKSPACE_ID })
+    mocks.checkUsageLimits.mockResolvedValue({ isExceeded: false })
+    environmentUtilsMockFns.mockGetEffectiveDecryptedEnv.mockResolvedValue({})
   })
 
   it.each<PersonalApiKeyPrincipal | SessionPrincipal>([principal, createSessionPrincipal()])(
@@ -498,6 +525,59 @@ describe('executeToolForCaller', () => {
     expect(mocks.recordUsage).not.toHaveBeenCalled()
   })
 
+  const referencedImageCall = {
+    toolId: 'image_generate',
+    input: { provider: '{{IMAGE_PROVIDER}}', apiKey: '{{IMAGE_KEY}}' },
+  }
+
+  it.each<[string, Parameters<typeof run>[0], Record<string, string>]>([
+    ['the key is omitted', { input: { url: 'https://a.co' } }, {}],
+    [
+      'the key references an empty variable',
+      { input: { url: 'https://a.co', apiKey: '{{FIRECRAWL_KEY}}' } },
+      { FIRECRAWL_KEY: ' ' },
+    ],
+    [
+      'a reference selects the hosted provider',
+      referencedImageCall,
+      { IMAGE_PROVIDER: 'falai', IMAGE_KEY: '' },
+    ],
+  ])('refuses a hosted-key call over the usage limit when %s', async (_case, input, env) => {
+    mocks.checkUsageLimits.mockResolvedValue({ isExceeded: true, message: 'Usage limit exceeded' })
+    environmentUtilsMockFns.mockGetEffectiveDecryptedEnv.mockResolvedValue(env)
+
+    await expect(run(input)).rejects.toBeInstanceOf(ToolUsageLimitExceededError)
+  })
+
+  it.each<[string, Parameters<typeof run>[0], Record<string, string>]>([
+    ['the caller brings their own key', { input: { url: 'https://a.co', apiKey: 'sk-own' } }, {}],
+    [
+      'the caller references a variable holding their own key',
+      { input: { url: 'https://a.co', apiKey: '{{FIRECRAWL_KEY}}' } },
+      { FIRECRAWL_KEY: 'fc-own' },
+    ],
+    [
+      'the reference pads the variable name',
+      { input: { url: 'https://a.co', apiKey: '{{ FIRECRAWL_KEY }}' } },
+      { FIRECRAWL_KEY: 'fc-own' },
+    ],
+    [
+      'a reference selects a provider Sim does not host',
+      referencedImageCall,
+      { IMAGE_PROVIDER: 'openai', IMAGE_KEY: '' },
+    ],
+    [
+      'the tool has no hosted key',
+      { toolId: 'zendesk_get_ticket', input: { ticketId: '4', subdomain: 'a', apiToken: 't' } },
+      {},
+    ],
+  ])('does not gate on usage when %s', async (_case, input, env) => {
+    mocks.checkUsageLimits.mockResolvedValue({ isExceeded: true, message: 'Usage limit exceeded' })
+    environmentUtilsMockFns.mockGetEffectiveDecryptedEnv.mockResolvedValue(env)
+
+    await expect(run(input)).resolves.toMatchObject({ status: 'succeeded' })
+  })
+
   /**
    * The provider already ran and already charged Sim's key, so losing the
    * ledger row must not also lose the caller's result.

apps/sim/lib/tool-execution/application/execute-tool.ts

@@ -2,6 +2,7 @@ import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
 import { resolveBillingAttribution, toBillingContext } from '@/lib/billing/core/billing-attribution'
+import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import { recordUsage } from '@/lib/billing/core/usage-log'
 import {
   isBlockTypeAllowed,
@@ -16,8 +17,11 @@ import { defineAuthorizedWorkspaceUseCase } from '@/lib/core/application'
 import { ForbiddenOperationError } from '@/lib/core/application/forbidden'
 import { isHosted } from '@/lib/core/config/env-flags'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
+import { getEffectiveDecryptedEnv } from '@/lib/environment/utils'
 import { principalUserId } from '@/lib/integrations/principal-scope.server'
 import { toolExecutionOperations } from '@/lib/tool-execution/application/operations'
+import { isEnvVarReference } from '@/executor/constants'
+import { resolveEnvVarReferences } from '@/executor/utils/reference-validation'
 import { executeTool as executeRegistryTool } from '@/tools'
 import type { ExecutableToolConfig } from '@/tools/types'
 import { getTool } from '@/tools/utils'
@@ -34,6 +38,13 @@ export interface ExecuteToolInput {
   timeoutSeconds?: number
 }
 
+export class ToolUsageLimitExceededError extends Error {
+  constructor(message: string) {
+    super(message)
+    this.name = 'ToolUsageLimitExceededError'
+  }
+}
+
 export interface ExecuteToolResult {
   toolId: string
   status: 'succeeded' | 'failed'
@@ -49,10 +60,10 @@ export interface ExecuteToolResult {
  * deployment hosts keys, any `enabled` predicate accepts these params, and the
  * caller has not brought a key of their own — which wins where present.
  *
- * Pre-dispatch only, for the required-input exemption: a parameter Sim will
- * fill is not missing. It is deliberately NOT the metering gate — it cannot see
- * a BYOK key, which the registry injects while reporting the call as *not*
- * hosted, so after dispatch the registry's own verdict is read instead.
+ * Pre-dispatch only: the required-input exemption (a parameter Sim will fill is
+ * not missing) and usage admission. It is deliberately NOT the metering gate —
+ * it cannot see a BYOK key, which the registry injects while reporting the call
+ * as *not* hosted, so after dispatch the registry's own verdict is read instead.
  */
 function hostedKeyParamFor(
   tool: ExecutableToolConfig,
@@ -65,6 +76,42 @@ function hostedKeyParamFor(
   return tool.hosting.apiKeyParam
 }
 
+/**
+ * `params` as the registry holds them when it decides on Sim's key.
+ *
+ * The registry resolves each whole-value `{{VAR}}` in a `user-only` parameter
+ * before `injectHostedKeyIfNeeded` runs, so a reference can supply the key, or
+ * the value an `enabled` predicate reads, and an empty variable leaves the key
+ * for Sim's to fill. Resolved the same way here — same environment, same
+ * options. A missing variable stays as written: the registry refuses the call
+ * on it before any key is spent.
+ */
+async function resolveUserOnlyReferences(
+  tool: ExecutableToolConfig,
+  params: Record<string, unknown>,
+  userId: string,
+  workspaceId: string
+): Promise<Record<string, unknown>> {
+  const referenced = Object.entries(tool.params ?? {})
+    .filter(([name, declaration]) => {
+      const value = params[name]
+      return (
+        declaration?.visibility === 'user-only' &&
+        typeof value === 'string' &&
+        isEnvVarReference(value)
+      )
+    })
+    .map(([name]) => name)
+  if (referenced.length === 0) return params
+
+  const env = await getEffectiveDecryptedEnv(userId, workspaceId)
+  const resolved = { ...params }
+  for (const name of referenced) {
+    resolved[name] = resolveEnvVarReferences(params[name], env, { allowEmbedded: false })
+  }
+  return resolved
+}
+
 /**
  * The three spellings the executor accepts for "which credential".
  *
@@ -302,6 +349,25 @@ export const executeToolForCaller = defineAuthorizedWorkspaceUseCase({
       workspaceId: context.workspaceId,
     })
 
+    /**
+     * Admission before Sim's key is spent: metering runs only after the provider
+     * has charged it. A BYOK workspace is gated too, since only the registry can
+     * see that key — the same standing every workflow run is held to.
+     */
+    if (
+      hostedKeyParamFor(
+        tool,
+        await resolveUserOnlyReferences(tool, callerParams, userId, context.workspaceId)
+      )
+    ) {
+      const usage = await checkExecutionUsageLimits(billingAttribution)
+      if (usage.isExceeded) {
+        throw new ToolUsageLimitExceededError(
+          usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.'
+        )
+      }
+    }
+
     const params: Record<string, unknown> = {
       ...callerParams,
       _context: {