340 lines
9.3 KiB
Markdown
340 lines
9.3 KiB
Markdown
|
# 自定义节点系统设计说明
|
||
|
|
|
||
|
|
## 🎯 设计目标
|
||
|
|
|
||
|
|
为 TraceStudio 提供一个**安全、易用、可扩展**的自定义节点系统，类似 ComfyUI 的插件机制，但增强了安全性和用户体验。
|
||
|
|
|
||
|
|
## 🏗️ 核心架构
|
||
|
|
|
||
|
|
### 1. 三层安全防护
|
||
|
|
|
||
|
|
```
|
||
|
|
┌─────────────────────────────────────┐
|
||
|
|
│ 前端编辑器 (UI层) │
|
||
|
|
│ - Monaco Editor 语法高亮 │
|
||
|
|
│ - 实时验证反馈 │
|
||
|
|
│ - 操作确认对话框 │
|
||
|
|
└──────────────┬──────────────────────┘
|
||
|
|
│
|
||
|
|
┌──────────────▼──────────────────────┐
|
||
|
|
│ 代码验证器 (node_validator.py) │
|
||
|
|
│ - AST 语法分析 │
|
||
|
|
│ - 黑名单检查 (os, eval, subprocess) │
|
||
|
|
│ - 节点规范验证 (继承、方法) │
|
||
|
|
└──────────────┬──────────────────────┘
|
||
|
|
│
|
||
|
|
┌──────────────▼──────────────────────┐
|
||
|
|
│ 沙箱执行器 (node_sandbox.py) │
|
||
|
|
│ - 超时限制 (30s) │
|
||
|
|
│ - 内存限制 (512MB) │
|
||
|
|
│ - 异常捕获 │
|
||
|
|
└─────────────────────────────────────┘
|
||
|
|
```
|
||
|
|
|
||
|
|
### 2. 动态加载机制
|
||
|
|
|
||
|
|
```python
|
||
|
|
# 工作流程
|
||
|
|
1. 用户保存节点文件 (.py)
|
||
|
|
2. 验证器检查代码安全性
|
||
|
|
3. 动态导入模块 (importlib)
|
||
|
|
4. 提取节点类 (继承TraceNode)
|
||
|
|
5. 注册到节点注册表
|
||
|
|
6. 立即可用，无需重启服务器
|
||
|
|
```
|
||
|
|
|
||
|
|
## 🔒 安全机制详解
|
||
|
|
|
||
|
|
### 危险操作黑名单
|
||
|
|
|
||
|
|
| 类别 | 禁止项 | 原因 |
|
||
|
|
|------|--------|------|
|
||
|
|
| **系统操作** | `os`, `subprocess`, `sys` | 防止执行系统命令 |
|
||
|
|
| **动态执行** | `eval`, `exec`, `compile` | 防止代码注入 |
|
||
|
|
| **模块导入** | `__import__`, `importlib` | 防止绕过限制 |
|
||
|
|
| **文件操作** | `open` (警告) | 防止任意文件读写 |
|
||
|
|
|
||
|
|
### AST 检查示例
|
||
|
|
|
||
|
|
```python
|
||
|
|
# ❌ 会被拒绝
|
||
|
|
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'])
|
||
|
|
```
|
||
|
|
|
||
|
|
### 执行沙箱
|
||
|
|
|
||
|
|
```python
|
||
|
|
# 超时保护
|
||
|
|
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: 创建节点文件
|
||
|
|
|
||
|
|
```python
|
||
|
|
# 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: 前端验证
|
||
|
|
|
||
|
|
```typescript
|
||
|
|
// 点击"验证代码"按钮
|
||
|
|
const response = await api.validateCustomNode(code)
|
||
|
|
|
||
|
|
if (response.validation.valid) {
|
||
|
|
// ✅ 通过：显示检测到的节点类
|
||
|
|
showSuccess(response.validation.node_classes)
|
||
|
|
} else {
|
||
|
|
// ❌ 失败：显示错误信息
|
||
|
|
showErrors(response.validation.errors)
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|
||
|
|
### Step 3: 保存并加载
|
||
|
|
|
||
|
|
```typescript
|
||
|
|
// 点击"保存"按钮
|
||
|
|
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. 安全测试
|
||
|
|
|
||
|
|
```python
|
||
|
|
# 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. 功能测试
|
||
|
|
|
||
|
|
```python
|
||
|
|
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']
|
||
|
|
```
|
||
|
|
|
||
|
|
## 📊 性能优化
|
||
|
|
|
||
|
|
### 缓存机制
|
||
|
|
|
||
|
|
```python
|
||
|
|
class NodeLoader:
|
||
|
|
def __init__(self):
|
||
|
|
self.loaded_nodes = {} # 类缓存
|
||
|
|
self.node_metadata = {} # 元数据缓存
|
||
|
|
|
||
|
|
def get_node_class(self, name):
|
||
|
|
# 直接从缓存返回，无需重新导入
|
||
|
|
return self.loaded_nodes.get(name)
|
||
|
|
```
|
||
|
|
|
||
|
|
### 热重载
|
||
|
|
|
||
|
|
```python
|
||
|
|
# 修改节点文件后
|
||
|
|
loader.load_node('my_node.py', force_reload=True)
|
||
|
|
|
||
|
|
# 自动执行:
|
||
|
|
# 1. 卸载旧模块
|
||
|
|
# 2. 重新导入模块
|
||
|
|
# 3. 更新缓存
|
||
|
|
# 4. 无需重启服务器
|
||
|
|
```
|
||
|
|
|
||
|
|
## 🔄 与文件管理器的区别
|
||
|
|
|
||
|
|
| 特性 | custom_nodes | 文件管理器 |
|
||
|
|
|------|-------------|-----------|
|
||
|
|
| **访问权限** | 所有用户共享 | 根目录共享，用户目录隔离 |
|
||
|
|
| **操作限制** | 需验证+确认 | 直接操作 |
|
||
|
|
| **文件类型** | 仅 `.py` | 所有类型 |
|
||
|
|
| **加载机制** | 动态导入 | 静态文件 |
|
||
|
|
| **安全检查** | AST验证 | 路径检查 |
|
||
|
|
|
||
|
|
## 🚀 未来扩展
|
||
|
|
|
||
|
|
### 1. 节点市场
|
||
|
|
|
||
|
|
```
|
||
|
|
- 节点分享平台
|
||
|
|
- 版本管理
|
||
|
|
- 依赖声明
|
||
|
|
- 一键安装
|
||
|
|
```
|
||
|
|
|
||
|
|
### 2. 更严格的沙箱
|
||
|
|
|
||
|
|
```python
|
||
|
|
# 使用 RestrictedPython
|
||
|
|
from RestrictedPython import compile_restricted
|
||
|
|
|
||
|
|
code = compile_restricted(user_code)
|
||
|
|
```
|
||
|
|
|
||
|
|
### 3. GPU 支持
|
||
|
|
|
||
|
|
```python
|
||
|
|
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. **生产使用时**
|
||
|
|
- 定期备份节点文件
|
||
|
|
- 记录节点版本号
|
||
|
|
- 监控执行时间
|
||
|
|
|
||
|
|
---
|
||
|
|
|
||
|
|
**设计理念**: 在保证安全性的前提下，最大化用户的创作自由度。
|