TraceStudio-dev/docs/server1.2/CUSTOM_NODE_DESIGN.md
2026-01-09 21:37:02 +08:00

9.3 KiB
Raw Blame History

自定义节点系统设计说明

🎯 设计目标

为 TraceStudio 提供一个安全、易用、可扩展的自定义节点系统,类似 ComfyUI 的插件机制,但增强了安全性和用户体验。

🏗️ 核心架构

1. 三层安全防护

┌─────────────────────────────────────┐
│      前端编辑器 (UI层)              │
│  - Monaco Editor 语法高亮          │
│  - 实时验证反馈                     │
│  - 操作确认对话框                   │
└──────────────┬──────────────────────┘
               │
┌──────────────▼──────────────────────┐
│    代码验证器 (node_validator.py)   │
│  - AST 语法分析                     │
│  - 黑名单检查 (os, eval, subprocess) │
│  - 节点规范验证 (继承、方法)         │
└──────────────┬──────────────────────┘
               │
┌──────────────▼──────────────────────┐
│  沙箱执行器 (node_sandbox.py)       │
│  - 超时限制 (30s)                   │
│  - 内存限制 (512MB)                 │
│  - 异常捕获                         │
└─────────────────────────────────────┘

2. 动态加载机制

# 工作流程
1. 用户保存节点文件 (.py)
2. 验证器检查代码安全性
3. 动态导入模块 (importlib)
4. 提取节点类 (继承TraceNode)
5. 注册到节点注册表
6. 立即可用无需重启服务器

🔒 安全机制详解

危险操作黑名单

类别 禁止项 原因
系统操作 os, subprocess, sys 防止执行系统命令
动态执行 eval, exec, compile 防止代码注入
模块导入 __import__, importlib 防止绕过限制
文件操作 open (警告) 防止任意文件读写

AST 检查示例

# ❌ 会被拒绝
import os
os.system('rm -rf /')

eval("malicious_code()")

# ✅ 允许
import pandas as pd
import numpy as np

data = pd.DataFrame(...)
result = np.mean(data['column'])

执行沙箱

# 超时保护
def execute_with_timeout(node, inputs, timeout=30):
    thread = Thread(target=node.execute, args=(inputs,))
    thread.start()
    thread.join(timeout)
    
    if thread.is_alive():
        raise TimeoutError("节点执行超时")

📝 节点开发流程

Step 1: 创建节点文件

# custom_nodes/my_node.py
from app.core.node_base import TraceNode

class MyNode(TraceNode):
    @staticmethod
    def get_metadata():
        return {
            "display_name": "我的节点",
            "category": "Custom",
            "inputs": [...],
            "outputs": [...],
            "params": [...]
        }
    
    def execute(self, inputs):
        # 实现逻辑
        return {...}

Step 2: 前端验证

// 点击"验证代码"按钮
const response = await api.validateCustomNode(code)

if (response.validation.valid) {
  // ✅ 通过:显示检测到的节点类
  showSuccess(response.validation.node_classes)
} else {
  // ❌ 失败:显示错误信息
  showErrors(response.validation.errors)
}

Step 3: 保存并加载

// 点击"保存"按钮
const response = await api.saveCustomNode(filename, code)

// 自动执行:
// 1. 二次确认(如果文件存在)
// 2. 备份旧文件
// 3. 保存新文件
// 4. 动态加载节点类
// 5. 注册到系统

🎨 用户界面设计

编辑器布局

┌────────────────────────────────────────────────┐
│  [新建] [示例] | [验证] [保存] [关闭] | [刷新]  │  <- 工具栏
├──────────────┬─────────────────────────────────┤
│  节点列表    │   代码编辑器                    │
│  ┌────────┐ │   1 from app.core.node_base... │
│  │ node.py│ │   2 class MyNode(TraceNode):   │
│  │ ✅     │ │   3     def execute(self,...   │
│  └────────┘ │   4         return {...}        │
│              │                                  │
│              ├─────────────────────────────────┤
│              │   验证结果                      │
│              │   ✅ 检测到节点: MyNode        │
│              │   ⚠️ 建议实现 get_metadata()   │
└──────────────┴─────────────────────────────────┘

操作确认对话框

┌──────────────────────────────────────┐
│  ⚠️ 覆盖确认                         │
├──────────────────────────────────────┤
│  文件 "my_node.py" 已存在            │
│                                      │
│  修改时间: 2026-01-07 14:30:00      │
│  包含节点: DataFilterNode           │
│                                      │
│  是否覆盖?                          │
│                                      │
│  [取消]           [确定覆盖]        │
└──────────────────────────────────────┘

🧪 测试用例

1. 安全测试

# test_security.py
def test_dangerous_imports():
    code = "import os\nos.system('ls')"
    result = validate_node_code(code)
    
    assert not result['valid']
    assert '禁止导入危险模块' in result['errors'][0]

2. 功能测试

def test_valid_node():
    code = """
from app.core.node_base import TraceNode

class TestNode(TraceNode):
    @staticmethod
    def get_metadata():
        return {"display_name": "Test"}
    
    def execute(self, inputs):
        return {"output": inputs['input']}
"""
    result = validate_node_code(code)
    
    assert result['valid']
    assert 'TestNode' in result['node_classes']

📊 性能优化

缓存机制

class NodeLoader:
    def __init__(self):
        self.loaded_nodes = {}  # 类缓存
        self.node_metadata = {}  # 元数据缓存
    
    def get_node_class(self, name):
        # 直接从缓存返回,无需重新导入
        return self.loaded_nodes.get(name)

热重载

# 修改节点文件后
loader.load_node('my_node.py', force_reload=True)

# 自动执行:
# 1. 卸载旧模块
# 2. 重新导入模块
# 3. 更新缓存
# 4. 无需重启服务器

🔄 与文件管理器的区别

特性 custom_nodes 文件管理器
访问权限 所有用户共享 根目录共享,用户目录隔离
操作限制 需验证+确认 直接操作
文件类型 .py 所有类型
加载机制 动态导入 静态文件
安全检查 AST验证 路径检查

🚀 未来扩展

1. 节点市场

- 节点分享平台
- 版本管理
- 依赖声明
- 一键安装

2. 更严格的沙箱

# 使用 RestrictedPython
from RestrictedPython import compile_restricted

code = compile_restricted(user_code)

3. GPU 支持

class GPUNode(TraceNode):
    @property
    def requires_gpu(self):
        return True
    
    def execute(self, inputs):
        # CUDA 加速
        import torch
        ...

📖 API 端点总结

GET    /api/custom-nodes/list           # 列出所有节点
POST   /api/custom-nodes/validate       # 验证代码
GET    /api/custom-nodes/read/{file}    # 读取代码
POST   /api/custom-nodes/save           # 保存节点
POST   /api/custom-nodes/action         # 操作节点
GET    /api/custom-nodes/loaded         # 已加载节点
POST   /api/custom-nodes/reload-all     # 重新加载
GET    /api/custom-nodes/example        # 示例代码

🎓 设计优势

  1. 安全性

    • 多层安全检查
    • 危险操作黑名单
    • 沙箱执行环境
  2. 易用性

    • 可视化编辑器
    • 实时验证反馈
    • 示例代码模板
  3. 可扩展性

    • 热重载机制
    • 插件化架构
    • 标准化接口
  4. 用户体验

    • 二次确认防误操作
    • 自动备份机制
    • 详细错误提示

💡 最佳实践建议

  1. 开发节点时

    • 先使用示例代码模板
    • 频繁点击"验证代码"
    • 编写完整的 docstring
  2. 测试节点时

    • 使用小数据集测试
    • 检查边界条件
    • 捕获异常情况
  3. 生产使用时

    • 定期备份节点文件
    • 记录节点版本号
    • 监控执行时间

设计理念: 在保证安全性的前提下,最大化用户的创作自由度。