kadevin-ComfyUI-Active-Node-Highlighter/js/active_node_highlighter.js

/**
 * ComfyUI Active Node Highlighter
 * 
 * Highlights the currently executing node by inverting its colors.
 * Compatible with both legacy LiteGraph frontend and new Vue frontend.
 * 
 * @author kadevin
 * @license MIT
 */

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

const EXTENSION_NAME = "ActiveNodeHighlighter";

/**
 * Inverts a hex color string
 * @param {string} color - Hex color like "#RRGGBB" or "#RGB"
 * @returns {string} Inverted hex color
 */
function invertColor(color) {
    if (!color || typeof color !== 'string') {
        return null;
    }

    // Handle hex colors
    if (color.startsWith('#')) {
        let hex = color.slice(1);

        // Expand shorthand (#RGB -> #RRGGBB)
        if (hex.length === 3) {
            hex = hex.split('').map(c => c + c).join('');
        }

        if (hex.length === 6) {
            const r = 255 - parseInt(hex.slice(0, 2), 16);
            const g = 255 - parseInt(hex.slice(2, 4), 16);
            const b = 255 - parseInt(hex.slice(4, 6), 16);
            return `#${r.toString(16).padStart(2, '0')}${g.toString(16).padStart(2, '0')}${b.toString(16).padStart(2, '0')}`;
        }
    }

    // Handle rgb/rgba colors
    const rgbMatch = color.match(/rgba?\s*\(\s*(\d+)\s*,\s*(\d+)\s*,\s*(\d+)\s*(?:,\s*[\d.]+)?\s*\)/i);
    if (rgbMatch) {
        const r = 255 - parseInt(rgbMatch[1]);
        const g = 255 - parseInt(rgbMatch[2]);
        const b = 255 - parseInt(rgbMatch[3]);
        return `rgb(${r}, ${g}, ${b})`;
    }

    return null;
}

/**
 * Extension state management - using a closure to avoid global pollution
 */
const createState = () => ({
    // Currently highlighted node ID (only one at a time)
    currentNodeId: null,
    // Original colors of the currently highlighted node
    originalColor: undefined,
    originalBgColor: undefined,
    // Flag to indicate if execution is ACTUALLY in progress
    isExecuting: false,
    // Track if we've been initialized
    initialized: false
});

let state = createState();

/**
 * Trigger canvas redraw
 */
function triggerRedraw() {
    requestAnimationFrame(() => {
        if (app.canvas) {
            app.canvas.setDirty(true, true);
        }
    });
}

/**
 * Restores the currently highlighted node to its original colors
 */
function restoreCurrentNode() {
    if (state.currentNodeId === null) {
        return;
    }

    const graph = app.graph;
    if (!graph) {
        state.currentNodeId = null;
        return;
    }

    const node = graph.getNodeById(state.currentNodeId);
    if (node) {
        // Restore original colors
        node.color = state.originalColor;
        node.bgcolor = state.originalBgColor;

        // Trigger redraw
        if (node.setDirtyCanvas) {
            node.setDirtyCanvas(true, true);
        }
    }

    // Clear highlighting state
    state.currentNodeId = null;
    state.originalColor = undefined;
    state.originalBgColor = undefined;
}

/**
 * Highlights a node by inverting its colors
 * @param {number} nodeId - The node ID to highlight
 */
function highlightNode(nodeId) {
    // CRITICAL: Only highlight if we're actually executing
    if (!state.isExecuting) {
        return;
    }

    const graph = app.graph;
    if (!graph) return;

    const node = graph.getNodeById(nodeId);
    if (!node) return;

    // Save original colors BEFORE modifying
    state.currentNodeId = nodeId;
    state.originalColor = node.color;
    state.originalBgColor = node.bgcolor;

    // Apply inverted colors
    // Use default ComfyUI node colors if undefined
    const baseColor = node.color || "#353535";
    const baseBgColor = node.bgcolor || "#232323";

    const invertedColor = invertColor(baseColor);
    const invertedBgColor = invertColor(baseBgColor);

    if (invertedColor) node.color = invertedColor;
    if (invertedBgColor) node.bgcolor = invertedBgColor;

    // Trigger redraw
    if (node.setDirtyCanvas) {
        node.setDirtyCanvas(true, true);
    }
}

/**
 * Handles node execution event
 * @param {string|number|null} nodeId - The ID of the currently executing node
 */
function handleExecuting(nodeId) {
    // CRITICAL: Ignore if not in execution mode
    if (!state.isExecuting) {
        return;
    }

    // First, always restore the previous node
    restoreCurrentNode();

    // If nodeId is null, execution has finished
    if (nodeId === null || nodeId === undefined) {
        state.isExecuting = false;
        triggerRedraw();
        return;
    }

    // Convert to number if string
    const numericNodeId = typeof nodeId === 'string' ? parseInt(nodeId, 10) : nodeId;

    // Highlight the new node
    highlightNode(numericNodeId);

    // Trigger canvas redraw
    triggerRedraw();
}

/**
 * Called when execution starts
 */
function handleExecutionStart() {
    // Restore any leftover state first
    restoreCurrentNode();
    // Now mark as executing
    state.isExecuting = true;
    state.initialized = true;
}

/**
 * Reset all state when execution ends or errors
 */
function handleExecutionEnd() {
    restoreCurrentNode();
    state.isExecuting = false;
    triggerRedraw();
}

// Register the extension
app.registerExtension({
    name: EXTENSION_NAME,

    async setup() {
        // Reset state on setup to ensure clean start
        state = createState();

        // Listen for execution start - THIS is when we enable highlighting
        api.addEventListener("execution_start", () => {
            handleExecutionStart();
        });

        // Listen for executing events - only processed if isExecuting is true
        api.addEventListener("executing", ({ detail }) => {
            // The detail can be {node: nodeId} or just nodeId depending on version
            let nodeId = null;
            if (detail && typeof detail === 'object') {
                nodeId = detail.node;
            } else if (detail !== undefined) {
                nodeId = detail;
            }
            handleExecuting(nodeId);
        });

        // Listen for execution error to restore colors
        api.addEventListener("execution_error", () => {
            handleExecutionEnd();
        });

        // Listen for execution interrupted
        api.addEventListener("execution_interrupted", () => {
            handleExecutionEnd();
        });

        console.log(`[${EXTENSION_NAME}] Extension loaded successfully`);
    }
});