diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts index e04b0aac04b..c8b4943b687 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.test.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.test.ts @@ -297,6 +297,20 @@ describe('POST /api/copilot/api-keys/validate billing protocols', () => { expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() }) + it('fails account admission closed when its subscription cannot be read', async () => { + mockGetHighestPrioritySubscription.mockImplementationOnce( + async (_userId: string, options?: { onError?: 'throw' | 'return-null' }) => { + if (options?.onError === 'throw') throw new Error('database unavailable') + return null + } + ) + + const res = await POST(request(SELF_HOSTED_WORKSPACELESS_VALIDATE_BODY), createRouteContext({})) + + expect(res.status).toBe(500) + expect(mockCheckServerSideUsageLimits).not.toHaveBeenCalled() + }) + it('rejects markerless admission on hosted Sim', async () => { setEnvFlags({ isHosted: true }) const res = await POST(request(SELF_HOSTED_VALIDATE_BODY)) diff --git a/apps/sim/app/api/copilot/api-keys/validate/route.ts b/apps/sim/app/api/copilot/api-keys/validate/route.ts index 1a159072395..77dccc60f97 100644 --- a/apps/sim/app/api/copilot/api-keys/validate/route.ts +++ b/apps/sim/app/api/copilot/api-keys/validate/route.ts @@ -296,7 +296,9 @@ async function checkAdmissionUsage(admission: AdmissionBillingDecision): Promise } } - const subscription = await getHighestPrioritySubscription(admission.userId) + const subscription = await getHighestPrioritySubscription(admission.userId, { + onError: 'throw', + }) const usage = await checkServerSideUsageLimits(admission.userId, subscription) return { isExceeded: usage.isExceeded, diff --git a/apps/sim/lib/billing/core/chat-access.integration.ts b/apps/sim/lib/billing/core/chat-access.integration.ts new file mode 100644 index 00000000000..804597d16d2 --- /dev/null +++ b/apps/sim/lib/billing/core/chat-access.integration.ts @@ -0,0 +1,95 @@ +/** + * The Chat plan verdict against real PostgreSQL credit balances: a hosted free payer is refused + * Chat until it holds prepaid credits, on both the organization and personal balance paths, and a + * paid payer is admitted without a balance. + */ +import { db } from '@sim/db' +import { organization, user, userStats } from '@sim/db/schema' +import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock' +import { generateId } from '@sim/utils/id' +import { eq, inArray } from 'drizzle-orm' +import { afterAll, describe, expect, it, vi } from 'vitest' +import { payerHasChatAccess } from '@/lib/billing/core/chat-access' + +vi.mock('@/lib/core/config/env-flags', () => ({ + ...envFlagsMock, + isHosted: true, + isBillingEnabled: true, +})) + +const userIds: string[] = [] +const organizationIds: string[] = [] + +async function createUser(): Promise { + const id = `chat-access-user-${generateId()}` + const now = new Date() + await db.insert(user).values({ + id, + name: 'Chat Access Test', + email: `${id}@chat-access.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(userStats).values({ id: generateId(), userId: id }) + userIds.push(id) + return id +} + +async function createOrganization(): Promise { + const id = `chat-access-org-${generateId()}` + await db.insert(organization).values({ id, name: 'Chat Access Org', slug: id }) + organizationIds.push(id) + return id +} + +afterAll(async () => { + if (organizationIds.length > 0) { + await db.delete(organization).where(inArray(organization.id, organizationIds)) + } + if (userIds.length > 0) await db.delete(user).where(inArray(user.id, userIds)) +}) + +describe('payerHasChatAccess', () => { + it('refuses a free organization payer until it holds prepaid credits', async () => { + const organizationId = await createOrganization() + const payer = { + billingEntity: { type: 'organization' as const, id: organizationId }, + payerSubscription: null, + } + + expect(await payerHasChatAccess(payer)).toBe(false) + + await db + .update(organization) + .set({ creditBalance: '25' }) + .where(eq(organization.id, organizationId)) + + expect(await payerHasChatAccess(payer)).toBe(true) + }) + + it('refuses a free personal payer until it holds prepaid credits', async () => { + const userId = await createUser() + const payer = { + billingEntity: { type: 'user' as const, id: userId }, + payerSubscription: { plan: 'free' }, + } + + expect(await payerHasChatAccess(payer)).toBe(false) + + await db.update(userStats).set({ creditBalance: '10' }).where(eq(userStats.userId, userId)) + + expect(await payerHasChatAccess(payer)).toBe(true) + }) + + it('admits a paid payer that holds no credits', async () => { + const userId = await createUser() + + expect( + await payerHasChatAccess({ + billingEntity: { type: 'user', id: userId }, + payerSubscription: { plan: 'pro_6000' }, + }) + ).toBe(true) + }) +}) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 46bab93eb0b..62d2ad4fe1c 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -24,7 +24,10 @@ import { import { generateId } from '@sim/utils/id' import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { scopeProviderToolCallId } from '@/lib/mothership/request/go/tool-call-identity' -import { handleBillingLimitResponse } from '@/lib/mothership/request/tools/billing' +import { + handleBillingLimitResponse, + handleChatPlanRequiredResponse, +} from '@/lib/mothership/request/tools/billing' import type { ExecutionContext, StreamingContext } from '@/lib/mothership/request/types' import { executeFunctionExecute } from '@/lib/mothership/tools/handlers/function-execute' import { executeRunCode } from '@/lib/mothership/tools/handlers/run-code' @@ -204,6 +207,11 @@ vi.mock('@/lib/mothership/request/tools/billing', () => ({ handleChatPlanRequiredResponse: vi.fn(), })) +const mockPayerHasChatAccess = vi.hoisted(() => vi.fn()) +vi.mock('@/lib/billing/core/chat-access', () => ({ + payerHasChatAccess: mockPayerHasChatAccess, +})) + const mockRequestExplicitStreamAbort = vi.hoisted(() => vi.fn()) vi.mock('@/lib/mothership/request/session/explicit-abort', () => ({ requestExplicitStreamAbort: mockRequestExplicitStreamAbort, @@ -276,6 +284,7 @@ describe('runCopilotLifecycle', () => { }) mockCreateRunSegment.mockResolvedValue({ status: 'active' }) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) + mockPayerHasChatAccess.mockResolvedValue(true) continuationAuth.mockResolvedValue(undefined) mockEnv.COPILOT_API_KEY = undefined mockEnv.MSHIP_SYSPROMPT_OVERRIDE = undefined @@ -2311,17 +2320,6 @@ describe('runCopilotLifecycle', () => { expect(capturedExecContext?.userPermission).toBe('read') }) - /** A paid organization payer, which the free-plan Chat gate admits. */ - const paidOrganizationPayer = { - id: 'sub-org-1', - referenceId: 'org-1', - plan: 'team', - status: 'active', - seats: 2, - periodStart: null, - periodEnd: null, - } - it.each([ { workspaceId: 'ws-1', organizationId: undefined }, { workspaceId: undefined, organizationId: 'org-1' }, @@ -2336,7 +2334,7 @@ describe('runCopilotLifecycle', () => { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z', }, - payerSubscription: paidOrganizationPayer, + payerSubscription: null, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockEnv.COPILOT_API_KEY = 'sim-agent-key' @@ -2476,7 +2474,7 @@ describe('runCopilotLifecycle', () => { start: '2026-07-01T00:00:00.000Z', end: '2026-08-01T00:00:00.000Z', }, - payerSubscription: paidOrganizationPayer, + payerSubscription: null, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) mockCheckAttributedUsageLimits.mockResolvedValue({ @@ -2505,6 +2503,43 @@ describe('runCopilotLifecycle', () => { expect(result.cancelled).not.toBe(true) }) + it('refuses dispatch at admission when the payer plan excludes Chat', async () => { + const billingAttribution = { + actorUserId: 'user-1', + workspaceId: 'ws-1', + organizationId: 'org-1', + billedAccountUserId: 'user-1', + billingEntity: { type: 'organization' as const, id: 'org-1' }, + billingPeriod: { + start: '2026-07-01T00:00:00.000Z', + end: '2026-08-01T00:00:00.000Z', + }, + payerSubscription: null, + } + setEnvFlags({ isHosted: true, isBillingEnabled: true }) + mockPayerHasChatAccess.mockResolvedValue(false) + + const result = await runCopilotLifecycle( + { message: 'hello', messageId: 'message-1' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'execution-1', + runId: 'run-1', + simRequestId: 'request-1', + billingAttribution, + } + ) + + expect(mockPayerHasChatAccess).toHaveBeenCalledWith(billingAttribution) + expect(handleChatPlanRequiredResponse).toHaveBeenCalledOnce() + expect(mockCheckAttributedUsageLimits).not.toHaveBeenCalled() + expect(handleBillingLimitResponse).not.toHaveBeenCalled() + expect(mockRunStreamLoop).not.toHaveBeenCalled() + expect(result.cancelled).not.toBe(true) + }) + it('stops the worker run when the worker itself refuses a leg at the usage limit', async () => { mockRequestExplicitStreamAbort.mockResolvedValue({ settled: true }) mockRunStreamLoop.mockRejectedValueOnce(new BillingLimitError('user-1')) @@ -2529,7 +2564,7 @@ describe('runCopilotLifecycle', () => { billedAccountUserId: 'user-1', billingEntity: { type: 'organization' as const, id: 'org-1' }, billingPeriod: { start: '2026-07-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, - payerSubscription: paidOrganizationPayer, + payerSubscription: null, } setEnvFlags({ isHosted: true, isBillingEnabled: true }) resetUsageGateCache()