refactor(billing): rename debitFlux to consumeFluxForLLM and enhance metadata handling

This commit is contained in:
RainbowBird
2026-03-28 02:25:44 +08:00
committed by RainbowBird
parent 750e48ad95
commit 87d3d65b6a
10 changed files with 193 additions and 66 deletions
+8 -4
View File
@@ -246,11 +246,13 @@ export function createV1CompletionsRoutes(fluxService: FluxService, billingServi
const requestId = nanoid()
let actualCharged = 0
try {
await billingService.debitFlux({
await billingService.consumeFluxForLLM({
userId: user.id,
amount: fluxConsumed,
requestId,
description: requestModel,
promptTokens: usage.promptTokens,
completionTokens: usage.completionTokens,
})
actualCharged = fluxConsumed
}
@@ -291,11 +293,13 @@ export function createV1CompletionsRoutes(fluxService: FluxService, billingServi
// Debit flux via DB transaction (source of truth)
// NOTICE: no try/catch — debit failure (e.g. insufficient balance) must block the response
const requestId = nanoid()
await billingService.debitFlux({
await billingService.consumeFluxForLLM({
userId: user.id,
amount: fluxConsumed,
requestId,
description: requestModel,
promptTokens: usage.promptTokens,
completionTokens: usage.completionTokens,
})
publishRequestLog({
@@ -355,7 +359,7 @@ export function createV1CompletionsRoutes(fluxService: FluxService, billingServi
}
const fluxPerRequest = await configKV.getOrThrow('FLUX_PER_REQUEST_TTS')
await billingService.debitFlux({
await billingService.consumeFluxForLLM({
userId: user.id,
amount: fluxPerRequest,
requestId: nanoid(),
@@ -425,7 +429,7 @@ export function createV1CompletionsRoutes(fluxService: FluxService, billingServi
}
const fluxPerRequest = await configKV.getOrThrow('FLUX_PER_REQUEST_ASR')
await billingService.debitFlux({
await billingService.consumeFluxForLLM({
userId: user.id,
amount: fluxPerRequest,
requestId: nanoid(),
@@ -24,7 +24,7 @@ function createMockFluxService(flux = 100): FluxService {
function createMockBillingService(flux = 100): BillingService {
let balance = flux
return {
debitFlux: vi.fn(async (input: { userId: string, amount: number }) => {
consumeFluxForLLM: vi.fn(async (input: { userId: string, amount: number }) => {
balance -= input.amount
return { userId: input.userId, flux: balance }
}),
@@ -168,7 +168,7 @@ describe('v1CompletionsRoutes', () => {
expect(data.id).toBe('chatcmpl-1')
// Verify flux was debited via billingService
expect(billingService.debitFlux).toHaveBeenCalledWith(
expect(billingService.consumeFluxForLLM).toHaveBeenCalledWith(
expect.objectContaining({ userId: 'user-1', amount: 1 }),
)
@@ -255,7 +255,7 @@ describe('v1CompletionsRoutes', () => {
expect(res.status).toBe(500)
// Post-billing: no charge on failed requests
expect(billingService.debitFlux).not.toHaveBeenCalled()
expect(billingService.consumeFluxForLLM).not.toHaveBeenCalled()
})
it('should return 503 when config keys are missing', async () => {
@@ -345,7 +345,7 @@ describe('v1CompletionsRoutes', () => {
await Promise.resolve()
expect(billingService.debitFlux).not.toHaveBeenCalled()
expect(billingService.consumeFluxForLLM).not.toHaveBeenCalled()
expect(billingMq.publish).not.toHaveBeenCalled()
})
})
@@ -29,7 +29,13 @@ export function createBillingConsumerHandler(db: Database) {
balanceBefore,
balanceAfter: event.payload.balanceAfter ?? balanceBefore - event.payload.amount,
requestId: event.requestId,
description: event.payload.source ?? 'LLM request',
description: event.payload.description ?? event.payload.source ?? 'LLM request',
metadata: event.payload.metadata != null || event.payload.source != null
? {
...(event.payload.metadata as Record<string, unknown>),
source: event.payload.source,
}
: undefined,
}).onConflictDoNothing()
logger.withFields({
@@ -30,6 +30,8 @@ const BalanceChangePayloadSchema = object({
amount: number(),
balanceAfter: optional(number()),
source: optional(pipe(string(), nonEmpty())),
description: optional(pipe(string(), nonEmpty())),
metadata: optional(unknown()),
})
const StripeCheckoutCompletedPayloadSchema = object({
@@ -56,68 +56,99 @@ export function createBillingService(
}
}
/**
* Debit flux from a user's balance within a DB transaction.
* The transaction ONLY locks the row and updates the balance.
* Ledger + audit entries are written by the billing-mq consumer
* after it processes the flux.debited event published post-commit.
*
* Private — call domain-specific wrappers (e.g. consumeFluxForLLM) instead.
*/
async function debitFlux(input: {
userId: string
amount: number
requestId?: string
description?: string
source: string
metadata?: Record<string, unknown>
}): Promise<{ userId: string, flux: number }> {
const result = await db.transaction(async (tx) => {
// 1. Lock the row and read current balance
const [row] = await tx
.select({ flux: fluxSchema.userFlux.flux })
.from(fluxSchema.userFlux)
.where(eq(fluxSchema.userFlux.userId, input.userId))
.for('update')
if (!row) {
throw new Error(`No flux record for user ${input.userId}`)
}
const balanceBefore = row.flux
if (balanceBefore < input.amount) {
metrics?.fluxInsufficientBalance.add(1)
throw createPaymentRequiredError('Insufficient flux')
}
const balanceAfter = balanceBefore - input.amount
// 2. Update balance
await tx.update(fluxSchema.userFlux)
.set({ flux: balanceAfter, updatedAt: new Date() })
.where(eq(fluxSchema.userFlux.userId, input.userId))
return { userId: input.userId, flux: balanceAfter, balanceBefore }
})
// 3. Update Redis cache after commit (best-effort)
await updateRedisCache(input.userId, result.flux)
// 4. Publish flux.debited event to stream; ledger + audit written by consumer
await publishEvent({
eventId: nanoid(),
eventType: 'flux.debited',
aggregateId: input.userId,
userId: input.userId,
requestId: input.requestId,
occurredAt: new Date().toISOString(),
schemaVersion: 1,
payload: {
amount: input.amount,
balanceAfter: result.flux,
source: input.source,
description: input.description,
metadata: input.metadata,
},
})
logger.withFields({ userId: input.userId, amount: input.amount, balance: result.flux }).log('Debited flux')
return { userId: result.userId, flux: result.flux }
}
return {
/**
* Debit flux from a user's balance within a DB transaction.
* The transaction ONLY locks the row and updates the balance.
* Ledger + audit entries are written by the billing-mq consumer
* after it processes the flux.debited event published post-commit.
* Debit flux for an LLM API request (chat, TTS, ASR).
* Passes token usage as opaque metadata carried through the flux.debited event
* so the billing-mq consumer can write it to the ledger.
*/
async debitFlux(input: {
async consumeFluxForLLM(input: {
userId: string
amount: number
requestId?: string
description?: string
promptTokens?: number
completionTokens?: number
}): Promise<{ userId: string, flux: number }> {
const result = await db.transaction(async (tx) => {
// 1. Lock the row and read current balance
const [row] = await tx
.select({ flux: fluxSchema.userFlux.flux })
.from(fluxSchema.userFlux)
.where(eq(fluxSchema.userFlux.userId, input.userId))
.for('update')
if (!row) {
throw new Error(`No flux record for user ${input.userId}`)
}
const balanceBefore = row.flux
if (balanceBefore < input.amount) {
metrics?.fluxInsufficientBalance.add(1)
throw createPaymentRequiredError('Insufficient flux')
}
const balanceAfter = balanceBefore - input.amount
// 2. Update balance
await tx.update(fluxSchema.userFlux)
.set({ flux: balanceAfter, updatedAt: new Date() })
.where(eq(fluxSchema.userFlux.userId, input.userId))
return { userId: input.userId, flux: balanceAfter, balanceBefore }
})
// 3. Update Redis cache after commit (best-effort)
await updateRedisCache(input.userId, result.flux)
// 4. Publish flux.debited event to stream; ledger + audit written by consumer
await publishEvent({
eventId: nanoid(),
eventType: 'flux.debited',
aggregateId: input.userId,
return debitFlux({
userId: input.userId,
amount: input.amount,
requestId: input.requestId,
occurredAt: new Date().toISOString(),
schemaVersion: 1,
payload: {
amount: input.amount,
balanceAfter: result.flux,
source: 'llm.request',
},
description: input.description,
source: 'llm.request',
metadata: input.promptTokens != null || input.completionTokens != null
? { promptTokens: input.promptTokens, completionTokens: input.completionTokens }
: undefined,
})
logger.withFields({ userId: input.userId, amount: input.amount, balance: result.flux }).log('Debited flux')
return { userId: result.userId, flux: result.flux }
},
/**
@@ -0,0 +1,68 @@
import type { Database } from '../../../libs/db'
import { eq } from 'drizzle-orm'
import { beforeAll, beforeEach, describe, expect, it } from 'vitest'
import { mockDB } from '../../../libs/mock-db'
import { createBillingConsumerHandler } from '../billing-consumer-handler'
import * as schema from '../../../schemas'
describe('billingConsumerHandler', () => {
let db: Database
beforeAll(async () => {
db = await mockDB(schema)
await db.insert(schema.user).values({
id: 'user-billing-handler-1',
name: 'Billing Handler User',
email: 'billing-handler@example.com',
})
})
beforeEach(async () => {
await db.delete(schema.fluxLedger).where(eq(schema.fluxLedger.userId, 'user-billing-handler-1'))
})
it('writes debit ledger metadata so token usage can be shown in the UI', async () => {
const handler = createBillingConsumerHandler(db)
await handler.handleMessage({
streamMessageId: '1-0',
event: {
eventId: 'evt-1',
eventType: 'flux.debited',
aggregateId: 'user-billing-handler-1',
userId: 'user-billing-handler-1',
requestId: 'req-1',
occurredAt: '2026-03-27T00:00:00.000Z',
schemaVersion: 1,
payload: {
amount: 3,
balanceAfter: 97,
source: 'llm.request',
description: 'gpt-5',
metadata: { promptTokens: 111, completionTokens: 222 },
},
},
})
const [ledgerRecord] = await db.select().from(schema.fluxLedger).where(eq(schema.fluxLedger.requestId, 'req-1'))
expect(ledgerRecord).toMatchObject({
userId: 'user-billing-handler-1',
type: 'debit',
amount: 3,
balanceBefore: 100,
balanceAfter: 97,
requestId: 'req-1',
description: 'gpt-5',
metadata: {
promptTokens: 111,
completionTokens: 222,
source: 'llm.request',
},
})
})
})
@@ -16,6 +16,8 @@ describe('billingEvents', () => {
amount: 12,
balanceAfter: 88,
source: 'llm',
description: 'gpt-5',
metadata: { promptTokens: 100, completionTokens: 200 },
},
}
@@ -32,6 +34,8 @@ describe('billingEvents', () => {
amount: 12,
balanceAfter: 88,
source: 'llm',
description: 'gpt-5',
metadata: { promptTokens: 100, completionTokens: 200 },
}),
})
@@ -16,6 +16,8 @@ function createEvent() {
amount: 5,
balanceAfter: 95,
source: 'llm',
description: 'gpt-5',
metadata: { promptTokens: 100, completionTokens: 200 },
},
}
}
@@ -58,6 +60,8 @@ describe('billingMqService', () => {
amount: 5,
balanceAfter: 95,
source: 'llm',
description: 'gpt-5',
metadata: { promptTokens: 100, completionTokens: 200 },
}),
)
})
@@ -124,6 +128,8 @@ describe('billingMqService', () => {
amount: 5,
balanceAfter: 95,
source: 'llm',
description: 'gpt-5',
metadata: { promptTokens: 100, completionTokens: 200 },
}),
],
]],
@@ -189,6 +195,8 @@ describe('billingMqService', () => {
amount: 5,
balanceAfter: 95,
source: 'llm',
description: 'gpt-5',
metadata: { promptTokens: 100, completionTokens: 200 },
}),
],
]],
@@ -153,16 +153,18 @@ describe('billingService', () => {
})
})
describe('debitFlux', () => {
describe('consumeFluxForLLM', () => {
it('deducts balance, publishes flux.debited event, updates Redis', async () => {
// Setup: give user some flux first
await db.insert(schema.userFlux).values({ userId: 'user-billing-1', flux: 100 })
const result = await billingService.debitFlux({
const result = await billingService.consumeFluxForLLM({
userId: 'user-billing-1',
amount: 30,
requestId: 'req-1',
description: 'gpt-4',
promptTokens: 120,
completionTokens: 80,
})
expect(result).toEqual({ userId: 'user-billing-1', flux: 70 })
@@ -179,6 +181,8 @@ describe('billingService', () => {
payload: expect.objectContaining({
amount: 30,
balanceAfter: 70,
description: 'gpt-4',
metadata: { promptTokens: 120, completionTokens: 80 },
}),
}))
@@ -189,7 +193,7 @@ describe('billingService', () => {
it('throws 402 when balance is insufficient', async () => {
await db.insert(schema.userFlux).values({ userId: 'user-billing-1', flux: 5 })
await expect(billingService.debitFlux({
await expect(billingService.consumeFluxForLLM({
userId: 'user-billing-1',
amount: 10,
})).rejects.toThrow('Insufficient flux')
@@ -49,7 +49,7 @@ async function fetchAuditHistory(loadMore = false) {
auditLoading.value = true
try {
const offset = loadMore ? auditOffset.value : 0
const res = await client.api.flux.history.$get({
const res = await client.api.v1.flux.history.$get({
query: { limit: String(AUDIT_PAGE_SIZE), offset: String(offset) },
})
if (res.ok) {
@@ -78,7 +78,7 @@ function formatDate(iso: string): string {
async function fetchPackages() {
try {
const res = await client.api.stripe.packages.$get()
const res = await client.api.v1.stripe.packages.$get()
if (res.ok)
packages.value = await res.json() as { amount: number, label: string, price: string }[]
}
@@ -104,7 +104,7 @@ async function handleBuy(amount: number) {
loadingAmount.value = amount
message.value = null
try {
const res = await client.api.stripe.checkout.$post({ json: { amount } })
const res = await client.api.v1.stripe.checkout.$post({ json: { amount } })
if (!res.ok) {
const data = await res.json() as { error?: string, message?: string }
message.value = { type: 'error', text: data.message || t('settings.pages.flux.checkout.error') }