9.3 KiB
9.3 KiB
自定义节点系统设计说明
🎯 设计目标
为 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 # 示例代码
🎓 设计优势
-
安全性 ✅
- 多层安全检查
- 危险操作黑名单
- 沙箱执行环境
-
易用性 ✅
- 可视化编辑器
- 实时验证反馈
- 示例代码模板
-
可扩展性 ✅
- 热重载机制
- 插件化架构
- 标准化接口
-
用户体验 ✅
- 二次确认防误操作
- 自动备份机制
- 详细错误提示
💡 最佳实践建议
-
开发节点时
- 先使用示例代码模板
- 频繁点击"验证代码"
- 编写完整的 docstring
-
测试节点时
- 使用小数据集测试
- 检查边界条件
- 捕获异常情况
-
生产使用时
- 定期备份节点文件
- 记录节点版本号
- 监控执行时间
设计理念: 在保证安全性的前提下,最大化用户的创作自由度。