Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6b586e0951 | |||
| 7c18c998bd |
@@ -194,6 +194,12 @@ export async function migrateSchema(pool) {
|
|||||||
`ALTER TABLE h5_usage_records ADD UNIQUE KEY uniq_h5_usage_request_id (request_id)`,
|
`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 = [
|
const assetForeignKeys = [
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -511,6 +511,7 @@ CREATE TABLE IF NOT EXISTS h5_usage_records (
|
|||||||
output_tokens INT NOT NULL DEFAULT 0,
|
output_tokens INT NOT NULL DEFAULT 0,
|
||||||
cost_cents BIGINT NOT NULL,
|
cost_cents BIGINT NOT NULL,
|
||||||
balance_after_cents BIGINT NOT NULL,
|
balance_after_cents BIGINT NOT NULL,
|
||||||
|
billing_source VARCHAR(16) NOT NULL DEFAULT 'wallet',
|
||||||
created_at BIGINT NOT NULL,
|
created_at BIGINT NOT NULL,
|
||||||
UNIQUE KEY uniq_h5_usage_request_id (request_id),
|
UNIQUE KEY uniq_h5_usage_request_id (request_id),
|
||||||
KEY idx_h5_usage_user_time (user_id, created_at),
|
KEY idx_h5_usage_user_time (user_id, created_at),
|
||||||
|
|||||||
+30
-6
@@ -1442,10 +1442,12 @@ export function createUserAuth(pool, options = {}) {
|
|||||||
);
|
);
|
||||||
|
|
||||||
// Subscription quota check: consume tokens from active plan before touching balance.
|
// Subscription quota check: consume tokens from active plan before touching balance.
|
||||||
|
let subscriptionCovered = false;
|
||||||
if (costCents > 0 && subscriptionService) {
|
if (costCents > 0 && subscriptionService) {
|
||||||
const coverage = await subscriptionService.consumeQuota(userId, deltaTokens, conn);
|
const coverage = await subscriptionService.consumeQuota(userId, deltaTokens, conn);
|
||||||
if (coverage.fullyCovers) {
|
if (coverage.fullyCovers) {
|
||||||
costCents = 0;
|
costCents = 0;
|
||||||
|
subscriptionCovered = true;
|
||||||
} else if (coverage.overageRate < 1.0) {
|
} else if (coverage.overageRate < 1.0) {
|
||||||
costCents = Math.max(1, Math.ceil(costCents * coverage.overageRate));
|
costCents = Math.max(1, Math.ceil(costCents * coverage.overageRate));
|
||||||
}
|
}
|
||||||
@@ -1482,8 +1484,8 @@ export function createUserAuth(pool, options = {}) {
|
|||||||
|
|
||||||
await conn.query(
|
await conn.query(
|
||||||
`INSERT INTO h5_usage_records
|
`INSERT INTO h5_usage_records
|
||||||
(user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, created_at)
|
(user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, billing_source, created_at)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
|
VALUES (?, ?, ?, ?, ?, ?, ?, 'wallet', ?)`,
|
||||||
[
|
[
|
||||||
userId,
|
userId,
|
||||||
agentSessionId,
|
agentSessionId,
|
||||||
@@ -1524,6 +1526,28 @@ export function createUserAuth(pool, options = {}) {
|
|||||||
userId,
|
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 {
|
} else {
|
||||||
const user = await getUserById(userId);
|
const user = await getUserById(userId);
|
||||||
balanceAfter = user ? Number(user.balance_cents) : null;
|
balanceAfter = user ? Number(user.balance_cents) : null;
|
||||||
@@ -1556,12 +1580,12 @@ export function createUserAuth(pool, options = {}) {
|
|||||||
if (userId) params.push(userId);
|
if (userId) params.push(userId);
|
||||||
const [rows] = await pool.query(
|
const [rows] = await pool.query(
|
||||||
`SELECT r.id, r.user_id, u.username, r.agent_session_id, r.request_id,
|
`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
|
FROM h5_usage_records r JOIN h5_users u ON u.id = r.user_id
|
||||||
${where} ORDER BY r.created_at DESC LIMIT ${safeLimit}`,
|
${where} ORDER BY r.created_at DESC LIMIT ${safeLimit}`,
|
||||||
params,
|
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 safePageSize = Math.min(Math.max(Number(pageSize) || 50, 1), 200);
|
||||||
const safePage = Math.max(Number(page) || 1, 1);
|
const safePage = Math.max(Number(page) || 1, 1);
|
||||||
@@ -1575,7 +1599,7 @@ export function createUserAuth(pool, options = {}) {
|
|||||||
);
|
);
|
||||||
const [rows] = await pool.query(
|
const [rows] = await pool.query(
|
||||||
`SELECT r.id, r.user_id, u.username, r.agent_session_id, r.request_id,
|
`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
|
FROM h5_usage_records r JOIN h5_users u ON u.id = r.user_id
|
||||||
${where}
|
${where}
|
||||||
ORDER BY r.created_at DESC
|
ORDER BY r.created_at DESC
|
||||||
@@ -1583,7 +1607,7 @@ export function createUserAuth(pool, options = {}) {
|
|||||||
params,
|
params,
|
||||||
);
|
);
|
||||||
return {
|
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),
|
total: Number(total),
|
||||||
page: safePage,
|
page: safePage,
|
||||||
pageSize: safePageSize,
|
pageSize: safePageSize,
|
||||||
|
|||||||
@@ -770,6 +770,114 @@ test('billSessionUsage auto gifts low-balance bonus once for eligible new users'
|
|||||||
assert.deepEqual(notificationTypes, ['low_balance_gift']);
|
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 () => {
|
test('updateUser rejects quota smaller than occupied bytes', async () => {
|
||||||
const userRow = {
|
const userRow = {
|
||||||
id: 'user-3',
|
id: 'user-3',
|
||||||
|
|||||||
Reference in New Issue
Block a user