Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
224 changes: 224 additions & 0 deletions apps/sim/app/api/billing/update-cost/replay.integration.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,224 @@
/**
* @vitest-environment node
*/
import { type ExecFileException, execFile } from 'node:child_process'
import { createServer } from 'node:http'
import { promisify } from 'node:util'
import { resetEnvFlagsMock, resetEnvMock, setEnv, setEnvFlags } from '@sim/testing'
import { NextRequest } from 'next/server'
import type { Sql } from 'postgres'
import { afterAll, describe, expect, it, vi } from 'vitest'

const state = vi.hoisted(() => ({
databaseUrl: process.env.BILLING_REPLAY_IT_DATABASE_URL,
copilotDirectory: process.env.BILLING_REPLAY_COPILOT_DIR,
client: null as Sql | null,
schema: `billing_callback_${process.pid}`,
temporaryFailures: 1,
}))

vi.unmock('drizzle-orm')
vi.unmock('@sim/db/schema')
vi.mock('@sim/db', async () => {
const { drizzle } = await import('drizzle-orm/postgres-js')
const { default: postgres } = await import('postgres')
const client = postgres(state.databaseUrl ?? 'postgres://127.0.0.1:1/unused', {
max: 2,
connection: { search_path: state.schema },
onnotice: () => {},
})
state.client = client
const db = drizzle(client)
return { db, dbReplica: db }
})

/** Subscription and payment-provider fixtures; callback, ledger and replay code are real. */
vi.mock('@/lib/billing/core/subscription', () => {
const subscription = async (userId: string) => {
if (userId === 'billing-replay-transient' && state.temporaryFailures-- > 0) {
throw new Error('Temporary subscription lookup failure')
}
return {
id: 'billing-replay-subscription',
referenceId: userId,
plan: 'pro',
status: 'active',
periodStart: new Date('2025-02-01T00:00:00.000Z'),
periodEnd: new Date('2025-03-01T00:00:00.000Z'),
}
}
return {
getHighestPrioritySubscription: subscription,
getHighestPriorityPersonalSubscription: subscription,
getOrganizationSubscriptionUsable: vi.fn(),
}
})
vi.mock('@/lib/billing/core/plan', () => ({
getHighestPrioritySubscription: vi.fn(),
getHighestPriorityPersonalSubscription: vi.fn(),
}))
vi.mock('@/lib/billing/core/access', () => ({
getEffectiveBillingStatus: async () => ({ billingBlocked: false }),
isOrganizationBillingBlocked: async () => false,
}))
vi.mock('@/lib/billing/core/billing', () => ({
calculateSubscriptionOverage: async () => 0,
computeOrgOverageAmount: vi.fn(),
getOrganizationSubscription: vi.fn(),
}))
vi.mock('@/lib/billing/cycle-close', () => ({ isSubscriptionCycleCloseCurrent: async () => true }))
vi.mock('@/lib/billing/plan-helpers', () => ({ isEnterprise: () => false, isFree: () => false }))
vi.mock('@/lib/billing/subscriptions/utils', () => ({
hasUsableSubscriptionAccess: () => true,
isOrgScopedSubscription: () => false,
}))
vi.mock('@/lib/billing/calculations/usage-monitor', () => ({
checkBillingBlocked: vi.fn(),
checkBillingEntityBlocked: vi.fn(),
checkOrganizationMemberUsageLimit: vi.fn(),
checkUsageStatus: vi.fn(),
}))
vi.mock('@/lib/billing/webhooks/outbox-handlers', () => ({
OUTBOX_EVENT_TYPES: { STRIPE_THRESHOLD_OVERAGE_INVOICE: 'stripe.threshold-overage-invoice' },
}))
vi.mock('@/lib/core/outbox/service', () => ({ enqueueOutboxEvent: vi.fn() }))
vi.mock('@sim/audit', () => ({ AuditAction: {}, AuditResourceType: {}, recordAudit: vi.fn() }))
vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: vi.fn() }))
vi.mock('@/lib/copilot/request/otel', () => ({
withIncomingGoSpan: (
_headers: unknown,
_name: unknown,
_attrs: unknown,
run: (span: { setAttribute: () => void; setAttributes: () => void }) => unknown
) => run({ setAttribute: vi.fn(), setAttributes: vi.fn() }),
}))

import { POST } from '@/app/api/billing/update-cost/route'

const run = promisify(execFile)

afterAll(async () => {
await state.client?.end()
resetEnvMock()
resetEnvFlagsMock()
})

/**
* Requires an isolated localhost PostgreSQL database and the Copilot checkout.
* Runs the actual Go client/reconciler/repository through HTTP into this route,
* with real header validation, cumulative ledger SQL and period classification.
*/
describe.skipIf(!state.databaseUrl || !state.copilotDirectory)(
'cross-service billing replay',
() => {
it('quarantines expired charges durably and recovers temporary failures without double billing', async () => {
const databaseUrl = new URL(state.databaseUrl as string)
expect(['127.0.0.1', 'localhost']).toContain(databaseUrl.hostname)
const client = state.client
if (!client) throw new Error('Test database client was not initialized')
setEnv({ INTERNAL_API_SECRET: 'billing-replay-local-secret' })
setEnvFlags({ isBillingEnabled: true, isHosted: true })
await client.unsafe(`CREATE SCHEMA "${state.schema}"`)
const statuses: number[] = []
const server = createServer(async (request, response) => {
try {
const chunks: Buffer[] = []
let bytes = 0
for await (const chunk of request) {
const buffer = Buffer.from(chunk)
bytes += buffer.length
if (bytes > 16384) throw new Error('Test request exceeds the callback fixture limit')
chunks.push(buffer)
}
const headers = new Headers()
for (const [key, value] of Object.entries(request.headers)) {
if (typeof value === 'string') headers.set(key, value)
}
const result = await POST(
new NextRequest(`http://127.0.0.1${request.url}`, {
method: 'POST',
headers,
body: Buffer.concat(chunks).toString(),
})
)
statuses.push(result.status)
response.writeHead(result.status, Object.fromEntries(result.headers))
response.end(await result.text())
} catch (error) {
response.writeHead(500)
response.end(String(error))
}
})
try {
await client.unsafe(`CREATE TABLE "user" (id text PRIMARY KEY);
INSERT INTO "user" (id) VALUES ('billing-replay-actor'), ('billing-replay-transient');
CREATE TABLE usage_log (
id text PRIMARY KEY, user_id text NOT NULL, category text NOT NULL, source text NOT NULL,
description text NOT NULL, metadata jsonb, cost numeric NOT NULL, event_key text,
billing_entity_type text, billing_entity_id text, billing_period_start timestamp,
billing_period_end timestamp, workspace_id text, workflow_id text, execution_id text,
created_at timestamp NOT NULL DEFAULT now(),
CONSTRAINT usage_log_user_id_user_id_fk FOREIGN KEY (user_id) REFERENCES "user"(id)
); CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key) WHERE event_key IS NOT NULL`)
await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
const address = server.address()
if (!address || typeof address === 'string') throw new Error('Expected HTTP server port')
const result = await run(
'go',
[
'test',
'./internal/analytics',
'-run',
'^TestBillingReplayCrossService$',
'-count=1',
'-v',
],
{
cwd: state.copilotDirectory,
env: {
...process.env,
INTERNAL_API_SECRET: 'billing-replay-local-secret',
BILLING_REPLAY_SIM_URL: `http://127.0.0.1:${address.port}`,
BILLING_REPLAY_IT_DATABASE_URL: state.databaseUrl,
},
timeout: 60_000,
maxBuffer: 1024 * 1024,
}
).catch((error: ExecFileException & { stdout?: string; stderr?: string }) => {
throw new Error([error.message, error.stdout, error.stderr].filter(Boolean).join('\n'), {
cause: error,
})
})
expect(result.stdout).toContain('--- PASS: TestBillingReplayCrossService')
expect(statuses.filter((status) => status === 503)).toHaveLength(1)
expect(statuses.filter((status) => status === 409)).toHaveLength(10)
const rows =
await client`SELECT user_id, cost, billing_period_start::text AS period_start, billing_period_end::text AS period_end FROM usage_log ORDER BY user_id`
expect(rows).toHaveLength(4)
for (const row of rows) {
expect(Number(row.cost)).toBe(1.25)
expect(row.period_start).toBe(
row.user_id === 'billing-replay-transient'
? '2025-02-01 00:00:00'
: '2025-01-01 00:00:00'
)
expect(row.period_end).toBe(
row.user_id === 'billing-replay-transient'
? '2025-03-01 00:00:00'
: '2025-02-01 00:00:00'
)
}
} finally {
try {
if (server.listening) {
await new Promise<void>((resolve, reject) =>
server.close((error) => (error ? reject(error) : resolve()))
)
}
} finally {
await client.unsafe(`DROP SCHEMA "${state.schema}" CASCADE`)
}
}
}, 90_000)
}
)
108 changes: 107 additions & 1 deletion apps/sim/app/api/billing/update-cost/route.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,9 @@ const {
MockCumulativeUsageContextMismatchError: class extends Error {},
MockThresholdSettlementError: class extends Error {
readonly code: string
readonly retryable = true
get retryable() {
return this.code !== 'billing_period_elapsed'
}

constructor(code: string) {
super('Billing settlement temporarily unavailable')
Expand Down Expand Up @@ -587,6 +589,110 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
)
})

it.each(['legacy-v0', 'attribution-v1', 'direct-v1'])(
'returns a distinct non-retryable conflict for an elapsed %s period and preserves usage attribution',
async (protocol) => {
const billingRequestId = '0190c03f-9f7d-4b79-8b58-e7f779fd29e1'
const direct = protocol === 'direct-v1'
setEnvFlags({ isBillingEnabled: true, isHosted: true })
mockCheckAndBillPayerOverageThreshold.mockRejectedValue(
new MockThresholdSettlementError('billing_period_elapsed')
)
mockRecordCumulativeUsage
.mockResolvedValueOnce({ billed: true, delta: 0.5, total: 0.5 })
.mockResolvedValueOnce({ billed: false, delta: 0, total: 0.5 })

for (let attempt = 0; attempt < 2; attempt++) {
const res = await POST(
createMockRequest(
'POST',
{
userId: 'user-1',
cost: 0.5,
model: 'claude-opus-4.8',
source: 'copilot',
idempotencyKey: billingRequestId,
...(direct ? {} : { workspaceId: 'ws-1' }),
},
{
'x-api-key': 'internal',
'x-sim-billing-protocol': protocol,
...(protocol === 'legacy-v0' ? {} : { 'x-sim-billing-request-id': billingRequestId }),
...(direct
? { 'x-sim-billing-account-decision': 'serialized-account-decision' }
: { 'x-sim-billing-attribution': 'serialized-attribution' }),
}
)
)
expect(res.status).toBe(409)
expect(res.headers.get('retry-after')).toBeNull()
await expect(res.json()).resolves.toMatchObject({
success: false,
code: 'BILLING_PERIOD_ELAPSED',
error: 'Billing period has elapsed; reconciliation required',
retryable: false,
})
}
expect(mockRecordCumulativeUsage).toHaveBeenCalledTimes(2)
expect(mockRecordCumulativeUsage).toHaveBeenLastCalledWith(
expect.objectContaining({
eventKey: `update-cost:${billingRequestId}`,
billingPeriod: {
start: new Date('2026-07-01T00:00:00.000Z'),
end: new Date('2026-08-01T00:00:00.000Z'),
...(direct ? { source: 'reporting' } : {}),
},
})
)
}
)

it.each([
['23503', 'usage_log_user_id_user_id_fk', false, 409],
['23503', 'usage_log_workspace_id_workspace_id_fk', false, 500],
['40001', 'usage_log_user_id_user_id_fk', false, 500],
['23503', 'usage_log_user_id_user_id_fk', true, 500],
])(
'classifies the exact missing-user constraint safely (%s, %s, markerless=%s)',
async (code, constraint, markerless, status) => {
mockRecordCumulativeUsage.mockRejectedValueOnce(
new Error('Insert failed', {
cause: { code, constraint_name: constraint },
})
)
const res = await POST(
createMockRequest('POST', SELF_HOSTED_UPDATE_COST_BODY, {
'x-api-key': 'internal',
...(markerless
? {}
: {
'x-sim-billing-protocol': 'legacy-v0',
'x-sim-billing-attribution': 'serialized-attribution',
}),
})
)
expect(res.status).toBe(status)
expect(mockCheckAndBillPayerOverageThreshold).not.toHaveBeenCalled()
expect(mockCheckAndBillOverageThreshold).not.toHaveBeenCalled()
if (status === 409) {
await expect(res.json()).resolves.toMatchObject({
code: 'BILLING_USER_NOT_FOUND',
retryable: false,
})
}
}
)

it('does not expose elapsed-period 409 to markerless clients that treat all conflicts as success', async () => {
mockCheckAndBillPayerOverageThreshold.mockRejectedValueOnce(
new MockThresholdSettlementError('billing_period_elapsed')
)
const res = await POST(
createMockRequest('POST', SELF_HOSTED_UPDATE_COST_BODY, { 'x-api-key': 'internal' })
)
expect(res.status).toBe(503)
})

it('returns a stable retryable 503 when modern threshold settlement fails', async () => {
const billingRequestId = '0190c03f-9f7d-4b79-8b58-e7f779fd29e1'
mockCheckAndBillPayerOverageThreshold.mockRejectedValueOnce(
Expand Down
34 changes: 32 additions & 2 deletions apps/sim/app/api/billing/update-cost/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,38 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
)
}

const pgCode = getPostgresErrorCode(error)
const pgConstraint = getPostgresConstraintName(error)
const reconciliationOutcome =
error instanceof ThresholdSettlementError && !error.retryable
? BILLING_CALLBACK_OUTCOME.billingPeriodElapsed
: pgCode === '23503' && pgConstraint === 'usage_log_user_id_user_id_fk'
? BILLING_CALLBACK_OUTCOME.billingUserNotFound
: undefined

/** Old markerless clients treat every 409 as a successful duplicate. */
if (reconciliationOutcome && !isMarkerlessLegacy) {
logger.warn(`[${requestId}] Billing callback requires reconciliation`, {
code: reconciliationOutcome.code,
duration,
billingProtocol:
req.headers.get(COPILOT_BILLING_PROTOCOL_HEADER) ?? COPILOT_BILLING_PROTOCOL.legacy,
})
span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.ReconciliationRequired)
span.setAttribute(TraceAttr.HttpStatusCode, 409)
span.setAttribute(TraceAttr.BillingDurationMs, duration)
return NextResponse.json(
{
success: false,
code: reconciliationOutcome.code,
error: reconciliationOutcome.message,
retryable: false,
requestId,
},
{ status: 409 }
)
}

if (error instanceof ThresholdSettlementError) {
logger.error(`[${requestId}] Retryable threshold settlement failure`, {
settlementErrorCode: error.code,
Expand Down Expand Up @@ -425,8 +457,6 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
// lock timeout) — Drizzle's "Failed query" wrapper alone cannot
// distinguish them, which made the dead-workspace incident undiagnosable
// from logs.
const pgCode = getPostgresErrorCode(error)
const pgConstraint = getPostgresConstraintName(error)
logger.error(`[${requestId}] Cost update failed`, {
error: toError(error).message,
...(pgCode && { pgCode }),
Expand Down
Loading
Loading