koishi-plugin-aka-ai-generator 0.6.8 → 0.6.10

package/lib/index.js

@@ -295,7 +295,7 @@ var GptGodProvider = class {
   constructor(config) {
     this.config = config;
   }
-  async generateImages(prompt, imageUrls, numImages) {
+  async generateImages(prompt, imageUrls, numImages, onImageGenerated) {
     const urls = Array.isArray(imageUrls) ? imageUrls : [imageUrls];
     const logger = this.config.logger;
     const ctx = this.config.ctx;
@@ -410,7 +410,18 @@ var GptGodProvider = class {
               throw new Error(`生成失败：${shortError}`);
             }
           }
-          allImages.push(...images);
+          for (let imgIdx = 0; imgIdx < images.length; imgIdx++) {
+            const imageUrl = images[imgIdx];
+            const currentIndex = allImages.length;
+            allImages.push(imageUrl);
+            if (onImageGenerated) {
+              try {
+                await onImageGenerated(imageUrl, currentIndex, numImages);
+              } catch (callbackError) {
+                logger.warn("图片生成回调函数执行失败", { error: sanitizeError(callbackError) });
+              }
+            }
+          }
           logger.success("GPTGod 图像编辑 API 调用成功", { current: i + 1, total: numImages });
           success = true;
           break;
@@ -609,7 +620,7 @@ var GeminiProvider = class {
   constructor(config) {
     this.config = config;
   }
-  async generateImages(prompt, imageUrls, numImages) {
+  async generateImages(prompt, imageUrls, numImages, onImageGenerated) {
     let urls = [];
     if (Array.isArray(imageUrls)) {
       urls = imageUrls.filter((url) => url && typeof url === "string" && url.trim());
@@ -682,8 +693,19 @@ var GeminiProvider = class {
           });
         } else {
           logger.success("Gemini API 调用成功", { current: i + 1, total: numImages, imagesCount: images.length });
+          for (let imgIdx = 0; imgIdx < images.length; imgIdx++) {
+            const imageUrl = images[imgIdx];
+            const currentIndex = allImages.length;
+            allImages.push(imageUrl);
+            if (onImageGenerated) {
+              try {
+                await onImageGenerated(imageUrl, currentIndex, numImages);
+              } catch (callbackError) {
+                logger.warn("图片生成回调函数执行失败", { error: sanitizeError(callbackError) });
+              }
+            }
+          }
         }
-        allImages.push(...images);
       } catch (error) {
         const sanitizedError = sanitizeError(error);
         const safeMessage = typeof error?.message === "string" ? sanitizeString(error.message) : "未知错误";
@@ -983,6 +1005,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 +1111,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 };
@@ -1318,34 +1370,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,
@@ -1359,6 +1413,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) {
@@ -1467,7 +1527,7 @@ function apply(ctx, config) {
     return { images: collectedImages, text: collectedText };
   }
   __name(getInputData, "getInputData");
-  async function requestProviderImages(prompt, imageUrls, numImages, requestContext) {
+  async function requestProviderImages(prompt, imageUrls, numImages, requestContext, onImageGenerated) {
     const providerType = requestContext?.provider || config.provider;
     const targetModelId = requestContext?.modelId;
     const providerInstance = getProviderInstance(providerType, targetModelId);
@@ -1478,7 +1538,7 @@ function apply(ctx, config) {
         numImages
       });
     }
-    return await providerInstance.generateImages(prompt, imageUrls, numImages);
+    return await providerInstance.generateImages(prompt, imageUrls, numImages, onImageGenerated);
   }
   __name(requestProviderImages, "requestProviderImages");
   async function processImageWithTimeout(session, img, prompt, styleName, requestContext, displayInfo, mode = "single") {
@@ -1493,7 +1553,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 !== "命令执行超时") {
@@ -1575,20 +1634,38 @@ ${infoParts.join("\n")}`;
       }
       statusMessage += "...";
       await session.send(statusMessage);
-      const images = await requestProviderImages(finalPrompt, imageUrls, imageCount, requestContext);
+      const generatedImages = [];
+      let creditDeducted = false;
+      const onImageGenerated = /* @__PURE__ */ __name(async (imageUrl, index, total) => {
+        if (checkTimeout && checkTimeout()) {
+          throw new Error("命令执行超时");
+        }
+        generatedImages.push(imageUrl);
+        if (!creditDeducted && generatedImages.length > 0) {
+          creditDeducted = true;
+          await recordUserUsage(session, styleName, total);
+          logger.info("流式处理：第一张图片生成，积分已扣除", {
+            userId,
+            totalImages: total,
+            currentIndex: index
+          });
+        }
+        await session.send(import_koishi2.h.image(imageUrl));
+        logger.debug("流式处理：图片已发送", { index: index + 1, total });
+        if (total > 1 && index < total - 1) {
+          await new Promise((resolve) => setTimeout(resolve, 1e3));
+        }
+      }, "onImageGenerated");
+      const images = await requestProviderImages(finalPrompt, imageUrls, imageCount, requestContext, onImageGenerated);
       if (checkTimeout && checkTimeout()) throw new Error("命令执行超时");
       if (images.length === 0) {
         return "图像处理失败：未能生成图片";
       }
-      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));
-        }
+      if (!creditDeducted) {
+        await recordUserUsage(session, styleName, images.length);
+        logger.warn("流式处理：积分在最后扣除（异常情况）", { userId, imagesCount: images.length });
       }
-      await recordUserUsage(session, styleName, images.length);
+      await session.send("图像处理完成！");
     } finally {
       userManager.endTask(userId);
     }
@@ -1610,7 +1687,8 @@ ${infoParts.join("\n")}`;
           }
           const userPromptText = userPromptParts.join(" - ");
           const numImages = options?.num || config.defaultNumImages;
-          const limitCheck = await userManager.checkDailyLimit(session.userId, config, numImages);
+          const userName = session.username || session.userId || "未知用户";
+          const limitCheck = await userManager.checkAndReserveQuota(session.userId, userName, numImages, config);
           if (!limitCheck.allowed) {
             return limitCheck.message;
           }
@@ -1646,7 +1724,8 @@ ${infoParts.join("\n")}`;
   ctx.command(`${COMMANDS.TXT_TO_IMG} [prompt:text]`, "根据文字描述生成图像").option("num", "-n <num:number> 生成图片数量 (1-4)").action(async ({ session, options }, prompt) => {
     if (!session?.userId) return "会话无效";
     const numImages = options?.num || config.defaultNumImages;
-    const limitCheck = await userManager.checkDailyLimit(session.userId, config, numImages);
+    const userName = session.username || session.userId || "未知用户";
+    const limitCheck = await userManager.checkAndReserveQuota(session.userId, userName, numImages, config);
     if (!limitCheck.allowed) {
       return limitCheck.message;
     }
@@ -1659,7 +1738,8 @@ ${infoParts.join("\n")}`;
     if (!session?.userId) return "会话无效";
     const numImages = options?.num || config.defaultNumImages;
     const mode = options?.multiple ? "multiple" : "single";
-    const limitCheck = await userManager.checkDailyLimit(session.userId, config, numImages);
+    const userName = session.username || session.userId || "未知用户";
+    const limitCheck = await userManager.checkAndReserveQuota(session.userId, userName, numImages, config);
     if (!limitCheck.allowed) {
       return limitCheck.message;
     }
@@ -1674,11 +1754,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 = [];
@@ -1723,7 +1801,8 @@ ${infoParts.join("\n")}`;
           if (imageCount < 1 || imageCount > 4) {
             return "生成数量必须在 1-4 之间";
           }
-          const limitCheck = await userManager.checkDailyLimit(userId, config, imageCount);
+          const userName = session.username || userId || "未知用户";
+          const limitCheck = await userManager.checkAndReserveQuota(userId, userName, imageCount, config);
           if (!limitCheck.allowed) {
             return limitCheck.message;
           }
@@ -1737,20 +1816,38 @@ ${infoParts.join("\n")}`;
           });
           await session.send(`开始合成图（${collectedImages.length}张）...
 Prompt: ${prompt}`);
-          const resultImages = await requestProviderImages(prompt, collectedImages, imageCount);
+          const generatedImages = [];
+          let creditDeducted = false;
+          const onImageGenerated = /* @__PURE__ */ __name(async (imageUrl, index, total) => {
+            if (isTimeout) {
+              throw new Error("命令执行超时");
+            }
+            generatedImages.push(imageUrl);
+            if (!creditDeducted && generatedImages.length > 0) {
+              creditDeducted = true;
+              await recordUserUsage(session, COMMANDS.COMPOSE_IMAGE, total);
+              logger.info("流式处理：第一张图片生成，积分已扣除", {
+                userId,
+                totalImages: total,
+                currentIndex: index
+              });
+            }
+            await session.send(import_koishi2.h.image(imageUrl));
+            logger.debug("流式处理：图片已发送", { index: index + 1, total });
+            if (total > 1 && index < total - 1) {
+              await new Promise((resolve) => setTimeout(resolve, 1e3));
+            }
+          }, "onImageGenerated");
+          const resultImages = await requestProviderImages(prompt, collectedImages, imageCount, void 0, onImageGenerated);
           if (isTimeout) throw new Error("命令执行超时");
           if (resultImages.length === 0) {
             return "图片合成失败：未能生成图片";
           }
-          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));
-            }
+          if (!creditDeducted) {
+            await recordUserUsage(session, COMMANDS.COMPOSE_IMAGE, resultImages.length);
+            logger.warn("流式处理：积分在最后扣除（异常情况）", { userId, imagesCount: resultImages.length });
           }
-          await recordUserUsage(session, COMMANDS.COMPOSE_IMAGE, resultImages.length);
+          await session.send("图片合成完成！");
         } finally {
           userManager.endTask(userId);
         }
@@ -1762,7 +1859,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/providers/gemini.d.ts

@@ -7,5 +7,5 @@ export interface GeminiConfig extends ProviderConfig {
 export declare class GeminiProvider implements ImageProvider {
     private config;
     constructor(config: GeminiConfig);
-    generateImages(prompt: string, imageUrls: string | string[], numImages: number): Promise<string[]>;
+    generateImages(prompt: string, imageUrls: string | string[], numImages: number, onImageGenerated?: (imageUrl: string, index: number, total: number) => void | Promise<void>): Promise<string[]>;
 }

package/lib/providers/gptgod.d.ts

@@ -6,5 +6,5 @@ export interface GptGodConfig extends ProviderConfig {
 export declare class GptGodProvider implements ImageProvider {
     private config;
     constructor(config: GptGodConfig);
-    generateImages(prompt: string, imageUrls: string | string[], numImages: number): Promise<string[]>;
+    generateImages(prompt: string, imageUrls: string | string[], numImages: number, onImageGenerated?: (imageUrl: string, index: number, total: number) => void | Promise<void>): Promise<string[]>;
 }

package/lib/providers/types.d.ts

@@ -4,9 +4,10 @@ export interface ImageProvider {
      * @param prompt 提示词
      * @param imageUrls 输入图片 URL 数组
      * @param numImages 需要生成的图片数量
+     * @param onImageGenerated 可选的回调函数，每生成一张图片时调用（用于流式处理）
      * @returns 生成的图片 URL 数组（data: URL 或 http URL）
      */
-    generateImages(prompt: string, imageUrls: string | string[], numImages: number): Promise<string[]>;
+    generateImages(prompt: string, imageUrls: string | string[], numImages: number, onImageGenerated?: (imageUrl: string, index: number, total: number) => void | Promise<void>): Promise<string[]>;
 }
 export interface ProviderConfig {
     apiTimeout: number;

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.8",
+  "version": "0.6.10",
   "main": "lib/index.js",
   "typings": "lib/index.d.ts",
   "files": [