WWWEN8-ComfyUI-CCXManager/web/queue_utils.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

class EventManager {
  constructor() {
    this.listeners = new Map();
  }

  addEventListener(event, callback) {
    if (!this.listeners.has(event)) {
      this.listeners.set(event, []);
    }
    this.listeners.get(event).push(callback);
  }

  removeEventListener(event, callback) {
    if (this.listeners.has(event)) {
      const callbacks = this.listeners.get(event);
      const index = callbacks.indexOf(callback);
      if (index > -1) {
        callbacks.splice(index, 1);
      }
    }
  }

  dispatchEvent(event, detail = {}) {
    if (this.listeners.has(event)) {
      const callbacks = this.listeners.get(event);
      callbacks.forEach(callback => {
        try {
          callback({ detail });
        } catch (error) {
          console.error(`Error in event listener for ${event}:`, error);
        }
      });
    }
  }
}

class QueueManager {
  constructor() {
    this.eventManager = new EventManager();
    this.queueNodeIds = null;
    this.processingQueue = false;
    this.lastAdjustedMouseEvent = null;
    this.initializeHooks();
  }

  initializeHooks() {
    const originalQueuePrompt = app.queuePrompt;
    const originalGraphToPrompt = app.graphToPrompt;
    const originalApiQueuePrompt = api.queuePrompt;

    app.queuePrompt = async function() {
      queueManager.processingQueue = true;
      queueManager.eventManager.dispatchEvent("queue");
      try {
        await originalQueuePrompt.apply(app, [...arguments]);
      } finally {
        queueManager.processingQueue = false;
        queueManager.eventManager.dispatchEvent("queue-end");
      }
    };

    app.graphToPrompt = async function() {
      queueManager.eventManager.dispatchEvent("graph-to-prompt");
      let promise = originalGraphToPrompt.apply(app, [...arguments]);
      await promise;
      queueManager.eventManager.dispatchEvent("graph-to-prompt-end");
      return promise;
    };

    api.queuePrompt = async function(index, prompt) {
      // 创建局部副本，避免异步执行时queueNodeIds被重置导致数据丢失
      const localQueueNodeIds = [...(queueManager.queueNodeIds || [])];
      
      // 修复：检测到有 CCXGroupExecutorSender 输出节点时，只保留它
      if (prompt.output && prompt.nodes) {
        let hasSenderNode = false;
        let senderNodeId = null;
        
        for (const outputNodeId in prompt.output) {
          const nodeData = prompt.nodes[outputNodeId];
          if (nodeData && nodeData.class_type === "CCXGroupExecutorSender") {
            hasSenderNode = true;
            senderNodeId = outputNodeId;
            break;
          }
        }

        // 只在没有明确指定输出节点ID时才进行此修改
        // 避免与CCXGroupExecutor的自定义执行逻辑冲突
        if (hasSenderNode && senderNodeId && !localQueueNodeIds.length) {
          console.log(`[QueueManager] 检测到 CCXGroupExecutorSender 节点，只保留它作为输出节点`);
          // 只保留 sender 节点
          const oldOutput = prompt.output;
          prompt.output = {};
          prompt.output[senderNodeId] = oldOutput[senderNodeId];
        }
      }

      if (localQueueNodeIds.length && prompt.output) {
        const oldOutput = prompt.output;
        let newOutput = {};
        for (const queueNodeId of localQueueNodeIds) {
          queueManager.recursiveAddNodes(String(queueNodeId), oldOutput, newOutput);
        }
        prompt.output = newOutput;
      }
      queueManager.eventManager.dispatchEvent("comfy-api-queue-prompt-before", {
        workflow: prompt.workflow,
        output: prompt.output,
      });
      const response = originalApiQueuePrompt.apply(api, [index, prompt]);
      queueManager.eventManager.dispatchEvent("comfy-api-queue-prompt-end");
      return response;
    };

    const originalProcessMouseDown = LGraphCanvas.prototype.processMouseDown;
    const originalAdjustMouseEvent = LGraphCanvas.prototype.adjustMouseEvent;
    const originalProcessMouseMove = LGraphCanvas.prototype.processMouseMove;

    LGraphCanvas.prototype.processMouseDown = function(e) {
      const result = originalProcessMouseDown.apply(this, [...arguments]);
      queueManager.lastAdjustedMouseEvent = e;
      return result;
    };

    LGraphCanvas.prototype.adjustMouseEvent = function(e) {
      originalAdjustMouseEvent.apply(this, [...arguments]);
      queueManager.lastAdjustedMouseEvent = e;
    };

    LGraphCanvas.prototype.processMouseMove = function(e) {
      const result = originalProcessMouseMove.apply(this, [...arguments]);
      if (e && !e.canvasX && !e.canvasY) {
        const canvas = app.canvas;
        const offset = canvas.convertEventToCanvasOffset(e);
        e.canvasX = offset[0];
        e.canvasY = offset[1];
      }
      queueManager.lastAdjustedMouseEvent = e;
      return result;
    };
  }
  recursiveAddNodes(nodeId, oldOutput, newOutput) {
    let currentId = nodeId;
    let currentNode = oldOutput[currentId];
    if (!currentNode) {
      console.warn(`[QueueManager] 节点 ${currentId} 在输出中不存在`);
      return newOutput;
    }
    if (newOutput[currentId] == null) {
      newOutput[currentId] = currentNode;
      for (const inputValue of Object.values(currentNode.inputs || [])) {
        if (Array.isArray(inputValue)) {
          this.recursiveAddNodes(inputValue[0], oldOutput, newOutput);
        }
      }
    }
    return newOutput;
  }
  async queueOutputNodes(nodeIds) {
    // 如果队列正在处理中，等待它完成后再继续
    while (this.processingQueue) {
      await new Promise(resolve => setTimeout(resolve, 100));
    }
    
    try {
      this.queueNodeIds = nodeIds;
      this.processingQueue = true;
      this.eventManager.dispatchEvent("queue");
      await app.queuePrompt();
    } catch (e) {
      console.error("队列节点时出错:", e);
    } finally {
      this.queueNodeIds = null;
      this.processingQueue = false;
      this.eventManager.dispatchEvent("queue-end");
    }
  }
  getLastMouseEvent() {
    return this.lastAdjustedMouseEvent;
  }
  addEventListener(event, callback) {
    this.eventManager.addEventListener(event, callback);
  }
  removeEventListener(event, callback) {
    this.eventManager.removeEventListener(event, callback);
  }
}

function getOutputNodes(nodes) {
  return (nodes?.filter((n) => {
    return (n.mode != LiteGraph.NEVER &&
      n.constructor.nodeData?.output_node);
  }) || []);
}
const queueManager = new QueueManager();
function queueSelectedOutputNodes() {
  const selectedNodes = app.canvas.selected_nodes;
  if (!selectedNodes || Object.keys(selectedNodes).length === 0) {
    console.log("[LG]队列: 没有选中的节点");
    return;
  }

  const outputNodes = getOutputNodes(Object.values(selectedNodes));
  if (!outputNodes || outputNodes.length === 0) {
    console.log("[LG]队列: 选中的节点中没有输出节点");
    return;
  }

  console.log(`[LG]队列: 执行 ${outputNodes.length} 个输出节点`);
  queueManager.queueOutputNodes(outputNodes.map((n) => n.id));
}

function queueGroupOutputNodes() {
  const lastMouseEvent = queueManager.getLastMouseEvent();
  if (!lastMouseEvent) {
    return;
  }

  let canvasX = lastMouseEvent.canvasX;
  let canvasY = lastMouseEvent.canvasY;
  
  if (!canvasX || !canvasY) {
    const canvas = app.canvas;
    const mousePos = canvas.getMousePos();
    canvasX = mousePos[0];
    canvasY = mousePos[1];
  }

  const group = app.graph.getGroupOnPos(canvasX, canvasY);

  if (!group) {
    return;
  }

  group.recomputeInsideNodes();

  if (!group._nodes || group._nodes.length === 0) {
    return;
  }
  
  const outputNodes = getOutputNodes(group._nodes);
  if (!outputNodes || outputNodes.length === 0) {
    return;
  }

  queueManager.queueOutputNodes(outputNodes.map((n) => n.id));
}

export { queueManager, getOutputNodes, queueSelectedOutputNodes, queueGroupOutputNodes };