/** * @vitest-environment node */ import { beforeEach, describe, expect, it, vi } from 'vitest' const { mockGetHighestPrioritySubscription, mockInsert, mockIsOrgScopedSubscription, mockOnConflictDoNothing, mockReturning, mockValues, mockTransaction, mockUpdate, } = vi.hoisted(() => ({ mockGetHighestPrioritySubscription: vi.fn(), mockInsert: vi.fn(), mockIsOrgScopedSubscription: vi.fn(), mockOnConflictDoNothing: vi.fn(), mockReturning: vi.fn(), mockValues: vi.fn(), mockTransaction: vi.fn(), mockUpdate: vi.fn(), })) vi.mock('@sim/db', () => { const instance = { insert: mockInsert, transaction: mockTransaction } return { db: instance, dbReplica: instance } }) vi.mock('@sim/db/schema', () => ({ usageLog: { billingEntityId: 'usageLog.billingEntityId', billingEntityType: 'usageLog.billingEntityType', billingPeriodEnd: 'usageLog.billingPeriodEnd', billingPeriodStart: 'usageLog.billingPeriodStart', category: 'usageLog.category', cost: 'usageLog.cost', createdAt: 'usageLog.createdAt', description: 'usageLog.description', eventKey: 'usageLog.eventKey', executionId: 'usageLog.executionId', id: 'usageLog.id', metadata: 'usageLog.metadata', source: 'usageLog.source', userId: 'usageLog.userId', workflowId: 'usageLog.workflowId', workspaceId: 'usageLog.workspaceId', }, })) vi.mock('@/lib/billing/core/plan', () => ({ getHighestPrioritySubscription: mockGetHighestPrioritySubscription, })) vi.mock('@/lib/billing/subscriptions/utils', () => ({ isOrgScopedSubscription: mockIsOrgScopedSubscription, })) vi.mock('@/lib/core/config/env-flags', () => ({ isBillingEnabled: true, })) import { CUMULATIVE_COST_EPSILON, recordCumulativeUsage, recordUsage, resolveCumulativeTopUp, } from '@/lib/billing/core/usage-log' describe('recordUsage', () => { beforeEach(() => { vi.clearAllMocks() mockReturning.mockResolvedValue([{ cost: '0.10' }, { cost: '0.20' }]) mockOnConflictDoNothing.mockReturnValue({ returning: mockReturning }) mockValues.mockReturnValue({ onConflictDoNothing: mockOnConflictDoNothing, }) mockInsert.mockReturnValue({ values: mockValues }) mockGetHighestPrioritySubscription.mockResolvedValue({ periodEnd: new Date('2026-06-01T00:00:00.000Z'), periodStart: new Date('2026-05-01T00:00:00.000Z'), referenceId: 'org-1', }) mockIsOrgScopedSubscription.mockReturnValue(true) }) it('commits canonical usage rows with deterministic event keys and billing scope', async () => { await recordUsage({ userId: 'user-1', workspaceId: 'workspace-1', workflowId: 'workflow-1', executionId: 'execution-1', entries: [ { category: 'fixed', source: 'workflow', description: 'execution_fee', cost: 0.1 }, { category: 'model', source: 'workflow', description: 'gpt-4', cost: 0.2, metadata: { inputTokens: 10, outputTokens: 20 }, }, ], }) const values = mockValues.mock.calls[0][0] expect(values).toHaveLength(2) expect(values[0]).toMatchObject({ billingEntityId: 'org-1', billingEntityType: 'organization', billingPeriodEnd: new Date('2026-06-01T00:00:00.000Z'), billingPeriodStart: new Date('2026-05-01T00:00:00.000Z'), }) expect(values[0].eventKey).toMatch(/^[a-f0-9]{64}$/) expect(values[1].eventKey).toMatch(/^[a-f0-9]{64}$/) expect(values[0].eventKey).not.toBe(values[1].eventKey) expect(mockOnConflictDoNothing).toHaveBeenCalledWith( expect.objectContaining({ target: 'usageLog.eventKey' }) ) }) it('uses pre-resolved billing context without loading subscriptions', async () => { await recordUsage({ userId: 'user-1', billingEntity: { type: 'user', id: 'user-1' }, billingPeriod: { start: new Date('2026-05-01T00:00:00.000Z'), end: new Date('2026-06-01T00:00:00.000Z'), }, entries: [{ category: 'fixed', source: 'workflow', description: 'execution_fee', cost: 0.1 }], }) expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled() expect(mockValues.mock.calls[0][0][0]).toMatchObject({ billingEntityId: 'user-1', billingEntityType: 'user', }) }) }) describe('resolveCumulativeTopUp', () => { it('bills the full amount on the first flush (nothing recorded yet)', () => { expect(resolveCumulativeTopUp(0, 0.3474447)).toEqual({ shouldBill: true, delta: 0.3474447, newTotal: 0.3474447, }) }) it('bills only the delta when the cumulative grows (recovered request)', () => { const result = resolveCumulativeTopUp(0.3474447, 0.4662453) expect(result.shouldBill).toBe(true) expect(result.newTotal).toBe(0.4662453) expect(result.delta).toBeCloseTo(0.1188006, 9) }) it('is a no-op when the cumulative is unchanged (abort-race duplicate)', () => { expect(resolveCumulativeTopUp(0.4662453, 0.4662453)).toEqual({ shouldBill: false, delta: 0, newTotal: 0.4662453, }) }) it('is a no-op when an out-of-order flush carries a lower cumulative', () => { expect(resolveCumulativeTopUp(0.4662453, 0.3)).toMatchObject({ shouldBill: false, delta: 0 }) }) it('ignores sub-epsilon increases from decimal round-trips', () => { expect( resolveCumulativeTopUp(0.4662453, 0.4662453 + CUMULATIVE_COST_EPSILON / 2) ).toMatchObject({ shouldBill: false }) }) }) describe('recordCumulativeUsage', () => { beforeEach(() => { vi.clearAllMocks() mockReturning.mockResolvedValue([{ cost: '0.3474447' }]) mockOnConflictDoNothing.mockReturnValue({ returning: mockReturning }) mockValues.mockReturnValue({ onConflictDoNothing: mockOnConflictDoNothing }) mockInsert.mockReturnValue({ values: mockValues }) mockGetHighestPrioritySubscription.mockResolvedValue({ periodEnd: new Date('2026-06-01T00:00:00.000Z'), periodStart: new Date('2026-05-01T00:00:00.000Z'), referenceId: 'org-1', }) mockIsOrgScopedSubscription.mockReturnValue(true) }) const setupTx = (existingRow: { id: string; cost: string } | null) => { const limit = vi.fn().mockResolvedValue(existingRow ? [existingRow] : []) const where = vi.fn().mockReturnValue({ limit }) const from = vi.fn().mockReturnValue({ where }) const select = vi.fn().mockReturnValue({ from }) const updateWhere = vi.fn().mockResolvedValue(undefined) const updateSet = vi.fn().mockReturnValue({ where: updateWhere }) mockUpdate.mockReturnValue({ set: updateSet }) const tx = { execute: vi.fn().mockResolvedValue(undefined), select, update: mockUpdate, insert: mockInsert, // recordUsage(tx) reuses the shared insert chain } mockTransaction.mockImplementation(async (fn: (t: typeof tx) => unknown) => fn(tx)) return { tx, select, updateSet } } /** True when any tx.execute call ran a sql`` template containing the substring. */ const executedSqlContaining = (tx: { execute: ReturnType }, substring: string) => tx.execute.mock.calls.some(([arg]) => { const strings = (arg as { strings?: readonly string[] } | null)?.strings return Array.isArray(strings) && strings.some((s) => s.includes(substring)) }) it('inserts the full cumulative on the first flush', async () => { setupTx(null) const result = await recordCumulativeUsage({ userId: 'user-1', workspaceId: 'ws-1', source: 'workspace-chat', model: 'claude-opus-4.8', cost: 0.3474447, eventKey: 'update-cost:msg-1-billing', metadata: { inputTokens: 100, outputTokens: 5 }, }) expect(result).toEqual({ billed: true, delta: 0.3474447, total: 0.3474447 }) expect(mockInsert).toHaveBeenCalledTimes(1) expect(mockUpdate).not.toHaveBeenCalled() }) it('tops up to the higher cumulative and bills only the delta', async () => { const { updateSet } = setupTx({ id: 'row-1', cost: '0.3474447' }) const result = await recordCumulativeUsage({ userId: 'user-1', source: 'workspace-chat', model: 'claude-opus-4.8', cost: 0.4662453, eventKey: 'update-cost:msg-1-billing', }) expect(result.billed).toBe(true) expect(result.total).toBe(0.4662453) expect(result.delta).toBeCloseTo(0.1188006, 9) expect(updateSet).toHaveBeenCalledWith(expect.objectContaining({ cost: '0.4662453' })) expect(mockInsert).not.toHaveBeenCalled() }) it('does not bill when the cumulative is not higher than recorded', async () => { const { updateSet } = setupTx({ id: 'row-1', cost: '0.4662453' }) const result = await recordCumulativeUsage({ userId: 'user-1', source: 'workspace-chat', model: 'claude-opus-4.8', cost: 0.4662453, eventKey: 'update-cost:msg-1-billing', }) expect(result).toEqual({ billed: false, delta: 0, total: 0.4662453 }) expect(updateSet).not.toHaveBeenCalled() expect(mockInsert).not.toHaveBeenCalled() }) it('resolves the billing context before opening the locked transaction, exactly once', async () => { setupTx(null) await recordCumulativeUsage({ userId: 'user-1', source: 'workspace-chat', model: 'claude-opus-4.8', cost: 0.3474447, eventKey: 'update-cost:msg-1-billing', }) // One lookup total: pre-resolved outside the tx, and the first-flush // insert reuses it instead of re-resolving on the pool inside the tx. expect(mockGetHighestPrioritySubscription).toHaveBeenCalledTimes(1) expect(mockGetHighestPrioritySubscription.mock.invocationCallOrder[0]).toBeLessThan( mockTransaction.mock.invocationCallOrder[0] ) }) it('stamps the pre-resolved billing context onto the first-flush insert', async () => { setupTx(null) await recordCumulativeUsage({ userId: 'user-1', source: 'workspace-chat', model: 'claude-opus-4.8', cost: 0.3474447, eventKey: 'update-cost:msg-1-billing', }) expect(mockValues.mock.calls[0][0][0]).toMatchObject({ billingEntityId: 'org-1', billingEntityType: 'organization', }) }) it('bounds the advisory-lock wait and locks on the 64-bit event-key hash', async () => { const { tx } = setupTx({ id: 'row-1', cost: '0.3474447' }) await recordCumulativeUsage({ userId: 'user-1', source: 'workspace-chat', model: 'claude-opus-4.8', cost: 0.4662453, eventKey: 'update-cost:msg-1-billing', }) expect(executedSqlContaining(tx, 'lock_timeout')).toBe(true) expect(executedSqlContaining(tx, 'pg_advisory_xact_lock')).toBe(true) expect(executedSqlContaining(tx, 'hashtextextended')).toBe(true) }) })