fix: improve generation task reliability

src/features/workbench/toolKeepalive.ts

@@ -3,6 +3,8 @@
  * Persists task state to localStorage so in-progress tasks survive page switches.
  */
 
+import { waitForTask } from "../../api/taskSubscription";
+
 const KEEPALIVE_PREFIX = "omniai:tool-task:";
 
 interface ToolTaskKeepalive {
@@ -59,38 +61,19 @@ export function clearToolTaskState(key: string): void {
   try { window.localStorage.removeItem(KEEPALIVE_PREFIX + key); } catch { /* ignore */ }
 }
 
-const TASK_POLL_INTERVAL = 3000;
-const TASK_POLL_TIMEOUT = 30 * 60 * 1000;
-
 export async function pollTaskUntilDone(
   taskId: string,
   onProgress?: (progress: number) => void,
   abortRef?: { current: boolean },
+  kind: "image" | "video" = "video",
 ): Promise<string | null> {
-  const startTime = Date.now();
-  const { aiGenerationClient } = await import("../../api/aiGenerationClient");
-
-  while (true) {
-    if (abortRef?.current) return null;
-    if (Date.now() - startTime > TASK_POLL_TIMEOUT) return null;
-
-    try {
-      const task = await aiGenerationClient.getTaskStatus(taskId);
-      if (!task) return null;
-
-      const progress = Math.min(99, task.progress || 0);
-      onProgress?.(progress);
-
-      if (task.status === "completed") {
-        return task.resultUrl || null;
-      }
-      if (task.status === "failed" || task.status === "cancelled") {
-        return null;
-      }
-    } catch {
-      // retry on next poll
-    }
-
-    await new Promise((r) => setTimeout(r, TASK_POLL_INTERVAL));
+  try {
+    return await waitForTask(taskId, {
+      kind,
+      abortRef,
+      onProgress: (event) => onProgress?.(Math.min(99, Number(event.progress || 0))),
+    });
+  } catch {
+    return null;
   }
 }