diff --git a/apps/sim/lib/billing/core/chat-access.ts b/apps/sim/lib/billing/core/chat-access.ts index 72645f0afd8..7e86759e023 100644 --- a/apps/sim/lib/billing/core/chat-access.ts +++ b/apps/sim/lib/billing/core/chat-access.ts @@ -5,23 +5,30 @@ import { isFree } from '@/lib/billing/plan-helpers' import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade' import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' -const logger = createLogger('ChatAccess') +const logger = createLogger('PaidAccess') + +interface PaidAccessOptions { + /** + * The verdict for a free payer whose credit balance cannot be read. Chat admits, because a + * database blip must not lock a paying customer out and every turn is bounded by the usage + * check that follows. A hosted key denies, because nothing after acquisition bounds it. + */ + onBalanceUnavailable: 'admit' | 'deny' +} /** - * Whether a payer may use Chat (talking to Sim). Every turn spends hosted model credits, so a - * hosted deployment that bills withholds Chat from a payer on the free plan — unless that payer - * holds prepaid credits, which an admin grants to customers billed outside a subscription. + * Whether a payer may use a feature a hosted deployment reserves for paying payers: Chat, and + * hosted keys a tool marks `requiresPaidPlan`. A hosted deployment that bills withholds these from + * a payer on the free plan, unless that payer holds prepaid credits, which an admin grants to + * customers billed outside a subscription. * * The plan is the payer's entitled subscription (active or past due), so a canceled * plan reads as free; a blocked payer is refused by the usage check every caller runs next. - * - * This is the single verdict behind every Chat plan gate: run admission, chat titles, scheduled - * job dispatch, Copilot key admission, and the workspace host context the composer reads, so the - * composer never offers a send that admission will refuse. A balance that cannot be read admits - * the payer: a database blip must not lock a paying customer out, and a free payer stays bounded - * by its usage limit. */ -export async function payerHasChatAccess(payer: UsageUpgradePayer): Promise { +export async function payerHasPaidAccess( + payer: UsageUpgradePayer, + options: PaidAccessOptions +): Promise { if (!isHosted || !isBillingEnabled) return true if (!isFree(payer.payerSubscription?.plan)) return true try { @@ -31,10 +38,23 @@ export async function payerHasChatAccess(payer: UsageUpgradePayer): Promise 0 } catch (error) { - logger.warn('Credit balance unavailable; admitting the payer to Chat', { + logger.warn('Credit balance unavailable for a free payer', { billingEntity: payer.billingEntity, + verdict: options.onBalanceUnavailable, error: getErrorMessage(error), }) - return true + return options.onBalanceUnavailable === 'admit' } } + +/** + * Whether a payer may use Chat (talking to Sim): {@link payerHasPaidAccess}, since every turn + * spends hosted model credits, admitting a payer whose balance cannot be read. + * + * This is the single verdict behind every Chat plan gate: run admission, chat titles, scheduled + * job dispatch, Copilot key admission, and the workspace host context the composer reads, so the + * composer never offers a send that admission will refuse. + */ +export function payerHasChatAccess(payer: UsageUpgradePayer): Promise { + return payerHasPaidAccess(payer, { onBalanceUnavailable: 'admit' }) +} diff --git a/apps/sim/tools/errors.ts b/apps/sim/tools/errors.ts index 8a8809258a7..316d103f17e 100644 --- a/apps/sim/tools/errors.ts +++ b/apps/sim/tools/errors.ts @@ -17,6 +17,19 @@ export class HostedKeyRateLimitedError extends HttpError { } } +/** + * The tool's hosted key is reserved for paying payers and this run's payer is on the free plan. + * Carries a 402 so the run fails as the caller's plan, not as a Sim fault. + */ +export class HostedKeyRequiresPaidPlanError extends HttpError { + readonly statusCode = 402 + + constructor(message: string) { + super(message) + this.name = 'HostedKeyRequiresPaidPlanError' + } +} + /** No hosted keys are configured or available for this provider. */ export class HostedKeyUnavailableError extends HttpError { readonly statusCode = 503 diff --git a/apps/sim/tools/index.test.ts b/apps/sim/tools/index.test.ts index b7bef8e6e80..fd6f413a055 100644 --- a/apps/sim/tools/index.test.ts +++ b/apps/sim/tools/index.test.ts @@ -3,6 +3,10 @@ import { createSessionPrincipal } from '@sim/testing/factories/principal.factory import { jsonResponse } from '@sim/testing/helpers/http' import { apiKeyByokMock, apiKeyByokMockFns } from '@sim/testing/mocks/api-key-byok.mock' import { authInternalMock, authInternalMockFns } from '@sim/testing/mocks/auth-internal.mock' +import { + billingAttributionMock, + billingAttributionMockFns, +} from '@sim/testing/mocks/billing-attribution.mock' import { billingUsageLogMock } from '@sim/testing/mocks/billing-usage-log.mock' import { executorPrincipalMock, @@ -111,7 +115,9 @@ const { mockExecuteChatFunction, mockGetInternalToolOperationHandler, mockExecuteInternalToolOperation, + mockPayerHasPaidAccess, } = vi.hoisted(() => ({ + mockPayerHasPaidAccess: vi.fn(), mockGetToolAsync: vi.fn(), mockRateLimiterFns: { acquireKey: vi.fn(), @@ -173,6 +179,12 @@ vi.mock('@/lib/core/rate-limiter/hosted-key', () => ({ getHostedKeyRateLimiter: () => mockRateLimiterFns, })) +vi.mock('@/lib/billing/core/chat-access', () => ({ + payerHasPaidAccess: mockPayerHasPaidAccess, +})) + +vi.mock('@/lib/billing/core/attribution', () => billingAttributionMock) + vi.mock('@/lib/uploads/contexts/workspace/file-manager', () => workspaceFileManagerMock) vi.mock( @@ -442,6 +454,7 @@ import { createToolConfig, getTool } from '@/tools/utils' import { getToolAsync } from '@/tools/utils.server' const mockGetBYOKKey = apiKeyByokMockFns.mockGetBYOKKey +const mockResolveWorkspaceBillingPayer = billingAttributionMockFns.mockResolveWorkspaceBillingPayer const mockGenerateInternalToken = authInternalMockFns.mockGenerateInternalToken const mockCreateExecutorPrincipalFromExecutionContext = executorPrincipalMockFns.mockCreateExecutorPrincipalFromExecutionContext @@ -6551,6 +6564,153 @@ describe('Hosted Key Injection', () => { } }) + describe('requiresPaidPlan', () => { + const PAID_ONLY_TOOL_ID = 'test_paid_only_hosting' + + beforeEach(() => { + setEnvFlags({ isHosted: true, isBillingEnabled: true }) + mockRateLimiterFns.acquireKey.mockResolvedValue({ + success: true, + key: 'platform-fixture-key', + keyIndex: 0, + envVarName: 'TEST_HOSTED_KEY', + }) + tools[PAID_ONLY_TOOL_ID] = { + id: PAID_ONLY_TOOL_ID, + name: 'Paid-only fixture', + description: 'Hosted key reserved for paying payers', + version: '1.0.0', + params: {}, + request: { url: 'https://example.com/fixture', method: 'GET', headers: () => ({}) }, + hosting: { + envKeyPrefix: 'TEST_HOSTED_KEY', + apiKeyParam: 'apiKey', + byokProviderId: 'falai', + requiresPaidPlan: true, + pricing: { type: 'per_request', cost: 0.5 }, + rateLimit: { mode: 'per_request', requestsPerMinute: 100 }, + }, + transformResponse: async () => ({ success: true, output: {} }), + } + }) + + afterEach(() => { + Reflect.deleteProperty(tools, PAID_ONLY_TOOL_ID) + resetEnvFlagsMock() + }) + + async function runPaidOnlyTool() { + mockValidateUrlWithDNS.mockResolvedValue({ isValid: true, resolvedIP: '93.184.216.34' }) + mockSecureFetchWithPinnedIP.mockResolvedValue({ + ok: true, + status: 200, + statusText: 'OK', + headers: { get: () => null, toRecord: () => ({}) }, + json: async () => ({ ok: true }), + text: async () => '{"ok":true}', + }) + const request = mockSecureFetchWithPinnedIP + const result = await executeTool( + PAID_ONLY_TOOL_ID, + {}, + { executionContext: createToolExecutionContext({ userId: 'user-1' }) } + ) + return { result, request } + } + + it('refuses the hosted key to a payer without paid access before acquiring one', async () => { + mockGetBYOKKey.mockResolvedValue(null) + mockPayerHasPaidAccess.mockResolvedValue(false) + + const { result, request } = await runPaidOnlyTool() + + expect(result.success).toBe(false) + expect(result.error).toContain('available on paid plans') + expect(mockPayerHasPaidAccess).toHaveBeenCalledWith(TEST_BILLING_ATTRIBUTION, { + onBalanceUnavailable: 'deny', + }) + expect(mockRateLimiterFns.acquireKey).not.toHaveBeenCalled() + expect(request).not.toHaveBeenCalled() + }) + + it('grants the hosted key to a payer with paid access', async () => { + mockGetBYOKKey.mockResolvedValue(null) + mockPayerHasPaidAccess.mockResolvedValue(true) + + const { result, request } = await runPaidOnlyTool() + + expect(result.success).toBe(true) + expect(mockRateLimiterFns.acquireKey).toHaveBeenCalledTimes(1) + expect(request).toHaveBeenCalledTimes(1) + }) + + async function runPaidOnlyToolWithoutAttribution() { + mockGetBYOKKey.mockResolvedValue(null) + mockValidateUrlWithDNS.mockResolvedValue({ isValid: true, resolvedIP: '93.184.216.34' }) + mockSecureFetchWithPinnedIP.mockResolvedValue({ + ok: true, + status: 200, + statusText: 'OK', + headers: { get: () => null, toRecord: () => ({}) }, + json: async () => ({ ok: true }), + text: async () => '{"ok":true}', + }) + const result = await executeTool(PAID_ONLY_TOOL_ID, { + _context: { workspaceId: 'workspace-without-run', userId: 'user-1' }, + }) + return { result, request: mockSecureFetchWithPinnedIP } + } + + it('judges the workspace payer when the run carries no attribution', async () => { + mockResolveWorkspaceBillingPayer.mockResolvedValue({ + billedAccountUserId: 'payer-1', + organizationId: 'org-1', + payerSubscription: { plan: 'team' }, + }) + mockPayerHasPaidAccess.mockResolvedValue(true) + + const { result, request } = await runPaidOnlyToolWithoutAttribution() + + expect(result.success).toBe(true) + expect(mockResolveWorkspaceBillingPayer).toHaveBeenCalledWith('workspace-without-run', { + onMissing: 'return-null', + }) + expect(mockPayerHasPaidAccess).toHaveBeenCalledWith( + { + billingEntity: { type: 'organization', id: 'org-1' }, + payerSubscription: { plan: 'team' }, + }, + { onBalanceUnavailable: 'deny' } + ) + expect(request).toHaveBeenCalledTimes(1) + }) + + it('refuses the hosted key when the workspace has no payer', async () => { + mockResolveWorkspaceBillingPayer.mockResolvedValue(null) + mockPayerHasPaidAccess.mockResolvedValue(true) + + const { result, request } = await runPaidOnlyToolWithoutAttribution() + + expect(result.success).toBe(false) + expect(result.error).toContain('available on paid plans') + expect(mockPayerHasPaidAccess).not.toHaveBeenCalled() + expect(mockRateLimiterFns.acquireKey).not.toHaveBeenCalled() + expect(request).not.toHaveBeenCalled() + }) + + it("runs a free payer's own workspace key without consulting the plan", async () => { + mockGetBYOKKey.mockResolvedValue({ apiKey: 'workspace-own-key' }) + mockPayerHasPaidAccess.mockResolvedValue(false) + + const { result, request } = await runPaidOnlyTool() + + expect(result.success).toBe(true) + expect(mockPayerHasPaidAccess).not.toHaveBeenCalled() + expect(mockRateLimiterFns.acquireKey).not.toHaveBeenCalled() + expect(request).toHaveBeenCalledTimes(1) + }) + }) + it('should not inject hosted key when tool has no hosting config', async () => { const mockTool = { id: 'test_no_hosting', diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 307dbc87f9f..05dff63008e 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -13,9 +13,12 @@ import type { InternalSandboxProfile } from '@/lib/auth/internal' import { BILLING_ATTRIBUTION_HEADER, type BillingAttributionSnapshot, + resolveWorkspaceBillingPayer, serializeBillingAttributionHeader, } from '@/lib/billing/core/attribution' -import { isHosted } from '@/lib/core/config/env-flags' +import { payerHasPaidAccess } from '@/lib/billing/core/chat-access' +import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade' +import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags' import { findDatabaseQueryError, getUserFacingErrorMessage, @@ -111,7 +114,11 @@ import { } from '@/executor/utils/resolved-secret-trace-registry' import type { ErrorInfo } from '@/tools/error-extractors' import { extractErrorMessage, redactErrorData } from '@/tools/error-extractors' -import { HostedKeyRateLimitedError, HostedKeyUnavailableError } from '@/tools/errors' +import { + HostedKeyRateLimitedError, + HostedKeyRequiresPaidPlanError, + HostedKeyUnavailableError, +} from '@/tools/errors' import { getOwnEnumerableDataEntries, prepareToolRequest, @@ -547,6 +554,40 @@ async function registerAssistantCredentialSecrets( } } +/** + * Refuses a `requiresPaidPlan` hosted key to a run whose payer lacks paid access. The payer is + * the run's admitted attribution, or the workspace's payer for a surface that carries none; a + * workspace without a payer, or a free payer whose credit balance cannot be read, is refused. + */ +async function assertPayerMayUseHostedKey( + tool: ToolDefinition, + workspaceId: string | undefined, + billingAttribution: BillingAttributionSnapshot | undefined +): Promise { + if (!isBillingEnabled) return + + let payer: UsageUpgradePayer | null = billingAttribution ?? null + if (!payer && workspaceId) { + const workspacePayer = await resolveWorkspaceBillingPayer(workspaceId, { + onMissing: 'return-null', + }) + if (workspacePayer) { + payer = { + billingEntity: workspacePayer.organizationId + ? { type: 'organization', id: workspacePayer.organizationId } + : { type: 'user', id: workspacePayer.billedAccountUserId }, + payerSubscription: workspacePayer.payerSubscription, + } + } + } + + if (payer && (await payerHasPaidAccess(payer, { onBalanceUnavailable: 'deny' }))) return + + throw new HostedKeyRequiresPaidPlanError( + `${tool.name} with Sim's hosted key is available on paid plans. Upgrade your plan, or add your own API key under Settings > BYOK.` + ) +} + /** Result from hosted key injection */ interface HostedKeyInjectionResult { isUsingHostedKey: boolean @@ -577,7 +618,10 @@ async function injectHostedKeyIfNeeded( return { isUsingHostedKey: false } } - const { workspaceId, userId, workflowId } = resolveToolScope(params, executionContext) + const { workspaceId, userId, workflowId, billingAttribution } = resolveToolScope( + params, + executionContext + ) if (byokProviderId && workspaceId) { const byokResult = await getBYOKKey(workspaceId, byokProviderId as BYOKProviderId, { signal }) @@ -588,6 +632,10 @@ async function injectHostedKeyIfNeeded( } } + if (tool.hosting.requiresPaidPlan) { + await assertPayerMayUseHostedKey(tool, workspaceId, billingAttribution) + } + const rateLimiter = getHostedKeyRateLimiter() const provider = byokProviderId || tool.id const billingActorId = workspaceId diff --git a/apps/sim/tools/types.ts b/apps/sim/tools/types.ts index 7f3eba3beb8..b3cd3a530df 100644 --- a/apps/sim/tools/types.ts +++ b/apps/sim/tools/types.ts @@ -499,6 +499,11 @@ export interface ToolHostingConfig

> { apiKeyParam: string /** BYOK provider ID for workspace key lookup */ byokProviderId?: BYOKProviderId + /** + * Withholds the hosted key from a free-plan payer without prepaid credits. A workspace BYOK key + * or a key the user passes still runs; only Sim's own key is refused. + */ + requiresPaidPlan?: boolean /** Pricing when using hosted key */ pricing: ToolHostingPricing

/** Hosted key rate limit configuration (required for hosted key distribution) */ diff --git a/apps/sim/tools/video/falai.ts b/apps/sim/tools/video/falai.ts index 224a97c4d2c..1c8907eb184 100644 --- a/apps/sim/tools/video/falai.ts +++ b/apps/sim/tools/video/falai.ts @@ -71,6 +71,7 @@ export const falaiVideoTool: InternalToolConfig = { envKeyPrefix: 'FALAI_API_KEY', apiKeyParam: 'apiKey', byokProviderId: 'falai', + requiresPaidPlan: true, pricing: { type: 'custom', getCost: (_params, output) => {