agens-studio 0.1.0

package/src/services/taskManager.js

@@ -0,0 +1,246 @@
+import { randomUUID } from 'node:crypto';
+import { config } from '../../config.js';
+import * as agens from './agensClient.js';
+import { downloadAndRegister } from './assetStore.js';
+import { reviewImage } from './mimoVision.js';
+import { addSample } from './promptMemory.js';
+
+/**
+ * 异步任务管理器
+ * - generate 接口提交后立即返回本地 taskId
+ * - 后台轮询远端状态，状态机：pending → processing → success / failed
+ * - 成功后自动下载生成物并登记到资产库
+ *
+ * 本地单机：用内存 Map 存任务。重启会丢，但够用。
+ */
+
+const tasks = new Map(); // localTaskId -> task object
+
+// 任务队列只保留最近 N 条（前端展示上限），超出的自动丢弃
+const MAX_TASKS_KEPT = 4;
+
+// 生成任务并发限制：agnes 同时只允许 1 个任务（视频报 Service busy (tasks:1)）
+// 用简单的"等待 + 重试"机制，而不是阻塞式 promise 链（避免卡死）
+const AGNES_MAX_CONCURRENT = 1;
+let agnesRunning = 0;
+const agnesWaitQueue = [];
+
+/** 获取 agnes 调用权（信号量），释放时调 release */
+async function acquireAgnes() {
+  if (agnesRunning < AGNES_MAX_CONCURRENT) {
+    agnesRunning++;
+    return () => { agnesRunning--; drainAgnesQueue(); };
+  }
+  // 排队等待
+  await new Promise((resolve) => agnesWaitQueue.push(resolve));
+  agnesRunning++;
+  return () => { agnesRunning--; drainAgnesQueue(); };
+}
+function drainAgnesQueue() {
+  if (agnesRunning < AGNES_MAX_CONCURRENT && agnesWaitQueue.length) {
+    const next = agnesWaitQueue.shift();
+    next();
+  }
+}
+
+function createTask({ type, mode, prompt, params, imageUrls }) {
+  const id = randomUUID();
+  const task = {
+    id,
+    type, // 'image' | 'video'
+    mode, // 视频：text2video/image2video/multi2video/keyframe；生图：'image'
+    prompt,
+    params,
+    imageUrls,
+    status: 'pending', // pending / processing / success / failed
+    progress: 0,
+    remoteTaskId: null,
+    resultAssetId: null, // 成功后对应资产 id
+    resultUrl: null, // 原始远端 URL（下载前）
+    error: null,
+    createdAt: Date.now(),
+    updatedAt: Date.now(),
+  };
+  tasks.set(id, task);
+  pruneOldTasks();
+  return task;
+}
+
+/** 保留最近 MAX_TASKS_KEPT 条任务，其余从内存移除 */
+function pruneOldTasks() {
+  if (tasks.size <= MAX_TASKS_KEPT) return;
+  // 按 createdAt 倒序，保留最新的 MAX_TASKS_KEPT 条
+  const sorted = Array.from(tasks.values()).sort((a, b) => b.createdAt - a.createdAt);
+  const keep = new Set(sorted.slice(0, MAX_TASKS_KEPT).map((t) => t.id));
+  for (const id of tasks.keys()) {
+    if (!keep.has(id)) tasks.delete(id);
+  }
+}
+
+/**
+ * 提交一个生成任务
+ */
+export async function submitGenerateJob({ type, mode, prompt, params = {}, imageUrls = [] }) {
+  const task = createTask({ type, mode, prompt, params, imageUrls });
+
+  // 异步执行（不阻塞请求），内部用信号量限制 agnes 并发
+  runJob(task).catch((err) => {
+    task.status = 'failed';
+    task.error = err?.message || String(err);
+    task.updatedAt = Date.now();
+  });
+
+  return task;
+}
+
+async function runJob(task) {
+  // 等待 agnes 调用权（排队期间是 pending，获取到才变 processing）
+  const release = await acquireAgnes();
+  let released = false;
+  const releaseOnce = () => { if (!released) { released = true; release(); } };
+  task.status = 'processing';
+  task.updatedAt = Date.now();
+
+  try {
+    if (task.type === 'image') {
+      const { url, mime, meta } = await agens.generateImage({
+        prompt: task.prompt,
+        params: { ...task.params, imageUrls: task.imageUrls },
+      });
+      const asset = await downloadAndRegister({
+        url, kind: 'image', mime,
+        meta: { prompt: task.prompt, mode: task.mode, params: task.params },
+      });
+      task.status = 'success';
+      task.resultUrl = url;
+      task.resultAssetId = asset.id;
+      task.updatedAt = Date.now();
+      // 生图 agnes 调用结束，释放信号量（复看不占 agnes）
+      releaseOnce();
+      task.reviewStatus = 'pending';
+      reviewImage(asset.path, task.prompt, task.mode)
+        .then(async (review) => {
+          if (review.improvedPrompt) {
+            try {
+              const sample = await addSample({ kind: 'positive', original: task.prompt, optimized: review.improvedPrompt, mode: task.mode, imageDescription: '' });
+              review.sampleId = sample.id;
+            } catch {}
+          }
+          task.review = review;
+          task.reviewStatus = 'done';
+          task.updatedAt = Date.now();
+        })
+        .catch((e) => { task.review = { error: e.message }; task.reviewStatus = 'failed'; task.updatedAt = Date.now(); });
+      return;
+    }
+
+    // 视频：信号量保持到整个任务结束
+    const { videoId } = await agens.generateVideo({ mode: task.mode, prompt: task.prompt, imageUrls: task.imageUrls, params: task.params });
+    task.remoteTaskId = videoId;
+    task.updatedAt = Date.now();
+
+    for (let i = 0; i < config.agens.pollMaxAttempts; i++) {
+      await sleep(config.agens.pollInterval);
+      let info;
+      try { info = await agens.getTaskStatus(videoId); }
+      catch (err) { task.error = `查询失败（重试中）：${err.message}`; task.updatedAt = Date.now(); continue; }
+      task.error = null;
+      task.updatedAt = Date.now();
+      if (typeof info.progress === 'number' && info.progress > 0) task.progress = Math.min(99, info.progress);
+      if (info.status === 'processing') { if (typeof info.progress !== 'number') task.progress = Math.min(95, 10 + i * 3); continue; }
+      if (info.status === 'pending') continue;
+      if (info.status === 'failed') { task.status = 'failed'; task.error = info.result?.error || '生成失败'; task.updatedAt = Date.now(); return; }
+      if (info.status === 'success') {
+        task.progress = 100;
+        const v = info.result?.videoUrl;
+        if (!v) { task.status = 'failed'; task.error = '成功但未返回视频 URL'; task.updatedAt = Date.now(); return; }
+        task.resultUrl = v;
+        let vasset;
+        try {
+          vasset = await downloadAndRegister({ url: v, kind: 'video', mime: 'video/mp4', meta: { prompt: task.prompt, mode: task.mode, params: { ...task.params, duration: info.result?.duration }, taskId: task.id } });
+          task.resultAssetId = vasset.id;
+        } catch (e) {
+          // 下载失败：标记失败（而不是 success 但没资产）
+          task.status = 'failed';
+          task.error = `视频已生成但下载失败：${e.message}`;
+          task.updatedAt = Date.now();
+          return;
+        }
+        task.status = 'success';
+        task.updatedAt = Date.now();
+        // 视频也触发 AI 复看（mimo-v2.5 能看视频）
+        releaseOnce(); // 视频任务结束，释放信号量
+        task.reviewStatus = 'pending';
+        task.updatedAt = Date.now();
+        if (vasset) {
+          reviewImage(vasset.path, task.prompt, task.mode)
+            .then(async (review) => {
+              if (review.improvedPrompt) {
+                try {
+                  const sample = await addSample({ kind: 'positive', original: task.prompt, optimized: review.improvedPrompt, mode: task.mode, imageDescription: '' });
+                  review.sampleId = sample.id;
+                } catch {}
+              }
+              task.review = review;
+              task.reviewStatus = 'done';
+              task.updatedAt = Date.now();
+            })
+            .catch((e) => { task.review = { error: e.message }; task.reviewStatus = 'failed'; task.updatedAt = Date.now(); });
+        }
+        return;
+      }
+    }
+    task.status = 'failed';
+    task.error = '轮询超时';
+    task.updatedAt = Date.now();
+  } catch (err) {
+    task.status = 'failed';
+    task.error = err?.message || String(err);
+    task.updatedAt = Date.now();
+  } finally {
+    releaseOnce();
+  }
+}
+
+/** 查询任务（对外给前端，已脱敏） */
+export function getTask(localId) {
+  const t = tasks.get(localId);
+  if (!t) return null;
+  return {
+    id: t.id,
+    type: t.type,
+    mode: t.mode,
+    prompt: t.prompt,
+    status: t.status,
+    progress: t.progress,
+    resultAssetId: t.resultAssetId,
+    resultUrl: t.resultUrl,
+    error: t.error,
+    reviewStatus: t.reviewStatus || null,
+    review: t.review || null,
+    createdAt: t.createdAt,
+    updatedAt: t.updatedAt,
+  };
+}
+
+/** 最近 N 条任务（默认与 MAX_TASKS_KEPT 对齐） */
+export function recentTasks(limit = MAX_TASKS_KEPT) {
+  return Array.from(tasks.values())
+    .sort((a, b) => b.createdAt - a.createdAt)
+    .slice(0, limit)
+    .map((t) => ({
+      id: t.id,
+      type: t.type,
+      mode: t.mode,
+      prompt: t.prompt,
+      status: t.status,
+      progress: t.progress,
+      resultAssetId: t.resultAssetId,
+      error: t.error,
+      createdAt: t.createdAt,
+    }));
+}
+
+function sleep(ms) {
+  return new Promise((r) => setTimeout(r, ms));
+}