From 49ff7c9badb29945b39700b93bfef89c7d2e8d3c Mon Sep 17 00:00:00 2001 From: tkmind Date: Thu, 6 Aug 2026 10:52:55 +0000 Subject: [PATCH] feat(billing): record subscription-covered token usage in usage records (#46) --- db.mjs | 6 +++ schema.sql | 1 + user-auth.mjs | 36 ++++++++++++--- user-auth.test.mjs | 108 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 145 insertions(+), 6 deletions(-) diff --git a/db.mjs b/db.mjs index f1a94ab..7e02fe3 100644 --- a/db.mjs +++ b/db.mjs @@ -194,6 +194,12 @@ export async function migrateSchema(pool) { `ALTER TABLE h5_usage_records ADD UNIQUE KEY uniq_h5_usage_request_id (request_id)`, ); } + if (!(await columnExists(pool, 'h5_usage_records', 'billing_source'))) { + await pool.query( + `ALTER TABLE h5_usage_records + ADD COLUMN billing_source VARCHAR(16) NOT NULL DEFAULT 'wallet' AFTER balance_after_cents`, + ); + } const assetForeignKeys = [ { diff --git a/schema.sql b/schema.sql index f6c1cb3..14f7e83 100644 --- a/schema.sql +++ b/schema.sql @@ -511,6 +511,7 @@ CREATE TABLE IF NOT EXISTS h5_usage_records ( output_tokens INT NOT NULL DEFAULT 0, cost_cents BIGINT NOT NULL, balance_after_cents BIGINT NOT NULL, + billing_source VARCHAR(16) NOT NULL DEFAULT 'wallet', created_at BIGINT NOT NULL, UNIQUE KEY uniq_h5_usage_request_id (request_id), KEY idx_h5_usage_user_time (user_id, created_at), diff --git a/user-auth.mjs b/user-auth.mjs index 8639b2c..6eedfbb 100644 --- a/user-auth.mjs +++ b/user-auth.mjs @@ -1442,10 +1442,12 @@ export function createUserAuth(pool, options = {}) { ); // Subscription quota check: consume tokens from active plan before touching balance. + let subscriptionCovered = false; if (costCents > 0 && subscriptionService) { const coverage = await subscriptionService.consumeQuota(userId, deltaTokens, conn); if (coverage.fullyCovers) { costCents = 0; + subscriptionCovered = true; } else if (coverage.overageRate < 1.0) { costCents = Math.max(1, Math.ceil(costCents * coverage.overageRate)); } @@ -1482,8 +1484,8 @@ export function createUserAuth(pool, options = {}) { await conn.query( `INSERT INTO h5_usage_records - (user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + (user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, billing_source, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, 'wallet', ?)`, [ userId, agentSessionId, @@ -1524,6 +1526,28 @@ export function createUserAuth(pool, options = {}) { userId, ]); } + } else if (subscriptionCovered && deltaTokens > 0) { + const [walletRows] = await conn.query( + `SELECT balance_cents, tokens_used FROM h5_user_wallets WHERE user_id = ?`, + [userId], + ); + balanceAfter = walletRows[0] ? Number(walletRows[0].balance_cents) : 0; + tokensUsedAfter = walletRows[0] ? Number(walletRows[0].tokens_used ?? 0) : null; + + await conn.query( + `INSERT INTO h5_usage_records + (user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, billing_source, created_at) + VALUES (?, ?, ?, ?, ?, 0, ?, 'subscription', ?)`, + [ + userId, + agentSessionId, + normalizedRequestId, + deltaIn, + deltaOut, + balanceAfter, + now, + ], + ); } else { const user = await getUserById(userId); balanceAfter = user ? Number(user.balance_cents) : null; @@ -1556,12 +1580,12 @@ export function createUserAuth(pool, options = {}) { if (userId) params.push(userId); const [rows] = await pool.query( `SELECT r.id, r.user_id, u.username, r.agent_session_id, r.request_id, - r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.created_at + r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.billing_source, r.created_at FROM h5_usage_records r JOIN h5_users u ON u.id = r.user_id ${where} ORDER BY r.created_at DESC LIMIT ${safeLimit}`, params, ); - return rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), createdAt: Number(row.created_at) })); + return rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), billingSource: row.billing_source ?? 'wallet', createdAt: Number(row.created_at) })); } const safePageSize = Math.min(Math.max(Number(pageSize) || 50, 1), 200); const safePage = Math.max(Number(page) || 1, 1); @@ -1575,7 +1599,7 @@ export function createUserAuth(pool, options = {}) { ); const [rows] = await pool.query( `SELECT r.id, r.user_id, u.username, r.agent_session_id, r.request_id, - r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.created_at + r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.billing_source, r.created_at FROM h5_usage_records r JOIN h5_users u ON u.id = r.user_id ${where} ORDER BY r.created_at DESC @@ -1583,7 +1607,7 @@ export function createUserAuth(pool, options = {}) { params, ); return { - records: rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), createdAt: Number(row.created_at) })), + records: rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), billingSource: row.billing_source ?? 'wallet', createdAt: Number(row.created_at) })), total: Number(total), page: safePage, pageSize: safePageSize, diff --git a/user-auth.test.mjs b/user-auth.test.mjs index bc3aca5..4c554f7 100644 --- a/user-auth.test.mjs +++ b/user-auth.test.mjs @@ -770,6 +770,114 @@ test('billSessionUsage auto gifts low-balance bonus once for eligible new users' assert.deepEqual(notificationTypes, ['low_balance_gift']); }); +test('billSessionUsage writes usage record when subscription fully covers tokens', async () => { + const userRow = { + id: 'user-sub-1', + username: 'pro_user', + slug: 'pro_user', + email: 'pro@example.com', + display_name: 'Pro User', + role: 'user', + status: 'active', + plan_type: 'pro', + workspace_root: '/tmp/pro-user', + balance_cents: 200, + tokens_used: 0, + spent_cents: 0, + }; + const stateBySession = new Map(); + let walletBalance = 200; + let tokensUsed = 0; + const usageRecords = []; + let ledgerCount = 0; + let consumeQuotaCalls = 0; + + const subscriptionService = { + async consumeQuota(userId, deltaTokens) { + consumeQuotaCalls += 1; + assert.equal(userId, userRow.id); + assert.equal(deltaTokens, 15_000); + return { fullyCovers: true, overageRate: 0.5 }; + }, + }; + + const pool = { + async query(sql) { + if (sql.includes('FROM h5_users u') && sql.includes('WHERE u.id = ?')) { + return [[{ ...userRow, balance_cents: walletBalance, tokens_used: tokensUsed }]]; + } + throw new Error(`unexpected pool query: ${sql}`); + }, + async getConnection() { + return { + async beginTransaction() {}, + async commit() {}, + async rollback() {}, + release() {}, + async query(sql, params = []) { + if (sql.includes('SELECT cost_cents FROM h5_usage_records WHERE request_id = ? LIMIT 1')) return [[]]; + if (sql.includes('INSERT INTO h5_session_billing_state') && sql.includes('agent_session_id = agent_session_id')) { + return [{ affectedRows: 1 }, []]; + } + if (sql.includes('FROM h5_session_billing_state') && sql.includes('FOR UPDATE')) { + const row = stateBySession.get(params[0]); + return [row ? [row] : []]; + } + if (sql.includes('INSERT INTO h5_session_billing_state') && sql.includes('ON DUPLICATE KEY UPDATE')) { + stateBySession.set(params[0], { + last_accumulated_cost: params[2], + last_input_tokens: params[3], + last_output_tokens: params[4], + }); + return [{ affectedRows: 1 }, []]; + } + if (sql.includes('SELECT balance_cents, tokens_used FROM h5_user_wallets WHERE user_id = ?')) { + return [[{ balance_cents: walletBalance, tokens_used: tokensUsed }]]; + } + if (sql.includes('INSERT INTO h5_usage_records')) { + const subscription = sql.includes("'subscription'"); + usageRecords.push({ + user_id: params[0], + agent_session_id: params[1], + request_id: params[2], + input_tokens: params[3], + output_tokens: params[4], + cost_cents: subscription ? 0 : params[5], + balance_after_cents: subscription ? params[5] : params[6], + billing_source: subscription ? 'subscription' : 'wallet', + }); + return [{ affectedRows: 1 }, []]; + } + if (sql.includes("INSERT INTO h5_billing_ledger") && sql.includes("'deduct'")) { + ledgerCount += 1; + return [{ affectedRows: 1 }, []]; + } + throw new Error(`unexpected connection query: ${sql}`); + }, + }; + }, + }; + + const auth = createUserAuth(pool, { persistSessions: false, subscriptionService }); + const result = await auth.billSessionUsage( + userRow.id, + 'session-sub-1', + { accumulatedOutputTokens: 15_000 }, + 'req-sub-1', + ); + + assert.equal(result.ok, true); + assert.equal(result.costCents, 0); + assert.equal(result.balanceCents, 200); + assert.equal(consumeQuotaCalls, 1); + assert.equal(ledgerCount, 0); + assert.equal(usageRecords.length, 1); + assert.equal(usageRecords[0].cost_cents, 0); + assert.equal(usageRecords[0].billing_source, 'subscription'); + assert.equal(usageRecords[0].output_tokens, 15_000); + assert.equal(walletBalance, 200); +}); + test('updateUser rejects quota smaller than occupied bytes', async () => { const userRow = { id: 'user-3',