ai-chatbot/docs/每月出图限额功能实施文档.md

每月出图限额功能实施文档

文档版本：v1.0
创建日期：2024-12-18
状态：待实施

---

一、功能概述

1.1 需求背景

为了更好地管理系统资源和控制成本，需要在系统设置的"用户管理"中增加限制每个用户每月出图数量的功能。

1.2 功能目标

• 管理员可为每个用户设置每月图片生成限额
• 用户达到限额后无法继续生成图片，给予友好提示
• 用户可查看自己的配额使用情况
• 每月自动重置配额（按自然月统计）

1.3 核心规则

规则	说明
默认值	null 或 0 表示无限制
统计周期	自然月（每月1日0点自动重置）
管理员	可设置为无限制
计数对象	GeneratedImage 表中的记录

---

二、数据库设计

2.1 表结构变更

修改 User 表，添加新字段：

-- 迁移文件：0022_add_monthly_image_limit.sql

-- 添加每月出图限额字段
ALTER TABLE "User" ADD COLUMN IF NOT EXISTS "monthlyImageLimit" integer DEFAULT NULL;

-- 添加字段注释
COMMENT ON COLUMN "User"."monthlyImageLimit" IS '每月图片生成限额，NULL或0表示无限制';

2.2 Schema 定义更新

文件：lib/db/schema.ts

export const user = pgTable("User", {
  id: uuid("id").primaryKey().notNull().defaultRandom(),
  email: varchar("email", { length: 64 }).notNull(),
  password: varchar("password", { length: 64 }),
  name: varchar("name", { length: 100 }),
  role: varchar("role", { enum: ["guest", "regular", "admin"] })
    .notNull()
    .default("regular"),
  isActive: boolean("isActive").notNull().default(true),
  monthlyImageLimit: integer("monthlyImageLimit"), // 🆕 每月出图限额
  createdAt: timestamp("createdAt").notNull().defaultNow(),
  updatedAt: timestamp("updatedAt").notNull().defaultNow(),
});

---

三、后端实现

3.1 配额查询函数

文件：lib/db/queries.ts

import { generatedImage } from "./schema";

/**
 * 获取用户当月已生成图片数量
 */
export async function getUserMonthlyImageCount(userId: string): Promise<number> {
  const now = new Date();
  const startOfMonth = new Date(now.getFullYear(), now.getMonth(), 1);
  
  const result = await db
    .select({ count: sql<number>`count(*)::int` })
    .from(generatedImage)
    .where(and(
      eq(generatedImage.userId, userId),
      gte(generatedImage.createdAt, startOfMonth)
    ));
  
  return result[0]?.count ?? 0;
}

/**
 * 检查用户图片生成配额
 * @returns 配额状态对象
 */
export async function checkUserImageQuota(userId: string): Promise<{
  allowed: boolean;      // 是否允许继续生成
  used: number;          // 本月已使用数量
  limit: number | null;  // 限额（null表示无限制）
  remaining: number | null; // 剩余数量
}> {
  // 获取用户信息
  const [userData] = await db
    .select({ monthlyImageLimit: user.monthlyImageLimit })
    .from(user)
    .where(eq(user.id, userId));
  
  const limit = userData?.monthlyImageLimit;
  const used = await getUserMonthlyImageCount(userId);
  
  // null 或 0 表示无限制
  if (!limit || limit === 0) {
    return { 
      allowed: true, 
      used, 
      limit: null, 
      remaining: null 
    };
  }
  
  return {
    allowed: used < limit,
    used,
    limit,
    remaining: Math.max(0, limit - used),
  };
}

3.2 更新用户函数

文件：lib/db/queries.ts - 修改 updateUser 函数

export async function updateUser({
  id,
  email,
  role,
  isActive,
  monthlyImageLimit, // 🆕 新增参数
}: {
  id: string;
  email?: string;
  role?: "guest" | "regular" | "admin";
  isActive?: boolean;
  monthlyImageLimit?: number | null; // 🆕 null 表示无限制
}) {
  try {
    const updateData: Record<string, unknown> = { updatedAt: new Date() };
    
    if (email !== undefined) updateData.email = email;
    if (role !== undefined) updateData.role = role;
    if (isActive !== undefined) updateData.isActive = isActive;
    if (monthlyImageLimit !== undefined) {
      // 0 或空字符串转为 null
      updateData.monthlyImageLimit = monthlyImageLimit || null;
    }

    return await db
      .update(user)
      .set(updateData)
      .where(eq(user.id, id))
      .returning();
  } catch (_error) {
    throw new ChatSDKError("bad_request:database", "Failed to update user");
  }
}

3.3 用户管理 API 更新

文件：app/api/admin/users/route.ts

// 更新 Schema 验证
const updateUserSchema = z.object({
  id: z.string().uuid(),
  role: z.enum(["guest", "regular", "admin"]).optional(),
  isActive: z.boolean().optional(),
  email: z.string().email().optional(),
  monthlyImageLimit: z.number().int().min(0).nullable().optional(), // 🆕
});

// GET 请求 - 返回用户列表时包含配额信息
export async function GET(request: Request) {
  // ... 现有代码 ...
  
  // 返回时包含 monthlyImageLimit 字段
  const users = await db
    .select({
      id: user.id,
      email: user.email,
      name: user.name,
      role: user.role,
      isActive: user.isActive,
      monthlyImageLimit: user.monthlyImageLimit, // 🆕
      createdAt: user.createdAt,
    })
    .from(user)
    // ...
}

3.4 用户配额查询 API

新建文件：app/api/user/image-quota/route.ts

import { NextResponse } from "next/server";
import { auth } from "@/app/(auth)/auth";
import { checkUserImageQuota } from "@/lib/db/queries";

export async function GET() {
  try {
    const session = await auth();

    if (!session?.user?.id) {
      return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
    }

    const quota = await checkUserImageQuota(session.user.id);

    return NextResponse.json({
      used: quota.used,
      limit: quota.limit,
      remaining: quota.remaining,
      unlimited: quota.limit === null,
      // 计算百分比（用于进度条显示）
      percentage: quota.limit ? Math.round((quota.used / quota.limit) * 100) : 0,
    });
  } catch (error) {
    console.error("[User Image Quota API] Error:", error);
    return NextResponse.json(
      { error: "Failed to fetch quota" },
      { status: 500 }
    );
  }
}

3.5 图片生成前配额检查

文件：app/(chat)/api/chat/route.ts

在调用生图模型之前添加配额检查：

import { checkUserImageQuota } from "@/lib/db/queries";

// 在处理图片生成请求时
if (isImageModel || model.supportsImageGen) {
  // 检查用户配额
  const quota = await checkUserImageQuota(userId);
  
  if (!quota.allowed) {
    // 返回友好的错误提示
    const errorMessage = `⚠️ 本月图片生成配额已用完\n\n` +
      `已使用：${quota.used} 张\n` +
      `月度限额：${quota.limit} 张\n\n` +
      `配额将于下月1日自动重置，如需更多配额请联系管理员。`;
    
    return new Response(errorMessage, { 
      status: 429,
      headers: { "Content-Type": "text/plain; charset=utf-8" }
    });
  }
}

---

四、前端实现

4.1 类型定义更新

文件：app/(admin)/settings/admin-settings-client.tsx

type User = {
  id: string;
  email: string;
  name: string | null;
  role: "guest" | "regular" | "admin";
  isActive: boolean;
  monthlyImageLimit: number | null; // 🆕
  createdAt: string;
};

4.2 用户列表显示配额

{/* 用户列表 */}
{usersData?.users?.map((user) => (
  <div key={user.id} className="flex items-center justify-between p-4 border rounded-lg">
    <div className="flex-1">
      <div className="font-medium">{user.name || user.email}</div>
      <div className="text-sm text-muted-foreground">
        邮箱: {user.email}
      </div>
      <div className="text-sm text-muted-foreground">
        角色: {ROLE_LABELS[user.role]} | 状态: {user.isActive ? "已激活" : "未激活"}
      </div>
      {/* 🆕 显示配额信息 */}
      <div className="text-sm text-muted-foreground">
        月度出图限额: {user.monthlyImageLimit ? `${user.monthlyImageLimit} 张/月` : "无限制"}
      </div>
    </div>
    {/* 操作按钮 */}
  </div>
))}

4.3 编辑用户对话框

{/* 编辑用户对话框 */}
<Dialog open={isEditUserOpen} onOpenChange={setIsEditUserOpen}>
  <DialogContent>
    <DialogHeader>
      <DialogTitle>编辑用户</DialogTitle>
    </DialogHeader>
    
    {/* 现有字段... */}
    
    {/* 🆕 每月出图限额 */}
    <div className="space-y-2">
      <Label htmlFor="monthlyImageLimit">每月出图限额</Label>
      <div className="flex items-center gap-2">
        <Input
          id="monthlyImageLimit"
          type="number"
          min="0"
          placeholder="留空或填0表示无限制"
          value={editingUser?.monthlyImageLimit || ""}
          onChange={(e) => setEditingUser({
            ...editingUser,
            monthlyImageLimit: e.target.value ? parseInt(e.target.value) : null
          })}
          className="flex-1"
        />
        <span className="text-sm text-muted-foreground">张/月</span>
      </div>
      <p className="text-xs text-muted-foreground">
        留空或填写 0 表示无限制
      </p>
    </div>
    
    <DialogFooter>
      <Button onClick={handleUpdateUser}>保存</Button>
    </DialogFooter>
  </DialogContent>
</Dialog>

4.4 用户首页配额显示（可选）

文件：components/greeting.tsx

// 获取配额信息
const [quota, setQuota] = useState<{
  used: number;
  limit: number | null;
  percentage: number;
} | null>(null);

useEffect(() => {
  const fetchQuota = async () => {
    try {
      const response = await fetch("/api/user/image-quota");
      if (response.ok) {
        const data = await response.json();
        setQuota(data);
      }
    } catch (error) {
      console.error("Failed to fetch quota:", error);
    }
  };
  
  if (session?.user?.id) {
    fetchQuota();
  }
}, [session?.user?.id]);

// 显示配额进度条
{quota && quota.limit && (
  <motion.div
    animate={{ opacity: 1, y: 0 }}
    initial={{ opacity: 0, y: 10 }}
    transition={{ delay: 0.7 }}
    className="mt-4"
  >
    <div className="text-sm text-zinc-500 mb-1">
      本月出图：{quota.used}/{quota.limit} 张
    </div>
    <div className="w-64 h-2 bg-zinc-200 rounded-full overflow-hidden">
      <div 
        className={`h-full rounded-full transition-all ${
          quota.percentage >= 90 ? 'bg-red-500' : 
          quota.percentage >= 70 ? 'bg-yellow-500' : 'bg-green-500'
        }`}
        style={{ width: `${Math.min(quota.percentage, 100)}%` }}
      />
    </div>
  </motion.div>
)}

---

五、实施步骤

阶段一：数据库迁移（预计 10 分钟）

1. 创建迁移文件 0022_add_monthly_image_limit.sql
2. 更新 lib/db/schema.ts
3. 在服务器执行迁移

阶段二：后端开发（预计 30 分钟）

1. 添加 getUserMonthlyImageCount 函数
2. 添加 checkUserImageQuota 函数
3. 修改 updateUser 函数
4. 更新 /api/admin/users 路由
5. 创建 /api/user/image-quota 接口
6. 在图片生成流程中添加配额检查

阶段三：前端开发（预计 45 分钟）

1. 更新 User 类型定义
2. 用户列表显示配额信息
3. 编辑用户对话框添加限额设置
4. （可选）首页显示配额进度

阶段四：测试与部署（预计 30 分钟）

1. 本地功能测试
2. 部署到服务器
3. 生产环境验证

---

六、测试用例

6.1 功能测试

测试场景	预期结果
设置限额为 10，生成第 11 张图	显示配额已用完提示
设置限额为 0	无限制，可无限生成
留空限额	无限制，可无限生成
月初查看配额	已使用数量重置为 0
管理员查看用户列表	显示每个用户的限额设置

6.2 边界测试

测试场景	预期结果
负数限额	表单验证拒绝
超大数字限额	正常保存
并发生成（接近限额时）	正确计数，不会超额

---

七、风险与注意事项

7.1 兼容性

• ✅ 现有用户默认无限制（null）
• ✅ 不影响现有功能
• ✅ 无需停机迁移

7.2 性能考虑

• 配额检查会增加一次数据库查询
• 可考虑使用 Redis 缓存配额信息（后续优化）

7.3 安全性

• 配额修改仅管理员可操作
• 配额检查在服务端执行，不可绕过

---

八、后续优化建议

1. 批量设置：支持批量为多个用户设置相同限额
2. 配额预警：用户达到 80% 限额时发送提醒
3. 历史统计：记录每月配额使用历史
4. 角色默认值：为不同角色设置默认限额模板
5. 临时额度：支持临时增加额度功能

---

九、相关文件清单

文件路径	操作
lib/db/schema.ts	修改
lib/db/queries.ts	修改
lib/db/migrations/0022_add_monthly_image_limit.sql	新建
app/api/admin/users/route.ts	修改
app/api/user/image-quota/route.ts	新建
app/(chat)/api/chat/route.ts	修改
app/(admin)/settings/admin-settings-client.tsx	修改
components/greeting.tsx	修改（可选）

---

文档结束