118 lines
4.8 KiB
TypeScript
118 lines
4.8 KiB
TypeScript
// Verification: docs/ai/context/verification-automation.md
|
|
// Source doc: apps/server/docs/ai-context/verifications/flux-unbilled-exploit-fix.md
|
|
//
|
|
// Covers the "user path" section of the verification doc:
|
|
// Scenario: user with 0 < balance < fallbackRate fires N concurrent LLM
|
|
// completion requests.
|
|
// Expected (post-patch): all N are rejected at pre-flight with 402; the
|
|
// upstream router is never invoked; flux_transaction ledger gains
|
|
// no 'debit' row; user_flux.flux stays at the seeded value.
|
|
|
|
import type { Harness } from './_harness'
|
|
|
|
import { and, eq } from 'drizzle-orm'
|
|
import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'
|
|
|
|
import { startVerificationContext } from './_harness'
|
|
|
|
describe('verification: flux-unbilled-exploit-fix', () => {
|
|
let ctx: Harness
|
|
|
|
beforeAll(async () => {
|
|
ctx = await startVerificationContext()
|
|
})
|
|
|
|
beforeEach(async () => {
|
|
// Clean ledger between cases so balanceBefore / balanceAfter assertions
|
|
// don't carry state across tests.
|
|
await ctx.db.delete(ctx.schema.fluxTransaction)
|
|
await ctx.db.delete(ctx.schema.userFlux)
|
|
ctx.redisStore.clear()
|
|
})
|
|
|
|
afterAll(async () => {
|
|
// PGlite is in-memory — letting it fall out of scope is enough.
|
|
})
|
|
|
|
it('rejects N concurrent partial-balance requests at pre-flight, never writes a debit ledger row', async () => {
|
|
// ROOT CAUSE:
|
|
//
|
|
// Before commit 7267b0d6b the pre-flight gate read `if (flux.flux <= 0)`,
|
|
// so users with `0 < balance < FLUX_PER_REQUEST` could spawn N concurrent
|
|
// requests, each pass pre-flight, complete the upstream LLM call, and
|
|
// race on a post-billing debit that almost always failed (insufficient
|
|
// funds for the requested amount, full-amount rollback). Result: N free
|
|
// LLM responses.
|
|
//
|
|
// After 7267b0d6b: pre-flight rejects when `flux.flux < fallbackRate`
|
|
// before the upstream call. This test pins the after-patch behavior so
|
|
// a regression that loosens the gate back to `<= 0` will fail here.
|
|
|
|
const userId = 'user-partial-balance'
|
|
const N = 5
|
|
|
|
ctx.setConfig({ FLUX_PER_REQUEST: 100, FLUX_PER_1K_TOKENS: 50 })
|
|
await ctx.seedUser({ id: userId, balance: 5 })
|
|
ctx.setSessionUser({ id: userId, email: `${userId}@example.com` })
|
|
|
|
const responses = await Promise.all(
|
|
Array.from({ length: N }, () => ctx.app.request('/api/v1/openai/chat/completions', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({ model: 'auto', messages: [{ role: 'user', content: 'hi' }] }),
|
|
})),
|
|
)
|
|
|
|
// All N must be rejected before the gateway is touched.
|
|
expect(responses).toHaveLength(N)
|
|
for (const res of responses) {
|
|
expect(res.status).toBe(402)
|
|
}
|
|
const bodies = await Promise.all(responses.map(r => r.json() as Promise<{ error: string, message: string }>))
|
|
for (const body of bodies) {
|
|
expect(body.error).toBe('PAYMENT_REQUIRED')
|
|
expect(body.message).toMatch(/Insufficient flux/i)
|
|
}
|
|
|
|
// Ledger must be untouched. The pre-flight reject path never reaches
|
|
// `consumeFluxForLLM`, so no `debit` rows should exist.
|
|
const ledger = await ctx.db.query.fluxTransaction.findMany({
|
|
where: eq(ctx.schema.fluxTransaction.userId, userId),
|
|
})
|
|
expect(ledger).toEqual([])
|
|
|
|
// Balance must be exactly the seeded value. No drain, no rollback.
|
|
const fluxRow = await ctx.db.query.userFlux.findFirst({
|
|
where: and(
|
|
eq(ctx.schema.userFlux.userId, userId),
|
|
// Active rows only — deletedAt IS NULL is fluxService's read guard
|
|
),
|
|
})
|
|
expect(fluxRow?.flux).toBe(5)
|
|
})
|
|
|
|
it('allows a request when balance >= fallbackRate (pre-flight pass-through smoke check)', async () => {
|
|
// Companion case: with sufficient balance the pre-flight gate must NOT
|
|
// block. We don't run the request to completion (no upstream gateway
|
|
// available in this harness), only assert that the failure mode here is
|
|
// upstream-fetch error (not 402). This protects against a regression
|
|
// that flips the comparison and rejects everyone.
|
|
const userId = 'user-funded'
|
|
|
|
ctx.setConfig({ FLUX_PER_REQUEST: 100, FLUX_PER_1K_TOKENS: 50 })
|
|
await ctx.seedUser({ id: userId, balance: 1000 })
|
|
ctx.setSessionUser({ id: userId, email: `${userId}@example.com` })
|
|
|
|
const res = await ctx.app.request('/api/v1/openai/chat/completions', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({ model: 'auto', messages: [{ role: 'user', content: 'hi' }] }),
|
|
})
|
|
|
|
// Pre-flight passed -> route delegated to the mock router. The harness
|
|
// wires `llmRouter.route` to a vi.fn that returns 200, so we only assert
|
|
// the pre-flight gate did NOT short-circuit with 402.
|
|
expect(res.status).not.toBe(402)
|
|
})
|
|
})
|