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

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

---

## 一、功能概述

### 1.1 需求背景

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

### 1.2 功能目标

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

### 1.3 核心规则

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

---

## 二、数据库设计

### 2.1 表结构变更

**修改 `User` 表，添加新字段：**

```sql
-- 迁移文件：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`**

```typescript
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`**

```typescript
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` 函数

```typescript
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`**

```typescript
// 更新 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`**

```typescript
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`**

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

```typescript
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`**

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

### 4.2 用户列表显示配额

```tsx
{/* 用户列表 */}
{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 编辑用户对话框

```tsx
{/* 编辑用户对话框 */}
<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`**

```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` | 修改（可选） |

---

*文档结束*