feat: auto grant one-time low balance gift for new users

2026-06-30 20:39:46 +08:00
parent ae6eaf8c3e
commit 9c653e9b16
4 changed files with 232 additions and 9 deletions
@@ -431,6 +431,130 @@ test('purchaseSpaceQuota deducts balance and increases quota', async () => {
  assert.equal(result.quota.availableBytes, 15 * 1024 * 1024 - 1024 - 2048);
 });

+test('billSessionUsage auto gifts low-balance bonus once for eligible new users', async () => {
+  const userRow = {
+    id: 'user-3',
+    username: 'bonus_user',
+    slug: 'bonus_user',
+    email: 'bonus@example.com',
+    display_name: 'Bonus User',
+    role: 'user',
+    status: 'active',
+    plan_type: 'free',
+    workspace_root: '/tmp/bonus-user',
+    balance_cents: 150,
+    tokens_used: 0,
+    spent_cents: 0,
+  };
+  const stateBySession = new Map();
+  let walletBalance = 150;
+  let tokensUsed = 0;
+  let lowBalanceGiftEligible = 1;
+  let lowBalanceGiftGrantedAt = null;
+  let ledgerAdjustCount = 0;
+  let notificationTypes = [];
+
+  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 w.balance_cents, w.tokens_used, u.status')) {
+            return [[{ balance_cents: walletBalance, tokens_used: tokensUsed, status: userRow.status }]];
+          }
+          if (sql.includes('UPDATE h5_user_wallets') && sql.includes('tokens_used = tokens_used + ?')) {
+            walletBalance = params[0];
+            tokensUsed += params[1];
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes('INSERT INTO h5_usage_records')) return [{ affectedRows: 1 }, []];
+          if (sql.includes("INSERT INTO h5_billing_ledger") && sql.includes("'deduct'")) return [{ affectedRows: 1 }, []];
+          if (sql.includes('SELECT low_balance_gift_eligible, low_balance_gift_granted_at')) {
+            return [[{
+              low_balance_gift_eligible: lowBalanceGiftEligible,
+              low_balance_gift_granted_at: lowBalanceGiftGrantedAt,
+            }]];
+          }
+          if (sql.includes('UPDATE h5_user_wallets') && sql.includes('SET balance_cents = ?, updated_at = ?')) {
+            walletBalance = params[0];
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes('UPDATE h5_users') && sql.includes('low_balance_gift_eligible = 0')) {
+            lowBalanceGiftEligible = 0;
+            lowBalanceGiftGrantedAt = params[0];
+            userRow.status = 'active';
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes("INSERT INTO h5_billing_ledger") && sql.includes("'adjust'")) {
+            ledgerAdjustCount += 1;
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes('INSERT INTO h5_user_notifications')) {
+            notificationTypes.push('low_balance_gift');
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes('UPDATE h5_users SET status = \'suspended\'')) {
+            userRow.status = 'suspended';
+            return [{ affectedRows: 1 }, []];
+          }
+          throw new Error(`unexpected connection query: ${sql}`);
+        },
+      };
+    },
+  };
+
+  const auth = createUserAuth(pool, { persistSessions: false });
+
+  const first = await auth.billSessionUsage(
+    userRow.id,
+    'session-bonus-1',
+    { accumulatedOutputTokens: 10_000 },
+    'req-bonus-1',
+  );
+  assert.equal(first.ok, true);
+  assert.equal(first.costCents, 60);
+  assert.equal(first.balanceCents, 1090);
+  assert.equal(ledgerAdjustCount, 1);
+  assert.deepEqual(notificationTypes, ['low_balance_gift']);
+
+  const second = await auth.billSessionUsage(
+    userRow.id,
+    'session-bonus-2',
+    { accumulatedOutputTokens: 10_000 },
+    'req-bonus-2',
+  );
+  assert.equal(second.ok, true);
+  assert.equal(second.costCents, 60);
+  assert.equal(second.balanceCents, 1030);
+  assert.equal(ledgerAdjustCount, 1);
+  assert.deepEqual(notificationTypes, ['low_balance_gift']);
+});
+
 test('updateUser rejects quota smaller than occupied bytes', async () => {
  const userRow = {
    id: 'user-3',