koishi-plugin-aka-ai-generator 0.6.7 → 0.6.9

package/lib/index.js

@@ -983,6 +983,36 @@ var UserManager = class {
     }
     return { allowed: true, isAdmin: false };
   }
+  // 原子性地检查并预留额度（防止并发绕过）
+  async checkAndReserveQuota(userId, userName, numImages, config) {
+    if (this.isAdmin(userId, config)) {
+      return { allowed: true, reservationId: "admin" };
+    }
+    const rateLimitCheck = this.checkRateLimit(userId, config);
+    if (!rateLimitCheck.allowed) {
+      return { ...rateLimitCheck };
+    }
+    this.updateRateLimit(userId);
+    return await this.dataLock.acquire(async () => {
+      const userData = await this.getUserData(userId, userName);
+      const today = (/* @__PURE__ */ new Date()).toDateString();
+      const lastReset = new Date(userData.lastDailyReset || userData.createdAt).toDateString();
+      let dailyCount = userData.dailyUsageCount;
+      if (today !== lastReset) {
+        dailyCount = 0;
+      }
+      const remainingToday = Math.max(0, config.dailyFreeLimit - dailyCount);
+      const totalAvailable = remainingToday + userData.remainingPurchasedCount;
+      if (totalAvailable < numImages) {
+        return {
+          allowed: false,
+          message: `生成 ${numImages} 张图片需要 ${numImages} 次可用次数，但您的可用次数不足（今日免费剩余：${remainingToday}次，充值剩余：${userData.remainingPurchasedCount}次，共${totalAvailable}次）`
+        };
+      }
+      const reservationId = `${userId}_${Date.now()}_${Math.random()}`;
+      return { allowed: true, reservationId };
+    });
+  }
   // 扣减额度并记录使用
   async consumeQuota(userId, userName, commandName, numImages, config) {
     return await this.dataLock.acquire(async () => {
@@ -1059,7 +1089,7 @@ var UserManager = class {
     if (blockCount >= config.securityBlockWarningThreshold && !hasWarning) {
       this.securityWarningMap.set(userId, true);
       shouldWarn = true;
-    } else if (hasWarning) {
+    } else if (blockCount > config.securityBlockWarningThreshold) {
       shouldDeduct = true;
     }
     return { shouldWarn, shouldDeduct, blockCount };
@@ -1182,16 +1212,13 @@ var Config = import_koishi2.Schema.intersect([
       import_koishi2.Schema.const("gptgod").description("GPTGod 服务"),
       import_koishi2.Schema.const("gemini").description("Google Gemini 原生")
     ]).default("yunwu").description("图像生成供应商"),
-    // 所有字段都定义，但根据 provider 在运行时验证必填
-    yunwuApiKey: import_koishi2.Schema.string().description("云雾API密钥").role("secret").default(""),
+    yunwuApiKey: import_koishi2.Schema.string().description("云雾API密钥").role("secret").required(),
     yunwuModelId: import_koishi2.Schema.string().default("gemini-2.5-flash-image").description("云雾图像生成模型ID"),
     gptgodApiKey: import_koishi2.Schema.string().description("GPTGod API 密钥").role("secret").default(""),
     gptgodModelId: import_koishi2.Schema.string().default("nano-banana").description("GPTGod 模型ID"),
     geminiApiKey: import_koishi2.Schema.string().description("Gemini API 密钥").role("secret").default(""),
     geminiModelId: import_koishi2.Schema.string().default("gemini-2.5-flash").description("Gemini 模型ID"),
-    geminiApiBase: import_koishi2.Schema.string().default("https://generativelanguage.googleapis.com").description("Gemini API 基础地址")
-  }),
-  import_koishi2.Schema.object({
+    geminiApiBase: import_koishi2.Schema.string().default("https://generativelanguage.googleapis.com").description("Gemini API 基础地址"),
     modelMappings: import_koishi2.Schema.array(import_koishi2.Schema.object({
       suffix: import_koishi2.Schema.string().required().description("指令后缀（例如 4K，对应输入 -4K）"),
       provider: import_koishi2.Schema.union([
@@ -1242,15 +1269,6 @@ var Config = import_koishi2.Schema.intersect([
 ]);
 function apply(ctx, config) {
   const logger = ctx.logger("aka-ai-generator");
-  if (config.provider === "yunwu" && !config.yunwuApiKey) {
-    throw new Error('当选择"云雾 Gemini 服务"时，云雾API密钥为必填项');
-  }
-  if (config.provider === "gptgod" && !config.gptgodApiKey) {
-    throw new Error('当选择"GPTGod 服务"时，GPTGod API 密钥为必填项');
-  }
-  if (config.provider === "gemini" && !config.geminiApiKey) {
-    throw new Error('当选择"Google Gemini 原生"时，Gemini API 密钥为必填项');
-  }
   const userManager = new UserManager(ctx.baseDir, logger);
   function getProviderInstance(providerType, modelId) {
     return createImageProvider({
@@ -1330,34 +1348,36 @@ function apply(ctx, config) {
     return input || null;
   }
   __name(getPromptInput, "getPromptInput");
-  async function recordUserUsage(session, commandName, numImages = 1) {
-    const userId = session.userId;
-    const userName = session.username || session.userId || "未知用户";
-    if (!userId) return;
-    const { userData, consumptionType, freeUsed, purchasedUsed } = await userManager.consumeQuota(userId, userName, commandName, numImages, config);
-    if (userManager.isAdmin(userId, config)) {
-      await session.send(`📊 使用统计 [管理员]
+  function buildStatsMessage(userData, numImages, consumptionType, freeUsed, purchasedUsed, config2) {
+    if (userManager.isAdmin(userData.userId, config2)) {
+      return `📊 使用统计 [管理员]
 用户：${userData.userName}
 总调用次数：${userData.totalUsageCount}次
-状态：无限制使用`);
+状态：无限制使用`;
+    }
+    const remainingToday = Math.max(0, config2.dailyFreeLimit - userData.dailyUsageCount);
+    let consumptionText = "";
+    if (consumptionType === "mixed") {
+      consumptionText = `每日免费次数 -${freeUsed}，充值次数 -${purchasedUsed}`;
+    } else if (consumptionType === "free") {
+      consumptionText = `每日免费次数 -${freeUsed}`;
     } else {
-      const remainingToday = Math.max(0, config.dailyFreeLimit - userData.dailyUsageCount);
-      let consumptionText = "";
-      if (consumptionType === "mixed") {
-        consumptionText = `每日免费次数 -${freeUsed}，充值次数 -${purchasedUsed}`;
-      } else if (consumptionType === "free") {
-        consumptionText = `每日免费次数 -${freeUsed}`;
-      } else {
-        consumptionText = `充值次数 -${purchasedUsed}`;
-      }
-      await session.send(`📊 使用统计
+      consumptionText = `充值次数 -${purchasedUsed}`;
+    }
+    return `📊 使用统计
 用户：${userData.userName}
 本次生成：${numImages}张图片
 本次消费：${consumptionText}
 总调用次数：${userData.totalUsageCount}次
 今日剩余免费：${remainingToday}次
-充值剩余：${userData.remainingPurchasedCount}次`);
-    }
+充值剩余：${userData.remainingPurchasedCount}次`;
+  }
+  __name(buildStatsMessage, "buildStatsMessage");
+  async function recordUserUsage(session, commandName, numImages = 1) {
+    const userId = session.userId;
+    const userName = session.username || session.userId || "未知用户";
+    if (!userId) return;
+    const { userData, consumptionType, freeUsed, purchasedUsed } = await userManager.consumeQuota(userId, userName, commandName, numImages, config);
     logger.info("用户调用记录", {
       userId,
       userName: userData.userName,
@@ -1371,6 +1391,12 @@ function apply(ctx, config) {
       remainingPurchasedCount: userData.remainingPurchasedCount,
       isAdmin: userManager.isAdmin(userId, config)
     });
+    try {
+      const statsMessage = buildStatsMessage(userData, numImages, consumptionType, freeUsed, purchasedUsed, config);
+      await session.send(statsMessage);
+    } catch (error) {
+      logger.warn("发送统计信息失败", { userId, error: sanitizeError(error) });
+    }
   }
   __name(recordUserUsage, "recordUserUsage");
   async function recordSecurityBlock(session, numImages = 1) {
@@ -1505,7 +1531,6 @@ function apply(ctx, config) {
         }, config.commandTimeout * 1e3)
       )
     ]).catch(async (error) => {
-      if (userId) userManager.endTask(userId);
       const sanitizedError = sanitizeError(error);
       logger.error("图像处理超时或失败", { userId, error: sanitizedError });
       if (error?.message !== "命令执行超时") {
@@ -1526,18 +1551,13 @@ function apply(ctx, config) {
     if (!userManager.startTask(userId)) {
       return "您有一个图像处理任务正在进行中，请等待完成";
     }
-    let taskLockReleased = false;
     try {
       const imageCount = requestContext?.numImages || config.defaultNumImages;
       if (imageCount < 1 || imageCount > 4) {
-        userManager.endTask(userId);
-        taskLockReleased = true;
         return "生成数量必须在 1-4 之间";
       }
       const inputResult = await getInputData(session, img, mode);
       if ("error" in inputResult) {
-        userManager.endTask(userId);
-        taskLockReleased = true;
         return inputResult.error;
       }
       if (checkTimeout && checkTimeout()) throw new Error("命令执行超时");
@@ -1551,23 +1571,17 @@ function apply(ctx, config) {
         await session.send("请发送画面描述");
         const promptMsg = await session.prompt(3e4);
         if (!promptMsg) {
-          userManager.endTask(userId);
-          taskLockReleased = true;
           return "未检测到描述，操作已取消";
         }
         const elements = import_koishi2.h.parse(promptMsg);
         const images2 = import_koishi2.h.select(elements, "img");
         if (images2.length > 0) {
-          userManager.endTask(userId);
-          taskLockReleased = true;
           return "检测到图片，本功能仅支持文字输入";
         }
         const text = import_koishi2.h.select(elements, "text").map((e) => e.attrs.content).join(" ").trim();
         if (text) {
           finalPrompt = text;
         } else {
-          userManager.endTask(userId);
-          taskLockReleased = true;
           return "未检测到有效文字描述，操作已取消";
         }
       }
@@ -1601,55 +1615,20 @@ ${infoParts.join("\n")}`;
       const images = await requestProviderImages(finalPrompt, imageUrls, imageCount, requestContext);
       if (checkTimeout && checkTimeout()) throw new Error("命令执行超时");
       if (images.length === 0) {
-        userManager.endTask(userId);
-        taskLockReleased = true;
         return "图像处理失败：未能生成图片";
       }
-      try {
-        await recordUserUsage(session, styleName, images.length);
-      } catch (err) {
-        logger.error(`记录用户使用统计失败 (用户: ${userId}):`, err);
-      }
-      const resultMessage = "图像处理完成！";
-      taskLockReleased = true;
-      setImmediate(async () => {
-        try {
-          await session.send(resultMessage);
-          for (let i = 0; i < images.length; i++) {
-            try {
-              await Promise.race([
-                session.send(import_koishi2.h.image(images[i])),
-                new Promise((_, reject) => setTimeout(() => reject(new Error("SendTimeout")), 2e4))
-              ]);
-            } catch (err) {
-              logger.warn(`图片发送可能超时 (用户: ${userId}, 图片 ${i + 1}/${images.length}): ${err instanceof Error ? err.message : String(err)}`);
-              try {
-                await session.send(`⚠️ 图片 ${i + 1} 发送失败，请重试`);
-              } catch (sendErr) {
-                logger.error(`发送错误通知失败 (用户: ${userId}):`, sendErr);
-              }
-            }
-            if (images.length > 1 && i < images.length - 1) {
-              await new Promise((resolve) => setTimeout(resolve, 1e3));
-            }
-          }
-        } catch (err) {
-          logger.error(`异步发送图片时出错 (用户: ${userId}):`, err);
-          try {
-            await session.send("⚠️ 图片发送过程中出现错误，但积分已扣除。如未收到图片，请联系管理员。");
-          } catch (sendErr) {
-            logger.error(`发送错误通知失败 (用户: ${userId}):`, sendErr);
-          }
-        } finally {
-          userManager.endTask(userId);
+      await recordUserUsage(session, styleName, images.length);
+      if (checkTimeout && checkTimeout()) throw new Error("命令执行超时");
+      await session.send("图像处理完成！");
+      for (let i = 0; i < images.length; i++) {
+        if (checkTimeout && checkTimeout()) break;
+        await session.send(import_koishi2.h.image(images[i]));
+        if (images.length > 1 && i < images.length - 1) {
+          await new Promise((resolve) => setTimeout(resolve, 1e3));
         }
-      });
-      return "";
-    } catch (error) {
-      if (!taskLockReleased) {
-        userManager.endTask(userId);
       }
-      throw error;
+    } finally {
+      userManager.endTask(userId);
     }
   }
   __name(processImage, "processImage");
@@ -1733,11 +1712,9 @@ ${infoParts.join("\n")}`;
     if (!userManager.startTask(userId)) {
       return "您有一个图像处理任务正在进行中，请等待完成";
     }
-    userManager.endTask(userId);
     let isTimeout = false;
     return Promise.race([
       (async () => {
-        if (!userManager.startTask(userId)) return "您有一个图像处理任务正在进行中";
         try {
           await session.send("多张图片+描述");
           const collectedImages = [];
@@ -1801,48 +1778,18 @@ Prompt: ${prompt}`);
           if (resultImages.length === 0) {
             return "图片合成失败：未能生成图片";
           }
-          try {
-            await recordUserUsage(session, COMMANDS.COMPOSE_IMAGE, resultImages.length);
-          } catch (err) {
-            logger.error(`记录用户使用统计失败 (用户: ${userId}):`, err);
-          }
-          const resultMessage = "图片合成完成！";
-          setImmediate(async () => {
-            try {
-              await session.send(resultMessage);
-              for (let i = 0; i < resultImages.length; i++) {
-                try {
-                  await Promise.race([
-                    session.send(import_koishi2.h.image(resultImages[i])),
-                    new Promise((_, reject) => setTimeout(() => reject(new Error("SendTimeout")), 2e4))
-                  ]);
-                } catch (err) {
-                  logger.warn(`图片合成发送可能超时 (用户: ${userId}, 图片 ${i + 1}/${resultImages.length}): ${err instanceof Error ? err.message : String(err)}`);
-                  try {
-                    await session.send(`⚠️ 图片 ${i + 1} 发送失败，请重试`);
-                  } catch (sendErr) {
-                    logger.error(`发送错误通知失败 (用户: ${userId}):`, sendErr);
-                  }
-                }
-                if (resultImages.length > 1 && i < resultImages.length - 1) {
-                  await new Promise((resolve) => setTimeout(resolve, 1e3));
-                }
-              }
-            } catch (err) {
-              logger.error(`异步发送合成图片时出错 (用户: ${userId}):`, err);
-              try {
-                await session.send("⚠️ 图片发送过程中出现错误，但积分已扣除。如未收到图片，请联系管理员。");
-              } catch (sendErr) {
-                logger.error(`发送错误通知失败 (用户: ${userId}):`, sendErr);
-              }
-            } finally {
-              userManager.endTask(userId);
+          await recordUserUsage(session, COMMANDS.COMPOSE_IMAGE, resultImages.length);
+          if (isTimeout) throw new Error("命令执行超时");
+          await session.send("图片合成完成！");
+          for (let i = 0; i < resultImages.length; i++) {
+            if (isTimeout) break;
+            await session.send(import_koishi2.h.image(resultImages[i]));
+            if (resultImages.length > 1 && i < resultImages.length - 1) {
+              await new Promise((resolve) => setTimeout(resolve, 1e3));
             }
-          });
-          return "";
-        } catch (error) {
+          }
+        } finally {
           userManager.endTask(userId);
-          throw error;
         }
       })(),
       new Promise(
@@ -1852,7 +1799,6 @@ Prompt: ${prompt}`);
         }, config.commandTimeout * 1e3)
       )
     ]).catch(async (error) => {
-      if (userId) userManager.endTask(userId);
       const sanitizedError = sanitizeError(error);
       logger.error("图片合成超时或失败", { userId, error: sanitizedError });
       if (error?.message !== "命令执行超时") {

package/lib/services/UserManager.d.ts

@@ -74,6 +74,11 @@ export declare class UserManager {
         message?: string;
         isAdmin?: boolean;
     }>;
+    checkAndReserveQuota(userId: string, userName: string, numImages: number, config: Config): Promise<{
+        allowed: boolean;
+        message?: string;
+        reservationId?: string;
+    }>;
     consumeQuota(userId: string, userName: string, commandName: string, numImages: number, config: Config): Promise<{
         userData: UserData;
         consumptionType: 'free' | 'purchased' | 'mixed';

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "koishi-plugin-aka-ai-generator",
   "description": "自用AI生成插件(GPTGod & Yunwu)",
-  "version": "0.6.7",
+  "version": "0.6.9",
   "main": "lib/index.js",
   "typings": "lib/index.d.ts",
   "files": [