bizydraft 0.1.29__py3-none-any.whl → 0.2.0__py3-none-any.whl

bizydraft/static/js/postEvent.js

@@ -1,4 +1,6 @@
 import { app } from "../../scripts/app.js";
+import { enableAIAppMode, disableAIAppMode, selectInputNode, deselectInputNode, updateInputNodeWidget, getSelectedInputNodes, clearSelectedInputNodes, toggleExportMode } from "./aiAppHandler.js";
+import { focusNodeOnly } from "./nodeFocusHandler.js";
 
 app.registerExtension({
     name: "comfy.BizyAir.Socket",
@@ -10,28 +12,28 @@ app.registerExtension({
     socket: null,
     isConnecting: false,
     taskRunning: false,
-    
+
     // 心跳检测
     pingInterval: 5000, // 5秒发送一次心跳
     pingTimer: null,
     pingTimeout: 3000,
-    pongReceived: false, 
+    pongReceived: false,
     pingTimeoutTimer: null, // ping超时计时器
-    
+
     /**
      * 开始心跳检测
      */
     startPing() {
         this.stopPing();
-        
+
         if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
             return;
         }
-        
+
         // 立即发送一次ping消息
         this.pongReceived = false;
         this.socket.send('ping');
-        
+
         // 设置ping超时检测
         this.pingTimeoutTimer = setTimeout(() => {
             if (!this.pongReceived) {
@@ -39,12 +41,12 @@ app.registerExtension({
                 this.reconnect();
             }
         }, this.pingTimeout);
-        
+
         // 设置定时发送ping
         this.pingTimer = setInterval(() => {
             if (this.socket && this.socket.readyState === WebSocket.OPEN) {
                 this.pongReceived = false;
-                this.socket.send('ping');             
+                this.socket.send('ping');
                 // 设置ping超时检测
                 this.pingTimeoutTimer = setTimeout(() => {
                     // 如果没有收到pong响应
@@ -59,7 +61,7 @@ app.registerExtension({
             }
         }, this.pingInterval);
     },
-    
+
     /**
      * 停止心跳检测
      */
@@ -68,25 +70,25 @@ app.registerExtension({
             clearInterval(this.pingTimer);
             this.pingTimer = null;
         }
-        
+
         if (this.pingTimeoutTimer) {
             clearTimeout(this.pingTimeoutTimer);
             this.pingTimeoutTimer = null;
         }
     },
-    
+
     /**
      * 重新连接
      */
     reconnect() {
         if (this.isConnecting) {
             return;
-        }    
-        const url = this.socket ? this.socket.url : app.api.socket.url;      
-        this.closeSocket();    
+        }
+        const url = this.socket ? this.socket.url : app.api.socket.url;
+        this.closeSocket();
         this.createSocket(url);
     },
-    
+
     /**
      * 创建新的WebSocket连接
      */
@@ -96,21 +98,21 @@ app.registerExtension({
             console.log('WebSocket连接已在创建中，避免重复创建');
             return null;
         }
-        
+
         // 标记为连接中
         this.isConnecting = true;
-        
+
         // 先关闭现有连接
         this.closeSocket();
-        
+
         const url = customUrl || app.api.socket.url;
         console.log('创建WebSocket连接:', url);
-        
+
         try {
             const socket = new WebSocket(url);
             const dispatchCustomEvent = this.dispatchCustomEvent;
             const self = this;
-            
+
             socket.onopen = function() {
                 console.log('WebSocket连接已打开');
                 // 清除连接中标志
@@ -122,19 +124,19 @@ app.registerExtension({
                 // 开始心跳检测
                 self.startPing();
             };
-    
+
             socket.onmessage = function (event) {
                 try {
                     // 处理心跳响应
                     if (event.data === 'pong') {
                         // 标记收到pong响应
                         self.pongReceived = true;
-                        return; 
-                    }               
+                        return;
+                    }
                     if (event.data instanceof ArrayBuffer) {
                         const view = new DataView(event.data);
                         const eventType = view.getUint32(0);
-    
+
                         let imageMime;
                         switch (eventType) {
                             case 3:
@@ -177,20 +179,27 @@ app.registerExtension({
                             self.closeSocket(1000);
                             return;
                         }
-                        
+
                         const msg = JSON.parse(event.data);
                         window.parent.postMessage({
                             type: 'functionResult',
                             method: 'progress_info_change',
                             result: msg.progress_info
                         }, '*');
-                        
                         switch (msg.type) {
+                            case 'load_start':
+                            case 'load_end':
+                            case 'prompt_id':
+                                window.parent.postMessage({
+                                    type: 'functionResult',
+                                    method: 'preparingStatus',
+                                    result: msg
+                                }, '*')
                             case 'status':
                                 if (msg.data.sid) {
                                     const clientId = msg.data.sid;
-                                    window.name = clientId; 
-                                    sessionStorage.setItem('clientId', clientId); 
+                                    window.name = clientId;
+                                    sessionStorage.setItem('clientId', clientId);
                                     socket.clientId = clientId;
                                 }
                                 dispatchCustomEvent('status', msg.data.status ?? null);
@@ -209,6 +218,17 @@ app.registerExtension({
                                 }
                                 dispatchCustomEvent(msg.type, msg.data);
                                 break;
+                            case 'error':
+                                self.taskRunning = false;
+                                dispatchCustomEvent('execution_error', {
+                                    exception_message: msg.data.error_message || 'Unknown error',
+                                    traceback: ['Manual error triggered'],
+                                    executed: [],
+                                    prompt_id: 'manual',
+                                    node_id: '0',
+                                    node_type: 'RunError',
+                                });
+                                break;
                             case 'execution_error':
                             case 'execution_interrupted':
                                 self.taskRunning = false;
@@ -229,7 +249,7 @@ app.registerExtension({
                             default:
                                 const registeredTypes = socket.registeredTypes || new Set();
                                 const reportedUnknownMessageTypes = socket.reportedUnknownMessageTypes || new Set();
-    
+
                                 if (registeredTypes.has(msg.type)) {
                                     app.dispatchEvent(
                                         new CustomEvent(msg.type, { detail: msg.data })
@@ -244,7 +264,7 @@ app.registerExtension({
                     console.warn('Unhandled message:', event.data, error);
                 }
             };
-    
+
             socket.onerror = function(error) {
                 console.log('WebSocket 错误:', error);
                 // 清除连接中标志
@@ -252,7 +272,7 @@ app.registerExtension({
                 // 停止心跳
                 self.stopPing();
             };
-    
+
             socket.onclose = function(event) {
                 console.log('WebSocket 连接已关闭, 状态码:', event.code, event.reason);
                 // 清除连接中标志
@@ -264,10 +284,10 @@ app.registerExtension({
                 // 停止心跳
                 self.stopPing();
             };
-    
+
             socket.registeredTypes = new Set();
             socket.reportedUnknownMessageTypes = new Set();
-    
+
             // 返回创建的socket，但不要立即使用，等待onopen
             return socket;
         } catch (error) {
@@ -288,7 +308,7 @@ app.registerExtension({
                 resolve(this.socket);
                 return;
             }
-            
+
             // 如果连接正在创建中，等待一段时间后检查
             if (this.isConnecting) {
                 console.log('WebSocket连接创建中，等待...');
@@ -303,26 +323,26 @@ app.registerExtension({
                 }, 100); // 每100ms检查一次
                 return;
             }
-            
+
             // 创建新连接
             const socket = this.createSocket(customUrl);
             if (!socket) {
                 reject(new Error('创建WebSocket连接失败'));
                 return;
             }
-            
+
             // 监听连接打开事件
             socket.addEventListener('open', () => {
                 resolve(socket);
             });
-            
+
             // 监听错误事件
             socket.addEventListener('error', (error) => {
                 reject(error);
             });
         });
     },
-    
+
     /**
      * 获取可用的socket连接，如果不存在则创建
      * 同步版本，可能返回尚未就绪的连接
@@ -332,7 +352,7 @@ app.registerExtension({
         if (this.socket && this.socket.readyState === WebSocket.OPEN) {
             return this.socket;
         }
-        
+
         // 创建新连接
         return this.createSocket(customUrl);
     },
@@ -344,7 +364,7 @@ app.registerExtension({
     closeSocket(code) {
         // 先停止心跳
         this.stopPing();
-        
+
         if (this.socket) {
             if (this.socket.readyState === WebSocket.OPEN || this.socket.readyState === WebSocket.CONNECTING) {
                 console.log('关闭WebSocket连接');
@@ -352,10 +372,10 @@ app.registerExtension({
             }
             this.socket = null;
         }
-        
+
         // 重置任务状态
         this.taskRunning = false;
-        
+
         return true;
     },
 
@@ -365,7 +385,7 @@ app.registerExtension({
     changeSocketUrl(newUrl) {
         const clientId = sessionStorage.getItem("clientId");
         const fullUrl = newUrl + "?clientId=" + clientId + "&a=1";
-        
+
         return this.createSocket(fullUrl);
     },
 
@@ -402,7 +422,7 @@ app.registerExtension({
             return false;
         }
     },
-    
+
     getCookie(name) {
         const value = `; ${document.cookie}`;
         const parts = value.split(`; ${name}=`);
@@ -411,14 +431,11 @@ app.registerExtension({
 
     async setup() {
         const createSocket = this.createSocket.bind(this);
-        const getSocket = this.getSocket.bind(this);
-        const getSocketAsync = this.getSocketAsync.bind(this);
         const closeSocket = this.closeSocket.bind(this);
-        const changeSocketUrl = this.changeSocketUrl.bind(this);
-        const sendSocketMessage = this.sendSocketMessage.bind(this);
-        const sendPrompt = this.sendPrompt.bind(this);
-        const startPing = this.startPing.bind(this);
-        const stopPing = this.stopPing.bind(this);
+
+        const customErrorStyles = new Map()
+
+
 
         // 方法映射
         const methods = {
@@ -432,26 +449,6 @@ app.registerExtension({
                 return socket;
             },
 
-            startSocket: async function (params) {
-                try {
-                    const socket = await getSocketAsync(params.url);
-                    window.parent.postMessage({
-                        type: 'functionResult',
-                        method: 'startSocket',
-                        result: 'Socket连接已启动'
-                    }, '*');
-                    return socket;
-                } catch (error) {
-                    window.parent.postMessage({
-                        type: 'functionResult',
-                        method: 'startSocket',
-                        result: 'Socket连接失败: ' + error.message,
-                        error: true
-                    }, '*');
-                    return null;
-                }
-            },
-
             closeSocket: function () {
                 const result = closeSocket();
                 window.parent.postMessage({
@@ -462,34 +459,6 @@ app.registerExtension({
                 return result;
             },
 
-            changeSocketUrl: function (params) {
-                if (!params.url) {
-                    console.error('缺少url参数');
-                    return false;
-                }
-                const socket = changeSocketUrl(params.url);
-                window.parent.postMessage({
-                    type: 'functionResult',
-                    method: 'changeSocketUrl',
-                    result: 'Socket URL已更改为' + params.url
-                }, '*');
-                return socket;
-            },
-
-            sendSocketMessage: async function (params) {
-                if (!params.message) {
-                    console.error('缺少message参数');
-                    return false;
-                }
-                const result = await sendSocketMessage(params.message);
-                window.parent.postMessage({
-                    type: 'functionResult',
-                    method: 'sendSocketMessage',
-                    result: result ? '消息发送成功' : '消息发送失败'
-                }, '*');
-                return result;
-            },
-
             clearCanvas: function () {
                 app.graph.clear();
                 window.parent.postMessage({
@@ -499,11 +468,10 @@ app.registerExtension({
                 }, '*');
                 return true;
             },
-
             loadWorkflow: function (params) {
                 app.graph.clear();
-                document.dispatchEvent(new CustomEvent('workflowLoaded', { 
-                    detail: params 
+                document.dispatchEvent(new CustomEvent('workflowLoaded', {
+                    detail: params.json
                 }));
                 if (params.json.version) {
                     app.loadGraphData(params.json);
@@ -537,6 +505,19 @@ app.registerExtension({
                 }, '*');
                 return graph.workflow;
             },
+            // 新增：获取 workflow 和 output
+            getWorkflowWithOutput: async function () {
+                const graph = await app.graphToPrompt();
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'getWorkflowWithOutput',
+                    result: {
+                        workflow: graph.workflow,
+                        output: graph.output
+                    }
+                }, '*');
+                return { workflow: graph.workflow, output: graph.output };
+            },
             saveApiJson: async function (params) {
                 const graph = await app.graphToPrompt();
                 window.parent.postMessage({
@@ -559,9 +540,10 @@ app.registerExtension({
                 try {
                     // 确保有连接
                     // await getSocketAsync();
-                    
+
                     const graph = await app.graphToPrompt();
                     const clientId = sessionStorage.getItem("clientId");
+                    await app.queuePrompt(graph.output);
                     const resPrompt = await fetch("api/prompt", {
                         method: "POST",
                         body: JSON.stringify({
@@ -577,8 +559,45 @@ app.registerExtension({
                     });
                     const resPromptJson = await resPrompt.json();
                     if (resPromptJson.error) {
-                        await app.queuePrompt(graph.output);
+                        this.openCustomError({
+                            nodeId: resPromptJson.node_id,
+                            nodeType: resPromptJson.node_type,
+                            errorMessage: resPromptJson.error,
+                            borderColor: '#FF0000'
+                        })
+                        return
+                    }
+
+                    for (const i in resPromptJson.node_errors) {
+
+                        if (resPromptJson.node_errors[i].errors) {
+                            const err = resPromptJson.node_errors[i].errors[0]
+                            if (err) return
+                            this.openCustomError({
+                                nodeId: i,
+                                nodeType: err.type,
+                                errorMessage: err.error,
+                                borderColor: '#FF0000'
+                            })
+                            return
+                        } else {
+                            console.log(resPromptJson.node_errors[i])
+                        }
                     }
+
+                    if (Object.keys(resPromptJson.node_errors).length) return
+                    graph.workflow.nodes.forEach(node => {
+                        for (const key in graph.output) {
+                            if (graph.output[key].class_type === node.type) {
+                                if (!graph.output[key]._meta) {
+                                    graph.output[key]._meta = {};
+                                }
+                                graph.output[key]._meta.id = node.id;
+                                break; // 找到匹配的就跳出循环
+                            }
+                        }
+                    });
+                    console.log(graph.output)
                     window.parent.postMessage({
                         type: 'functionResult',
                         method: 'runWorkflow',
@@ -611,7 +630,8 @@ app.registerExtension({
                     }
                     document.cookie = name + "=" + (value || "") + expires + "; path=/";
                 };
-                console.log("-----------setCookie-----------", params)
+                // console.log("-----------setCookie-----------", params)
+                // console.log("-----------setCookie-----------", params)
                 setCookie(params.name, params.value, params.days);
 
 
@@ -675,11 +695,13 @@ app.registerExtension({
                 return graph.workflow;
             },
 
-            graphToPrompt: async function () {
+            graphToPrompt: async function (params) {
+                console.log('postEvent.js - graphToPrompt被调用，参数:', params);
                 const graph = await app.graphToPrompt();
                 window.parent.postMessage({
                     type: 'functionResult',
                     method: 'graphToPrompt',
+                    params: params, // 传递原始参数
                     result: {
                         workflow: graph.workflow,
                         output: graph.output
@@ -703,8 +725,204 @@ app.registerExtension({
                     }, '*');
                     return true;
             },
+
+            // AI应用相关方法
+            toggleAIAppMode: function(params) {
+                const enable = params.enable === true;
+                const result = enable ? enableAIAppMode() : disableAIAppMode();
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'toggleAIAppMode',
+                    result: result
+                }, '*');
+                return result;
+            },
+
+            selectInputNode: function(params) {
+                if (!params.nodeId) return false;
+                const result = selectInputNode(params.nodeId);
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'selectInputNode',
+                    result: result
+                }, '*');
+                return result;
+            },
+
+            selectExportNode: function(params) {
+                if (!params.nodeId) return false;
+                const result = selectInputNode(params.nodeId);
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'selectExportNode',
+                    result: result
+                }, '*');
+                return result;
+            },
+
+            deselectInputNode: function(params) {
+                if (!params.nodeId) return false;
+                const result = deselectInputNode(params.nodeId);
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'deselectInputNode',
+                    result: result
+                }, '*');
+                return result;
+            },
+
+            updateInputNodeWidget: function(params) {
+                if (!params.nodeId || params.widgetName === undefined) return false;
+                const result = updateInputNodeWidget(params.nodeId, params.widgetName, params.value);
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'updateInputNodeWidget',
+                    result: result
+                }, '*');
+                return result;
+            },
+
+            getInputNodes: function() {
+                const result = getSelectedInputNodes();
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'getInputNodes',
+                    result: result
+                }, '*');
+                return result;
+            },
+
+            clearInputNodes: function() {
+                const result = clearSelectedInputNodes();
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'clearInputNodes',
+                    result: result
+                }, '*');
+                return result;
+            },
+            toggleExportMode: function(params) {
+                const result = toggleExportMode(params);
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'toggleExportMode',
+                    result: result
+                }, '*');
+                return result;
+            },
+            openCustomError: function (params) {
+                const { nodeId, nodeType, errorMessage, borderColor='#FF0000' } = params;
+                const nodeIds = Array.isArray(nodeId) ? nodeId : [nodeId];
+                function injectErrorDialogStyles() {
+                    const styleId = 'custom-error-dialog-styles';
+                    if (document.getElementById(styleId)) {
+                        return; // 样式已经存在
+                    }
+
+                    const style = document.createElement('style');
+                    style.id = styleId;
+                    style.textContent = `
+                        .comfy-error-report .no-results-placeholder p {
+                            text-align: left;
+                        }
+                    `;
+                    document.head.appendChild(style);
+                };
+                injectErrorDialogStyles()
+                function simulateExecutionError(nodeId, nodeType, errorMessage, borderColor) {
+                    // const originalNodeErrorStyle = node.strokeStyles?.['nodeError']
+                    const node = app.graph.getNodeById(nodeId)
+                    if (!node) return
+                    if (!customErrorStyles.has(nodeId)) {
+                        customErrorStyles.set(nodeId, {
+                          originalStyle: node.strokeStyles?.['nodeError'],
+                          customColor: borderColor,
+                          nodeId: nodeId
+                        })
+                    }
+                    node.strokeStyles = node.strokeStyles || {}
+                    node.strokeStyles['nodeError'] = function () {
+                        // if (this.id === nodeId) {
+                        return { color: borderColor, lineWidth: 2 } // 自定义颜色和线宽
+                        // }
+                    }
+                    const mockErrorEvent = {
+                        detail: {
+                            node_id: nodeId,
+                            node_type: nodeType,
+                            exception_message: errorMessage,
+                            exception_type: 'ManualError',
+                            traceback: ['Manual error triggered'],
+                            executed: [],
+                            prompt_id: 'manual',
+                            timestamp: Date.now()
+                        }
+                    }
+
+                    // 手动触发事件监听器
+                    app.api.dispatchCustomEvent('execution_error', mockErrorEvent.detail)
+                }
+
+                nodeIds.forEach(id => {
+                    simulateExecutionError(id, nodeType, errorMessage, borderColor)
+                })
+
+                app.canvas.draw(true, true)
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'openCustomError',
+                    result: true
+                }, '*');
+            },
+            clearAllCustomStyles: function () {
+                customErrorStyles.forEach((styleInfo, nodeId) => {
+                  const node = app.graph.getNodeById(nodeId)
+                  if (!node) return
+                  console.log(node)
+                  // 恢复原始样式
+                  if (styleInfo.originalStyle) {
+                    node.strokeStyles['nodeError'] = styleInfo.originalStyle
+                  } else {
+                    delete node.strokeStyles['nodeError']
+                  }
+
+                  // 从映射中移除
+                  customErrorStyles.delete(nodeId)
+                })
+
+                // 重绘画布
+                app.canvas.draw(true, true)
+            },
+            focusNodeOnly: function(params) {
+                if (!params.nodeId) return false;
+                const result = focusNodeOnly(params.nodeId);
+                window.parent.postMessage({
+                    type: 'functionResult',
+                    method: 'focusNodeOnly',
+                    result: result
+                }, '*');
+                return result;
+            },
         };
 
+        methods.deselectExportNode = function(params) {
+            if (params && params.nodeId !== undefined) {
+                if (typeof window.deselectInputNode === 'function') {
+                    window.deselectInputNode(params.nodeId);
+                }
+            }
+        };
+        methods.clearExportNodes = function() {
+            if (typeof window.clearExportNodes === 'function') {
+                window.clearExportNodes();
+            }
+        };
+        // 保存工作流的原始节点颜色信息
+        methods.saveOriginalNodeColors = function(params) {
+            if (typeof window.saveOriginalNodeColors === 'function') {
+                window.saveOriginalNodeColors(params.workflowId);
+            }
+        };
         window.addEventListener('message', function (event) {
             if (event.data && event.data.type === 'callMethod') {
                 const methodName = event.data.method;
@@ -732,5 +950,17 @@ app.registerExtension({
             }
         });
         window.parent.postMessage({ type: 'iframeReady' }, '*');
+        app.api.addEventListener('graphChanged', (e) => {
+            console.log('Graph 发生变化，当前 workflow JSON:', e.detail)
+            window.parent.postMessage({
+                type: 'functionResult',
+                method: 'workflowChanged',
+                result: e.detail
+            }, '*');
+
+            document.dispatchEvent(new CustomEvent('workflowLoaded', {
+                detail: e.detail
+            }));
+        })
     }
 });