From fe5a839b37dc392852fb547b1121a406b29666b7 Mon Sep 17 00:00:00 2001
From: Stringadmin <admin@omniai.net.cn>
Date: Mon, 8 Jun 2026 14:47:27 +0800
Subject: [PATCH] fix: harden generation task polling fallback

---
 src/App.tsx                      |  5 +++++
 src/api/generationConcurrency.ts | 17 ++++++++++++++---
 src/api/keyServerClient.ts       |  1 +
 src/api/taskSubscription.ts      |  4 +---
 src/types.ts                     |  1 +
 5 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/App.tsx b/src/App.tsx
index ba85629..77b1cce 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -8,6 +8,7 @@ import ToastContainer from "./components/toast/ToastContainer";
 import { toast } from "./components/toast/toastStore";
 import { aiGenerationClient } from "./api/aiGenerationClient";
 import { keyServerClient } from "./api/keyServerClient";
+import { setUserMaxConcurrency } from "./api/generationConcurrency";
 import { notificationClient } from "./api/notificationClient";
 import {
   SERVER_SESSION_REPLACED_EVENT,
@@ -466,6 +467,7 @@ function App() {
   const clearAuthenticatedState = useCallback((options?: { resetView?: boolean }) => {
     clearAllUserStorage();
     clearSessionState();
+    setUserMaxConcurrency(null);
     setProjects([]);
     setProjectsLoaded(true);
     setUsage(emptyUsageSummary);
@@ -574,6 +576,7 @@ function App() {
       const nextSession = await keyServerClient.getCurrentSession();
       if (cancelled) return;
       setSession(nextSession);
+      setUserMaxConcurrency(nextSession?.user?.maxConcurrency);
       await hydrateAccountData(nextSession);
     };
 
@@ -606,6 +609,7 @@ function App() {
         if (cancelled) return;
         if (nextSession) {
           setSession(nextSession);
+          setUserMaxConcurrency(nextSession?.user?.maxConcurrency);
         } else {
           clearAuthenticatedState({ resetView: true });
         }
@@ -943,6 +947,7 @@ function App() {
     async (nextSession: WebUserSession) => {
       hideSessionReplaced();
       setSession(nextSession);
+      setUserMaxConcurrency(nextSession?.user?.maxConcurrency);
       await hydrateAccountData(nextSession);
 
       if (nextSession.user.email && !nextSession.user.emailVerified) {
diff --git a/src/api/generationConcurrency.ts b/src/api/generationConcurrency.ts
index ed9df37..a4e3b56 100644
--- a/src/api/generationConcurrency.ts
+++ b/src/api/generationConcurrency.ts
@@ -7,10 +7,20 @@ interface GenerationSlot {
   createdAt: number;
 }
 
-const MAX_ACTIVE_GENERATION_TASKS = 3;
+const DEFAULT_MAX_ACTIVE_GENERATION_TASKS = 3;
 const STALE_SLOT_MS = 6 * 60 * 60 * 1000;
 const activeSlots = new Map<string, GenerationSlot>();
 
+let userMaxConcurrency: number | null = null;
+
+export function setUserMaxConcurrency(limit: number | null | undefined): void {
+  userMaxConcurrency = typeof limit === "number" && limit > 0 ? limit : null;
+}
+
+function getEffectiveLimit(): number {
+  return userMaxConcurrency ?? DEFAULT_MAX_ACTIVE_GENERATION_TASKS;
+}
+
 export function getGenerationUserKey(userId?: string | number | null): string {
   return userId === undefined || userId === null || userId === "" ? "anonymous" : String(userId);
 }
@@ -39,8 +49,9 @@ export function claimGenerationSlot(input: {
 }): () => void {
   pruneStaleSlots();
   const activeCount = getActiveGenerationTaskCount(input.userKey);
-  if (activeCount >= MAX_ACTIVE_GENERATION_TASKS) {
-    throw new Error("当前账号同时最多生成 3 个图片/视频任务，请等待已有任务完成后再提交。");
+  const effectiveLimit = getEffectiveLimit();
+  if (activeCount >= effectiveLimit) {
+    throw new Error(`当前账号同时最多生成 ${effectiveLimit} 个图片/视频任务，请等待已有任务完成后再提交。`);
   }
 
   const id = input.id || `generation-slot-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
diff --git a/src/api/keyServerClient.ts b/src/api/keyServerClient.ts
index 777d46e..4653047 100644
--- a/src/api/keyServerClient.ts
+++ b/src/api/keyServerClient.ts
@@ -434,6 +434,7 @@ function normalizeUser(raw: unknown): WebUserSession["user"] | null {
         candidate.enterpriseBalance ??
         candidate.enterprise_balance,
     ),
+    maxConcurrency: toNumber(candidate.maxConcurrency ?? candidate.max_concurrency),
     activePackages: toActivePackages(candidate.activePackages ?? candidate.active_packages),
   };
 }
diff --git a/src/api/taskSubscription.ts b/src/api/taskSubscription.ts
index ae702ec..3084581 100644
--- a/src/api/taskSubscription.ts
+++ b/src/api/taskSubscription.ts
@@ -44,7 +44,6 @@ export function waitForTask(
     let settled = false;
     let cleanup: (() => void) | null = null;
     let timeoutId: ReturnType<typeof setTimeout> | null = null;
-    let sseConnected = false;
     let fallbackTimerId: ReturnType<typeof setTimeout> | null = null;
     let lastProgress = 0;
     let lastProgressAt = startedAt;
@@ -83,10 +82,9 @@ export function waitForTask(
     };
 
     cleanup = aiGenerationClient.subscribeTaskStatus(taskId, handleUpdate);
-    sseConnected = true;
 
     fallbackTimerId = setTimeout(() => {
-      if (settled || !sseConnected) return;
+      if (settled) return;
       if (cleanup) cleanup();
       startPolling();
     }, 5000);
diff --git a/src/types.ts b/src/types.ts
index 9040a64..776e73f 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -73,6 +73,7 @@ export interface WebUserSession extends WebApiResultMeta {
     enterpriseAdminUserId?: number | string | null;
     balanceCents?: number;
     enterpriseBalanceCents?: number;
+    maxConcurrency?: number;
     activePackages?: Array<{
       name: string;
       expiresAt: string;