omniai-web/src/api/taskSubscription.ts

import { aiGenerationClient } from "./aiGenerationClient";

export interface TaskProgressEvent {
  taskId: string;
  status: string;
  progress: number;
  resultUrl?: string | null;
  error?: string | null;
}

export interface WaitForTaskOptions {
  onProgress?: (event: TaskProgressEvent) => void;
  abortRef?: { current: boolean };
  timeoutMs?: number;
}

const POLL_INTERVAL = 3000;
const DEFAULT_TIMEOUT = 30 * 60 * 1000;

export function waitForTask(
  taskId: string,
  options: WaitForTaskOptions = {},
): Promise<string | null> {
  const { onProgress, abortRef, timeoutMs = DEFAULT_TIMEOUT } = options;

  return new Promise((resolve, reject) => {
    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;

    const settle = (fn: () => void) => {
      if (settled) return;
      settled = true;
      if (timeoutId) clearTimeout(timeoutId);
      if (fallbackTimerId) clearTimeout(fallbackTimerId);
      if (cleanup) cleanup();
      fn();
    };

    timeoutId = setTimeout(
      () => settle(() => reject(new Error("等待任务结果超时，请稍后在任务历史中查看"))),
      timeoutMs,
    );

    const handleUpdate = (event: TaskProgressEvent) => {
      if (settled) return;
      if (abortRef?.current) {
        settle(() => resolve(null));
        return;
      }
      onProgress?.(event);
      if (event.status === "completed") {
        settle(() => resolve(event.resultUrl || null));
      } else if (event.status === "failed" || event.status === "cancelled") {
        settle(() => reject(new Error(event.error || "任务失败，请稍后重试")));
      }
    };

    cleanup = aiGenerationClient.subscribeTaskStatus(taskId, handleUpdate);
    sseConnected = true;

    fallbackTimerId = setTimeout(() => {
      if (settled || !sseConnected) return;
      if (cleanup) cleanup();
      startPolling();
    }, 5000);

    function startPolling() {
      const poll = async () => {
        while (!settled) {
          if (abortRef?.current) {
            settle(() => resolve(null));
            return;
          }
          await new Promise((r) => setTimeout(r, POLL_INTERVAL));
          if (settled || abortRef?.current) return;
          try {
            const task = await aiGenerationClient.getTaskStatus(taskId);
            handleUpdate({
              taskId,
              status: task.status,
              progress: task.progress || 0,
              resultUrl: task.resultUrl,
              error: task.error,
            });
          } catch (e) {
            if (!settled) settle(() => reject(e));
          }
        }
      };
      poll();
    }
  });
}