shuzhiren-comfyui/任务队列后端/worker_threads/wait/generatTask.js

import { parentPort } from 'worker_threads';
import redis from '../../redis/index.js';
import initQueue from '../../redis/initQueue.js';
import { externalPostRequest } from '../../outside/generat.js';
import { platformData } from '../../config/Config.js';

const REDIS_KEYS = {
  JWT: `${process.env.PROJECT_PREFIX}:md:jwt`
};

async function getJwtTokenFromRedis() {
  try {
    const token = await redis.get(REDIS_KEYS.JWT);
    return token;
  } catch (error) {
    console.error('[generatTask] 从 Redis 获取 JWT Token 失败:', error);
    return null;
  }
}

async function generatTask(tasksData) {
  const jwtToken = await getJwtTokenFromRedis();
  
  const generatTasks = []
  for (const task of tasksData) {
    const generatTaskPromise = externalPostRequest(task, jwtToken)
    generatTasks.push(generatTaskPromise)
  }

  try {
    const responseTasks = await Promise.all(generatTasks)
    return responseTasks
  } catch (error) {
    console.error('[generatTask] 批量请求出错:', error);
    return [];
  }
}

async function storeGeneratTasks(tasks) {
  if (!tasks || !Array.isArray(tasks)) {
    console.error('storeGeneratTasks函数接收到无效的tasks参数:', tasks);
    return;
  }
  
  const multi = redis.multi();
  let errorCount = 0;
  const taskErrorCountMap = new Map();
  const taskCountMap = new Map();
  
  for (const task of tasks) {
    if(task.remoteTaskId?.type === 2){
      console.error(`[generatTask] 任务失败: taskId=${task.taskId}, 错误:`, task.remoteTaskId.message);
      const aigc = task.AIGC || task.aigc;
      const platform = task.platform || task.platformName;
      
      multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', JSON.stringify(task.remoteTaskId.message));
      multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
      multi.lPush(initQueue.errorList, task.taskId);
      
      errorCount++;
      const key = `${aigc}:${platform}`;
      if(taskErrorCountMap.has(key)){
        taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
      } else {
        taskErrorCountMap.set(key, 1);
      }
      continue;
    }
    
    let externalTaskId;
    if (task.remoteTaskId?.type === 1 && task.remoteTaskId?.data) {
      const responseData = task.remoteTaskId.data;
      
      if (typeof responseData === 'string') {
        externalTaskId = responseData;
      } else if (typeof responseData === 'object' && responseData !== null) {
        try {
          const platform = task.platform || task.platformName;
          if ((responseData.msg === 'success' || platform === 'coze') && responseData.code === 0) {
            if (platform === 'coze') {
              externalTaskId = responseData.execute_id;
            } else {
              externalTaskId = responseData.data?.taskId;
            }
            
            if (!externalTaskId) {
              console.error(`[generatTask] 无法提取任务ID: taskId=${task.taskId}`);
              const errorMessage = JSON.stringify({ message: '无法从响应中提取外部平台任务ID', response: responseData });
              multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
              multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
              multi.lPush(initQueue.errorList, task.taskId);
              errorCount++;
              const key = `${task.AIGC}:${task.platform}`;
              if(taskErrorCountMap.has(key)){
                taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
              } else {
                taskErrorCountMap.set(key, 1);
              }
              continue;
            }
          } else {
            console.error(`[generatTask] 平台返回错误: taskId=${task.taskId}, 响应:`, responseData);
            const aigc = task.AIGC || task.aigc;
            const platform = task.platform || task.platformName;
            const errorMessage = JSON.stringify(responseData);
            multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
            multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
            multi.lPush(initQueue.errorList, task.taskId);
            errorCount++;
            const key = `${aigc}:${platform}`;
            if(taskErrorCountMap.has(key)){
              taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
            } else {
              taskErrorCountMap.set(key, 1);
            }
            continue;
          }
        } catch (extractError) {
          console.error(`[generatTask] 提取任务ID失败: taskId=${task.taskId}`, extractError);
          const aigc = task.AIGC || task.aigc;
          const platform = task.platform || task.platformName;
          const errorMessage = JSON.stringify({ message: '提取外部平台任务ID失败', error: extractError.message });
          multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
          multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
          multi.lPush(initQueue.errorList, task.taskId);
          errorCount++;
          const key = `${aigc}:${platform}`;
          if(taskErrorCountMap.has(key)){
            taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
          } else {
            taskErrorCountMap.set(key, 1);
          }
          continue;
        }
      } else {
        console.error(`[generatTask] remoteTaskId.data 类型异常: taskId=${task.taskId}, type=${typeof responseData}`);
        const aigc = task.AIGC || task.aigc;
        const platform = task.platform || task.platformName;
        const errorMessage = JSON.stringify({ message: 'remoteTaskId.data 类型异常', type: typeof responseData });
        multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'resultData', errorMessage);
        multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'status', 'failed');
        multi.lPush(initQueue.errorList, task.taskId);
        errorCount++;
        const key = `${aigc}:${platform}`;
        if(taskErrorCountMap.has(key)){
          taskErrorCountMap.set(key, taskErrorCountMap.get(key) + 1);
        } else {
          taskErrorCountMap.set(key, 1);
        }
        continue;
      }
    } else {
      externalTaskId = task.remoteTaskId;
    }
    
    const aigc = task.AIGC || task.aigc;
    const platform = task.platform || task.platformName;
    
    // console.log(`[generatTask] 任务映射: taskId=${task.taskId}, externalTaskId=${externalTaskId}, platform=${platform}`);
    
    if(platformData.callback.includes(platform)) {
      multi.set(`${initQueue.callback}:${externalTaskId}`, task.taskId, { EX: 7200 });
      await initQueue.addCallbackPendingTask(externalTaskId, task.taskId, aigc, platform);
    } else {
      const pollingKey = `${initQueue.prefix}:processPolling:${aigc}:${platform}`;
      let workflowId = task.workflowId || '';
      try {
        if (!workflowId && task.taskData) {
          const taskDataParsed = JSON.parse(task.taskData);
          workflowId = taskDataParsed.workflow_id || '';
        }
      } catch (e) {
        console.error('[generatTask] 解析taskData获取workflow_id失败:', e);
      }
      const pollingData = {
        taskId: task.taskId, 
        platform: platform, 
        AIGC: aigc,
        workflowId: workflowId
      };
      multi.hSet(pollingKey, externalTaskId, JSON.stringify(pollingData));
    }
    
    multi.hSet(`${initQueue.prefix}:task:${task.taskId}`, 'remoteTaskId', externalTaskId);
    multi.expire(`${initQueue.prefix}:task:${task.taskId}`, 7200);
    
    const key = `${aigc}:${platform}`;
    if(taskCountMap.has(key)){
      taskCountMap.set(key, taskCountMap.get(key) + 1);
    } else {
      taskCountMap.set(key, 1);
    }
  }

  if(errorCount > 0){
    initQueue.addEQtaskALL(errorCount);
  }
  await multi.exec();
  
  console.log(`[generatTask] 批次处理完成: 成功=${taskCountMap.size} 个平台, 错误=${errorCount} 个`);
}

parentPort.on('message', async (tasksData) => {
    await generatTask(tasksData)
      .then (tasks => storeGeneratTasks(tasks))
    parentPort.postMessage({ status: 'completed' });
});