24.05-额度系统

-- 额度使用记录表
CREATE TABLE quota_usage (
  id TEXT PRIMARY KEY,
  tenant_id TEXT NOT NULL,
  user_id TEXT NOT NULL,
  quota_type TEXT NOT NULL,        -- tokens/requests/storage/knowledge_bases/members
  amount INTEGER NOT NULL,         -- 使用量
  model TEXT,                      -- 模型名称（仅 tokens 类型需要）
  metadata TEXT,                   -- 额外信息（JSON）
  created_at INTEGER NOT NULL,
  FOREIGN KEY (tenant_id) REFERENCES tenants(id)
)
 
-- 额度快照表（用于快速查询当前用量）
CREATE TABLE quota_snapshots (
  id TEXT PRIMARY KEY,
  tenant_id TEXT NOT NULL,
  quota_type TEXT NOT NULL,
  period_start INTEGER NOT NULL,   -- 周期开始时间
  total_used INTEGER NOT NULL DEFAULT 0,
  updated_at INTEGER NOT NULL,
  UNIQUE(tenant_id, quota_type, period_start)
)
 
-- 索引：加速周期内查询
CREATE INDEX idx_quota_usage_period
ON quota_usage(tenant_id, quota_type, created_at)

// src/types/quota.ts
export type QuotaType =
  | 'tokens'
  | 'requests'
  | 'storage'
  | 'knowledge_bases'
  | 'members'
 
export interface QuotaUsage {
  id: string
  tenantId: string
  userId: string
  quotaType: QuotaType
  amount: number
  model?: string
  metadata?: Record<string, any>
  createdAt: number
}
 
export interface QuotaSnapshot {
  id: string
  tenantId: string
  quotaType: QuotaType
  periodStart: number
  totalUsed: number
  updatedAt: number
}
 
export interface QuotaCheck {
  type: QuotaType
  used: number
  limit: number
  remaining: number
  percentage: number
}

// src/lib/quota.ts
import { QuotaType } from '../types/quota'
 
export class QuotaManager {
  constructor(private db: D1Database, private env: Env) {}
 
  /**
   * 检查并扣减额度
   * 返回是否成功，以及当前用量
   */
  async consume(
    tenantId: string,
    userId: string,
    type: QuotaType,
    amount: number,
    model?: string,
    metadata?: Record<string, any>
  ): Promise<{ success: boolean; used: number; limit: number }> {
    // 1. 获取套餐配额
    const limit = await this.getQuotaLimit(tenantId, type)
 
    // 2. 获取当前用量
    const used = await this.getCurrentUsage(tenantId, type)
 
    // 3. 检查是否超额
    if (limit > 0 && used + amount > limit) {
      return { success: false, used, limit }
    }
 
    // 4. 记录使用（原子操作）
    const now = Date.now()
    const id = crypto.randomUUID()
 
    await this.db.prepare(`
      INSERT INTO quota_usage (id, tenant_id, user_id, quota_type, amount, model, metadata, created_at)
      VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    `).bind(
      id,
      tenantId,
      userId,
      type,
      amount,
      model || null,
      metadata ? JSON.stringify(metadata) : null,
      now
    ).run()
 
    // 5. 更新快照（使用 UPSERT）
    const periodStart = this.getPeriodStart(type)
 
    await this.db.prepare(`
      INSERT INTO quota_snapshots (id, tenant_id, quota_type, period_start, total_used, updated_at)
      VALUES (?, ?, ?, ?, ?, ?)
      ON CONFLICT(tenant_id, quota_type, period_start)
      DO UPDATE SET total_used = total_used + ?, updated_at = ?
    `).bind(
      crypto.randomUUID(),
      tenantId,
      type,
      periodStart,
      amount,
      now,
      amount,
      now
    ).run()
 
    return { success: true, used: used + amount, limit }
  }
 
  /**
   * 获取当前周期内的使用量
   */
  async getCurrentUsage(tenantId: string, type: QuotaType): Promise<number> {
    const periodStart = this.getPeriodStart(type)
 
    const snapshot = await this.db.prepare(`
      SELECT total_used FROM quota_snapshots
      WHERE tenant_id = ? AND quota_type = ? AND period_start = ?
    `).bind(tenantId, type, periodStart).first()
 
    return snapshot?.total_used || 0
  }
 
  /**
   * 获取配额上限
   */
  private async getQuotaLimit(tenantId: string, type: QuotaType): Promise<number> {
    const tenant = await this.db.prepare(`
      SELECT plan FROM tenants WHERE id = ?
    `).bind(tenantId).first()
 
    if (!tenant) return 0
 
    const plan = getPlan(tenant.plan)
 
    switch (type) {
      case 'tokens':
        return plan.tokenQuota
      case 'requests':
        return plan.requestQuota || -1
      case 'storage':
        return plan.storageQuota || -1
      case 'knowledge_bases':
        return plan.maxKnowledgeBases
      case 'members':
        return plan.maxTeamMembers
      default:
        return -1
    }
  }
 
  /**
   * 获取周期开始时间
   */
  private getPeriodStart(type: QuotaType): number {
    const now = new Date()
 
    switch (type) {
      case 'tokens':
      case 'requests':
        // 月度周期：当月 1 号 00:00:00
        return new Date(now.getFullYear(), now.getMonth(), 1).getTime()
      case 'storage':
      case 'knowledge_bases':
      case 'members':
        // 累计型：从 0 开始
        return 0
      default:
        return 0
    }
  }
}

// src/routes/chat.ts
import { Hono } from 'hono'
import { QuotaManager } from '../lib/quota'
 
const app = new Hono()
 
app.post('/chat', async (c) => {
  const tenantId = c.get('tenantId')
  const userId = c.get('userId')
  const { messages, model } = c.req.valid('json')
 
  const quotaManager = new QuotaManager(c.env.DB, c.env)
 
  // 1. 预估 token 消耗（简单估算）
  const estimatedTokens = messages.reduce((sum, msg) => {
    return sum + Math.ceil(msg.content.length / 4)
  }, 0) + 1000  // 预留 1000 token 给回复
 
  // 2. 检查并预扣额度
  const preConsume = await quotaManager.consume(
    tenantId,
    userId,
    'tokens',
    estimatedTokens,
    model,
    { type: 'estimate' }
  )
 
  if (!preConsume.success) {
    return c.json({
      error: 'Token 额度不足',
      quota: {
        used: preConsume.used,
        limit: preConsume.limit,
      },
    }, 429)
  }
 
  // 3. 调用 LLM
  const response = await callLLM(messages, model)
  const actualTokens = response.usage.total_tokens
 
  // 4. 修正额度（多退少补）
  const diff = actualTokens - estimatedTokens
  if (diff !== 0) {
    await quotaManager.consume(
      tenantId,
      userId,
      'tokens',
      diff,
      model,
      { type: 'correction' }
    )
  }
 
  return c.json({
    message: response.message,
    usage: {
      promptTokens: response.usage.prompt_tokens,
      completionTokens: response.usage.completion_tokens,
      totalTokens: response.usage.total_tokens,
    },
  })
})

// src/middleware/quota-check.ts
import { Context, Next } from 'hono'
import { QuotaManager } from '../lib/quota'
 
export function requireQuota(type: QuotaType, amount: number = 1) {
  return async (c: Context, next: Next) => {
    const tenantId = c.get('tenantId')
    const userId = c.get('userId')
 
    const quotaManager = new QuotaManager(c.env.DB, c.env)
 
    // 获取当前用量
    const used = await quotaManager.getCurrentUsage(tenantId, type)
    const limit = await quotaManager.getQuotaLimit(tenantId, type)
 
    // 检查是否超额
    if (limit > 0 && used + amount > limit) {
      return c.json({
        error: '额度不足',
        quota: {
          type,
          used,
          limit,
          remaining: Math.max(0, limit - used),
        },
      }, 429)
    }
 
    // 存储到上下文
    c.set('quotaUsed', used)
    c.set('quotaLimit', limit)
 
    await next()
  }
}
 
// 使用
app.post('/knowledge-bases',
  requireQuota('knowledge_bases', 1),
  async (c) => {
    // 创建知识库
    // 在成功后扣减额度
    const quotaManager = new QuotaManager(c.env.DB, c.env)
    await quotaManager.consume(
      c.get('tenantId'),
      c.get('userId'),
      'knowledge_bases',
      1
    )
  }
)

// src/routes/quota.ts
import { Hono } from 'hono'
import { QuotaManager } from '../lib/quota'
 
const app = new Hono()
 
// 获取所有额度信息
app.get('/', async (c) => {
  const tenantId = c.get('tenantId')
  const quotaManager = new QuotaManager(c.env.DB, c.env)
 
  const quotaTypes: QuotaType[] = [
    'tokens',
    'requests',
    'storage',
    'knowledge_bases',
    'members',
  ]
 
  const quotas = await Promise.all(
    quotaTypes.map(async (type) => {
      const used = await quotaManager.getCurrentUsage(tenantId, type)
      const limit = await quotaManager.getQuotaLimit(tenantId, type)
 
      return {
        type,
        used,
        limit,
        remaining: limit > 0 ? Math.max(0, limit - used) : -1,
        percentage: limit > 0 ? Math.round((used / limit) * 100) : 0,
      }
    })
  )
 
  return c.json({ quotas })
})
 
// 获取使用历史
app.get('/history', async (c) => {
  const tenantId = c.get('tenantId')
  const type = c.req.query('type')
  const days = parseInt(c.req.query('days') || '30')
 
  const startTime = Date.now() - days * 24 * 60 * 60 * 1000
 
  let query = `
    SELECT quota_type, SUM(amount) as total, created_at
    FROM quota_usage
    WHERE tenant_id = ? AND created_at > ?
  `
 
  const params: any[] = [tenantId, startTime]
 
  if (type) {
    query += ` AND quota_type = ?`
    params.push(type)
  }
 
  query += ` GROUP BY DATE(created_at / 1000, 'unixepoch'), quota_type
             ORDER BY created_at DESC`
 
  const history = await c.env.DB.prepare(query).bind(...params).all()
 
  return c.json({ history: history.results })
})

// src/lib/quota-alert.ts
export async function checkQuotaAlerts(
  tenantId: string,
  type: QuotaType,
  used: number,
  limit: number,
  env: Env
) {
  if (limit <= 0) return
 
  const percentage = (used / limit) * 100
 
  // 预警阈值：80%, 90%, 100%
  const thresholds = [80, 90, 100]
 
  for (const threshold of thresholds) {
    if (percentage >= threshold) {
      // 检查是否已发送过该阈值的预警
      const alertKey = `quota_alert:${tenantId}:${type}:${threshold}`
      const sent = await env.QUOTA_ALERT.get(alertKey)
 
      if (!sent) {
        // 发送预警
        await sendQuotaAlert(tenantId, type, used, limit, threshold, env)
 
        // 标记已发送（24 小时内不重复）
        await env.QUOTA_ALERT.put(alertKey, '1', {
          expirationTtl: 24 * 60 * 60,
        })
      }
    }
  }
}
 
async function sendQuotaAlert(
  tenantId: string,
  type: QuotaType,
  used: number,
  limit: number,
  threshold: number,
  env: Env
) {
  const tenant = await env.DB.prepare(`
    SELECT name FROM tenants WHERE id = ?
  `).bind(tenantId).first()
 
  const message = `租户 ${tenant.name} 的 ${type} 额度已达到 ${threshold}%`
 
  // 发送邮件
  await sendEmail({
    to: await getTenantAdminEmail(tenantId, env),
    subject: '额度预警',
    body: `${message}\n\n已使用：${used} / ${limit}`,
  })
 
  // 记录日志
  console.warn(`[Quota Alert] ${message}`, {
    tenantId,
    type,
    used,
    limit,
    percentage: Math.round((used / limit) * 100),
  })
}

// 在 QuotaManager.consume 中添加
async consume(...) {
  // ... 扣减逻辑
 
  // 检查预警
  await checkQuotaAlerts(tenantId, type, used + amount, limit, this.env)
 
  return { success: true, used: used + amount, limit }
}

// src/cron/quota-reset.ts
export async function resetMonthlyQuota(env: Env) {
  const now = new Date()
  const currentPeriodStart = new Date(
    now.getFullYear(),
    now.getMonth(),
    1
  ).getTime()
 
  // 获取所有活跃租户
  const tenants = await env.DB.prepare(`
    SELECT id FROM tenants WHERE status = 'active'
  `).all()
 
  for (const tenant of tenants.results) {
    // 为新周期创建初始快照
    const quotaTypes: QuotaType[] = ['tokens', 'requests']
 
    for (const type of quotaTypes) {
      await env.DB.prepare(`
        INSERT OR IGNORE INTO quota_snapshots (id, tenant_id, quota_type, period_start, total_used, updated_at)
        VALUES (?, ?, ?, ?, 0, ?)
      `).bind(
        crypto.randomUUID(),
        tenant.id,
        type,
        currentPeriodStart,
        Date.now()
      ).run()
    }
  }
 
  console.log(`[Cron] Monthly quota reset completed for ${tenants.results.length} tenants`)
}
 
// wrangler.toml
// [triggers]
// crons = ["0 0 1 * *"]  // 每月 1 号 00:00